// Copyright (c) 2013-2014 IONU Security, Inc. All rights reserved.
//
#include <cstdio>
#include <cstring>
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/hmac.h>

#include "eyeutils.h"
#include "eyering.h"



static unsigned char salt[] = {'I', 's', 'O', 'a', 'n', 'l', 'U', 't'};
static unsigned char iv[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 'a', 'b', 'c', 'd', 'e', 'f'};
static const char* PERMS = "{\"doorway\":[],\"writeAssign\":[],\"admin\":[\"owner\"],\"owner\":\"00000002\",\"write\":[],\"read\":[\"owner\"],\"readAssign\":[\"owner\"],\"adminAssign\":[],\"doorwayAssign\":[]}";


void usage()
{
   printf ("performance algorithm iterations|password\n");
   printf ("   algorithm = md5, sha1, sha256, sha512, hmac, pbkdf2, rsa-encrypt, rsa-decrypt, ec-encrypt, ec-decrypt, aes256\n");
   printf ("   iteration = [1-10000000]\n");
   printf ("   password = [A-z]+[any ascii]\n");
}

int main (int argc, char* argv[])
{
   if (argc != 3) {
      usage();
      exit (1);
   }
   char hex [256];
   unsigned char key[IONU_AES_KEY_LEN];
   RAND_bytes (key, IONU_AES_KEY_LEN);
   unsigned int i;
   unsigned int rounds = 0;
   unsigned int len = 0;
   size_t ebytes = 0;
   size_t ecbytes = 0;
   Keyring* ring = new Keyring ("testring", "Test Key Ring");
   ring->GenerateRSAKey("RSA", "RSA keypair");
   const Key* privateK = ring->GetKey("RSA");
   char* pk = ring->GetPubKey("RSA");
   Key* publicK = new Key ("RSAPublic", "CloudGuard public key", (int)strlen(pk), (const unsigned char*) pk, Key::RSA_PUBLIC);
   unsigned char* ct = publicK->PublicKeyEncrypt ((unsigned char*)PERMS, strlen (PERMS), &ebytes);

   ring->GenerateECKey("me", "mine");
   ring->GenerateECKey("peer", "theirs");
   const Key* me = ring->GetKey ("me");
   const Key* peer = ring->GetKey ("peer");
   unsigned char* ect = me->ECDHEncrypt (peer, (unsigned char*)PERMS, strlen (PERMS), &ecbytes);

   ring->GenerateAESKey("AES", "AES key");
   const Key* aesKey =  ring->GetKey ("AES");

   char* password = NULL;
   char c = argv[2][0];
   if (c >= '0' && c <= '9')
      rounds = atoi (argv[2]);
   else {
      password = argv[2];
      len = static_cast<unsigned int>(strlen (password));
   }
   const EVP_MD* md = NULL;
   if (strcmp ("md5", argv[1]) == 0)
      md = EVP_md5();
   else if (strcmp ("sha1", argv[1]) == 0)
      md = EVP_sha1();
   else if (strcmp ("sha256", argv[1]) == 0)
      md = EVP_sha256();
   else if (strcmp ("sha512", argv[1]) == 0)
      md = EVP_sha512();

   long ms; // System time in milliseconds
   ms =  ionu::GetSystemTime();
   if (md) {
      unsigned char md_value[EVP_MAX_MD_SIZE];
      unsigned int md_len;
      EVP_MD_CTX* mdctx = EVP_MD_CTX_create();
      EVP_DigestInit_ex (mdctx, md, NULL);
      EVP_DigestUpdate (mdctx, password, len);
      EVP_DigestFinal_ex (mdctx, md_value, &md_len);
      for (i = 1; i < rounds; ++i) {
         EVP_DigestInit_ex (mdctx, md, NULL);
         EVP_DigestUpdate (mdctx, md_value, md_len);
         EVP_DigestFinal_ex (mdctx, md_value, &md_len);
      }
      EVP_MD_CTX_destroy (mdctx);
      BinaryToHex (md_value, md_len, hex);
      ms =  ionu::GetSystemTime() - ms;
      printf ("%s: rounds = %d, time = %ld ms\n", argv[1], rounds, ms);

   }
   else if (strcmp ("hmac", argv[1]) == 0) {
      unsigned char md_value[EVP_MAX_MD_SIZE];
      unsigned int md_len;
      HMAC_CTX hmctx;
      HMAC_CTX_init (&hmctx);
      HMAC_Init (&hmctx, key, IONU_AES_KEY_LEN, EVP_sha1());
      HMAC_Update (&hmctx, key, len);
      HMAC_Final (&hmctx, md_value, &md_len);
      for (i = 1; i < rounds; ++i) {
         HMAC_Init (&hmctx, key, IONU_AES_KEY_LEN, EVP_sha1());
         HMAC_Update (&hmctx, md_value, md_len);
         HMAC_Final (&hmctx, md_value, &md_len);
      }
      HMAC_CTX_cleanup (&hmctx);
      HMAC_cleanup (&hmctx);

      BinaryToHex (md_value, md_len, hex);
      ms =  ionu::GetSystemTime() - ms;
      printf ("hmac: rounds = %d, time = %ld ms\n", rounds, ms);

   }
   else if (strcmp ("pbkdf2", argv[1]) == 0) {
      unsigned char pbkey[IONU_AES_KEY_LEN];
      if (password)
         ionu::DeriveKey (password, pbkey);
      else
         PKCS5_PBKDF2_HMAC_SHA1 (password, len, salt, sizeof(salt), rounds, IONU_AES_KEY_LEN, pbkey);
      BinaryToHex (pbkey, IONU_AES_KEY_LEN, hex);
      ms =  ionu::GetSystemTime() - ms;
      if (password)
         printf ("pbkdf2: password = %s, time = %ld ms\n", password, ms);
      else
         printf ("pbkdf2: rounds = %d, time = %ld ms\n", rounds, ms);
   }
   else if (strcmp ("rsa-encrypt", argv[1]) == 0) {
      for (i = 0; i < rounds; ++i) {
         ct = publicK->PublicKeyEncrypt ((unsigned char*)PERMS, strlen (PERMS), &ebytes);
         delete[] ct;
      }
      ms =  ionu::GetSystemTime() - ms;
      printf ("rsa-encrypt: rounds = %d, time = %ld ms\n", rounds, ms);
   }
   else if (strcmp ("rsa-decrypt", argv[1]) == 0) {
      size_t bytes = 0;
      for (i = 0; i < rounds; ++i) {
         unsigned char* back = privateK->PrivateKeyDecrypt (ct, ebytes, &bytes);
         delete[] back;
      }
      ms =  ionu::GetSystemTime() - ms;
      printf ("rsa-decrypt: rounds = %d, time = %ld ms\n", rounds, ms);
   }
   else if (strcmp ("ec-encrypt", argv[1]) == 0) {
      for (i = 0; i < rounds; ++i) {
         ect = me->ECDHEncrypt (peer, (unsigned char*)PERMS, strlen (PERMS), &ecbytes);
         delete[] ect;
      }
      ms =  ionu::GetSystemTime() - ms;
      printf ("ec-encrypt: rounds = %d, time = %ld ms\n", rounds, ms);
   }
   else if (strcmp ("ec-decrypt", argv[1]) == 0) {
      size_t bytes = 0;
      for (i = 0; i < rounds; ++i) {
         unsigned char* back = peer->ECDHDecrypt (me, ect, ecbytes, &bytes);
         delete[] back;
      }
      ms =  ionu::GetSystemTime() - ms;
      printf ("ec-decrypt: rounds = %d, time = %ld ms\n", rounds, ms);
   }
   else if (strcmp ("aes256", argv[1]) == 0) {
      for (i = 0; i < rounds; ++i) {
         ct = aesKey->SymmetricEncryptBuffer ((unsigned char*)PERMS, strlen (PERMS), &ebytes, iv);
         delete[] ct;
      }
      ms =  ionu::GetSystemTime() - ms;
      printf ("aes: rounds = %d, time = %ld ms\n", rounds, ms);
   }
   else {
      usage();
      exit (1);
   }

   exit (0);
}