#include "utils.h"

#define HASH_LENGTH 32
#define NONCE_LENGTH 16
#define KEY_LENGTH 16
#define MILLION 1E6
#define IV_LENGTH 12
#define EC_KEY_SIZE 32
#define KEY_LENGTH 16
#define TAG_SIZE 16
#define CLOCKS_PER_MS (CLOCKS_PER_SEC/1000)
#define AES_GCM_BLOCK_SIZE_IN_BYTES 16
#define PRIME256V1_KEY_SIZE 32

int AES_GCM_128_encrypt (unsigned char *plaintext, int plaintext_len, unsigned char *aad,
  int aad_len, unsigned char *key, unsigned char *iv, int iv_len,
  unsigned char *ciphertext, unsigned char *tag)
{
  EVP_CIPHER_CTX *ctx;

  int len;
  int ciphertext_len;

  /* Create and initialise the context */
  if(!(ctx = EVP_CIPHER_CTX_new())) {
    printf("Failed context intialization for OpenSSL EVP\n");
  }

  /* Initialise the encryption operation. */
  if(1 != EVP_EncryptInit_ex(ctx, EVP_aes_128_gcm(), NULL, NULL, NULL)){
    printf("Failed AES_GCM_128 intialization for OpenSSL EVP\n");	
  }

  /* Set IV length if default 12 bytes (96 bits) is not appropriate */
  if(1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, iv_len, NULL)){
    printf("Failed IV config\n");
  }

  /* Initialise key and IV */
  if(1 != EVP_EncryptInit_ex(ctx, NULL, NULL, key, iv)) {
    printf("Failed intialization for key and IV for AES_GCM\n");	
  }

  /* Provide any AAD data. This can be called zero or more times as
   * required
   */
  //if(1 != EVP_EncryptUpdate(ctx, NULL, &len, aad, aad_len))
  //	printf("Failed AAD\n");

  /* Provide the message to be encrypted, and obtain the encrypted output.
   * EVP_EncryptUpdate can be called multiple times if necessary
   */
  //printf("Error code = %d\n\n", EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len));
  if(1 != EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len)) {
    printf("Failed AES_GCM encrypt\n");
  }
  ciphertext_len = len;

  /* Finalise the encryption. Normally ciphertext bytes may be written at
   * this stage, but this does not occur in GCM mode
   */
  if(1 != EVP_EncryptFinal_ex(ctx, ciphertext + len, &len)){
    printf("Failed Finalizing ciphertext\n");	
  }
  ciphertext_len += len;

  /* Get the tag */
  if(1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, TAG_SIZE, tag))
    printf("Failed tag for AES_GCM_encrypt\n");

  /* Clean up */
  EVP_CIPHER_CTX_free(ctx);

  return ciphertext_len;
}

int AES_GCM_128_decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char *aad,
  int aad_len, unsigned char *tag, unsigned char *key, unsigned char *iv,
  int iv_len, unsigned char *plaintext)
{
  EVP_CIPHER_CTX *ctx;
  int len;
  int plaintext_len;
  int ret;

  /* Create and initialise the context */
  if(!(ctx = EVP_CIPHER_CTX_new())) 
    printf("Failed context intialization for OpenSSL EVP\n");

  /* Initialise the decryption operation. */
  if(!EVP_DecryptInit_ex(ctx, EVP_aes_128_gcm(), NULL, NULL, NULL))
    printf("Failed AES_GCM_128 intialization for OpenSSL EVP\n");	

  /* Set IV length. Not necessary if this is 12 bytes (96 bits) */
  if(!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, iv_len, NULL))
    printf("Failed IV config\n");

  /* Initialise key and IV */
  if(!EVP_DecryptInit_ex(ctx, NULL, NULL, key, iv))
    printf("Failed intialization for key and IV for AES_GCM_128\n");	
  /* Provide any AAD data. This can be called zero or more times as
   * required
   */
  //if(!EVP_DecryptUpdate(ctx, NULL, &len, aad, aad_len))
  //	printf("Failed AAD\n");

  /* Provide the message to be decrypted, and obtain the plaintext output.
   * EVP_DecryptUpdate can be called multiple times if necessary
   */
  if(!EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len))
    printf("Failed AES_GCM decrypt\n");
  plaintext_len = len;

  /* Set expected tag value. Works in OpenSSL 1.0.1d and later */
  if(!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, TAG_SIZE, tag))
    printf("Failed tag for AES_GCM_decrypt\n");

  /* Finalise the decryption. A positive return value indicates success,
   * anything else is a failure - the plaintext is not trustworthy.
   */
  ret = EVP_DecryptFinal_ex(ctx, plaintext + len, &len);

  /* Clean up */
  EVP_CIPHER_CTX_free(ctx);

  if(ret > 0)
  {
    /* Success */
    plaintext_len += len;
    return plaintext_len;
  }
  else
  {
    /* Verify failed */
    return -1;
  }
}


int serializedSizeLSORAMRequest(uint32_t key_size, uint32_t value_size) {
  uint32_t serialized_request_size = key_size + value_size + 2*sizeof(uint32_t);
  return serialized_request_size;
}

/*

Inputs: a target pub key, a seriailzed request and request size.
Outputs: instantiates and populates serialized_request

*/
int32_t serializeLSORAMRequest(unsigned char *key, uint32_t key_size,
         unsigned char *value, uint32_t value_size, unsigned char** serialized_request) {
  // SerializedRequest:
  // key_size, key bytes, value_size, value bytes
  uint32_t serialized_request_size = serializedSizeLSORAMRequest(key_size, value_size);
  *serialized_request = (unsigned char*) malloc(serialized_request_size);
  
  unsigned char* req_ptr = *serialized_request;
  memcpy(req_ptr, &key_size, sizeof(uint32_t));
  req_ptr+=sizeof(uint32_t);
  memcpy(req_ptr, key, key_size);
  req_ptr+=key_size;
  
  memcpy(req_ptr, &value_size, sizeof(uint32_t));
  req_ptr+=sizeof(uint32_t);
  memcpy(req_ptr, value, value_size);
  req_ptr+=value_size;
  
  return serialized_request_size; 
}

int ECDH_encrypt(const string &request, uint32_t request_size, const string &params, string &shared_secret, string &pirquery){ 
  /* 
  1) Extracts Enclave public key from params
  2) Samples a public key for itself
  3) Performs ECDH with enclave public key and stores to encrypted_request: 
     <client_sampled_public_key, request> encrypted with shared_secret
  */

  int ret;
  EC_KEY *ephemeral_key = NULL;
  BIGNUM *x, *y;
  x = BN_new();
  y = BN_new();
  BN_CTX *bn_ctx = BN_CTX_new();
  const EC_GROUP *curve = NULL;

  if(NULL == (curve = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1)))
    printf("Setting EC_GROUP failed \n");

  ephemeral_key = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
  if(ephemeral_key==NULL)
    printf("Client: EC_KEY_new_by_curve_name Fail\n");

  ret = EC_KEY_generate_key(ephemeral_key);
  if(ret!=1)
    printf("Client: EC_KEY_generate_key Fail\n");

  const EC_POINT *pub_point;
  pub_point = EC_KEY_get0_public_key((const EC_KEY *) ephemeral_key);
  if(pub_point == NULL)
    printf("Client: EC_KEY_get0_public_key Fail\n");
      
  ret = EC_POINT_get_affine_coordinates_GFp(curve, pub_point, x, y, bn_ctx);
  if(ret==0)
    printf("Client: EC_POINT_get_affine_coordinates_GFp Failed \n");
      
  unsigned char *bin_x, *bin_y;
  uint32_t size_bin_x = BN_num_bytes(x);
  uint32_t size_bin_y = BN_num_bytes(y);
  bin_x = (unsigned char*) malloc(size_bin_x);
  bin_y = (unsigned char*) malloc(size_bin_y);
  BN_bn2bin(x, bin_x);
  BN_bn2bin(y, bin_y);
  unsigned char *serialized_client_public_key = (unsigned char*) malloc(size_bin_x + size_bin_y);
  memcpy(serialized_client_public_key, bin_x, size_bin_x);
  memcpy(serialized_client_public_key + size_bin_x, bin_y, size_bin_y);
  //Done with serializeing sampled client public key
  /*TEST snippet for key comparison at client and enclave
  const EC_POINT *point = EC_KEY_get0_public_key(ENCLAVE_PUBLIC_KEY);
  BIGNUM *x1, *y1;
  x1 = BN_new();
  y1 = BN_new();
  ret = EC_POINT_get_affine_coordinates_GFp(curve, point, x1, y1, bn_ctx);
  unsigned char *bin_point = (unsigned char*) malloc(32*2);
  BN_bn2bin(x1,bin_point);
  BN_bn2bin(y1,bin_point+32);	

  printf("Serialized Client's Public Key at Client :\n");
  for(int t = 0; t < size_bin_x+size_bin_y; t++)
	  printf("%02X", (*serialized_client_public_key)[t]);
  printf("\n");

  printf("Serialized Enclave's Public Key at Client :\n");
  for(int t = 0; t < size_bin_x+size_bin_y; t++)
	  printf("%02X", bin_point[t]);
  printf("\n");
  */


  //Load Enclave public key from string params
  EC_KEY *ENCLAVE_PUBLIC_KEY;
  ENCLAVE_PUBLIC_KEY = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
  BIGNUM *x_enclave, *y_enclave;
  EC_POINT *pub_point_enclave = EC_POINT_new(curve);
  x_enclave = BN_new();
  y_enclave = BN_new();

  unsigned char *bin_x_enclave = (unsigned char*) malloc(PRIME256V1_KEY_SIZE);
  unsigned char *bin_y_enclave = (unsigned char*) malloc(PRIME256V1_KEY_SIZE);
  for(uint32_t i=0; i<PRIME256V1_KEY_SIZE; i++){
    bin_x_enclave[i] = params[i];
    bin_y_enclave[i] = params[PRIME256V1_KEY_SIZE+i];
  }
  x_enclave = BN_bin2bn(bin_x_enclave, PRIME256V1_KEY_SIZE, NULL);
  y_enclave = BN_bin2bn(bin_y_enclave, PRIME256V1_KEY_SIZE, NULL);

  if(EC_POINT_set_affine_coordinates_GFp(curve, pub_point_enclave, x_enclave, y_enclave, bn_ctx)==0)
    printf("EC_POINT_set_affine_coordinates FAILED \n");

  if(EC_KEY_set_public_key(ENCLAVE_PUBLIC_KEY, pub_point_enclave)==0)
    printf("EC_KEY_set_public_key FAILED \n");


  //Do the ECDH And encrypt request
  uint32_t field_size = EC_GROUP_get_degree(EC_KEY_get0_group(ENCLAVE_PUBLIC_KEY));
  uint32_t secret_len = (field_size+7)/8;
  unsigned char *secret = (unsigned char*) malloc(secret_len);
  //Returns a 32 byte secret	
  secret_len = ECDH_compute_key(secret, secret_len, EC_KEY_get0_public_key(ENCLAVE_PUBLIC_KEY),
					  ephemeral_key, NULL);

  unsigned char *ecdh_shared_aes_key = (unsigned char *) malloc(KEY_LENGTH);
  unsigned char *ecdh_iv = (unsigned char*) malloc(IV_LENGTH);
  memcpy(ecdh_shared_aes_key, secret, KEY_LENGTH);
  memcpy(ecdh_iv, secret + KEY_LENGTH, IV_LENGTH);
 
  //
  /*	
  printf("Secret computed by Client :\n");
  for(int t = 0; t < secret_len; t++)
	  printf("%02X", secret[t]);
  printf("\n");
  */

  int encrypted_request_size = request_size;
  unsigned char *request_cstr = (unsigned char*) malloc (request_size);
  for(uint32_t i = 0; i<request_size; i++)
    request_cstr[i] = request[i];
  unsigned char *encrypted_request_cstr = (unsigned char*) malloc (request_size);	
  unsigned char *tag_cstr = (unsigned char*) malloc (TAG_SIZE);
  encrypted_request_size = AES_GCM_128_encrypt(request_cstr, request_size, NULL, 0, (unsigned char*) ecdh_shared_aes_key, ecdh_iv, IV_LENGTH, encrypted_request_cstr, tag_cstr);
  
  //pirquery = aes_key||iv||plainquery_{aes_key}||tag
  pirquery.assign((const char*) ecdh_shared_aes_key, PRIME256V1_KEY_SIZE);
  pirquery.append((const char*) ecdh_iv, IV_LENGTH);
  pirquery.append((const char*) encrypted_request_cstr, encrypted_request_size);
  pirquery.append((const char*) tag_cstr, TAG_SIZE);

  BN_CTX_free(bn_ctx);
  free(request_cstr);
  free(ecdh_shared_aes_key);
  free(ecdh_iv);
  free(encrypted_request_cstr);
  free(tag_cstr);
  return encrypted_request_size;
}

//hybridEncrypt();