/*
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "mbedtls/config.h"

#include "mbedtls/aes.h"
#include "mbedtls/bignum.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ecp.h"
#include "mbedtls/ecdh.h"
#include "mbedtls/ecdsa.h"
#include "mbedtls/hmac_drbg.h"
#include "mbedtls/hkdf.h"
#include "mbedtls/md.h"
#include "mbedtls/pk.h"
#include "mbedtls/rsa.h"
#include "mbedtls/sha1.h"
#include "mbedtls/sha256.h"

#include "defines.h"
*/
#include "KeyGen.h"


/**

This is a Asymmetric key gen app which takes KeyDrv_context as input
derives parameters and creates a symmetric key 


Every key derivation should start with a new KD contxt	

**/

/*

/////////////////////////////////////TODO////////////////////
void cleanup()
{
	mbedtls_entropy_free(&entropyCtx);
   	mbedtls_ctr_drbg_free(&drbgCtx);
}
*/

DIMASTATUS use_dev_random(void *data, unsigned char *output,
                             size_t len, size_t *olen )
{
	//initiate HW TRNG 

    FILE *file;
    size_t ret, left = len;
    unsigned char *p = output;
    ((void) data);

    *olen = 0;

    file = fopen( "/dev/random", "rb" );
    if( file == NULL )
        return( MBEDTLS_ERR_ENTROPY_SOURCE_FAILED );

    while( left > 0 )
    {
        /* /dev/random can return much less than requested. If so, try again */
        ret = fread( p, 1, left, file );
        if( ret == 0 && ferror( file ) )
        {
            fclose( file );
            return( MBEDTLS_ERR_ENTROPY_SOURCE_FAILED );
        }

        p += ret;
        left -= ret;
        sleep( 1 );
    }
    fclose( file );
    *olen = len;

    return( DIMASUCCESS );
}

DIMASTATUS seedRNGSource(void *data, unsigned char *output, size_t len)
{
  //srand(); lib fun call //https://stackoverflow.com/questions/55927662/generate-every-time-same-rsa-key-with-c
  //ctr_drbg    //programs/test/benchmark.c:705
  //hmac_drbg

	//Fill entropy accum with CDI and pass to DRBG


  mbedtls_entropy_context * p_ent = data;
  printf("manual update entropy with CDI\n");
      
  if(memcpy(output, p_ent -> accumulator.buffer + ACCUM_BUFF_OFFSET , ENTROPY_LEN) < 0)
  	return( MBEDTLS_ERR_ENTROPY_SOURCE_FAILED );

  len = ENTROPY_LEN;


  return DIMASUCCESS;
}



DIMASTATUS AsymmKeyGen(KeyDrv_context * KD_ctx)
{

/*

	KeyDrv_context DID_ctx;
	DID_ctx.ENT_MODE = DFL_ENT;			//SW_PRNG,HW_TRNG,DETERM,
	DID_ctx.PKC_MODE = DFL_PKC;			//isRSA, isECC
	DID_ctx.seed = CDIKEY;
	DID_ctx.phrase = IDENTITY;			//IDENTITY,ALIAS,SESSION
	DID_ctx.KEY_FORM = DFL_FORM;		//BIN,PEM,DER
	DID_ctx.pub_file = "DID_pub.pem";
	DID_pub.priv_file = "DID_priv.pem";

*/
	DIMASTATUS ret = 0;
	int len = 0;

	unsigned char pubkeybuf[1024];
	size_t pubkeybufsize;
	unsigned char privkeybuf[1024];
	size_t	privkeybufsize;

	mbedtls_entropy_context entropyCtx;
	mbedtls_entropy_init(&entropyCtx);

	mbedtls_ctr_drbg_context drbgCtx;
	mbedtls_ctr_drbg_init(&drbgCtx);


	mbedtls_pk_context pkey_ctx;
	mbedtls_pk_init(&pkey_ctx);

	//assign entropy source and derive seed for key gen

	if (KD_ctx -> ENT_MODE == HW_TRNG)					// use TRNG
	{	
		if(DEBUG)
			printf("using /dev/random...this may take a moment\n");
		
		mbedtls_entropy_add_source( &entropyCtx, use_dev_random,
                               NULL, ENTROPY_LEN, MBEDTLS_ENTROPY_SOURCE_STRONG );
		mbedtls_ctr_drbg_seed(&drbgCtx, mbedtls_entropy_func, &entropyCtx,
                             (const unsigned char *) KD_ctx->phrase, strlen(KD_ctx->phrase) );
	}

	else if (KD_ctx -> ENT_MODE == SW_PRNG)				// use mbedtls sw prng
	{
		if(DEBUG)
			printf("Accumulating entropy...\n");
		
		//mbedtls_entropy_update_manual(&entropyCtx, KD_ctx->seed, ENTROPY_LEN);
		mbedtls_ctr_drbg_seed(&drbgCtx, mbedtls_entropy_func, &entropyCtx,
                              (const unsigned char *) KD_ctx->phrase, strlen(KD_ctx->phrase) );
	}
	
	else if (KD_ctx -> ENT_MODE == DETERM)				// used seed provided in the KD ctx
	{
		if(DEBUG)
			printf("Seeding entropy accumulator...\n");
		
		ret = mbedtls_entropy_update_manual(&entropyCtx, KD_ctx->seed, ENTROPY_LEN);
		if(ret < DIMASUCCESS)
		{
			perror("Seeding entropy failed\n");
			//cleanup();
			exit(DIMADRBGFAILURE);
		}

		ret = mbedtls_ctr_drbg_seed(&drbgCtx, seedRNGSource, &entropyCtx,
    				(const unsigned char *) KD_ctx->phrase, sizeof(&KD_ctx->phrase));
		if(ret < DIMASUCCESS)
		{
			perror("DRBG failed\n");
			//cleanup();
			exit(DIMADRBGFAILURE);
		}
	}

	else
	{
		perror("DIMAINVALIDSTATE : Check entropy source mode in Key Derivation struct\n");
		//cleanup();
		exit(DIMAINVALIDSTATE);
	}

	if(DEBUG) printf("Generating ECC asymmetric key pair...\n");

	//////ECC implementation

	if(KD_ctx->PKC_MODE == isECC)
	{
		ret = mbedtls_pk_setup(&pkey_ctx,mbedtls_pk_info_from_type((mbedtls_pk_type_t)MBEDTLS_PK_ECKEY));
		if(ret < DIMASUCCESS)
		{
			perror("DIMAPKFAILURE\n");
			//cleanup();
			exit(DIMAPKFAILURE);
		}

		ret = mbedtls_ecp_gen_key( (mbedtls_ecp_group_id) ECC_CURVE, mbedtls_pk_ec( pkey_ctx ), 
									mbedtls_ctr_drbg_random, &drbgCtx );
		if(ret < DIMASUCCESS)
		{
			perror("DIMAECCFAILURE\n");
			//cleanup();
			exit(DIMAECCFAILURE);
		}

		if( mbedtls_pk_get_type( &pkey_ctx ) == MBEDTLS_PK_ECKEY )
    	{
        	mbedtls_ecp_keypair *ecp = mbedtls_pk_ec( pkey_ctx );
        	if(DEBUG)
        	{
				printf( "key info : \ncurve: %s\n", mbedtls_ecp_curve_info_from_grp_id( ecp->grp.id )->name );
        		mbedtls_mpi_write_file( "X_Q:   ", &ecp->Q.X, 16, NULL );
        		mbedtls_mpi_write_file( "Y_Q:   ", &ecp->Q.Y, 16, NULL );
        		mbedtls_mpi_write_file( "D:     ", &ecp->d  , 16, NULL );
        	}

        	mbedtls_ecp_keypair_free(ecp);
        }
        else
        {
			perror("DIMAPKFAILURE\n");
			//cleanup();
			exit(DIMAPKFAILURE);
        }

        
    }

	/////RSA implementation
    else if(KD_ctx->PKC_MODE == isRSA)
    {

    	mbedtls_mpi N, P, Q, D, E, DP, DQ, QP;
		mbedtls_mpi_init( &N ); 
		mbedtls_mpi_init( &P ); 
		mbedtls_mpi_init( &Q );
		mbedtls_mpi_init( &D ); 
		mbedtls_mpi_init( &E ); 
		mbedtls_mpi_init( &DP );
		mbedtls_mpi_init( &DQ ); 
		mbedtls_mpi_init( &QP );
    	
    	ret = mbedtls_pk_setup(&pkey_ctx, mbedtls_pk_info_from_type((mbedtls_pk_type_t) MBEDTLS_PK_RSA));
    	if(ret < DIMASUCCESS)
		{
			perror("DIMAPKFAILURE\n");
			//cleanup();
			exit(DIMAPKFAILURE);
		}

		ret = mbedtls_rsa_gen_key( mbedtls_pk_rsa( pkey_ctx ), mbedtls_ctr_drbg_random, &drbgCtx, RSA_SIZE, RSA_EXP );
		if(ret < DIMASUCCESS)
		{
			perror("DIMARSAFAILURE\n");
			//cleanup();
			exit(DIMARSAFAILURE);
		}

		if( mbedtls_pk_get_type( &pkey_ctx ) == MBEDTLS_PK_RSA )
    	{
        	mbedtls_rsa_context *rsa = mbedtls_pk_rsa( pkey_ctx );

        	if( ( ret = mbedtls_rsa_export    ( rsa, &N, &P, &Q, &D, &E ) ) != 0 ||
            	( ret = mbedtls_rsa_export_crt( rsa, &DP, &DQ, &QP ) )      != 0 )
        	{
            	printf( " DIMARSAFAILURE:could not export RSA parameters\n\n" );
        	}

        	if(DEBUG)
        	{
        		printf("RSA Key pair info :\n");
        		mbedtls_mpi_write_file( "N:  ",  &N,  16, NULL );
       			mbedtls_mpi_write_file( "E:  ",  &E,  16, NULL );
    		    mbedtls_mpi_write_file( "D:  ",  &D,  16, NULL );
    		    mbedtls_mpi_write_file( "P:  ",  &P,  16, NULL );
    		    mbedtls_mpi_write_file( "Q:  ",  &Q,  16, NULL );
    		    mbedtls_mpi_write_file( "DP: ",  &DP, 16, NULL );
    		    mbedtls_mpi_write_file( "DQ:  ", &DQ, 16, NULL );
    		    mbedtls_mpi_write_file( "QP:  ", &QP, 16, NULL );
    		}
    	}

    	mbedtls_mpi_free(&N);
    	mbedtls_mpi_free(&P);
    	mbedtls_mpi_free(&Q);
    	mbedtls_mpi_free(&D);
    	mbedtls_mpi_free(&E);
    	mbedtls_mpi_free(&DP);
    	mbedtls_mpi_free(&DQ);
    	mbedtls_mpi_free(&QP);

	}
	else
    {
        perror("DIMAINVALIDSTATE : Check PKC mode in Key Derivation struct\n");
		//cleanup();
		exit(DIMAINVALIDSTATE);
    }

    ret = WritePrivKey(KD_ctx,&pkey_ctx);


    ret = WritePubKey(KD_ctx, &pkey_ctx);

    ////////////////////////////TODO///////////////////////////
    //free block
    mbedtls_entropy_free(&entropyCtx);
    mbedtls_ctr_drbg_free(&drbgCtx);
	mbedtls_pk_free(&pkey_ctx);



    return DIMASUCCESS;
}

//#define DFL_PUB		"keys/DID_pub." DFL_FORM
//#define DFL_PRIV		"SecureStorage/DID_priv" DFL_FORM

DIMASTATUS WritePrivKey(KeyDrv_context * KD_ctx, mbedtls_pk_context * pkey_ctx)
{
	DIMASTATUS ret = 0;
	int i = 0;
	size_t	len;
	FILE *fp;
	unsigned char dest_file[50];
	unsigned char * outbuf = calloc(1,sizeof(unsigned char)*KEY_BUF_SIZE);

	if(strcmp(KD_ctx -> phrase, IDENTITY) == 0)
	{
		if(DEBUG)
		{
			printf("Private DID key should be stored only inside secure storage.\n");
		}
		//codeblock to check existance of SS. 
		//return with warning otherwise
		//exit?
		//////////////////////////////TODO///////////////////


		strcpy(dest_file,"SecureStorage/");
		}
	else
	{
		strcpy(dest_file,"keys/");
	}


	strcat(dest_file, KD_ctx->phrase);
	strcat(dest_file, "_priv");

	if(KD_ctx->KEY_FORM == PEM)
	{
		strcat(dest_file, ".pem");
		//write pem does not return no of bytes written....stupid
		len = mbedtls_pk_write_key_pem(pkey_ctx, outbuf, KEY_BUF_SIZE);	
		if(len < 0)
		{
			perror("DIMAOUTPUTERROR:");
			return(DIMAOUTPUTERROR);
		}
	}
	else if(KD_ctx->KEY_FORM == DER)
	{
		strcat(dest_file, ".der");
		len = mbedtls_pk_write_key_der(pkey_ctx, outbuf, KEY_BUF_SIZE);
		if(len < 0)
		{
			perror("DIMAOUTPUTERROR:");
			return(DIMAOUTPUTERROR);
		}
	}
	else
	{
		perror("DIMAINVALIDSTATE : Check key format in Key Derivation struct\n");
		//cleanup();
		exit(DIMAINVALIDSTATE);
	}

	len = strlen( (char *) outbuf );

	if( ( fp = fopen( dest_file, "w" ) ) == NULL )
	{
		perror("DIMAOUTPUTERROR: ");
		return(DIMAOUTPUTERROR);
	}

    fwrite( outbuf, 1, len, fp );
    fclose( fp );
    free(outbuf);
    return(DIMASUCCESS);

}


DIMASTATUS WritePubKey(KeyDrv_context * KD_ctx, mbedtls_pk_context * pkey_ctx)
{
	DIMASTATUS ret = 0;
	int i = 0;
	size_t len;
	FILE * fp;
	unsigned char dest_file[50];
	unsigned char * outbuf = calloc(1,sizeof(unsigned char)*KEY_BUF_SIZE);

	strcpy(dest_file,"keys/");
	strcat(dest_file, KD_ctx->phrase);
	strcat(dest_file, "_pub");

	
	if(KD_ctx->KEY_FORM == PEM)
	{
		strcat(dest_file, ".pem");
		//write pem does not return no of bytes written....stupid
		len = mbedtls_pk_write_pubkey_pem(pkey_ctx, outbuf, KEY_BUF_SIZE);	
		if(len < 0)
		{
			perror("DIMAOUTPUTERROR:");
			return(DIMAOUTPUTERROR);
		}
	}
	else if(KD_ctx->KEY_FORM == DER)
	{
		strcat(dest_file, ".der");
		len = mbedtls_pk_write_pubkey_der(pkey_ctx, outbuf, KEY_BUF_SIZE);
		if(len < 0)
		{
			perror("DIMAOUTPUTERROR:");
			return(DIMAOUTPUTERROR);
		}
	}
	else
	{
		perror("DIMAINVALIDSTATE : Check key format in Key Derivation struct\n");
		//cleanup();
		exit(DIMAINVALIDSTATE);
	}

	len = strlen( (char *) outbuf );
	
	if( ( fp = fopen( dest_file, "w" ) ) == NULL )
	{
		perror("DIMAOUTPUTERROR: ");
		return(DIMAOUTPUTERROR);
	}

    fwrite( outbuf, 1, len, fp );
    fclose( fp );
    free(outbuf);
    return(DIMASUCCESS);

}