#include <stdio.h>
#include <stdlib.h>
#include <string.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"

/*

void main()
{
  int ret = 0;

  unsigned char pubkeybuf[100];
  size_t pubkeysize;
  char privkeybuf[100];
  size_t privkeysize;
  printf("Pass 1\n");
  mbedtls_entropy_context entropyCtx;
  mbedtls_entropy_init(&entropyCtx);

  mbedtls_ctr_drbg_context drbgCtx;
  mbedtls_ctr_drbg_init(&drbgCtx);
  printf("Pass 2\n");
  //mbedtls_entropy_add_source( &entropyCtx, use_dev_random,NULL, ENTROPY_LEN, MBEDTLS_ENTROPY_SOURCE_STRONG );


  mbedtls_pk_context pk_ctx;
  mbedtls_pk_init(&pk_ctx);
  if (mbedtls_pk_setup(&pk_ctx, mbedtls_pk_info_from_type (MBEDTLS_PK_ECKEY)) < 0)
    {
      perror("pk setup failed \n");
      return ;
    }
    printf("Pass 3\n");
  //if ECC or RSA
  
  mbedtls_ecp_keypair *ec_key = mbedtls_pk_ec( pk_ctx );
  if (ec_key == NULL )
    {
      perror("pk to ec failed \n");
      return ;
    }
  printf("Pass 4\n");
  mbedtls_ecp_group_load(&ec_key -> grp, MBEDTLS_ECP_DP_SECP256R1);
  printf("Pass 5\n");
  ret = mbedtls_ctr_drbg_seed(&drbgCtx, mbedtls_entropy_func, &entropyCtx,
                             (const unsigned char *) "seed", strlen("seed") );
  if(ret <0)
  {
    printf("%d\n", ret);
    perror("DRBG seed failure\n");
    return;
  }
  printf("Pass 5.5\n");
  if (mbedtls_ecp_gen_key (MBEDTLS_ECP_DP_SECP256R1, mbedtls_pk_ec( pk_ctx ), mbedtls_ctr_drbg_random, &drbgCtx) < 0)
    {
      perror("key gen failed \n");
      return ;
    }
  printf("Pass 6\n");
  ret = mbedtls_pk_write_pubkey_pem (&pk_ctx, pubkeybuf, sizeof(pubkeybuf));
  if(ret < 0)
  {
    perror("ECP write public failure\n");
    return;
  }

  for(int i = 0; i < sizeof(pubkeybuf); i++)
    printf("%hhx",pubkeybuf[i]);
  printf(" :  PubKey\n");

  printf("Pass 7\n");
  ret = mbedtls_mpi_write_string(&ec_key->d, 16, privkeybuf, sizeof(privkeybuf), &privkeysize);
  if(ret < 0)
  {
    printf("%d\n", ret);
    perror("MPI write point to string failure\n");
    return;
  }

  printf("%s :  PrivKey\n",privkeybuf);

  printf("Pass 100\n");


return;
}

*/




/**/



void main()
{
  int ret = 1;
  int len = 0;
  char buf[1024];

  unsigned char pubkeybuf[1024];
  size_t pubkeysize;
  char privkeybuf[1024];
  size_t privkeysize;

  const char * seed = "randomstring";

  mbedtls_entropy_context entropyCtx;
  mbedtls_entropy_init(&entropyCtx);

  mbedtls_ctr_drbg_context drbgCtx;
  mbedtls_ctr_drbg_init(&drbgCtx);

  printf("PASS 1\n");

  mbedtls_pk_context key_ctx;
  mbedtls_pk_init(&key_ctx);

  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_ctr_drbg_seed(&drbgCtx, mbedtls_entropy_func, &entropyCtx,
                             (const unsigned char *) seed, strlen(seed) );

  if(ret <0)
  {
    printf("%d\n", ret);
    perror("DRBG seed failure\n");
    return;
  }
printf("PASS 2\n");

/**
////////////////// ECC Curve

  ret = mbedtls_pk_setup( &key_ctx,
            mbedtls_pk_info_from_type( (mbedtls_pk_type_t) MBEDTLS_PK_ECKEY ) );

  if(ret <0)
  {
    printf("%d\n", ret);
    perror("PK typeset failure\n");
    return;
  }
printf("PASS 3\n");

  ret = mbedtls_ecp_gen_key( (mbedtls_ecp_group_id) MBEDTLS_ECP_DP_SECP256R1,
            mbedtls_pk_ec( key_ctx ), mbedtls_ctr_drbg_random, &drbgCtx );

printf("PASS 4\n");

  if( mbedtls_pk_get_type( &key_ctx ) == MBEDTLS_PK_ECKEY )
    {
        mbedtls_ecp_keypair *ecp = mbedtls_pk_ec( key_ctx );
        printf( "curve: %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 );
    }

printf("PASS 5\n");

**/
///////////////// RSA 

ret = mbedtls_pk_setup( &key_ctx, mbedtls_pk_info_from_type( (mbedtls_pk_type_t) MBEDTLS_PK_RSA ) );

ret = mbedtls_rsa_gen_key( mbedtls_pk_rsa( key_ctx ), mbedtls_ctr_drbg_random, &drbgCtx, 4096, 65537 );

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

        if( ( ret = mbedtls_rsa_export    ( rsa, &N, &P, &Q, &D, &E ) ) != 0 ||
            ( ret = mbedtls_rsa_export_crt( rsa, &DP, &DQ, &QP ) )      != 0 )
        {
            printf( " failed\n  ! could not export RSA parameters\n\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 );
    }

printf("PASS 6\n");

///////////////// Private key

len = mbedtls_pk_write_key_pem( &key_ctx, privkeybuf, sizeof(privkeybuf) );

printf("\nPrivate key : ");
for(int i = 0; i < sizeof(privkeybuf); i++)
    printf(" %hhx",privkeybuf[i]);


///////////////// Public key

len = mbedtls_pk_write_pubkey_pem( &key_ctx, pubkeybuf, sizeof(pubkeybuf) );
printf("\nPublic key : ");
for(int i = 0; i < sizeof(pubkeybuf); i++)
    printf(" %hhx",pubkeybuf[i]);



printf("\nPASS 100\n");

return;
}
/**/