/*
 * The Initial Developer of the Original Code is International
 * Business Machines Corporation. Portions created by IBM
 * Corporation are Copyright (C) 2007 International Business
 * Machines Corporation. All Rights Reserved.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the Common Public License as published by
 * IBM Corporation; either version 1 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * Common Public License for more details.
 *
 * You should have received a copy of the Common Public License
 * along with this program; if not, a copy can be viewed at
 * http://www.opensource.org/licenses/cpl1.0.php.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>		// getpass

#include <tss/platform.h>
#include <tss/tss_defines.h>
#include <tss/tss_typedef.h>
#include <tss/tss_structs.h>
#include <tss/tss_error.h>
#include <tss/tspi.h>

// for UTF16 PW
#include <trousers/tss.h>
#include <trousers/trousers.h>

#include <getopt.h>


// Local TCSD
#define SERVER    NULL


/*

 Create Sign Key under SRK

 Usage:
  tpm_createkey --uuid UUID --type sign
 

 */


/* options */
const struct option long_option[] = {
	{"uuid", required_argument, NULL, 'u'},
	{"force", no_argument, NULL, 'f'},
	{"list", no_argument, NULL, 'l'},
	{"noauth", no_argument, NULL, 'N'},
	{"auth", required_argument, NULL, 'a'},
	{"popup", no_argument, NULL, 'P'},
	{"type", required_argument, NULL, 't'},
	{"help", no_argument, NULL, 'h'},
	{"system", no_argument, NULL, 'S'},
	{"user", no_argument, NULL, 'U'},
	{"blob", required_argument, NULL, 'B'},
	{0, 0, 0, 0}
};
const char short_option[] = "u:flNPt:a:hSUB:";

int hex2bin(void *dest, const void *src, size_t n);
void printhex(char *str, unsigned char *buf, int len);

void usage() {
	printf("Usage: tpm_createkey [options]\n");
	printf("\t-h, --help\n");
	printf("\t\tDisplay command usage info.\n");
	
	printf("\t-u, --uuid UUID\n");
	printf("\t\tSet UUID of key. Default is randum number\n");
	
	printf("\t-N, --noauth\n");
	printf("\t\tCreate key without auth secret\n");
	
	printf("\t-a, --auth PASSWORD\n");
	printf("\t\tCreate key with auth secret, PASSWORD\n");
	
	printf("\t-p, --popup\n");
	printf("\t\tUse TSS diaglog to set the authsecret, PASSWORD\n");
	
	printf("\t-f, --force\n");
	printf("\t\tUpdate the key\n");
	
	/* Key storage */
	printf("\t-S, --system\n");
	printf("\t\tUse SYSTEM_PS\n");
	printf("\t-U, --user\n");
	printf("\t\tUse USER_PS\n");
	printf("\t-B, --blob FILENAME\n");
	printf("\t\tUse blob file\n");
	
	// TODO  Segmentation fault at Tspi_Context_GetRegisteredKeysByUUID
	//printf("\t-l, --list\n");
	//printf("\t\tList existing keys (in tcsd's persistence storage)\n");
}

int main(int argc, char *argv[])
{
	TSS_RESULT result;
	TSS_HCONTEXT hContext;
	TSS_HTPM hTPM;

	TSS_UUID SRK_UUID = TSS_UUID_SRK;
	TSS_HKEY hSRK;
	TSS_HPOLICY hSRKPolicy;
	BYTE srk_auth[] = "";	// No Auth

	TSS_HKEY hKey;
	TSS_HPOLICY hKeyPolicy;
	//BYTE key_auth[] = ""; // No Auth
	TSS_UUID uuid;
	int createUuid = 1;

	TSS_FLAG initFlag = TSS_KEY_SIZE_2048 | TSS_KEY_TYPE_SIGNING;
	BYTE *buf;
	int so;
	int force = 0;
	int list = 0;
	int noauth = 0;
	int popup = 0;
	char *auth = NULL;
	
	UINT32 ps_type = TSS_PS_TYPE_SYSTEM; 
	UINT32 keyLength;
	BYTE *keyBlob;
	int i;
	
	char * filename;
	FILE *fp;

	// PW
	BYTE *string = NULL;
	unsigned len = 0;

	while (1) {
		so = getopt_long(argc, argv, short_option, long_option, 0);
		if (so == -1)
			break;	// END

		switch (so) {
		case 'u':	/* UUID of AIK/SignKey */
			if (strlen(optarg) != 32) {
				printf("ERROR invalid UUID size, %s\n",
				       optarg);
				goto close;
			}
			hex2bin(&uuid, optarg, 32);
			createUuid = 0;
			break;

		case 'f':	/* force */
			force = 1;
			break;
		case 'l':	/* list */
			list = 1;
			break;
		case 'N':	/* noauth */
			noauth = 1;
			break;
		case 'P':	/* popup */
			popup = 1;
			break;
		case 'a':	/* auth */
			noauth = 0;
			auth = optarg;
			break;
		case 't':	/* type */
			// TODO 
			// Sign key
			initFlag =
			    TSS_KEY_SIZE_2048 | TSS_KEY_TYPE_SIGNING;

			break;
		case 'S':	/* SYSTEM_PS */
			ps_type = TSS_PS_TYPE_SYSTEM;
			break;
		case 'U':	/* USER_PS */
			ps_type = TSS_PS_TYPE_USER;
			break;
		case 'B':	/* BLOB */
			ps_type = 0;
			filename = optarg;
			break;
		case 'h':	/* Help */
			usage();
			return 0;
		default:
			usage();
			return -1;
		}
	}

	if (noauth != 1) {
		/* key needs authorization */
		initFlag |= TSS_KEY_AUTHORIZATION;
	}


	/* Open TSS and get TPM/SRK handles */

	result = Tspi_Context_Create(&hContext);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_Context_Create failed rc=0x%x\n",
		       result);
		goto close;
	}

	result = Tspi_Context_Connect(hContext, SERVER);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_Context_Connect failed rc=0x%x\n",
		       result);
		goto close;
	}

	/* Key List (for DEBUG) */

	if (list == 1) {
		UINT32 ulKeyHierarchySize;
		TSS_KM_KEYINFO **ppKeyHierarchy;

		buf = (BYTE *) & SRK_UUID;
		printhex("SRK uuid: ", buf, 16);
		
		result = Tspi_Context_GetRegisteredKeysByUUID(hContext, TSS_PS_TYPE_SYSTEM, 
								  NULL,	//&SRK_UUID,
							      &ulKeyHierarchySize,
							      ppKeyHierarchy);

		if (result != TSS_SUCCESS) {
			printf
			    ("ERROR: Tspi_Context_GetRegisteredKeysByUUID failed rc=0x%x\n",
			     result);
		} else {
			int i;
			TSS_KM_KEYINFO *info = ppKeyHierarchy[0];
			printf("Key number   : %d\n", ulKeyHierarchySize);
			for (i = 0; i < ulKeyHierarchySize; i++) {
				printf("Key %d\n", i);
				buf = (BYTE *) & info->versionInfo;
				printhex(" version     : ", buf, 4);
				buf = (BYTE *) & info->keyUUID;
				printhex(" uuid        : ", buf, 16);
				buf = (BYTE *) & info->parentKeyUUID;
				printhex(" parents uuid: ", buf, 16);

				info = info + 1;
			}
		}
		goto close;
	}

	/* TPM */

	result = Tspi_Context_GetTpmObject(hContext, &hTPM);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_Context_GetTpmObject failed rc=0x%x\n",
		       result);
		goto close;
	}

	/* SRK */

      loadsrk:
	result = Tspi_Context_LoadKeyByUUID(hContext,
					    TSS_PS_TYPE_SYSTEM, SRK_UUID,
					    &hSRK);

	if (result != TSS_SUCCESS) {
		printf
		    ("ERROR: Tspi_Context_LoadKeyByUUID (SRK) failed rc=0x%x\n",
		     result);
		if (result == 0x2020) {
			printf
			    ("Your key storage of tcsd is damaged or missing. \n");
		}
		goto close;
	}

	/* SRK Policy objects */

	result = Tspi_GetPolicyObject(hSRK, TSS_POLICY_USAGE, &hSRKPolicy);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_GetPolicyObject failed rc=0x%x\n",
		       result);
		goto close;
	}

	result = Tspi_Policy_SetSecret(hSRKPolicy, TSS_SECRET_MODE_PLAIN,	// TODO
				       strlen((char *) srk_auth),
				       srk_auth);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_Policy_SetSecret failed rc=0x%x\n",
		       result);
		goto close;
	}


	/* UUID  */
	if (createUuid == 1) {
		result = Tspi_TPM_GetRandom(hTPM, sizeof(TSS_UUID), &buf);
		if (result != TSS_SUCCESS) {
			printf
			    ("ERROR: Tspi_TPM_GetRandom failed rc=0x%x\n",
			     result);
			Tspi_Context_FreeMemory(hContext, NULL);
			goto close;
		}
		memcpy(&uuid, buf, sizeof(TSS_UUID));
		Tspi_Context_FreeMemory(hContext, buf);
	}

	/* */


	/* Create New Key object */

	result = Tspi_Context_CreateObject(hContext,
					   TSS_OBJECT_TYPE_RSAKEY,
					   initFlag, &hKey);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_Context_CreateObject failed rc=0x%x\n",
		       result);
		goto close;
	}

	/* Sign Key Policy objects */

	if (noauth == 0) {
		/* Needs auth */
		// TODO UTF??
		char *ps;
		char *ps0;
		char *ps1;
		int size0, size1;

		//result = Tspi_GetPolicyObject(hKey, TSS_POLICY_USAGE, &hKeyPolicy);
		result =
		    Tspi_Context_CreateObject(hContext,
					      TSS_OBJECT_TYPE_POLICY,
					      TSS_POLICY_USAGE,
					      &hKeyPolicy);

		if (result != TSS_SUCCESS) {
			printf
			    ("ERROR: Tspi_GetPolicyObject failed rc=0x%x\n",
			     result);
			goto close;
		}

		if (popup == 1) {
			/* popup - set message */
			// TODO did not work???
			char *popupMsg = "Signature Key Password";
			//printf("DEBUG popupMsg %s\n",popupMsg);
			result = Tspi_SetAttribData(hKeyPolicy,
						    TSS_TSPATTRIB_POLICY_POPUPSTRING,
						    0,
						    strlen(popupMsg),
						    (BYTE *) popupMsg);

			if (result != TSS_SUCCESS) {
				printf
				    ("ERROR: Tspi_SetAttribData failed rc=0x%x\n",
				     result);
				goto close;
			}

			/* popup - go */
			result = Tspi_Policy_SetSecret(hKeyPolicy,
						       TSS_SECRET_MODE_POPUP,
						       0, NULL);
		} else {
			// PW 

			if (auth == NULL) {
				// ask 
				ps = getpass("Enter Key password: ");
				size0 = strlen(ps);
				ps0 = malloc(size0 + 1);
				ps0[size0] = 0;
				memcpy(ps0, ps, size0);
				ps1 = getpass("Confirm password: ");
				size1 = strlen(ps1);
				printf("PW %s %s\n", ps0, ps1);

				if (size0 != size1) {
					printf
					    ("Passwords didn't match %d %d\n",
					     size0, size1);
					free(ps0);
					goto close;
				}

				if (strncmp(ps0, ps1, size0) != 0) {
					printf
					    ("Passwords didn't match %d\n",
					     strncmp(ps0, ps1, size0));
					free(ps0);
					goto close;
				}
				
				len = strlen(ps1);
				string =
				    (BYTE *)
				    Trspi_Native_To_UNICODE((BYTE *) ps1,
							    &len);

				/* flash */
				memset(ps0,0,size0);
				memset(ps1,0,size1);
				free(ps0);
			} else {
				// commandine 
				int len2;
				len = strlen(auth);
				len2=len;
				string =
				    (BYTE *)
				    Trspi_Native_To_UNICODE((BYTE *) auth,
							    &len);
				/* flash */
				memset(auth,0,len2);
			}

			result = Tspi_Policy_SetSecret(hKeyPolicy,
						       TSS_SECRET_MODE_PLAIN,
						       len,
						       (BYTE *) string);
			// TODO UTF?                    
			//result = Tspi_Policy_SetSecret(hKeyPolicy, 
			//                                      TSS_SECRET_MODE_PLAIN,  // TODO
			//                             size0,
			//                             (BYTE *) auth);
		}

		if (result != TSS_SUCCESS) {
			printf
			    ("ERROR: Tspi_Policy_SetSecret failed rc=0x%x\n",
			     result);
			goto close;
		}

		result = Tspi_Policy_AssignToObject(hKeyPolicy, hKey);
		if (result != TSS_SUCCESS) {
			printf
			    ("ERROR: Tspi_Policy_SetSecret failed rc=0x%x\n",
			     result);
			goto close;
		}
	}

	result = Tspi_Key_CreateKey(hKey, hSRK, 0);

	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_Key_CreateKey failed rc=0x%04x\n",
		       result);
		goto close;
	}

	/* RegisterKey */
	
	if (ps_type == 0) {
		/* save as blob */
		fp = fopen(filename, "w");
		
		result = Tspi_GetAttribData(
			hKey, 
			TSS_TSPATTRIB_KEY_BLOB, 
			TSS_TSPATTRIB_KEYBLOB_BLOB,
            &keyLength, 
            &keyBlob);
		
		if (result != TSS_SUCCESS) {
			printf("ERROR: Tspi_GetAttribData failed rc=0x%04x\n",
		       result);
			fclose(fp);
			goto close;
		}
		
		//printf(" keyblob len = %d \n",keyLength);		
		
		for (i=0;i<keyLength;i++) {
			fprintf(fp,"%c",keyBlob[i]);
		}
		
		fclose(fp);
		
	}
	else {
		/* managed by TSS  */
	
      regkey:
		result = Tspi_Context_RegisterKey(hContext,
					  hKey,
					  ps_type,//TSS_PS_TYPE_SYSTEM,
					  uuid, 
					  TSS_PS_TYPE_SYSTEM,
					  SRK_UUID);

		if (result != TSS_SUCCESS) {
			if (result == 0x2008) {
				if (force == 1) {
					/* delete key */
					TSS_HKEY hKey;
					result =
					    Tspi_Context_UnregisterKey(hContext,
								       ps_type,//TSS_PS_TYPE_SYSTEM,
								       uuid,
								       &hKey);
					if (result != TSS_SUCCESS) {
						printf
						    ("ERROR: Tspi_Context_UnregisterKey failed rc=0x%x\n",
						     result);
					} else {
						/* try again */
						goto regkey;
					}
				} else {
					printf
					    ("ERROR: Tspi_Context_RegisterKey failed rc=0x%x\n",
					     result);
					printf
					    ("       TSS_E_KEY_ALREADY_REGISTERED\n");
					buf = (BYTE *) & uuid;
					printhex("       uuid=", buf, 16);
				}
			} else {
				printf
				    ("ERROR: Tspi_Context_RegisterKey failed rc=0x%x\n",
				     result);
			}
	
	
			goto close;
		} else {
			// OK
		}
	} // ps_type
	
	/* Close TSS/TPM */

      close:
	Tspi_Context_Close(hContext);
	memset(string,0,len);
	
	return result;
}
