#include "security.hpp"

#include <openssl/conf.h>
#include <openssl/err.h>
#include <openssl/evp.h>
// AES
#include <openssl/aes.h>
// ECDH
#include <openssl/ec.h>
#include <openssl/pem.h>

#include <stdexcept>

using namespace binom;


void Security::handleErrors() {
  ERR_print_errors_fp(stderr);
  throw std::runtime_error("Security error");
}

EVP_PKEY* Security::genKey() {
  EVP_PKEY* key_pair = nullptr;
  EVP_PKEY_CTX* param_gen_ctx = nullptr;
  EVP_PKEY_CTX* key_gen_ctx = nullptr;
  EVP_PKEY* params= nullptr;

  if(!(param_gen_ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_EC, NULL))) handleErrors();
  if(!EVP_PKEY_paramgen_init(param_gen_ctx)) handleErrors();

  if(!EVP_PKEY_CTX_set_ec_paramgen_curve_nid(param_gen_ctx, NID_X9_62_prime256v1))
    handleErrors();

  if(!EVP_PKEY_paramgen(param_gen_ctx, &params)) handleErrors();

  if(!(key_gen_ctx = EVP_PKEY_CTX_new(params, nullptr))) handleErrors();
  if(!EVP_PKEY_keygen_init(key_gen_ctx)) handleErrors();
  if(!EVP_PKEY_keygen(key_gen_ctx, &key_pair)) handleErrors();

  EVP_PKEY_CTX_free(param_gen_ctx);
  EVP_PKEY_CTX_free(key_gen_ctx);
  return key_pair;
}

void Security::freeKey(EVP_PKEY* key) {EVP_PKEY_free(key);}

ByteArray Security::extractPublicKey(EVP_PKEY* key_pair) {
  ByteArray data(i2d_PublicKey(key_pair, nullptr));
  unsigned char* ptr = data.begin();
  i2d_PublicKey(key_pair, &ptr);
  return data;
}

ByteArray Security::extractPrivateKey(EVP_PKEY* key_pair) {
  EC_KEY* ec_key = EVP_PKEY_get1_EC_KEY(key_pair);
  const BIGNUM* ec_priv = EC_KEY_get0_private_key(ec_key);
  int length = BN_bn2mpi(ec_priv, nullptr);
  ByteArray data(length);
  BN_bn2mpi(ec_priv, data.begin());
  return data;
}

EVP_PKEY* Security::getKeyPair(ByteArray priv_key_raw, ByteArray pub_key_raw) {
  EVP_PKEY* key_pair = EVP_PKEY_new();
  EC_KEY *ec_key = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);

  const EC_GROUP* ec_group = EC_KEY_get0_group(ec_key);
  EC_POINT* ec_point = EC_POINT_new(ec_group);
  EC_POINT_oct2point(ec_group, ec_point, pub_key_raw.begin(), pub_key_raw.length(), nullptr);
  EC_KEY_set_public_key(ec_key, ec_point);
  EC_POINT_free(ec_point);

  BIGNUM* priv = BN_mpi2bn(priv_key_raw.begin(), priv_key_raw.length(), nullptr);
  EC_KEY_set_private_key(ec_key, priv);
  BN_free(priv);

  EVP_PKEY_set1_EC_KEY(key_pair, ec_key);
  EC_KEY_free(ec_key);
  return key_pair;
}

Security::AES_t Security::getSecret(ByteArray peer_key, EVP_PKEY* key_pair) {
  EC_KEY *temp_ec_key = nullptr;
  EVP_PKEY *peerkey = nullptr;

  temp_ec_key = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
  if(temp_ec_key == nullptr)
    handleErrors();
  if(EC_KEY_oct2key(temp_ec_key, peer_key.begin(), peer_key.length(), NULL) != 1)
    handleErrors();
  if(EC_KEY_check_key(temp_ec_key) != 1) handleErrors();
  peerkey = EVP_PKEY_new();
  if(peerkey == NULL)
    handleErrors();
  if(EVP_PKEY_assign_EC_KEY(peerkey, temp_ec_key)!= 1)
    handleErrors();

  // Getting secret
  EVP_PKEY_CTX *derivation_ctx = EVP_PKEY_CTX_new(key_pair, NULL);
  EVP_PKEY_derive_init(derivation_ctx);
  EVP_PKEY_derive_set_peer(derivation_ctx, peerkey);
  size_t lenght;
  void* ptr;
  if(1 != EVP_PKEY_derive(derivation_ctx, NULL, &lenght)) handleErrors();
  if(NULL == (ptr = OPENSSL_malloc(lenght))) handleErrors();
  if(1 != (EVP_PKEY_derive(derivation_ctx, (unsigned char*)ptr, &lenght))) handleErrors();
  EVP_PKEY_CTX_free(derivation_ctx);
  EVP_PKEY_free(peerkey);

  AES_t aes_key;
  EVP_MD_CTX *mdctx;
  if((mdctx = EVP_MD_CTX_new()) == NULL)
    handleErrors();
  if(1 != EVP_DigestInit_ex(mdctx, EVP_sha384(), NULL))
    handleErrors();
  if(1 != EVP_DigestUpdate(mdctx, ptr, lenght))
    handleErrors();
  unsigned int length;
  if(1 != EVP_DigestFinal_ex(mdctx, (unsigned char*)&aes_key, &length))
    handleErrors();
  EVP_MD_CTX_free(mdctx);
  OPENSSL_free(ptr);
  return aes_key;
}

Security::AES_t Security::createAES256Key(ByteArray data) {
  AES_t aes_key;
  EVP_MD_CTX *mdctx;
  if((mdctx = EVP_MD_CTX_new()) == NULL)
    handleErrors();
  if(1 != EVP_DigestInit_ex(mdctx, EVP_sha3_384(), NULL))
    handleErrors();
  if(1 != EVP_DigestUpdate(mdctx, data.begin(), data.length()))
    handleErrors();
  unsigned int length;
  if(1 != EVP_DigestFinal_ex(mdctx, (unsigned char*)&aes_key, &length))
    handleErrors();
  EVP_MD_CTX_free(mdctx);
  return aes_key;
}

ByteArray Security::getHashOf(ByteArray data, HashMethod hash_method) {return getHashOf(std::move(data), ByteArray(), hash_method);}
ByteArray Security::getHashOf(ByteArray data, ByteArray salt, Security::HashMethod hash_method) {
  auto getHashMethod = [](HashMethod hash_method)->const EVP_MD* {
    switch (hash_method) {
      case Security::HashMethod::sha1: return EVP_sha1();
      case Security::HashMethod::sha224: return EVP_sha224();
      case Security::HashMethod::sha256: return EVP_sha256();
      case Security::HashMethod::sha384: return EVP_sha384();
      case Security::HashMethod::sha512: return EVP_sha512();
      case Security::HashMethod::sha512_224: return EVP_sha512_224();
      case Security::HashMethod::sha512_256: return EVP_sha512_256();
      case Security::HashMethod::sha3_224: return EVP_sha3_224();
      case Security::HashMethod::sha3_256: return EVP_sha3_256();
      case Security::HashMethod::sha3_384: return EVP_sha3_384();
      case Security::HashMethod::sha3_512: return EVP_sha3_512();
      case Security::HashMethod::shake128: return EVP_shake128();
      case Security::HashMethod::shake256: return EVP_shake256();
      default: return EVP_sha3_512();
    }
  };

  auto getHashSize = [](HashMethod hash_method)->size_t {
    switch (hash_method) {
      case Security::HashMethod::sha1: return 20;
      case Security::HashMethod::sha224: return 28;
      case Security::HashMethod::sha256: return 32;
      case Security::HashMethod::sha384: return 48;
      case Security::HashMethod::sha512: return 64;
      case Security::HashMethod::sha512_224: return 28;
      case Security::HashMethod::sha512_256: return 32;
      case Security::HashMethod::sha3_224: return 28;
      case Security::HashMethod::sha3_256: return 32;
      case Security::HashMethod::sha3_384: return 48;
      case Security::HashMethod::sha3_512: return 64;
      case Security::HashMethod::shake128: return 16;
      case Security::HashMethod::shake256: return 32;
      default: return 64;
    }
  };

  data.pushBack(std::move(salt));
  ByteArray hash(getHashSize(hash_method));
  EVP_MD_CTX *mdctx;
  if((mdctx = EVP_MD_CTX_new()) == NULL)
    handleErrors();
  if(1 != EVP_DigestInit_ex(mdctx, getHashMethod(hash_method), NULL))
    handleErrors();
  if(1 != EVP_DigestUpdate(mdctx, data.begin(), data.length()))
    handleErrors();
  unsigned int length;
  if(1 != EVP_DigestFinal_ex(mdctx, hash.begin(), &length))
    handleErrors();
  EVP_MD_CTX_free(mdctx);
  return hash;
}

ByteArray Security::encodeBase64(ByteArray decoded) {
  ByteArray encoded((4*((decoded.length()+2)/3)) + 1);
  EVP_EncodeBlock(encoded.begin(), decoded.begin(), decoded.length());
  return encoded;
}

ByteArray Security::decodeBase64(ByteArray encoded) {
  ByteArray decoded((3*encoded.length()/4) + 1);
  size_t recived_data_size = EVP_DecodeBlock(decoded.begin(), encoded.begin(), encoded.length());
  if(recived_data_size < decoded.length())
    decoded.resize(recived_data_size);
  return decoded;
}



binom::ByteArray Security::encrypt(binom::ByteArray plain_text, AES_t aes_struct) {
  /* max ciphertext len for a n bytes of plaintext is n + AES_BLOCK_SIZE -1 bytes */
  ByteArray ciphertext(plain_text.length() % AES_BLOCK_SIZE == 0
                       ? plain_text.length()
                       : (plain_text.length() / AES_BLOCK_SIZE + 1) * AES_BLOCK_SIZE);

  EVP_CIPHER_CTX *ctx;
  if(!(ctx = EVP_CIPHER_CTX_new()))
    handleErrors();
  if(1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), nullptr, aes_struct.key, aes_struct.init_vector))
    handleErrors();

  int f_length, s_length;
  if(1 != EVP_EncryptUpdate(ctx, ciphertext.begin(), &f_length, plain_text.begin(), plain_text.length()))
    handleErrors();

  if(ui64(f_length) == ciphertext.length())
    ciphertext.addSize(AES_BLOCK_SIZE);
  elif(ui64(f_length) > ciphertext.length())
    throw std::runtime_error("Predicted ciphertext size lower then actual!");

  if(1 != EVP_EncryptFinal_ex(ctx, ciphertext.begin() + f_length, &s_length))
    handleErrors();
  if(ui64 reuired_length = f_length + s_length;reuired_length < ciphertext.length())
    ciphertext.resize(f_length + s_length);
  elif(reuired_length > ciphertext.length())
    throw std::runtime_error("Predicted ciphertext size lower then actual!");

  EVP_CIPHER_CTX_free(ctx);

  return ciphertext;
}

ByteArray Security::decrypt(ByteArray ciphertext, Security::AES_t aes_struct) {
  EVP_CIPHER_CTX *ctx;
  ByteArray plain_text(ciphertext.length());

  if(!(ctx = EVP_CIPHER_CTX_new()))
    handleErrors();

  if(1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, aes_struct.key, aes_struct.init_vector))
    handleErrors();

  int f_length, s_length;
  if(1 != EVP_DecryptUpdate(ctx, plain_text.begin(), &f_length, ciphertext.begin(), ciphertext.length()))
    handleErrors();

  if(1 != EVP_DecryptFinal_ex(ctx, plain_text.begin() + f_length, &s_length))
    handleErrors();
  plain_text.resize(f_length + s_length);

  EVP_CIPHER_CTX_free(ctx);

  return plain_text;
}

std::string uuid::generateUUIDv4() {
  std::stringstream ss;
  int i;
  ss << std::hex;
  for (i = 0; i < 8; i++) {
    ss << dis(gen);
  }
  ss << "-";
  for (i = 0; i < 4; i++) {
    ss << dis(gen);
  }
  ss << "-4";
  for (i = 0; i < 3; i++) {
    ss << dis(gen);
  }
  ss << "-";
  ss << dis2(gen);
  for (i = 0; i < 3; i++) {
    ss << dis(gen);
  }
  ss << "-";
  for (i = 0; i < 12; i++) {
    ss << dis(gen);
  };
  return ss.str();
}
