/*
 * 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


/*

 TPM Quote

 Usage:
  tpm_quote --uuid UUID 
 

 TODO
 - No(Null) auth for SRK 
 - No(Null) auth for AIK/SignKey 

 */

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

/* options */
const struct option long_option[] = {
	{"help", no_argument, NULL, 'h'},
	{"uuid", required_argument, NULL, 'u'},
	{"nonce", required_argument, NULL, 'n'},
	{"pcrindex", required_argument, NULL, 'p'},
	{"output", required_argument, NULL, 'o'},
	{"popup", no_argument, NULL, 'P'},
	{"noauth", no_argument, NULL, 'N'},
	{"auth", required_argument, NULL, 'a'},
	{"system", no_argument, NULL, 'S'},
	{"user", no_argument, NULL, 'U'},
	{"blob", required_argument, NULL, 'B'},
	{0, 0, 0, 0}
};
const char short_option[] = "u:p:n:o:Pa:NSUB:";

void usage() {
	printf("Usage: tpm_quote [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\n");

	printf("\t-n, --nonce NONCE\n");
	printf("\t\tSet NONCE\n");

	printf("\t-p, --pcr NUMBER\n");
	printf("\t\tPCR to quote to.  Default is none.  This option can be specified multiple times to choose more than one PCR.\n");

	printf("\t-o, --output FILE\n");
	printf("\t\tFilename to write quote result to.  Default is STDOUT.\n");

	/* Key Auth */
	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");
	
	/* 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");
}

#define KEY_BLOB_SIZE 1024

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

	TSS_HTPM hTPM;

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


	TSS_HKEY hKey;
	TSS_HPOLICY hKeyPolicy;
	TSS_UUID uuid;

	TSS_HPCRS hPcrComposite;
	TSS_VALIDATION validationData;

	BYTE *nonce = NULL;
	int i;

	BYTE *blob;
	UINT32 blobLength;

	BYTE pcr[24];		// TODO

	UINT32 ulSubCapLength;
	UINT32 rgbSubCap;
	UINT32 pulRespDataLength;
	BYTE *prgbRespData;
	UINT32 pcrnum;

	int so;
	int pcrindex;
	int pcrSelectCount = 0;
	char *filename = NULL;
	FILE *fp = stdout;
	int noauth = 1;
	int popup = 0;
	char *auth = NULL;
	
	// Key 
	UINT32 ps_type = TSS_PS_TYPE_SYSTEM; 
	UINT32 keyLength;
	BYTE keyBlob[KEY_BLOB_SIZE];
	
	char *keyFilename;
	FILE *keyFp;
	
	// PW
	BYTE *string = NULL;
	unsigned len = 0;

	memset(pcr, 0, 24);

	/* Connect to TCSD */

	result = Tspi_Context_Create(&hContext);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_Context_Create failed rc=0x%x\n",
		       result);
		if (result == 0x3011 ) {
			printf( " TSS_E_COMM_FAILURE. tcsd is not running?\n");
		}
		
		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;
	}

	/* Get TPM handle */

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

	/* Get PCR Num */

	ulSubCapLength = sizeof(UINT32);
	rgbSubCap = TSS_TPMCAP_PROP_PCR;

	result = Tspi_TPM_GetCapability(hTPM,
					TSS_TPMCAP_PROPERTY,
					ulSubCapLength,
					(BYTE *) & rgbSubCap,
					&pulRespDataLength, &prgbRespData);

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

	pcrnum = (UINT32) * prgbRespData;
	//fprintf(fp,"pcrnums=%d\n",pcrnum);


	/* PCR Composite */

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


	/* OK, now we parse the option args */

	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);
			break;

		case 'n':	/* nonce */
			if (strlen(optarg) != 40) {
				printf("ERROR invalid nonce size, %s\n",
				       optarg);
				goto close;
			}
			nonce = malloc(20);
			if (nonce == NULL) {
				printf("ERROR no memory \n");
				goto close;
			}
			hex2bin(nonce, optarg, 40);
			break;

		case 'p':	/* PCR */
			pcrindex = atoi(optarg);
			result =
			    Tspi_PcrComposite_SelectPcrIndex(hPcrComposite,
							     pcrindex);
			if (result != TSS_SUCCESS) {
				printf("ERROR: failed rc=0x%x\n", result);
				goto close;
			}
			pcr[pcrindex] = 1;
			pcrSelectCount ++;
			break;

		case 'o':	/* output file name */
			filename = optarg;
			fp = fopen(filename, "w");
			break;
		case 'P':	/* popup */
			noauth = 0;
			popup = 1;
			break;
		case 'N':	/* noauth */
			noauth = 1;
			break;
		case 'a':	/* auth */
			noauth = 0;
			auth = optarg;
			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;
			keyFilename = optarg;
			break;
		case 'h':	/* help*/
			usage();
			goto close;
		default:
			usage();
			goto close;
		}
	}
	
	if (pcrSelectCount == 0 ) {
		// No PCR is selected 
		printf("ERROR: No PCR is selected for quote\n");
		usage();
		goto close;
	}

	fprintf(fp, "pcrnums=%d\n", pcrnum);



	/* Get SRK handles */

	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(" TSS_E_PS_KEY_NOT_FOUND.\n");
			printf("  check system_ps_file setting in /etc/tcsd.conf. (default is /var/lib/tpm/system.data)\n");
			printf(" If system_ps_file size is zero. it does not contains SRK info \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;
	}


	/* Load AIK or Sign key */
	if (ps_type == 0 ) {
		/* Blob file */
		keyFp = fopen(keyFilename, "r");
		keyLength = fread( keyBlob, 1, KEY_BLOB_SIZE, keyFp );		
		fclose(keyFp);
		
		
		/* Load */
		result = Tspi_Context_LoadKeyByBlob(
			hContext, 
			hSRK, 
			keyLength, 
			keyBlob, 
			&hKey);
		if (result != TSS_SUCCESS) {
			printf
		    	("ERROR: Tspi_Context_LoadKeyByBlob (Key) failed rc=0x%x\n",
			     result);
			goto close;
		}
			
	}
	else {
		/* TSS PS*/
		result = Tspi_Context_LoadKeyByUUID(hContext,
					    ps_type,//TSS_PS_TYPE_SYSTEM,
					    uuid, 
					    &hKey);
		if (result != TSS_SUCCESS) {
			printf
		    	("ERROR: Tspi_Context_LoadKeyByUUID (Key) failed rc=0x%x\n",
			     result);
			goto close;
		}
	}


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

	/* needs auth? */
	//noauth = 0;
	if (noauth != 1) {
		/* we needs auth */
		//printf("Auth?\n");
		if (popup == 1) {
			/* popup - set message */
			// TODO did not work???
			//char *popupMsg = "Signature Key Password";
			// Trousers 0.3 (TSS1.2?) has TSS_UNICODE definition. 
			uint16_t popupMsg[] = {
				'E','n','t','e','r',' ',
				'Y','o','u','r',' ',
				'A','t','t','e','s','t','a','t','i','o','n',' ',
				'K','e','y',' ',
				'P','a','s','s','w','o','r','d'};
			//printf("DEBUG popupMsg %s\n",popupMsg);
			result = Tspi_SetAttribData(hKeyPolicy,
						    TSS_TSPATTRIB_POLICY_POPUPSTRING,
						    0,
						    sizeof(popupMsg),//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 {

			if (auth != NULL) {
				int len2;
				len = strlen(auth);
				len2=len;
				string =
				    (BYTE *)
				    Trspi_Native_To_UNICODE((BYTE *) auth,
							    &len);
				/* flash */
				memset(auth,0,len2);

			} else {
				// Ask by cmdline
				char *ps;
				char *ps0;
				char *ps1;
				int size0, size1;
				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);
			}
			result = Tspi_Policy_SetSecret(hKeyPolicy,
						       TSS_SECRET_MODE_PLAIN,
						       len,
						       (BYTE *) string);
		}

	} else {
		// noauth
		printf("No Auth?\n");
		result = Tspi_Policy_SetSecret(hKeyPolicy, TSS_SECRET_MODE_PLAIN, strlen((char *) srk_auth),	// TODO
					       srk_auth);
	}

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


	/* Setup Validation Data Structure */

	if (nonce == NULL) {
		result = Tspi_TPM_GetRandom(hTPM, 20, &blob);
		if (result != TSS_SUCCESS) {
			printf
			    ("ERROR: Tspi_TPM_GetRandom failed rc=0x%x\n",
			     result);
			goto free;
		}
		nonce = malloc(20);
		memcpy(nonce, blob, 20);
		Tspi_Context_FreeMemory(hContext, blob);
	}
	// BUG validationData.ulDataLength = 20;
	validationData.ulExternalDataLength = 20;
	validationData.rgbExternalData = nonce;

	fprinthex(fp, "nonce=", nonce, 20);


	/* Quote */

	result = Tspi_TPM_Quote(hTPM,
				hKey, hPcrComposite, &validationData);
	if (result != TSS_SUCCESS) {
		if (result == 0x01) {
			printf("ERROR: Tspi_TPM_Quote failed rc=0x%x\n",
			       result);
			printf
			    ("       Authorization faild, needs valid password\n");
		} else {
			printf("ERROR: Tspi_TPM_Quote failed rc=0x%x\n",
			       result);
		}
		goto free;
	}

	for (i = 0; i < pcrnum; i++) {
		if (pcr[i] == 1) {
			UINT32 length;	//1,length2;
			BYTE *data;	//1, *data2;
			result =
			    Tspi_PcrComposite_GetPcrValue(hPcrComposite, i,
							  &length, &data);
			if (result != TSS_SUCCESS) {
				printf
				    ("ERROR: Tspi_PcrComposite_GetPcrValue failed rc=0x%x\n",
				     result);
				goto free;
			}

			fprintf(fp, "pcr.%d=", i);
			fprinthex(fp, "", data, length);

			Tspi_Context_FreeMemory(hContext, data);

#if 0
			/* Check with current PCR value */
			result =
			    Tspi_TPM_PcrRead(hTPM, i, &length2, &data2);
			if (result != TSS_SUCCESS) {
				printf
				    ("ERROR: Tspi_TPM_PcrRead failed rc=0x%x\n",
				     result);
				goto free;
			}
			fprintf(fp, "pcr.%d=", i);
			fprinthex(fp, "", data2, length2);
			Tspi_Context_FreeMemory(hContext, data2);
#endif
		}
	}


	fprinthex(fp, "quoteinfo=",
		  validationData.rgbData, validationData.ulDataLength);
	fprinthex(fp, "signature=",
		  validationData.rgbValidationData,
		  validationData.ulValidationDataLength);

	//  PubKey ?
	result = Tspi_GetAttribData(hKey,	// PcrComposite, 
				    TSS_TSPATTRIB_KEY_BLOB,
				    TSS_TSPATTRIB_KEYBLOB_PUBLIC_KEY,
				    &blobLength, &blob);
	if (result != TSS_SUCCESS) {
		printf("ERROR: Tspi_GetAttribData failed rc=0x%x\n",
		       result);
		goto free;
	}

	fprinthex(fp, "pubkey=", blob, blobLength);


	/* Validation */
	// TODO 


      free:
	Tspi_Context_FreeMemory(hContext, NULL);
	Tspi_Context_CloseObject(hContext, hPcrComposite);
	Tspi_Context_CloseObject(hContext, hKey);
	free(nonce);

	/* Close TSS/TPM */
      close:
	Tspi_Context_Close(hContext);
	memset(string,0,len);

	fclose(fp);

	return result;
}
