#include "ROMprotocol.h"


DIMASTATUS ROMprotocol()
{

	printf("Start DICE protocol...\n");

	DIMASTATUS ret = 0;
	int len = 0;

	if(DEBUG)
	{printf("Calculating device secrets ...\n");}


	//Calculate DIMA RTM hash
	//RTM, i.e this file measures itself 

	uint8_t* RTM_ID = calloc(1,sizeof(uint8_t)*SHA256_DGST_SIZE);
   	FILE *rtm = NULL;
   	rtm = fopen("ROMprotocol.c", "rb");
   	if(!rtm)
	{
		perror("DIMAFILENOTFOUND: Unable to find RTM(?)\n");
		
		fclose(rtm);
		exit(DIMAFILENOTFOUND);
	}

	mbedtls_sha256_context RTM_ctx;
	mbedtls_sha256_init(&RTM_ctx);

	ret = mbedtls_sha256_starts_ret(&RTM_ctx,0);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE : Failed to create hash context\n");
		mbedtls_sha256_free(&RTM_ctx);
		exit(DIMASHAFAILURE);
	}

	uint8_t *RTMbuf = calloc(1, sizeof(uint8_t)*1024);

	while( (ret = fread(RTMbuf, 1024, 1, rtm))&& (ret != EOF))
	{

		ret = mbedtls_sha256_update_ret(&RTM_ctx, RTMbuf, 1024);
		if(ret < DIMASUCCESS)
		{
			perror("DIMASHAFAILURE : Failed to measure RTM\n");
			mbedtls_sha256_free(&RTM_ctx);
			exit(DIMASHAFAILURE);
		}
		memset(RTMbuf,0,1024);

	}

	ret = mbedtls_sha256_update_ret(&RTM_ctx, RTMbuf, ret);
		if(ret < DIMASUCCESS)
		{
			perror("DIMASHAFAILURE : Failed to measure RTM\n");
			mbedtls_sha256_free(&RTM_ctx);
			exit(DIMASHAFAILURE);
		}

	ret = mbedtls_sha256_finish_ret(&RTM_ctx,RTM_ID);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE : Failed to measure RTM\n");
		mbedtls_sha256_free(&RTM_ctx);
		exit(DIMASHAFAILURE);
	}

	fclose(rtm);
	free(RTMbuf);
	mbedtls_sha256_free(&RTM_ctx);

	if(DEBUG)
	{
  		for (int i = 0; i < SHA256_DGST_SIZE; i++)
  			printf("%hhx", RTM_ID[i]);
		printf(" :  RTM ID\n" );
	}


	//Calculating UDS hash

	FILE *fp = NULL;
	fp = fopen("SecureStorage/RANDFILE", "rb");
	if(!fp)
	{
		perror("DIMAFILENOTFOUND: Unable to access UDS\n");
		printf("For the same of testing, please run >> \
			\n dd if=/dev/urandom of=SecureStorage/RANDFILE bs=256 count=1 \n and retry ..\n");
		fclose(fp);
		exit(DIMAFILENOTFOUND);
	}



	uint8_t *UDSbuf = calloc(1, sizeof(uint8_t)*SHA256_DGST_SIZE);
	fread(UDSbuf,UDS_SIZE,1,fp);
	fclose(fp);

	uint8_t* UDS_ID = calloc(1,sizeof(uint8_t)*SHA256_DGST_SIZE);
	ret = mbedtls_sha256_ret( UDSbuf,UDS_SIZE,UDS_ID,0 );
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILUE:\n");
		
		free(UDSbuf);
		free(UDS_ID);
		exit(DIMAFAILUREUNKWN);
	}

	free(UDSbuf);
	//UDS_ID contains the UDS hash
	//Printer
if(DEBUG)
{
		for (int i = 0; i < UDS_SIZE; i++)
  			printf("%hhx", UDSbuf[i]);
		printf(" :	fuse secret\n" );  	

  		for (int i = 0; i < SHA256_DGST_SIZE; i++)
  			printf("%hhx", UDS_ID[i]);
		printf(" :  UDS ID\n" );
}

	//Calculating CDI SHA256

	uint8_t* CDI = calloc(1,sizeof(uint8_t)*SHA256_DGST_SIZE);
	mbedtls_sha256_context CDI_ctx;
	mbedtls_sha256_init(&CDI_ctx);

	ret = mbedtls_sha256_starts_ret(&CDI_ctx,0);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE\n");
		mbedtls_sha256_free(&CDI_ctx);
		exit(DIMASHAFAILURE);
	}

	ret = mbedtls_sha256_update_ret(&CDI_ctx, UDS_ID, SHA256_DGST_SIZE);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE\n");
		mbedtls_sha256_free(&CDI_ctx);
		exit(DIMASHAFAILURE);
	}

	ret = mbedtls_sha256_update_ret(&CDI_ctx, RTM_ID, SHA256_DGST_SIZE);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE\n");
		mbedtls_sha256_free(&CDI_ctx);
		exit(DIMASHAFAILURE);
	}

	ret = mbedtls_sha256_finish_ret(&CDI_ctx,CDI);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE\n");
		mbedtls_sha256_free(&CDI_ctx);
		exit(DIMASHAFAILURE);
	}

	free(UDS_ID);
	free(RTM_ID);
	mbedtls_sha256_free(&CDI_ctx);

	if(DEBUG)
	{
		for(int i = 0; i < SHA256_DGST_SIZE; i++)
	    	printf("%hhx",CDI[i]);
	    printf(" :  CDI\n");
	}

	//Calculating CDI key, HKDF

	const mbedtls_md_info_t * md_info;
	if(!(md_info = mbedtls_md_info_from_type(HKDF_ALG)))
	{
		perror("DIMAHKDFFAILURE: MD alg type def failed\n");
		exit(DIMAHKDFFAILURE);
	}

	uint8_t salt[32] = {	0x30,0xe2,0x3e,0xcc,0x28,0xc5,0x7b,0xbb,
							0x38,0x7d,0xe6,0x66,0xbb,0xf1,0xd9,0x1e,
    						0xbe,0x67,0x0a,0xf8,0xf3,0x92,0x0e,0xba,
    						0x68,0xd1,0x56,0xea,0x34,0x3f,0xbc,0x4f		};

    uint8_t * CDIKEY = calloc(1, sizeof(uint8_t)*HKDF_KEY_SIZE);
	ret = mbedtls_hkdf(md_info, salt, sizeof(salt), CDI, SHA256_DGST_SIZE,
              IDENTITY, sizeof(IDENTITY), CDIKEY, HKDF_KEY_SIZE);

	if(ret < DIMASUCCESS)
	{
		perror("DIMAHKDFFAILURE\n");
		free(CDI);
		free(CDIKEY);
		exit(DIMAHKDFFAILURE);
	}

	if(DEBUG)
	{
		for(int i = 0; i < HKDF_KEY_SIZE; i++)
	      printf("%hhx",CDIKEY[i]);
		printf(" :  CDIKEY\n");
	}
	
	//Every key derivation should start with a new KD contxt
	//setting context for DID

	//creating scope to reduce access to key derivation context
	//is there a better way to do this?
	{
	KeyDrv_context DID_ctx;
	DID_ctx.ENT_MODE = DETERM;
	DID_ctx.PKC_MODE = DFL_PKC;
	DID_ctx.seed = CDIKEY;
	DID_ctx.phrase = IDENTITY;
	DID_ctx.KEY_FORM = DFL_FORM;
	//DID_ctx.pub_file = DFL_PUB;
	//DID_ctx.priv_file = DFL_PRIV;			//Dont save DID priv outside SS

	//Deriving and storing DID
	ret = AsymmKeyGen(&DID_ctx);
	if(ret < DIMASUCCESS)
	{
		perror("DIMAFAILURE : DID key gen failed\n");
		free(CDIKEY);
		exit(DIMAFAILURE);
	}
	}

	free(CDIKEY);			//But CDIKey is needed to gen AliasKP 
	//delete DID_ctx??

	//setting context for Alias keys

//////////////////////////deriving alisa keys


	//FW_ID = #(MututalAuthentication.c)
	//conf_ID = #(define.h)
	if(DEBUG) printf("Measuring Firmware and config files ...\n");
	uint8_t * FW_ID = calloc(1,sizeof(uint8_t)*SHA256_DGST_SIZE);

	mbedtls_sha256_context FW_ctx;
	mbedtls_sha256_init(&FW_ctx);

	FILE *fw = NULL;
   	fw = fopen("MutualAttestation.c", "rb");
   	if(!fw)
	{
		perror("DIMAFILENOTFOUND: Unable to find stage 2 files(?)\n");
		
		fclose(fw);
		exit(DIMAFILENOTFOUND);
	}
	
	ret = mbedtls_sha256_starts_ret(&FW_ctx,0);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE : Failed to create hash context\n");
		mbedtls_sha256_free(&FW_ctx);
		exit(DIMASHAFAILURE);
	}

	uint8_t *FWbuf = calloc(1, sizeof(uint8_t)*1024);

	while( (ret = fread(FWbuf, 1024, 1, fw))&& (ret != EOF))
	{

		ret = mbedtls_sha256_update_ret(&FW_ctx, FWbuf, 1024);
		if(ret < DIMASUCCESS)
		{
			perror("DIMASHAFAILURE : Failed to measure Stage2 files\n");
			mbedtls_sha256_free(&FW_ctx);
			exit(DIMASHAFAILURE);
		}
		memset(FWbuf,0,1024);

	}

	ret = mbedtls_sha256_update_ret(&FW_ctx, FWbuf, ret);
		if(ret < DIMASUCCESS)
		{
			perror("DIMASHAFAILURE : Failed to measure RTM\n");
			mbedtls_sha256_free(&FW_ctx);
			exit(DIMASHAFAILURE);
		}

	ret = mbedtls_sha256_finish_ret(&FW_ctx,FW_ID);
	if(ret < DIMASUCCESS)
	{
		perror("DIMASHAFAILURE : Failed to measure RTM\n");
		mbedtls_sha256_free(&FW_ctx);
		exit(DIMASHAFAILURE);
	}

	fclose(fw);
	free(FWbuf);
	mbedtls_sha256_free(&FW_ctx);

	if(DEBUG)
	{
  		for (int i = 0; i < SHA256_DGST_SIZE; i++)
  			printf("%hhx", FW_ID[i]);
		printf(" :  FW ID\n" );
	}



	//defines.h is volatile right now! 
	//changing define, hence alias key adds 3-4 steps to per round of testing 
	//include conf measurement only after everything else is fixed
	//should be performed exactly as line 242, above

	//calc conf_ID
	//calc composite hash CDI2 with CDI, conf_ID nad fw_ID
	//derive FWKEY from CDI2


    uint8_t * FWKEY = calloc(1, sizeof(uint8_t)*HKDF_KEY_SIZE);

    //LAZY METHOD - create Alias key derivation material using CDI and seed and FW_ID as salt

    ret = mbedtls_hkdf(md_info, FW_ID, sizeof(FW_ID), CDI, SHA256_DGST_SIZE,
              ALIAS, sizeof(ALIAS), FWKEY, HKDF_KEY_SIZE);

    if(ret < DIMASUCCESS)
	{
		perror("DIMAHKDFFAILURE\n");
		exit(DIMAHKDFFAILURE);
	}

	//creating scope to reduce access to key derivation context
	//is there a better way to do this?
	{
	KeyDrv_context ALIAS_ctx;
	ALIAS_ctx.ENT_MODE = DETERM;
	ALIAS_ctx.PKC_MODE = DFL_PKC;
	ALIAS_ctx.seed = FWKEY;
	ALIAS_ctx.phrase = ALIAS;
	ALIAS_ctx.KEY_FORM = DFL_FORM;

	printf("Generating Alias keys\n");

	ret = AsymmKeyGen(&ALIAS_ctx);
	if(ret < DIMASUCCESS)
	{
		perror("DIMAFAILURE : ALIAS key gen failed\n");
		exit(DIMAFAILURE);
	}
	}

	//session keys?

	free(CDI);
	free(FWKEY);

	return DIMASUCCESS;

}



/////////////////////////