#include "Stage2.h"


DIMASTATUS load_nodes(Node_info * nodelist)
{

	DIMASTATUS ret = 0;
	
	printf("load node info... \n");

	//client nodes
	Node_info		Cl1;
	Cl1.ID = 		CLIENT1_ID;
	Cl1.status =	DFL_CL_STAT;
	Cl1.pubKey_file="clientkeys/Client1_pub.pem";

	nodelist[1] =	Cl1;

	//more nodes ....


	return(ret);
}



DIMASTATUS challenge_client(Node_info * Client)
{
	DIMASTATUS ret = 0;

	//if (DEBUG)
		printf("challence node %s.... \n", Client->ID);

	Chall_context Ch;
	Ch.siglen = 0;
	//Ch.sign = calloc(1, sizeof(uint8_t)*SIGN_SIZE);
	
	mbedtls_entropy_context entropy;
	mbedtls_entropy_init(&entropy);

	//uint8_t * RAND = calloc(1, sizeof(uint8_t)*NONCE_SIZE);
	ret = mbedtls_entropy_func(&entropy, Ch.NONCE, NONCE_SIZE);
	if(ret < DIMASUCCESS)
	{
		printf("DIMADRBGFAILURE : Failed to gen NONCE-0x%04x\n", (unsigned int) -ret );
		exit(DIMADRBGFAILURE);
	}

	//Ch.NONCE 		= RAND;
	Client->NONCE 	= Ch.NONCE;
	//memcpy(Client->NONCE, Ch.NONCE, NONCE_SIZE);


	//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 < NONCE_SIZE; i++)
	    	printf("%hhx",Client->NONCE[i]);
	    printf(" :  NONCE client\n");
	}
	
	//loading private key
		//if (DEBUG)
		printf("load private key\n");

	mbedtls_pk_context priv_ctx;
	mbedtls_pk_init(&priv_ctx);

	ret = mbedtls_pk_parse_keyfile(&priv_ctx, "keys/Alias_priv.pem", NULL);
	if(ret < DIMASUCCESS)
	{
		printf("DIMAPKFAILURE : Failed to load signing key -0x%04x\n", (unsigned int) -ret);
		return DIMAPKFAILURE;
	}

	//measuring NONCe
	//if(DEBUG)
		printf("measuring nonce\n");

	uint8_t* dgst = calloc(1, sizeof(uint8_t)*SHA256_DGST_SIZE);
	//uint8_t* sign = calloc(1, sizeof(uint8_t)*SIGN_SIZE);
	size_t siglen = 0;


	ret = mbedtls_sha256_ret(Ch.NONCE, NONCE_SIZE, dgst, 0);
	if(ret < DIMASUCCESS)
	{
		printf("DIMASHAFAILURE : could not perfom SHA digest -0x%04x\n", (unsigned int) -ret);
		return DIMASHAFAILURE;
	}
	//if(DEBUG)
	{
		for(int i = 0; i < SHA256_DGST_SIZE; i++)
	    	printf("%hhx",dgst[i]);
	    printf(" :  NONCE dgst\n");
	}

	//signing NONCE
	//if(DEBUG)
	printf("SIGNING nonce\n");
	ret = mbedtls_pk_sign(&priv_ctx, MBEDTLS_MD_SHA256, dgst, SHA256_DGST_SIZE, Ch.sign, &siglen, NULL, NULL);
	if(ret < DIMASUCCESS)
	{
		printf("DIMASIGNFAILURE: could not perform nonce signing...-0x%04x\n", (unsigned int) -ret);
		return DIMASIGNFAILURE;
	}
	Ch.siglen = siglen;

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

	//sending challenge
	//if(DEBUG)
		printf("sending challenge...\n");


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

	printf("%ld sig\n", Ch.siglen);

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

	sleep(1);
	ret = mbedtls_net_send(Client->fd, netBuff, NONCE_SIZE+Ch.siglen+PAD );
	if(ret < DIMASUCCESS)
	{
    	printf("DIMANETWORKFAILURE : Failed to send-0x%04x\n", (unsigned int) -ret);
    	return DIMANETWORKFAILURE;
 	}

 	//if(DEBUG)
 	printf("sent %d bytes\n", ret);

 	return ret;
}


DIMASTATUS verify_client(Node_info * Client, Resp_context * Rp)
{

	DIMASTATUS ret = 0;

	if(strcmp(Client->ID, Rp->ID) != DIMASUCCESS)
	{
		printf("Client ID mismatch\nAborting...\n");
		return DIMAFAILURE;
	}

	mbedtls_pk_context Cli_ctx;
	mbedtls_pk_init(&Cli_ctx);

	ret = mbedtls_pk_parse_public_keyfile(&Cli_ctx, Client->pubKey_file);
	if(ret < DIMASUCCESS)
	{
		printf("DIMAPKFAILURE : Failed to load Client's public key -0x%04x\n", (unsigned int) -ret);
		return DIMAPKFAILURE;
	}

	uint8_t* dgst = calloc(1,sizeof(uint8_t)*SHA256_DGST_SIZE);

	ret = mbedtls_sha256_ret( Client->NONCE, NONCE_SIZE, dgst, 0);
	if(ret < DIMASUCCESS)
	{
		printf("DIMASHAFAILURE : could not perfom SHA digest -0x%04x\n", (unsigned int) -ret);
		return DIMASHAFAILURE;
	}

	for (int i = 0; i < SHA256_DGST_SIZE; i++)
        printf("%hhx", dgst[i]);
    printf(" :  sha dgst\n" );

    ret = mbedtls_pk_verify(&Cli_ctx, MBEDTLS_MD_SHA256, dgst, SHA256_DGST_SIZE, Rp->sign, Rp->siglen);
	if(ret < DIMASUCCESS)
	{	
		printf("DIMASIGNFAILURE : could not verify key signature %d , -0x%04x,\n",ret, (unsigned int) -ret);
		return DIMASIGNFAILURE;
	}

	return DIMASUCCESS;
}


//Client side functions


DIMASTATUS verify_master(Node_info * Server, Chall_context * Ch)
{
	DIMASTATUS ret = 0;

	mbedtls_pk_context Serv_ctx;
	mbedtls_pk_init(&Serv_ctx);

	ret = mbedtls_pk_parse_public_keyfile(&Serv_ctx, Server->pubKey_file);
	if(ret < DIMASUCCESS)
	{
		printf("DIMAPKFAILURE : Failed to load server's public key -0x%04x\n", (unsigned int) -ret);
		return DIMAPKFAILURE;
	}

	uint8_t* dgst = calloc(1,sizeof(uint8_t)*SHA256_DGST_SIZE);

	ret = mbedtls_sha256_ret( Ch->NONCE, NONCE_SIZE, dgst, 0);
	if(ret < DIMASUCCESS)
	{
		printf("DIMASHAFAILURE : could not perfom SHA digest -0x%04x\n", (unsigned int) -ret);
		return DIMASHAFAILURE;
	}

	for (int i = 0; i < SHA256_DGST_SIZE; i++)
        printf("%hhx", dgst[i]);
    printf(" :  sha dgst\n" );

    ret = mbedtls_pk_verify(&Serv_ctx, MBEDTLS_MD_SHA256, dgst, SHA256_DGST_SIZE, Ch->sign, Ch->siglen);
	if(ret < DIMASUCCESS)
	{	
		printf("DIMASIGNFAILURE : could not verify key signature %d , -0x%04x,\n",ret, (unsigned int) -ret);
		return DIMASIGNFAILURE;
	}

	return DIMASUCCESS;
}


DIMASTATUS response_master(Node_info * Server, Resp_context  * Rp, unsigned char * NONCE)
{
	DIMASTATUS ret = 0;

	mbedtls_pk_context priv_ctx;
	mbedtls_pk_init(&priv_ctx);

	ret = mbedtls_pk_parse_keyfile(&priv_ctx, "keys/Alias_priv.pem", NULL);
	if(ret < DIMASUCCESS)
	{
		printf("DIMAPKFAILURE : Failed to load signing key -0x%04x\n", (unsigned int) -ret);
		return DIMAPKFAILURE;
	}

	//measuring NONCe
	//if(DEBUG)
		printf("measuring nonce\n");

	uint8_t* dgst = calloc(1, sizeof(uint8_t)*SHA256_DGST_SIZE);
	//uint8_t* sign = calloc(1, sizeof(uint8_t)*SIGN_SIZE);
	size_t siglen = 0;


	ret = mbedtls_sha256_ret(NONCE, NONCE_SIZE, dgst, 0);
	if(ret < DIMASUCCESS)
	{
		printf("DIMASHAFAILURE : could not perfom SHA digest -0x%04x\n", (unsigned int) -ret);
		return DIMASHAFAILURE;
	}
	//if(DEBUG)
	{
		for(int i = 0; i < SHA256_DGST_SIZE; i++)
	    	printf("%hhx",dgst[i]);
	    printf(" :  NONCE dgst\n");
	}

	//signing NONCE
	//if(DEBUG)
	printf("SIGNING nonce\n");
	ret = mbedtls_pk_sign(&priv_ctx, MBEDTLS_MD_SHA256, dgst, SHA256_DGST_SIZE, Rp->sign, &siglen, NULL, NULL);
	if(ret < DIMASUCCESS)
	{
		printf("DIMASIGNFAILURE: could not perform nonce signing...-0x%04x\n", (unsigned int) -ret);
		return DIMASIGNFAILURE;
	}
	Rp->siglen = siglen;

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

	//sending challenge
	//if(DEBUG)
		printf("sending response...\n");


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

	printf("%ld sig\n", Rp->siglen);

	memcpy(netBuff, Rp->ID, NODE_ID_SIZE);
	memcpy(&netBuff[NODE_ID_SIZE + PAD], Rp->sign, Rp->siglen);

	sleep(1);
	ret = mbedtls_net_send(Server->fd, netBuff, NODE_ID_SIZE + Rp->siglen + PAD );
	if(ret < DIMASUCCESS)
	{
    	printf("DIMANETWORKFAILURE : Failed to send-0x%04x\n", (unsigned int) -ret);
    	return DIMANETWORKFAILURE;
 	}

 	//if(DEBUG)
 	for(int i; i < NODE_ID_SIZE + Rp->siglen + PAD; i++)
 		printf("%hhx", netBuff[i]);
 	printf(" : sent %d bytes\n", ret);

 	return ret;
}