#include <stdio.h>
#include <string.h>

#include "include/trm_helper.h"


TSS_RESULT createKey(char* id1, char* id2){

TSS_HCONTEXT hContext;
TSS_HKEY hSRK;
TSS_HTPM hTPM;
TSS_HKEY TRMkey;
TSS_RESULT result;

	TSS_HPOLICY srkUsagePolicy, TRMkeyUsagePolicy;
	BYTE *rgbPcrValue;
	UINT32 ulPcrLen;
	TSS_HPCRS hPcrComposite;
	TSS_FLAG initFlags;
	int i;
	initFlags = TSS_KEY_TYPE_BIND | TSS_KEY_SIZE_512 | TSS_KEY_VOLATILE | 
			TSS_KEY_AUTHORIZATION | TSS_KEY_NOT_MIGRATABLE;
	UINT32 pulPubKeyLength;
	BYTE *prgbPubKey;
	FILE *fp;

	unsigned int hex1, hex2;

	printf("\n***********Entering createKey************\n");

	sscanf(id1,"%x",&hex1);
	sscanf(id2,"%x",&hex2);
/*

	TRMkey_UUID.ulTimeLow = 0x00000000;
	TRMkey_UUID.usTimeMid = 0x0000;
	TRMkey_UUID.usTimeHigh = 0x0000;
	TRMkey_UUID.bClockSeqHigh = 0x00;
	TRMkey_UUID.bClockSeqLow = 0x00;
	TRMkey_UUID.rgbNode[4]=hex1;
	TRMkey_UUID.rgbNode[5]=hex2;
	printf("@CreateKey@ *TRM key will be registered in system store as: %x %x\n",TRMkey_UUID.rgbNode[4], TRMkey_UUID.rgbNode[5]);
*/
	//Create context
	result = Tspi_Context_Create(&hContext);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ context create error\n");
		return -1;
	}

	//Connect context
	result = Tspi_Context_Connect(hContext, NULL);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ context connect error\n");
		return -1;
	}

	//Create key object to hold TRMkey
	result = Tspi_Context_CreateObject(hContext, TSS_OBJECT_TYPE_RSAKEY,initFlags, &TRMkey);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ create object error\n");
		goto out_close;
	}

	//Load SRK
	result = Tspi_Context_LoadKeyByUUID(hContext,TSS_PS_TYPE_SYSTEM,SRKUUID,&hSRK);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ load key error\n");
		goto out_close;
	}

	//Get SRK usage policy object
	result = Tspi_GetPolicyObject(hSRK, TSS_POLICY_USAGE, &srkUsagePolicy);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ get policy object error\n");
		goto out_close;
	}

	//Set SRK usage policy
	result = Tspi_Policy_SetSecret(srkUsagePolicy,TSS_SECRET_MODE_PLAIN,strlen(SRK_AUTH) ,SRK_AUTH);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ set secret error\n");
		goto out_close;
	}

	//Create policy object for TRMkey
	result = Tspi_Context_CreateObject(hContext, TSS_OBJECT_TYPE_POLICY, TSS_POLICY_USAGE, &TRMkeyUsagePolicy);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ create object error\n");
		goto out_close;
	}

	//Set TRMkey usage policy
	result = Tspi_Policy_SetSecret(TRMkeyUsagePolicy,TSS_SECRET_MODE_PLAIN,strlen(TRM_KEY_AUTH),TRM_KEY_AUTH);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ set secret error\n");
		goto out_close;
	}

	//Assign policy to TRMkey
	result = Tspi_Policy_AssignToObject(TRMkeyUsagePolicy, TRMkey);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ Assign policy to key error\n");
		goto out_close;
	}

	//Create pcr composite object to bind the TRMkey to
	result=Tspi_Context_CreateObject(hContext, TSS_OBJECT_TYPE_PCRS, 0, &hPcrComposite);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ Create Pcr composite object error\n");
		goto out_close;
	}

	//Get TPM handle, read pcr, set pcr value in pcr composite object
	result=Tspi_Context_GetTpmObject(hContext, &hTPM);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ Get TPM object error\n");
		goto out_close;
	}
	result=Tspi_TPM_PcrRead (hTPM, TRM_KEY_PCR_NUM, &ulPcrLen, &rgbPcrValue);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ Pcr read error\n");
		goto out_close;
	}
	result=Tspi_PcrComposite_SetPcrValue(hPcrComposite, TRM_KEY_PCR_NUM, ulPcrLen, rgbPcrValue);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ Pcr composite error\n");
		goto out_close;
	}

	printf("@CreateKey@ *TRMkey will be bound to Pcr number: %d\n",TRM_KEY_PCR_NUM);
	printf("@CreateKey@ *value at Pcr number: ");
	for(i=0;i<ulPcrLen;i++)
		printf("%02x ", rgbPcrValue[i]);
	printf("\n");

	//Create TRMkey and wrap it to TRM_KEY_PCR_NUM under SRK
	result = Tspi_Key_CreateKey(TRMkey, hSRK, hPcrComposite);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ create key error\n");		
		goto out_close;
	}

/*
	result=Tspi_Context_RegisterKey(hContext,TRMkey,TSS_PS_TYPE_SYSTEM,TRMkey_UUID,TSS_PS_TYPE_SYSTEM,SRKUUID);
	if(result!=TSS_SUCCESS){
		printf("@CreateKey@ Register key error\n");
		goto out_close;
	}
*/

	//get the public part of TRM key
	result = Tspi_GetAttribData(TRMkey,TSS_TSPATTRIB_KEY_BLOB, TSS_TSPATTRIB_KEYBLOB_PUBLIC_KEY,&pulPubKeyLength,&prgbPubKey);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ get TRM public key error\n");
		goto out_close;
	}


	fp=fopen(TRM_PUB_KEY_trm,"wb+");
	fwrite(prgbPubKey,pulPubKeyLength,1,fp);
	fflush(fp);
	fclose(fp);

	printf("@CreateKey@ *TRM pub key & length: %d\n",pulPubKeyLength);
	for(i=0;i<pulPubKeyLength;i++)
		printf("%x",prgbPubKey[i]);
	printf("\n");

	printf("***********Leaving createKey************\n");
	return TSS_SUCCESS;

out_close:
	Tspi_Context_FreeMemory(hContext, NULL);
	Tspi_Context_CloseObject(hContext, TRMkey);
	Tspi_Context_Close(hContext);
	printf("\n***********Leaving createKey***************\n");
	return -1;
}

TSS_RESULT Trspi_UnloadBlob_CERTIFY_INFO(UINT64 *offset, BYTE *blob, TCPA_CERTIFY_INFO* out){
	Trspi_UnloadBlob_TSS_VERSION(offset, blob, &out->version);
	Trspi_UnloadBlob_UINT16(offset, &out->keyUsage, blob);
	Trspi_UnloadBlob_KEY_FLAGS(offset, blob, &out->keyFlags);
	Trspi_UnloadBlob_BYTE(offset, &out->authDataUsage, blob);
	Trspi_UnloadBlob_KEY_PARMS(offset, blob, &out->algorithmParms);
	//Trspi_UnloadBlob_DIGEST(offset, blob, &out->pubkeyDigest);
	//Bug in Trspi_UnloadBlob_DIGEST; using unloadblob
	Trspi_UnloadBlob(offset, TCPA_SHA1_160_HASH_LEN, blob, &out->pubkeyDigest);
	Trspi_UnloadBlob(offset, sizeof(TCPA_NONCE), blob, &out->data);
	Trspi_UnloadBlob_BOOL(offset, &out->parentPCRStatus, blob);
	Trspi_UnloadBlob_UINT32(offset, &out->PCRInfoSize, blob);
	out->PCRInfo = malloc(sizeof(TCPA_PCR_INFO)); 
	if(Trspi_UnloadBlob_PCR_INFO(offset, blob, out->PCRInfo)!=TSS_SUCCESS)
		return TSS_E_FAIL;
	return TSS_SUCCESS;
}

TSS_RESULT certifyKey(char* gxgy){

TSS_HCONTEXT hContext;
TSS_HKEY hSRK;
TSS_HKEY TRMkey;
TSS_RESULT result;

	TSS_HKEY signKey;
	TSS_HPOLICY keyUsagePolicy, srkUsagePolicy;
	TSS_VALIDATION pValidationData;
	TSS_HHASH hHash;

	TCPA_CERTIFY_INFO *tci;
	TCPA_PCR_INFO* pcr;

	UINT32 pulSignatureLength;
	BYTE *prgbSignature;
	UINT32 pulPubKeyLength;
	BYTE *prgbPubKey;

	unsigned char digest[20];
	char certify_key_nonce[20]="";
	char nonce[20]="";
	UINT64 offset=0;
	int i;
	FILE *fp;

	printf("***********Entering certifyKey************\n");

	//use the same SIGMA nonce for certify key
	strcpy(certify_key_nonce, CERTIFY_KEY_NONCE);
	for(i=0;i<strlen(gxgy);i++)
		nonce[i]=gxgy[i];
	for(i=strlen(gxgy);i<strlen(CERTIFY_KEY_NONCE);i++)
		nonce[i]=certify_key_nonce[i];


	//Create context
	result = Tspi_Context_Create(&hContext);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ context create error\n");
		return -1;
	}

	//Connect context
	result = Tspi_Context_Connect(hContext, NULL);
	if (result != TSS_SUCCESS) {
		printf("@CreateKey@ context connect error\n");
		return -1;
	}


	//Load SRK By UUID
	result = Tspi_Context_LoadKeyByUUID(hContext,TSS_PS_TYPE_SYSTEM,SRKUUID,&hSRK);
	if (result != TSS_SUCCESS) {
		printf("SRK load key error: %x\n", result);
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		goto out_close;
	}

	//Get Policy Object
	result =
	    Tspi_GetPolicyObject(hSRK, TSS_POLICY_USAGE, &srkUsagePolicy);
	if (result != TSS_SUCCESS) {
		printf("get policy object error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		goto out_close;
	}

	//Set Secret
	result =
	    Tspi_Policy_SetSecret(srkUsagePolicy,TSS_SECRET_MODE_PLAIN,strlen(SRK_AUTH),SRK_AUTH);
	if (result != TSS_SUCCESS) {
		printf("set secret error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		goto out_close;
	}



	//Load the key used to certify
	result = Tspi_Context_GetKeyByUUID(hContext,TSS_PS_TYPE_SYSTEM,signKeyUUID,&signKey);
	if (result != TSS_SUCCESS) {
		printf("@certifyKey@ sign load key error\n");
		goto out_close;
	}
	result = Tspi_Key_LoadKey(signKey, hSRK);
	if (result != TSS_SUCCESS) {
		printf("@certifyKey@ Tspi_Key_LoadKey (signKey), %x\n", result);
		Tspi_Context_FreeMemory(hContext, NULL);
		Tspi_Context_Close(hContext);
		goto out_close;
	}

	//Get policy object
	result=Tspi_GetPolicyObject(signKey, TSS_POLICY_USAGE, &keyUsagePolicy);
	if (result != TSS_SUCCESS) {
		printf("@certifyKey@ get policy object error\n");
		goto out_close;
	}

	//Set Secret
	result=Tspi_Policy_SetSecret(keyUsagePolicy, TSS_SECRET_MODE_PLAIN, 4, "temp");
	if (result != TSS_SUCCESS) {
		printf("@certifyKey@ set secret error\n");
		goto out_close;
	}


	printf("TRMkey_UUID: %x, %x, %x, %x, %x, %x, %x\n", TRMkey_UUID.ulTimeLow, TRMkey_UUID.usTimeMid, TRMkey_UUID.usTimeHigh, TRMkey_UUID.bClockSeqHigh, TRMkey_UUID.bClockSeqLow, TRMkey_UUID.rgbNode[4], TRMkey_UUID.rgbNode[5]);



	//Load SRK By UUID
	result = Tspi_Context_LoadKeyByUUID(hContext,TSS_PS_TYPE_SYSTEM,SRKUUID,&hSRK);
	if (result != TSS_SUCCESS) {
		printf("load key error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		goto out_close;
	}

	//Get Policy Object
	result =
	    Tspi_GetPolicyObject(hSRK, TSS_POLICY_USAGE, &srkUsagePolicy);
	if (result != TSS_SUCCESS) {
		printf("get policy object error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		goto out_close;
	}

	//Set Secret
	result =
	    Tspi_Policy_SetSecret(srkUsagePolicy,TSS_SECRET_MODE_PLAIN,strlen(SRK_AUTH),SRK_AUTH);
	if (result != TSS_SUCCESS) {
		printf("set secret error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		goto out_close;
	}



	//Load the key that is to be certified
	result = Tspi_Context_GetKeyByUUID(hContext,TSS_PS_TYPE_SYSTEM,TRMkey_UUID,&TRMkey);
	if (result != TSS_SUCCESS) {
		printf("@certifyKey@ GetKeyByUUID error: bindKey\n");
		goto out_close;
	}
	result = Tspi_Key_LoadKey(TRMkey, hSRK);
	if (result != TSS_SUCCESS) {
		printf("@certifyKey@ Tspi_Key_LoadKey Error (TRMkey): %4x\n", result);
		Tspi_Context_FreeMemory(hContext, NULL);
		Tspi_Context_Close(hContext);
		goto out_close;
	}

	//Certify TRMkey using AIK
	pValidationData.ulExternalDataLength = strlen(CERTIFY_KEY_NONCE);
	pValidationData.rgbExternalData = nonce;
	result = Tspi_Key_CertifyKey(TRMkey,signKey,&pValidationData);
	if (result != TSS_SUCCESS) {
		printf("@certifyKey@ certify key error: %x\n", result);
		goto out_close;
	}
	
	//Create certified key file
	fp=fopen(CERTIFIED_KEY_trm,"wb+");
	fwrite(pValidationData.rgbValidationData,pValidationData.ulValidationDataLength,1,fp);
	fclose(fp);
	printf("*@certifyKey@ Certify key validation data (signature) %d:\n", pValidationData.ulValidationDataLength);
	for(i=0;i<pValidationData.ulValidationDataLength;i++)
		printf("%x", pValidationData.rgbValidationData[i]);
	printf("\n");

	//Create verification data file
	fp=fopen(VERIFICATION_DATA_trm,"wb+");
	fwrite(pValidationData.rgbData, pValidationData.ulDataLength, 1, fp);
	fclose(fp);

	printf("*@certifyKey@ Certify info for signature verification: %d:\n", pValidationData.ulDataLength);
	for(i=0;i<pValidationData.ulDataLength;i++)
		printf("%x", pValidationData.rgbData[i]);
	printf("\n");

	//Unload TCPA_CERTIFY_INFO structure
	tci = (TCPA_CERTIFY_INFO*)malloc(sizeof(TCPA_CERTIFY_INFO));
	Trspi_UnloadBlob_CERTIFY_INFO(&offset, pValidationData.rgbData, tci);

	//display contents of TCPA_CERTIFY_INFO
	printf("Certified key info:\n");
	printf("@certifyKey@ *version maj: %d, version min: %d, version revmaj: %d, version revmin: %d\n", tci->version.major, tci->version.minor, tci->version.revMajor, tci->version.revMinor);
	printf("\n@certifyKey@ *key usage: %x\n",(UINT32)tci->keyUsage);
	printf("\n@certifyKey@ *key flags: %x\n",(UINT32)tci->keyFlags);
	printf("\n@certifyKey@ *auth data usage: %x\n", tci->authDataUsage);
	printf("\n@certifyKey@ *Alg params:: \n");
	printf("@certifyKey@ *Alg id:  %x, enc scheme: %x, sig scheme: %x; parm size: %d\n",tci->algorithmParms.algorithmID, tci->algorithmParms.encScheme, tci->algorithmParms.sigScheme, tci->algorithmParms.parmSize);

	printf("\n@certifyKey@ *pub key digest:\n");
	for(i=0;i<DIGEST_SIZE;i++)
		printf("%x ",tci->pubkeyDigest.digest[i]);
	printf("\n");

	printf("\n@certifyKey@ *nonce:\n");
	for(i=0;i<strlen(CERTIFY_KEY_NONCE);i++)
		printf("%c ",tci->data.nonce[i]);
	printf("\n");
	
	printf("\n@certifyKey@ *parent pcr status: %d\n",tci->parentPCRStatus);
	printf("\n@certifyKey@ *pcr info for wrapped key::\n");
	printf("\n@certifyKey@ *pcr info size: %d\n",tci->PCRInfoSize);

	pcr=(TCPA_PCR_INFO*) tci->PCRInfo;
	printf("\n@certifyKey@ *PCR digest at release: \n");
	for(i=0;i<DIGEST_SIZE;i++)
		printf("%02x ",pcr->digestAtRelease.digest[i]);
	printf("\n");

	printf("\n@certifyKey@ *PCR selection size: %d, selected pcrs: %x\n", pcr->pcrSelection.sizeOfSelect,*((UINT16*)pcr->pcrSelection.pcrSelect));

	Tspi_Key_UnloadKey(TRMkey);
	Tspi_Key_UnloadKey(signKey);
	Tspi_Context_FreeMemory(hContext, NULL);
	Tspi_Context_CloseObject(hContext, TRMkey);
	Tspi_Context_CloseObject(hContext, signKey);
	Tspi_Context_Close(hContext);
	printf("***********Leaving certifyKey************\n");
	return TSS_SUCCESS;

out_close:
	Tspi_Context_FreeMemory(hContext, NULL);
	Tspi_Context_CloseObject(hContext, TRMkey);
	Tspi_Context_CloseObject(hContext, signKey);
	Tspi_Context_Close(hContext);
	printf("\n***********Leaving certifyKey***************\n");
	return -1;
}

TSS_RESULT unBind(int size){

TSS_HCONTEXT hContext;
TSS_HKEY hSRK;
TSS_HKEY TRMkey;
TSS_RESULT result;

	TSS_HENCDATA hEncData;
	TSS_HPOLICY TRMkeyUsagePolicy,srkUsagePolicy;

	BYTE *encGROUPkey;
	UINT32 pulUnboundDataLength;
	BYTE *prgbUnboundData;
	FILE *fp;
	int i;

	printf("\n***********Entering unBind***************\n");

	//Read encrypted group key from file
	fp=fopen(ENC_GROUP_KEY_trm,"rb");
	encGROUPkey=malloc(size);
	fread(encGROUPkey,size,1,fp);
	fclose(fp);

	//Create context
	result = Tspi_Context_Create(&hContext);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ context create error\n");
		return -1;
	}

	//Connect context
	result = Tspi_Context_Connect(hContext, NULL);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ context connect error\n");
		goto out_close;
	}
	
	//Create and set hEncData to hold encGROUPKey
	result = Tspi_Context_CreateObject(hContext,TSS_OBJECT_TYPE_ENCDATA,TSS_ENCDATA_BIND,&hEncData);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ hEncData object creation error\n");
		goto out_close;
	}
	result = Tspi_SetAttribData(hEncData,TSS_TSPATTRIB_ENCDATA_BLOB, TSS_TSPATTRIB_ENCDATABLOB_BLOB, size, encGROUPkey);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ set attrib error\n");
		goto out_close;
	}


	//Load SRK By UUID
	result = Tspi_Context_LoadKeyByUUID(hContext,TSS_PS_TYPE_SYSTEM,SRKUUID,&hSRK);
	if (result != TSS_SUCCESS) {
		printf("load key error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		exit(1);
	}

	//Get Policy Object
	result =
	    Tspi_GetPolicyObject(hSRK, TSS_POLICY_USAGE, &srkUsagePolicy);
	if (result != TSS_SUCCESS) {
		printf("get policy object error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		exit(1);
	}

	//Set Secret
	result =
	    Tspi_Policy_SetSecret(srkUsagePolicy,TSS_SECRET_MODE_PLAIN,strlen(SRK_AUTH),SRK_AUTH);
	if (result != TSS_SUCCESS) {
		printf("set secret error\n");
		Tspi_Context_CloseObject(hContext, hSRK);
		Tspi_Context_Close(hContext);
		exit(1);
	}

	//Load the TRMkey
	result = Tspi_Context_GetKeyByUUID(hContext,TSS_PS_TYPE_SYSTEM,TRMkey_UUID,&TRMkey);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ load key error: bindKey\n");
		goto out_close;
	}
	result = Tspi_Key_LoadKey(TRMkey, hSRK);
	if (result != TSS_SUCCESS) {
		printf("Tspi_Key_LoadKey, %x\n", result);
		Tspi_Context_FreeMemory(hContext, NULL);
		Tspi_Context_Close(hContext);
		exit(result);
	}
/*
	//Get policy object
	result=Tspi_GetPolicyObject(TRMkey, TSS_POLICY_USAGE, &TRMkeyUsagePolicy);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ get policy object error\n");
		goto out_close;
	}

	//Set trm key Secret
	result=Tspi_Policy_SetSecret(TRMkeyUsagePolicy, TSS_SECRET_MODE_PLAIN, strlen(TRM_KEY_AUTH), TRM_KEY_AUTH);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ set secret error\n");
		goto out_close;
	}
*/
	//Unbind encrypted group key
	result=Tspi_Data_Unbind(hEncData,TRMkey,&pulUnboundDataLength,&prgbUnboundData);
	if (result != TSS_SUCCESS) {
		printf("@unBind@ Unbind error: %x\n",result);
		goto out_close;
	}
	printf("@unBind@ Unbound data (GROUP secret): ");
	for(i=0;i < pulUnboundDataLength;i++)
		printf("%c",prgbUnboundData[i]);
	printf("\n");

	printf("@unBind@ *TRM key id: %x %x\n",TRMkey_UUID.rgbNode[4], TRMkey_UUID.rgbNode[5]);

	Tspi_Key_UnloadKey(TRMkey);
	Tspi_Context_FreeMemory(hContext, NULL);
	Tspi_Context_CloseObject(hContext, TRMkey);
	Tspi_Context_Close(hContext);
	printf("***********Leaving unBind************\n");
	return TSS_SUCCESS;

out_close:
	Tspi_Context_FreeMemory(hContext, NULL);
	Tspi_Context_CloseObject(hContext, TRMkey);
	Tspi_Context_Close(hContext);
	printf("\n***********Leaving unBind***************\n");
	return -1;

}