#include <stdio.h>
#include <string.h>
#ifdef WIN32
#include <Windows.h>
#endif // WIN32

#include "mbedtls/base64.h"
#include "mbedtls/x509_crt.h"
#include "mbedtls/pk.h"
#include "ECP_DP_CURVE25519.h"
////////////////////////////////////////////////////////////////////////////////


static inline void vHEXDUMP(uint8_t* pu8Banner, uint8_t* pu8String, size_t szLength)
{
   printf("HEXDUMP %s", pu8Banner);
   for (size_t i = 0; i < szLength; printf("%02X ", pu8String[i++]))
   {
      if (i % 16 == 0) printf("\r\n");
   }
   printf("\r\n\r\n");
}

static inline void vBASE64(uint8_t* pu8Banner, uint8_t* pu8String, size_t szLength)
{
   uint8_t* pu8Base64String = NULL;
   size_t   szBase64StringLength = 0;

   mbedtls_base64_encode(pu8Base64String, 0, &szBase64StringLength, pu8String, szLength);

   pu8Base64String = (uint8_t*)malloc(szBase64StringLength);
   if (pu8Base64String == NULL)
   {
      printf("Memory allocation failed\r\n");
      return;
   }
   int r = mbedtls_base64_encode(pu8Base64String, szBase64StringLength, &szBase64StringLength, pu8String, szLength);
   if (r)
   {
      printf("failed! mbedtls_base64_encode: 0x%04x", -r);
      return;
   }


   printf("B64Encoded %s\r\n%s\r\n\r\n", pu8Banner, pu8Base64String);
   free(pu8Base64String);
}

int32_t ECDHE_Init(ECDHE_t* ctx, ECDH_Endpoint_t Endpoint)
{
   int32_t r = EXIT_SUCCESS;

   mbedtls_ecdh_init(&ctx->ECDH);
   mbedtls_ctr_drbg_init(&ctx->DRBG);
   mbedtls_entropy_init(&ctx->ENTROPY);

   do
   {
      uint8_t  pu8SeedString[] = "HelloWorld!-G&zaMy$v=mLfRM%7jq4BYgPL";
      size_t   szLength = strlen(pu8SeedString);
 
      /* On target platform utlizse Rand context to Seed CTR_DRBG */
      if (EXIT_SUCCESS !=
         (r = mbedtls_ctr_drbg_seed(&ctx->DRBG, mbedtls_entropy_func,
            &ctx->ENTROPY,
            pu8SeedString, szLength)))
      {
         printf("failed! mbedtls_ctr_drbg_seed: 0x%04x", -r);
         break;
      }
      uint8_t u8DrbgBuffer[32];
      memset(u8DrbgBuffer, 0x00, 32);
      if (r = mbedtls_ctr_drbg_random(&(ctx->DRBG), u8DrbgBuffer, 32))
      {
          printf("failed! mbedtls_ctr_drbg_random: 0x%04x", -r);
          break;
     }
      printf("\n Random Number : \r\n");
      for (int i = 0; i < 32; printf("%02x ", u8DrbgBuffer[i++]));
      /* Setup ECDH with curve SECP256R1(NSA Suite B Standard) */
      if (EXIT_SUCCESS !=
         (r = mbedtls_ecdh_setup(&ctx->ECDH, MBEDTLS_ECP_DP_SECP256R1 )))
         /* (r = mbedtls_ecdh_setup(&ctx->ECDH, MBEDTLS_ECP_DP_CURVE25519 ))) */
      {
         printf("failed! mbedtls_ecdh_setup: 0x%04x", -r);
         break;
      }

   } while (0);

   return(r);
}

int32_t ECDHE_WriteOwnPublic(ECDHE_t* ctx, uint8_t* pu8SigningKey, size_t szKeyLength)
{
   int r = EXIT_SUCCESS;

   /* ECDH Endpoint generate public key to be shared with Peer */
   if (EXIT_SUCCESS !=
      (r = mbedtls_ecdh_make_public(&ctx->ECDH, &ctx->szPublicLength,
         ctx->pu8Public, ECDHE_PUBLIC_SIZE,
         mbedtls_ctr_drbg_random, &ctx->DRBG)))
   {
      printf("failed! mbedtls_ecdh_make_public: 0x%04x", -r);
      return(r);
   }

   /* Hexdump of ECDH->Public */
   vHEXDUMP("ECDHE->Public", ctx->pu8Public, ctx->szPublicLength);
   vBASE64("ECDHE->Public", ctx->pu8Public, ctx->szPublicLength);

   /* ECDH Endpoint has a ECP256 Private Key to sign the ECDH->Public */
   if (pu8SigningKey != NULL)
   {
      uint8_t pu8Digest[32];
      mbedtls_pk_context pk;

      mbedtls_pk_init(&pk);

      /* Parse the private key to be used for signing */
      if (EXIT_SUCCESS !=
         (r = mbedtls_pk_parse_key(&pk, pu8SigningKey, szKeyLength, NULL, 0)))
      {
         printf("failed! mbedtls_pk_parse_key: 0x%04x\r\n", -r);
         return(r);
      }

      /* Generate SHA256 Digest of ECDH->Public before signing it */
      if (EXIT_SUCCESS !=
         (r = mbedtls_md(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 
              ctx->pu8Public, ctx->szPublicLength, pu8Digest)))
      {
         printf("failed! mbedtls_md: 0x%04x\r\n", -r);
         return(r);
      }

      /* Sign ECDH->Public using the Private Key provided */
      if (EXIT_SUCCESS !=
         (r = mbedtls_pk_sign(&pk, MBEDTLS_MD_SHA256, pu8Digest, 32, 
              ctx->pu8Signature, &ctx->szSignatureLength,
              mbedtls_ctr_drbg_random, &ctx->DRBG)))
      {
         printf("failed! mbedtls_pk_sign: 0x%04x\r\n", -r);
         return(r);
      }

      /* Hexdump of Signature */
      vHEXDUMP("ECDHE->Public->Signature", ctx->pu8Signature, ctx->szSignatureLength);
      vBASE64("ECDHE->Public->Signature", ctx->pu8Signature, ctx->szSignatureLength);
   }
   return(r);
}


int32_t ECDHE_ReadPeerPublic(ECDHE_t* ctx, uint8_t* pu8PublicKey, size_t szKeyLength,
                             uint8_t* pu8Signature, size_t szSignatureLength,
                             uint8_t* pu8VerificationCert, size_t szCertLength)
{
   int r = EXIT_SUCCESS;

   /*
    * ECDH Endpoint has X.509 Certificate of peer to verify the Signature
    * of ECDH->Public read from peer
    */
   if (pu8Signature != NULL && pu8VerificationCert != NULL)
   {
      uint8_t pu8Digest[32];
      mbedtls_x509_crt Cert;
      mbedtls_x509_crt_init(&Cert);
      
      /* Parse the X.509 Certificate of peer to be used for verification */
      if (EXIT_SUCCESS !=
         (r = mbedtls_x509_crt_parse(&Cert, pu8VerificationCert, szCertLength)))
      {
         printf("failed! mbedtls_x509_crt_parse: 0x%04x\r\n", -r);
         return(r);
      }

      /* Generate SHA256 digest of Peer's ECDH->Public for verification */
      if (EXIT_SUCCESS !=
         (r = mbedtls_md(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256),
            pu8PublicKey, szKeyLength, pu8Digest)))
      {
         printf("failed! mbedtls_md: 0x%04x\r\n", -r);
         return(r);
      }

      /* Verify the signature using parsed X.509 Certificate */
      if (EXIT_SUCCESS !=
         (r = mbedtls_pk_verify(&Cert.pk, MBEDTLS_MD_SHA256, pu8Digest,
              32, pu8Signature, szSignatureLength)))
      {
         printf("failed! mbedtls_pk_verify: 0x%04x\r\n", -r);
         return(r);
      }

      printf("\r\nSignature Verification Success!\r\n\r\n");
   }

   /* ECDH Endpoint consumes ECDH->Public of peer */
   if (EXIT_SUCCESS !=
      (r = mbedtls_ecdh_read_public(&ctx->ECDH, pu8PublicKey, szKeyLength)))
   {
      printf("failed! mbedtls_ecdh_read_public: 0x%04x", -r);
      return(r);
   }

   /* ECDH Endpoint computes shared secret DH Math */
   if (EXIT_SUCCESS !=
      (r = mbedtls_ecdh_calc_secret(&ctx->ECDH, &ctx->szSecretLength, ctx->pu8Secret,
         ECDHE_PUBLIC_SIZE, mbedtls_ctr_drbg_random, &ctx->DRBG)))
   {
      printf("failed! mbedtls_ecdh_calc_secret: 0x%04x", -r);
      return(r);
   }
   printf("\r\nECHDE_SHARED_SECRET : \r\n");
   for (int i = 0; i < ctx->szSecretLength; printf("%02x ", ctx->pu8Secret[i++]));

   /* ECDH Endpoint computes Symmetric Key from SHA256 digest of shared secret */
   if (EXIT_SUCCESS !=
      (r = mbedtls_md(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256),
                      ctx->pu8Secret, ctx->szSecretLength, ctx->pu8SymmetricKey)))
   {
      printf("failed! mbedtls_md: 0x%04x\r\n", -r);
      return(r);
   }

   /* Hexdump of Computed Symmetric Key */
   vHEXDUMP("Established Symmetric Key ", ctx->pu8SymmetricKey, SHA256_DIGEST_SIZE);

   return(r);
}

void ECDHE_SelfTest(void)
{ 
   ECDHE_t Node1, Node2;

   /*
    * Creating ECDH Nodes resembling DTU and GGCPU 
    */
   printf("Creating ECDHE Node-1\r\n");
   ECDHE_Init(&Node1, ECDH_SERVER);
   printf("Creating ECDHE Node-2\r\n");
   ECDHE_Init(&Node2, ECDH_CLIENT);

   /*
    * ECDH Nodes generate ECDH->Public for sharing;
    * DTU Node: signs it using ECP256 Private Key.
    * GGCPU Node: does NOT signs as it does NOT under PKI.
    */
   printf("Node-1 Generates & Signs it's ECDH->Public\r\n");
   ECDHE_WriteOwnPublic(&Node1, TEST_EC_PRIVATE_KEY, sizeof(TEST_EC_PRIVATE_KEY));

   printf("Node-2 Generates it's ECDH->Public\r\n");
   ECDHE_WriteOwnPublic(&Node2, NULL, 0);

   /*
    * ECDH Nodes reads ECDH->Public from peer;
    * DTU Node: does NOT verifies peer ECDH->Public.
    * GGCPU Node: Verifies ECDH->Public of peer using X.509 Certificate of peer.
    *           In actual handshake, the DTU's certificate is already verified
    *           in certificate transaction phase.
    */
   printf("Node-1 Reads ECDH->Public from Node-2\r\n"); //DTU
   ECDHE_ReadPeerPublic(&Node1, Node2.pu8Public, Node2.szPublicLength, NULL, 0, NULL, 0);

   printf("Node-2 Reads & Verifies ECDH->Public from Node-1\r\n"); //GGCPU

   ECDHE_ReadPeerPublic(&Node2, Node1.pu8Public, Node1.szPublicLength,
                                Node1.pu8Signature, Node1.szSignatureLength,
                                TEST_EC_CERTIFICATE, sizeof(TEST_EC_CERTIFICATE));

}