#include <stdio.h>
#include <string.h>
#include "include/cc_helper.h"

TSS_HCONTEXT hContext;

TSS_RESULT tssConnect(){

	TSS_RESULT result;
	//Create context
	result = Tspi_Context_Create(&hContext);
	if (result != TSS_SUCCESS) {
		printf("@tssConnect@ context create error\n");
		return -1;
	}
	//Connect context
	result = Tspi_Context_Connect(hContext, NULL);
	if (result != TSS_SUCCESS) {
		printf("@tssConnect@ context connect error\n");
		return -1;
	}
	return TSS_SUCCESS;
}

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);
	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 verifyPubKeyDigest(BYTE* pubKeyDigest){
	FILE *fp;
	int size,i;
	int flag=-1;
	unsigned char pdmPubKey[MAX_SIZE], md[DIGEST_SIZE];

	TSS_HKEY encKey;
	TSS_RESULT result;
	UINT32 pubKeyLength;
	BYTE *pubKeyModulus;

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

	strcpy(pdmPubKey,"");
	strcpy(md,"");

	//read pdm public key
	fp=fopen(PDM_PUB_KEY_cc,"rb");
	size=fread(pdmPubKey,1,300,fp);

	//Create Object
	result = Tspi_Context_CreateObject(hContext, TSS_OBJECT_TYPE_RSAKEY, TSS_KEY_TYPE_BIND, &encKey);

	if (result != TSS_SUCCESS) {
		printf("@verifyPubKeyDigest@ create pdm key object error\n");
		return -1;
	}

	//load the public key into encKey
	result = Tspi_SetAttribData(encKey,TSS_TSPATTRIB_KEY_BLOB, TSS_TSPATTRIB_KEYBLOB_PUBLIC_KEY, size, pdmPubKey);
	if (result != TSS_SUCCESS) {
		printf("@verifyPubKeyDigest@ set pdm key attribute error\n");
		return -1;
	}

	//get the modulus
	//note that pub key digest is calculated only on modulus as per TSS spec
	result= Tspi_GetAttribData(encKey,TSS_TSPATTRIB_RSAKEY_INFO, TSS_TSPATTRIB_KEYINFO_RSA_MODULUS, &pubKeyLength, &pubKeyModulus);
	if (result != TSS_SUCCESS) {
		printf("@verifyPubKeyDigest@ get pdm key modulus attribute error\n");
		return -1;
	}

	//calculate digest
	SHA1(pubKeyModulus,pubKeyLength,md);

	//verify calculated digest with the expected digest from certify info
	printf("@verifyPubKeyDigest@ digest of pdm pub key modulus: ");
	for(i=0;i<DIGEST_SIZE;i++){
		printf("%02x ", md[i]);
		if(md[i]!=pubKeyDigest[i])
			flag=0;
	}
	printf("\n");

	if (flag==0){
		printf("@verifyPubKeyDigest@ PDM pub key digest mismatch...exiting...\n");
		free(pubKeyModulus);
		return -1;
	}

	free(pubKeyModulus);
	printf("***********Leaving verifyPubKeyDigest*************\n");
	return TSS_SUCCESS;

}

TSS_RESULT verifyCertifyInfo(TCPA_CERTIFY_INFO* tci, BYTE* vd, char* gxgy){

	TSS_RESULT result;
	TCPA_PCR_INFO* pcr;

	UINT64 offset=0;
	int i;
	int flag=0;
	char expectedNonce[20]="";
	char certify_key_nonce[20]="";

	printf("****************Entering verifyCertifyInfo*****************\n");
	strcpy(expectedNonce,gxgy);
	strcpy(certify_key_nonce, CERTIFY_KEY_NONCE);
	for(i=0;i<strlen(gxgy);i++)
		expectedNonce[i]=gxgy[i];

	for(i=strlen(gxgy);i<NONCE_LENGTH;i++)
		expectedNonce[i]=certify_key_nonce[i];

	//Unload TCPA_CERTIFY_INFO structure from TPM --need to unserialize and translate endianess
	result=Trspi_UnloadBlob_CERTIFY_INFO(&offset, vd, tci);
	if(result!=TSS_SUCCESS){
		printf("@verifyCertifyInfo@ Unload certify info failed...exiting...\n");
		return -1;
	}

	//display contents of TCPA_CERTIFY_INFO
	printf("@verifyCertifyInfo@ *TCPA_CERTIFY_INFO contents...\n");
	printf("@verifyCertifyInfo@ *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("@verifyCertifyInfo@ *key usage: %x\n",(UINT32)tci->keyUsage);
	printf("@verifyCertifyInfo@ *key flags: %x\n",(UINT32)tci->keyFlags);
	printf("@verifyCertifyInfo@ *auth data usage: %x\n", tci->authDataUsage);
	printf("@verifyCertifyInfo@ *Alg params:: \n");
	printf("@verifyCertifyInfo@ *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("@verifyCertifyInfo@ *pub key digest: ");
	for(i=0;i<DIGEST_SIZE;i++)
		printf("%02x ",tci->pubkeyDigest.digest[i]);
	printf("\n");

	printf("@verifyCertifyInfo@ *nonce: ");
	for(i=0;i<NONCE_LENGTH;i++)
		printf("%c ",tci->data.nonce[i]);
	printf("\n");
	
	printf("@verifyCertifyInfo@ *parent pcr status: %d\n",tci->parentPCRStatus);

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

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


	//Verify certified PDM pub key freshness
	for(i=0;i<strlen(CERTIFY_KEY_NONCE);i++)
		if(tci->data.nonce[i]!=expectedNonce[i]){
			printf("@verifyCertifyInfo@ certify key replay...exiting...\n");
			return -1;
		}
	printf("@verifyCertifyInfo@ *fresh certified PDM key received...\n");

	//Verify PDM pub key properties
	printf("@verifyCertifyInfo@ *verifying PDM key properties...\n");

	//Verify if PCR is trustworthy (as listed in cc.h)
	for(i=0;i<DIGEST_SIZE;i++){
		if(trustedPCR_01[i]!=pcr->digestAtRelease.digest[i]){
			flag=-1;
		}
	}

	if(flag==-1){
		flag=0;
		for(i=0;i<DIGEST_SIZE;i++){
			if(trustedPCR_02[i]!=pcr->digestAtRelease.digest[i]){
				flag=-2;
			}
		}
	}

	if(flag!=0){
		printf("@verifyCertifyInfo@ untrustworthy PCR...exiting...\n");
		return -1;
	}
	else printf("@verifyCertifyInfo@ *trustworthy PCR...\n");

	//Verify if selected pcrs are appropriate
	if (*((UINT16*)pcr->pcrSelection.pcrSelect)!=TRUSTED_PCR_SELECTION){
		printf("@verifyCertifyInfo@ *untrustworthy pcr selection...exiting...\n");
		return -1;
	}
	else printf("@verifyCertifyInfo@ *trustworthy pcr selection...\n");

	//Verify PDM key usage
	if((UINT32)tci->keyUsage != TRUSTED_KEY_USAGE){
		printf("@verifyCertifyInfo@ *untrustworthy key usage...exiting...\n");
		return -1;
	}
	else printf("@verifyCertifyInfo@ *trustworthy key usage...\n");

/*
	//Verify PDM key flags (migratable/not)
	if((UINT32)tci->keyFlags != TRUSTED_KEY_FLAGS){
		printf("@verifyCertifyInfo@ *untrustworthy key flags...exiting...\n");
		printf("@verifyCertifyInfo@ *keyFlags read: %x\n",(UINT32)tci->keyFlags);
		return -1;
	}
	else printf("@verifyCertifyInfo@ *trustworthy key flag...\n");
*/
	//Verify PDM key flags (migratable/not)
	printf("@verifyCertifyInfo@ *keyFlags read: %x\n",(UINT32)tci->keyFlags);
	if((UINT32)tci->keyFlags == 0x00000002 | (UINT32)tci->keyFlags == 0x00000003 | (UINT32)tci->keyFlags == 0x00000006){
		printf("@verifyCertifyInfo@ *untrustworthy key flags...exiting...\n");
		return -1;
	}
	else printf("@verifyCertifyInfo@ *trustworthy key flag...\n");

	//Verify digest of pdm pub key
	result=verifyPubKeyDigest(tci->pubkeyDigest.digest);
	if(result!=TSS_SUCCESS){
		printf("@verifyCertifyInfo@ verify pub key failed...exiting...\n");
		return -1;
	}
	else printf("@verifyCertifyInfo@ *valid pub key digest...\n");

	printf("@verifyCertifyInfo@ *everything looks good so far...\n");

	printf("****************Leaving verifyCertifyInfo*****************\n");

	return TSS_SUCCESS;

}

TSS_RESULT verifyCertifiedKey(BYTE* vd, int vd_size){

	TSS_HHASH hHash;
	TSS_HKEY verifyKey;
	TSS_RESULT result;
	BYTE signKey[MAX_SIZE], ck[MAX_SIZE];
	unsigned char digest[DIGEST_SIZE];

	FILE *fp;
	int i,c,size;

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

	//open received verification data of the certified key
	fp=fopen(VERIFICATION_DATA_cc,"rb");
	fread(vd,vd_size,1,fp);
	fclose(fp);
		
	//calculate hash of TCPA_CERTIFY_INFO (vd)
	SHA1(vd, vd_size, digest);
	printf("@verifyCertifiedKey@ digest of tcpa_certify_info: ");
	for (i=0;i<DIGEST_SIZE;i++)
		printf("%02x ", digest[i]);
	printf("\n");

	// create hash object
	result = Tspi_Context_CreateObject(hContext, TSS_OBJECT_TYPE_HASH, TSS_HASH_SHA1, &hHash);
	if (result != TSS_SUCCESS) {
		printf("@verifyCertifiedKey@ create hash object error\n");
		return -1;
	}
	//Set digest as the hash value on hHash
	result = Tspi_Hash_SetHashValue(hHash, DIGEST_SIZE, digest);
	if (result != TSS_SUCCESS) {
		printf("@verifyCertifiedKey@ set hash value error\n");
		return -1;
	}

	//get the public part of signKey (AIK)
	fp=fopen(SIGN_KEY_cc,"rb");
	size=fread(signKey,1,500,fp);
	fclose(fp);
	
	//Create Object
	result = Tspi_Context_CreateObject(hContext, TSS_OBJECT_TYPE_RSAKEY, TSS_KEY_TYPE_SIGNING, &verifyKey);
	if (result != TSS_SUCCESS) {
		printf("@verifyCertifiedKey@ create sign key object error\n");
		return -1;
	}

	//load the public key into verifyKey
	result = Tspi_SetAttribData(verifyKey,TSS_TSPATTRIB_KEY_BLOB, TSS_TSPATTRIB_KEYBLOB_PUBLIC_KEY, size, signKey);
	if (result != TSS_SUCCESS) {
		printf("@verifyCertifiedKey@ set sign key attribute error\n");
		return -1;
	}

	//read the certified key
	fp=fopen(CERTIFIED_KEY_cc,"rb");
	size=fread(ck,1,300,fp);
	fclose(fp);

	//verify signature on certified key
	result = Tspi_Hash_VerifySignature(hHash, verifyKey, size, ck);
	if (result != TSS_SUCCESS) {
		printf("@verifyCertifiedKey@ invalid signature!\n");
		return -1;
	}
	else printf("@verifyCertifiedKey@ valid Signature on certified PDM key...\n");

	printf("***********Leaving verifyCertifiedKey*************\n");
	return TSS_SUCCESS;
}

TSS_RESULT Bind(){
	TSS_HKEY encKey;
	TSS_RESULT result;
	TSS_HENCDATA hEncData;

	UINT32 GROUPkeylength=strlen(GROUP_KEY);

	BYTE *pk;
	BYTE *BLOB;
	UINT32 BlobLength;

	FILE *fp;
	int i,c;
	int size=0;

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

	//read pdm pub key
	fp=fopen(PDM_PUB_KEY_cc,"rb");
	pk=malloc(MAX_SIZE);
	size=fread(pk,1,MAX_SIZE,fp);
	fclose(fp);

	//Create Object
	result = Tspi_Context_CreateObject(hContext, TSS_OBJECT_TYPE_RSAKEY, TSS_KEY_TYPE_BIND, &encKey);

	if (result != TSS_SUCCESS) {
		printf("@Bind@ create pdm key object error\n");
		return -1;
	}

	//load the public key into encKey
	result = Tspi_SetAttribData(encKey,TSS_TSPATTRIB_KEY_BLOB, TSS_TSPATTRIB_KEYBLOB_PUBLIC_KEY, size, pk);
	if (result != TSS_SUCCESS) {
		printf("@Bind@ set pdm key attribute error\n");
		return -1;
	}

	//Create hEncData
	result = Tspi_Context_CreateObject(hContext,TSS_OBJECT_TYPE_ENCDATA,TSS_ENCDATA_BIND,&hEncData);
	if (result != TSS_SUCCESS) {
		printf("@Bind@ create bind object creation error\n");
		return -1;
	}

	//Bind GROUP secret
	result = Tspi_Data_Bind(hEncData,encKey,GROUPkeylength,GROUP_KEY);
	if (result != TSS_SUCCESS) {
		printf("@Bind@ bind error...exiting...: %x\n", result);
		return -1;
	}

	printf("@Bind@ Bind TSK key Success!\n");

	//display encrypted data
	result=Tspi_GetAttribData(hEncData,TSS_TSPATTRIB_ENCDATA_BLOB,TSS_TSPATTRIB_ENCDATABLOB_BLOB,&BlobLength,&BLOB);
	if (result != TSS_SUCCESS) {
		printf("@Bind@ get bind(GROUP) key attribute error\n");
		return -1;
	}
	printf("@Bind@ Encrypted GROUP key: %d\n",BlobLength);
	for(i=0;i<BlobLength;i++){
		printf("%x",BLOB[i]);
	}
	printf("\n");

	//write encrypted GROUP key to file
	fp=fopen(ENC_GROUP_KEY,"wb+");
	fwrite(BLOB,BlobLength,1,fp);
	fclose(fp);

	printf("***********Leaving Bind*************\n");

	return TSS_SUCCESS;
}

int verifyAttestaion(int vd_size, char* gxgy){

	TSS_RESULT result;
	TCPA_CERTIFY_INFO *tci;
	BYTE* vd;

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

	//connect to TCSD
	result=tssConnect();
	if(result!=TSS_SUCCESS){
		printf("@verifyAttestaion@ failed to connect to TSS...exiting...\n");
		goto cleanup;
	}

	//verify signature on certified key (PDM key)
	vd=malloc(vd_size);
	result=verifyCertifiedKey(vd, vd_size);
	if(result!=TSS_SUCCESS){
		printf("@verifyAttestaion@ invalid signature on certified key...exiting...\n");
		goto cleanup;
	}

	//verify properties of certified key (PDM key)
	tci = (TCPA_CERTIFY_INFO*)malloc(sizeof(TCPA_CERTIFY_INFO));
	result= verifyCertifyInfo(tci, vd, gxgy);
	if(result!=TSS_SUCCESS){
		printf("@verifyAttestaion@ error in certify info...exiting...\n");
		goto cleanup;
	}

	//Encrypt GROUP secret
	result=Bind();
	if(result!=TSS_SUCCESS){
		printf("@verifyAttestaion@ bind Error...exiting...\n");
		goto cleanup;
	}
	else printf("@verifyAttestaion@ bind Success...\n");

	//cleanup
	Tspi_Context_Close(hContext);
	free(vd);

	printf("***********Leaving verifyAttestation*************\n");
	return 1;

cleanup:
	Tspi_Context_Close(hContext);
	free(vd);
	printf("***********Leaving verifyAttestation*************\n");
	return -1;

}