#include <openssl/ec.h>
#include <openssl/bn.h>
#include <openssl/rsa.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/rand.h>
#include <openssl/ecdh.h>
#include <string.h>
EVP_CIPHER_CTX* ctx;

int generate_sha256_hash(const unsigned char *message, size_t message_len, unsigned char *digest)
{
	EVP_MD_CTX *mdctx; unsigned int digest_len;

	if((mdctx = EVP_MD_CTX_create()) == NULL)
		return 0x1;

	if(EVP_DigestInit_ex(mdctx, EVP_sha256(), NULL) != 1)
		return 0x2;

  if(EVP_DigestUpdate(mdctx, message, message_len) != 1)
		return 0x3;

	if(EVP_DigestFinal_ex(mdctx, digest, &digest_len) != 1)
		return 0x4;

	if(digest_len != 32)
		return 0x5;

	EVP_MD_CTX_destroy(mdctx);
	return 0;
}

int ecdh_key_gen(unsigned char* public_key_x, unsigned char* public_key_y, unsigned char* private_key)
{
	EC_KEY * keypair = NULL;
	EC_GROUP* ecdh_group = NULL;
	const EC_POINT *public_key_point = NULL;
	const BIGNUM *private_key_bignum = NULL;
	BIGNUM *public_key_x_bignum = NULL;
	BIGNUM *public_key_y_bignum = NULL;
	int ret;

	keypair = EC_KEY_new();
	if (NULL == keypair) {
		return 0x42;
	}

	ecdh_group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
	if (NULL == ecdh_group) {
		EC_KEY_free(keypair);
		return 0x0f;
	}

	if (0 == EC_KEY_set_group (keypair, ecdh_group)) {
		EC_KEY_free(keypair);
		EC_GROUP_free(ecdh_group);
		return 0x66;
	}

	// TODO: Check return values.
	EC_KEY_set_asn1_flag(keypair, OPENSSL_EC_NAMED_CURVE);

	// Generating keypair
	if (EC_KEY_generate_key(keypair) == 0) {
		EC_KEY_free(keypair);
		EC_GROUP_free(ecdh_group);
    return 0x01;
	}

	// Extracting public key point from keypair. TODO: Note that it seems that this pointer does not need to be freed separately from the keypair pointer.
	public_key_point = EC_KEY_get0_public_key(keypair);
	if (NULL == public_key_point)
	{
		EC_KEY_free(keypair);
		EC_GROUP_free(ecdh_group);
		return 0x04;
	}

	// Extracting private key bignum from keypair. TODO: Note that it seems that this pointer does not need to be freed separately from the keypair pointer.
	private_key_bignum = EC_KEY_get0_private_key(keypair);
	if (NULL == private_key_bignum)
	{
		EC_KEY_free(keypair);
		EC_GROUP_free(ecdh_group);
		return 0x05;
	}

	// Converting private key bignum to string (in big-endian format).
  if(!BN_bn2bin(private_key_bignum, private_key))
  {
    EC_KEY_free(keypair);
		EC_GROUP_free(ecdh_group);
		return 0x06;
  }

  public_key_x_bignum = BN_new();
  public_key_y_bignum = BN_new();
  if (NULL ==  public_key_x_bignum || NULL == public_key_y_bignum) {
		EC_KEY_free(keypair);
		EC_GROUP_free(ecdh_group);
    return 0x42;//SGX_ERROR_OUT_OF_MEMORY;
  }

	// Extracting public key bignums from public key points.
	if (!EC_POINT_get_affine_coordinates_GFp(ecdh_group, public_key_point, public_key_x_bignum, public_key_y_bignum, NULL))
	{
		BN_clear_free(public_key_x_bignum);
		BN_clear_free(public_key_y_bignum);
		EC_KEY_free(keypair);
		EC_GROUP_free(ecdh_group);
		return 0x07;
	}

	// Converting public key x bignum to string (in big-endian format) and setting first output argument.
	ret = BN_bn2bin(public_key_x_bignum, public_key_x);
	BN_clear_free(public_key_x_bignum);
	if(ret == 0)
	{
		BN_clear_free(public_key_y_bignum);
    EC_KEY_free(keypair);
	  EC_GROUP_free(ecdh_group);
		return 0x08;
	}

	// Converting public key y bignum to string (in big-endian format) and setting second output argument.
	ret = BN_bn2bin(public_key_y_bignum, public_key_y);
	BN_clear_free(public_key_y_bignum);
	EC_KEY_free(keypair);
	EC_GROUP_free(ecdh_group);

	if(ret == 0)
  		return 0x09;

	return 0;
}

unsigned long check_ecdh_public_key(unsigned char* public_key_x, unsigned char* public_key_y) //, unsigned char* priv_key, unsigned char* shared_key_x, unsigned char* shared_key_y)
{
	EC_GROUP *ec_group = NULL;
	EC_POINT *public_key_point = NULL;
	BIGNUM *public_key_x_bignum = NULL;
	BIGNUM *public_key_y_bignum = NULL;

	ec_group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
  if (NULL == ec_group) {
    return 0x0f;
  }

	public_key_point = EC_POINT_new(ec_group);
	if (NULL == public_key_point) {
		EC_GROUP_free(ec_group);
		return 0x42;
	}

	public_key_x_bignum = BN_bin2bn(public_key_x, 32, NULL);
	if (NULL == public_key_x_bignum) {
		EC_POINT_free(public_key_point);
		EC_GROUP_free(ec_group);
		return 0xfd;
	}

	// converts the y value of the point, represented as positive integer in big-endian into a BIGNUM
	public_key_y_bignum = BN_bin2bn(public_key_y, 32, NULL);
	if (NULL == public_key_y_bignum) {
		BN_clear_free(public_key_x_bignum);
		EC_POINT_free(public_key_point);
		EC_GROUP_free(ec_group);
		return 0xfe;
	}

	// sets point based on x,y coordinates
	int ret = EC_POINT_set_affine_coordinates_GFp(ec_group, public_key_point, public_key_x_bignum, public_key_y_bignum, NULL);
	BN_clear_free(public_key_x_bignum);
	BN_clear_free(public_key_y_bignum);
	if (1 != ret) {
		EC_POINT_free(public_key_point);
		EC_GROUP_free(ec_group);
		return ERR_get_error();
	}

	// checks if point is on curve
	//
	int ret_point_on_curve = EC_POINT_is_on_curve(ec_group, public_key_point, NULL);
	EC_POINT_free(public_key_point);
	EC_GROUP_free(ec_group);
	if (1 != ret_point_on_curve) {
		return ret_point_on_curve << 4; //ret_point_on_curve;//	break;
	}

  return 0;
}

unsigned long compute_ecdh_shared_key(unsigned char* given_public_key_x, unsigned char* given_public_key_y, unsigned char* own_private_key, unsigned char* derived_key)
{
	unsigned long long_ret = 0;
	EC_GROUP *ec_group = NULL;
	EC_POINT *given_public_key_point = NULL;
	BIGNUM *given_public_key_x_bignum = NULL;
	BIGNUM *given_public_key_y_bignum = NULL;
	EC_KEY *own_keypair = NULL;
	BIGNUM *own_private_key_bignum = NULL;
	int shared_key_len = 32;
	unsigned char shared_key[32];

	ec_group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
  if (NULL == ec_group) {
    return 0x0f;
  }

	given_public_key_point = EC_POINT_new(ec_group);
	if (NULL == given_public_key_point) {
		EC_GROUP_free(ec_group);
		return 0x42;
	}

	given_public_key_x_bignum = BN_bin2bn(given_public_key_x, 32, NULL);
	if (NULL == given_public_key_x_bignum) {
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		return 0xfd;
	}

	// converts the y value of the point, represented as positive integer in big-endian into a BIGNUM
	given_public_key_y_bignum = BN_bin2bn(given_public_key_y, 32, NULL);
	if (NULL == given_public_key_y_bignum) {
		BN_clear_free(given_public_key_x_bignum); // TODO: You should probably not be freeing it here - since ya dont own it?
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		return 0xfe;
	}

	// sets point based on x,y coordinates
	int ret = EC_POINT_set_affine_coordinates_GFp(ec_group, given_public_key_point, given_public_key_x_bignum, given_public_key_y_bignum, NULL);
	BN_clear_free(given_public_key_x_bignum);
	BN_clear_free(given_public_key_y_bignum);
	if (1 != ret) {
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		return ERR_get_error();
	}

	// checks if point is on curve
	//
	int ret_point_on_curve = EC_POINT_is_on_curve(ec_group, given_public_key_point, NULL);
	if (1 != ret_point_on_curve) {
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		return ret_point_on_curve << 4; //ret_point_on_curve;//	break;
	}

	// create empty shared key BN
	//
	own_keypair = EC_KEY_new();
	if (own_keypair == NULL) {
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		return 0x42;//ret = SGX_ERROR_OUT_OF_MEMORY;
	}

	// init private key group (set curve)
	//
	if (EC_KEY_set_group (own_keypair, ec_group) != 1) {
		EC_KEY_free(own_keypair);
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		return ERR_get_error(); //break;
	}

	// init private key with BN value
	//

	own_private_key_bignum = BN_bin2bn(own_private_key, 32, NULL);
	if(own_private_key_bignum == NULL) {
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		EC_KEY_free(own_keypair);
		return 0x44;
	}

	long_ret = EC_KEY_set_private_key(own_keypair, own_private_key_bignum);
	BN_clear_free(own_private_key_bignum);
	if (long_ret != 1) {
		EC_POINT_free(given_public_key_point);
		EC_GROUP_free(ec_group);
		EC_KEY_free(own_keypair);
		return ERR_get_error();
	}

	// calculate shared dh key
	//
	shared_key_len = ECDH_compute_key(shared_key, 32, given_public_key_point, own_keypair, NULL);
	if (shared_key_len <= 0) {
		long_ret = ERR_get_error();
	}
	else
		long_ret = generate_sha256_hash(shared_key, 32, derived_key);

	EC_POINT_free(given_public_key_point);
	EC_GROUP_free(ec_group);
	EC_KEY_free(own_keypair);
	return long_ret;
}

int compute_ecdsa_signature(unsigned char* signature_data, uint32_t signature_data_length, unsigned char* own_private_key, unsigned char* signature)
{
		unsigned long long_ret = 0;
		EC_GROUP *ec_group = NULL;
		EC_POINT *given_public_key_point = NULL;
		EC_KEY *own_keypair = NULL;
		BIGNUM *own_private_key_bignum = NULL;
		unsigned char effective_signature_data[32];
		ECDSA_SIG *openssl_signature;
		unsigned char** temp = NULL;
		uint32_t signature_length;
		// &signature_length will not be NULL

		ec_group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
	  if (NULL == ec_group) {
	    return 0x0f;
	  }

		own_keypair = EC_KEY_new();
		if (own_keypair == NULL) {
			EC_GROUP_free(ec_group);
			return 0x42;//ret = SGX_ERROR_OUT_OF_MEMORY;
		}

		// init private key group (set curve)
		//
		if (EC_KEY_set_group (own_keypair, ec_group) != 1) {
			EC_KEY_free(own_keypair);
			EC_GROUP_free(ec_group);
			return ERR_get_error(); //break;
		}

		// init private key with BN value
		//
		own_private_key_bignum = BN_bin2bn(own_private_key, 32, NULL);
		if(own_private_key_bignum == NULL) {
			EC_GROUP_free(ec_group);
			EC_KEY_free(own_keypair);
			return 0x44;
		}

		long_ret = EC_KEY_set_private_key(own_keypair, own_private_key_bignum);
		BN_clear_free(own_private_key_bignum);
		if (long_ret != 1) {
			EC_GROUP_free(ec_group);
			EC_KEY_free(own_keypair);
			return ERR_get_error();
		}

		long_ret = generate_sha256_hash(signature_data, signature_data_length, effective_signature_data);
		if(long_ret != 0)
		{
			EC_GROUP_free(ec_group);
			EC_KEY_free(own_keypair);
			return long_ret;
		}

		openssl_signature = ECDSA_do_sign(effective_signature_data, 32, own_keypair);
		if(openssl_signature==NULL)
		{
			EC_GROUP_free(ec_group);
			EC_KEY_free(own_keypair);
			return 0x64;
		}

		signature_length = i2d_ECDSA_SIG(openssl_signature, temp);
		if(signature_length==0)
		{
			EC_GROUP_free(ec_group);
			EC_KEY_free(own_keypair);
			return 0x77;
		}

		signature_length=i2d_ECDSA_SIG(openssl_signature, &signature);
		if(signature_length==0)
		{
			EC_GROUP_free(ec_group);
			EC_KEY_free(own_keypair);
			return 0x79;
		}
		return 0;
}

// plaintext in case of decryption, should always contain the actual ciphertext length, that is, minus the tag
// ciphertext in case of encryption consists of ciphertext plus tag and op_ciphertext_len should similarly also consist of ciphertext_length plus 16.
// Code adapted from here: https://wiki.openssl.org/index.php/EVP_Authenticated_Encryption_and_Decryption
int aes_gcm_128(int enc, unsigned char *key, unsigned char *iv, unsigned char* plaintext, uint32_t plaintext_len, unsigned char *ciphertext,  uint32_t* op_ciphertext_len, unsigned char* tag)
{
	int len;
	int ciphertext_len;
	int reset_return;
	if(ctx == NULL)
	{
		// Create and initialise the context //
		if(!(ctx = EVP_CIPHER_CTX_new())) { return 0x1; }
	}

	
	// Initialise the encryption operation. //
	if(1 != EVP_CipherInit_ex(ctx, EVP_aes_128_gcm(), NULL, key, iv, enc))
	{
		reset_return = EVP_CIPHER_CTX_reset(ctx);
		if(reset_return != 1)
			return 0xf2;
		return 0x2;
	}
	
	// Provide the message to be encrypted, and obtain the encrypted output.	 * EVP_EncryptUpdate can be called multiple times if necessary
	 //
	if(1 != EVP_CipherUpdate(ctx, ciphertext, &len, plaintext, plaintext_len))
	{
                reset_return = EVP_CIPHER_CTX_reset(ctx);
	        if(1 != reset_return)
			return 0xF3;
		return ERR_get_error();
	}
	ciphertext_len = len;
	
	if(enc == 0)
        {
                if(1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, 16, tag))
                {
                       reset_return = EVP_CIPHER_CTX_reset(ctx);
                       if(1 != reset_return)
                                return 0xF5;
                        return 0x5;
                }
        }

	// Finalise the encryption. Normally ciphertext bytes may be written at  this stage, but this does not occur in GCM mode
	 //
	// TODO: ^^^ Why the heck does it not occur in GCM mode ?
	if(1 != EVP_CipherFinal_ex(ctx, ciphertext + len, &len))
	{
		reset_return = EVP_CIPHER_CTX_reset(ctx);
                if(1 != reset_return)
                        return 0xF4;
                return 0x43;
	}
	ciphertext_len += len;

	// Get the tag
	if(enc == 1)
	{
		if(1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, 16, tag))
		{
                	reset_return = EVP_CIPHER_CTX_reset(ctx);
                	if(1 != reset_return)
                        	return 0xF5;
	                return 0x5;
		}
		//ciphertext_len += 16;
	}

	// Clean up //
	if(1 != EVP_CIPHER_CTX_reset(ctx))
	{
		return 0xF0;
	}

	*op_ciphertext_len=ciphertext_len;
	return 0;
}

uint32_t base64_decoding_wrapper(unsigned char* src, unsigned char* dest, uint32_t length)
{
	int length_with_padding = EVP_DecodeBlock(dest, src, length);
	if(length_with_padding == -1)
		return length_with_padding;
	char* first_equals_character = strstr((char*)src, "=");
	if(first_equals_character != NULL)
	{
		if(first_equals_character == (char*)src + length - 1) // the first equals character is also the last character in the string ==> Only one equals ==> 2 valid bytes, only 1 padding 0
			length_with_padding -= 1;
		else //if(first_equals_character == src + 38) // assuming that the base64 string is valid (EVP_DecodeBlock would have thrown an error in that case), then first_equals_character == src + 38 ==> Another equals at +29 ==> 1 valid bytes, 2 padding 0s
			length_with_padding -= 2;
	}

	return length_with_padding;
}