#include "ROMprotocol.h"
#include "Stage2.h"

int main()
{
	printf("Hello World!\n");
	
	DIMASTATUS ret = 0;

	/*lock all resources, no interruptons*/ 
	/*unlock Secure Storage to read UDS and write DID Priv key*/

	ret = ROMprotocol();
	
	/*lock Secure Storage*/
	

	//ALIAS and session keys, firmware execution

	//measure firmware.bin
	//measure firmware.conf
	//measure ....
	//calculate FW_ID - composite hash of all above measurements
	//create ALIAS KD_ctx
	//keygen
	//publish keys



	//Session key creation
	//receive session nonce from verifier?
	//placeholder 
	char * PERS = "session1";

	KeyDrv_context SSN1_ctx;
	SSN1_ctx.ENT_MODE = SW_PRNG;				//non determ, gen new key for every session
	SSN1_ctx.PKC_MODE = DFL_PKC;
	SSN1_ctx.seed = NULL;					
	SSN1_ctx.phrase = PERS;
	SSN1_ctx.KEY_FORM = DFL_FORM;

	printf("Generating Session keys\n");

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


	/////////////////KEY gen protocol ends here
	/////////////////DID Priv key is purged/secure, other keys are available for use

	/////////////////////DICE END//////////////////////////


	//////////////////////Now STAGE 2 protocols///////////////////


	//Load self info
	//read from xml config file?

	Node_info 		Self;
	Self.ID =		DFL_ID;
	Self.status = 	DFL_CL_STAT;

	printf("size of ID %ld\n", sizeof(Self.ID) );

	//Verifier Node code
	if(Self.ID == MASTER_ID)
	{
		//load node info
		//bind socket
		//accept conn
		//verify cli ID
		//store FD

		//chll(node_info cli)
			//chll ctx
			//cacl NONCE  -> chll, node_info
			//load self.priv key, check
			//sign nonce
			//send chll_ctx

		//recv resp
		//update self.status
		//verify(cli)
			//load cli pub key, check
			//verify
			//update cli.status
			//send ack


		printf("Running as Master node....\n");
		Node_info * nodelist = calloc(1, sizeof(Node_info)*CL_NOS);
		ret = load_nodes(nodelist);

		mbedtls_net_context fd_bind;
		mbedtls_net_init(&fd_bind);

		unsigned char netBuff[NETBUFSIZE] = {0};
		unsigned char netbuf2[2];


		ret = mbedtls_net_bind(&fd_bind, NULL, SERVER_PORT, MBEDTLS_NET_PROTO_TCP);
		if(ret < DIMASUCCESS)
		{
			printf("DIMANETWORKFAILURE : Failed to bind port -0x%04x\n", (unsigned int) -ret);
			return DIMANETWORKFAILURE;
		}

	//loop till all clients verified:
		mbedtls_net_context fd_client;
		mbedtls_net_init(&fd_client);

		printf("Waiting for client ...\n");

		ret = mbedtls_net_accept(&fd_bind, &fd_client, NULL,0,NULL);
		if(ret < DIMASUCCESS)
		{
			printf("DIMANETWORKFAILURE : Failed to connect client -0x%04x\n", (unsigned int) -ret);
			return DIMANETWORKFAILURE;
		}

		printf("Client connected ...\n");

		ret = mbedtls_net_recv(&fd_client, netBuff, sizeof(Self.ID));
		if(ret < DIMASUCCESS)
		{
			printf("DIMANETWORKFAILURE : Failed to recv ID -0x%04x\n", (unsigned int) -ret);
			return DIMANETWORKFAILURE;
		}
		//if(DEBUG)
			printf("recvd %d bytes ID %s \n", ret, netBuff);

		if(strcmp(netBuff, nodelist[1].ID) == DIMASUCCESS)
		{
			printf("Client1 identified. Verifying ...\n");
			nodelist[1].fd = &fd_client;
		}
		else
		{
			printf("Unknown client. abort client. retrying ...\n");
			mbedtls_net_free(&fd_client);
			return DIMAFAILUREUNKWN;
	//continue loop
		}

		//challenge client
		ret = challenge_client(&nodelist[1]);

		if(ret < DIMASUCCESS)
		{
			printf("Challenge failed %d\n", ret );
			exit DIMAFAILURE;
		}


		//response from client
		memset(netBuff,0,NETBUFSIZE);

		ret = mbedtls_net_recv(nodelist[1].fd, netBuff, 100);
		if(ret <= DIMASUCCESS)
		{
			printf("DIMANETWORKFAILURE : Failed to recv response -0x%04x\n", (unsigned int) -ret);
			return DIMANETWORKFAILURE;
		}

		for(int i; i < 100; i++)
			printf("%hhx",netBuff[i] );
		printf("Received data of len... %d\n Unpacking...\n", ret );


		Resp_context Rp;
		Rp.siglen = (ret - NODE_ID_SIZE - PAD);
		memcpy(Rp.ID, netBuff, NODE_ID_SIZE);
		memcpy(Rp.sign, &netBuff[NODE_ID_SIZE + PAD], Rp.siglen);

	
		//if(DEBUG)
		{
			for(int i=0; i< Rp.siglen; i++)
				printf("%hhx",Rp.sign[i] );
			printf(" : signature of size %ld\n", Rp.siglen);	
		}

		//Verify client sig
		printf("Verifying Client%s's public signature....\n", Rp.ID );
		ret = verify_client(&nodelist[1], &Rp);
		if(ret < DIMASUCCESS)
		{
			printf("DIMASIGNFAILURE : could not verify key signature %d , -0x%04x,\n",ret, (unsigned int) -ret);
			return DIMASIGNFAILURE;
		}

		printf("Signature verified \n");


	}
	//Client node code
	else
	{
		//load server info
		//connect to server
		//store FD
		//send cli ID
		//recv chll
		//load Master pub
		//verify Master sign
		//sign nonce
		//resp ctx -> self id, siglen, sign
		//send resp
		//recv ack


		printf("Running as Slave node....\n");
		Node_info 		Server;
		Server.ID 	=	MASTER_ID;
		Server.pubKey_file =	"clientkeys/Server_pub.pem";


		mbedtls_net_context fd_server;
  		mbedtls_net_init(&fd_server);

		unsigned char netBuff[NETBUFSIZE] = {0};
		unsigned char netbuf2[2];


		ret = mbedtls_net_connect(&fd_server, SERVER_ADD, SERVER_PORT, MBEDTLS_NET_PROTO_TCP);
		if(ret < DIMASUCCESS)
		{
			printf("DIMANETWORKFAILURE : Failed to connect to server -0x%04x\n", (unsigned int) -ret);
			return DIMANETWORKFAILURE;
		}

		printf("Connected to server ...\n");
		Server.fd = &fd_server;

		sleep(1);
		ret = mbedtls_net_send(Server.fd, Self.ID, sizeof(Self.ID));
		if(ret < DIMASUCCESS)
		{
			printf("DIMANETWORKFAILURE : Failed to send ID -0x%04x\n", (unsigned int) -ret);
			return DIMANETWORKFAILURE;
		}
		printf("Sent %d bytes ID\n", ret);

		//recv, unpack challenge contxt

		ret = mbedtls_net_recv(Server.fd, netBuff, 100);
		if(ret <= DIMASUCCESS)
		{
			printf("DIMANETWORKFAILURE : Failed to recv challenge -0x%04x\n", (unsigned int) -ret);
			if(ret == 0)
				printf("Server closed connection!\n");
			return DIMANETWORKFAILURE;
		}

		for(int i; i < 100; i++)
			printf("%hhx",netBuff[i] );
		printf("Received data of len... %d\n Unpacking...\n", ret );

		Chall_context Ch;
		Ch.siglen = (ret - NONCE_SIZE - PAD);
		memcpy(Ch.NONCE, netBuff, NONCE_SIZE);
		memcpy(Ch.sign, &netBuff[NONCE_SIZE + PAD], Ch.siglen);

		//if(DEBUG)
		{
			for(int i = 0; i < NONCE_SIZE; i++)
				printf("%hhx",Ch.NONCE[i]);
			printf(" :  NONCE chall\n");
		}

		//if(DEBUG)
		{
			for(int i=0; i< Ch.siglen; i++)
				printf("%hhx",Ch.sign[i] );
			printf(" : signature of size %ld\n", Ch.siglen);	
		}

		//Verify server sig
		ret = verify_master(&Server, &Ch);
		if(ret < DIMASUCCESS)
		{
			printf("DIMASIGNFAILURE : could not verify key signature %d , -0x%04x,\n",ret, (unsigned int) -ret);
			return DIMASIGNFAILURE;
		}

		printf("Signature verified \n");

		//Respond to Server
		Resp_context Rp;
		memcpy(Rp.ID, Self.ID, NODE_ID_SIZE);

		printf("Responding to Challenge...\n");
		ret = response_master(&Server, &Rp, Ch.NONCE);


	}


}

	// const uint8_t *			ID;
	// const uint8_t *			status;
	// const uint8_t *			NONCE;
	// const char *			pubKey_file;
	// mbedtls_pk_context *	pub_key;
	// mbedtls_pk_context *	priv_key;
	// mbedtls_net_context *	fd;