#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 <stdio.h>

int generate_sha256_hash(const unsigned char *message, size_t message_len, unsigned char *digest);

int ec_key_gen(unsigned char* pub_key_x, unsigned char* pub_key_y, unsigned char* priv_key)
{
//	unsigned char entropy_buf[ADD_ENTROPY_SIZE] = {0};

//	RAND_add(entropy_buf, sizeof(entropy_buf), ADD_ENTROPY_SIZE);
//	RAND_seed(entropy_buf, sizeof(entropy_buf));

	EC_KEY * ec_key = NULL;
    
		ec_key = EC_KEY_new();
		if (NULL == ec_key) {
			return 0x42;
		}

	EC_GROUP* ec_group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);	
	if (NULL == ec_group) {
		//EC_KEY_free(ec_key); 
		return 0x0f;
	}

			if (0 == EC_KEY_set_group (ec_key, ec_group)) {
			EC_KEY_free(ec_key); 
			EC_GROUP_free(ec_group);
			return 0x43; //break;
		}


	EC_KEY_set_asn1_flag(ec_key, OPENSSL_EC_NAMED_CURVE);

	int ret = EC_KEY_generate_key(ec_key);
	if (!ret) {
//        printf("EC_KEY_generate_key failure\n");
	    return 0x01;
	}
	///////////////////////// Openssl code ////////////
	const EC_POINT *public_k = NULL;
	const BIGNUM *private_k = NULL;
	public_k = EC_KEY_get0_public_key(ec_key);
	if (NULL == public_k) 
	{
		EC_KEY_free(ec_key);
		return 0x04;
	}

	private_k = EC_KEY_get0_private_key(ec_key);
	if (NULL == private_k) 
	{
		EC_KEY_free(ec_key);
		return 0x05;
	}

        if(!BN_bn2bin(private_k, priv_key))
        {
                EC_KEY_free(ec_key); 
		return 0x06;
        }


        BIGNUM *pub_k_x = NULL;
        BIGNUM *pub_k_y = NULL;
        pub_k_x = BN_new();
        pub_k_y = BN_new();
        if (NULL ==  pub_k_x || NULL == pub_k_y) {
                return 0x42;//SGX_ERROR_OUT_OF_MEMORY;
        }

	// extract two BNs representing the public key
	//
	if (!EC_POINT_get_affine_coordinates_GFp(ec_group, public_k, pub_k_x, pub_k_y, NULL)) 
	//EC_KEY_free(ec_key); 
	{
		//printf("%ld\n", ERR_get_error());
		BN_clear_free(pub_k_x); BN_clear_free(pub_k_y);	
		EC_KEY_free(ec_key);
		EC_GROUP_free(ec_group);
		return 0x07;
	}

	ret = BN_bn2bin(pub_k_x, pub_key_x); 
	BN_clear_free(pub_k_x); 
	if(ret == 0)
	{
		BN_clear_free(pub_k_y); 
        	EC_KEY_free(ec_key);
	        EC_GROUP_free(ec_group);
		return 0x08; 
	}

	ret = BN_bn2bin(pub_k_y, pub_key_y);

	BN_clear_free(pub_k_y);
	EC_KEY_free(ec_key);        
	EC_GROUP_free(ec_group);

	if(ret == 0)
  		return 0x09; 

	return 0;
}

unsigned long check_key(unsigned char* given_key_x, unsigned char* given_key_y) //, unsigned char* priv_key, unsigned char* shared_key_x, unsigned char* shared_key_y)
{
        EC_GROUP* ec_group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);  
        if (NULL == ec_group) {
                //EC_KEY_free(ec_key); 
                return 0x0f;
        }
	EC_POINT *ec_point = NULL;
	BIGNUM *b_x = NULL;
	BIGNUM *b_y = NULL;

	ec_point = EC_POINT_new(ec_group);
	if (NULL == ec_point) {
			return 0x42; //retval = SGX_ERROR_OUT_OF_MEMORY;
//			break;
	}

	// converts the x value of the point, represented as positive integer in big-endian into a BIGNUM
	b_x = BN_bin2bn(given_key_x, 32, NULL);
	if (NULL == b_x) {
		EC_POINT_free(ec_point);
		// return ERR_get_error(); 
		return 0xfd;
	}

	// converts the y value of the point, represented as positive integer in big-endian into a BIGNUM
	b_y = BN_bin2bn(given_key_y, 32, NULL);
	if (NULL == b_y) {
		BN_clear_free(b_x); 
		EC_POINT_free(ec_point); 
		return 0xfe;
	}

	// sets point based on x,y coordinates
	int ret = EC_POINT_set_affine_coordinates_GFp(ec_group, ec_point, b_x, b_y, NULL);
	if (1 != ret) {
		BN_clear_free(b_x);
		BN_clear_free(b_y);
		EC_POINT_free(ec_point);
		return ERR_get_error();	
	}

	// checks if point is on curve
	//
	int ret_point_on_curve = EC_POINT_is_on_curve(ec_group, ec_point, NULL);
	if (1 != ret_point_on_curve) {
	        EC_POINT_free(ec_point);
		BN_clear_free(b_x);
		BN_clear_free(b_y);
		return ret_point_on_curve << 4; //ret_point_on_curve;//	break;
	}
	
        BN_clear_free(b_x);
        BN_clear_free(b_y);
	EC_POINT_free(ec_point);
	return 0;
}

unsigned long compute_shared_ECDHE_key(unsigned char* given_key_x, unsigned char* given_key_y, unsigned char* priv_key, unsigned char* shared_key, unsigned char* derived_key)
{
//	int valid_point = check_key(given_key_x, given_key_y);
//	if(valid_point != 0)
//		return valid_point;
	unsigned long long_ret = 0; 
	EC_GROUP* ec_group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);  
        if (NULL == ec_group) {
                return 0x0f;
        }
        EC_POINT *given_point = NULL;
        BIGNUM *b_x = NULL;
        BIGNUM *b_y = NULL;
	EC_KEY *private_key = NULL;

        given_point = EC_POINT_new(ec_group);
        if (NULL == given_point) {
                        return 0x42; //retval = SGX_ERROR_OUT_OF_MEMORY;
//                      break;
        }

        // converts the x value of the point, represented as positive integer in big-endian into a BIGNUM
        b_x = BN_bin2bn(given_key_x, 32, NULL);
        if (NULL == b_x) {
                EC_POINT_free(given_point);
                return 0xfd;
        }

        // converts the y value of the point, represented as positive integer in big-endian into a BIGNUM
        b_y = BN_bin2bn(given_key_y, 32, NULL);
        if (NULL == b_y) {
                EC_POINT_free(given_point); 
	        BN_clear_free(b_x); 
                return 0xfe;
        }

        // sets point based on x,y coordinates
        int ret = EC_POINT_set_affine_coordinates_GFp(ec_group, given_point, b_x, b_y, NULL);
        if (1 != ret) {
                EC_POINT_free(given_point);
                BN_clear_free(b_x); 
                BN_clear_free(b_y);
		return ERR_get_error(); 
        }

        // checks if point is on curve
        //
        int ret_point_on_curve = EC_POINT_is_on_curve(ec_group, given_point, NULL);
        if (1 != ret_point_on_curve) {
                EC_POINT_free(given_point);
		BN_clear_free(b_x);
		BN_clear_free(b_y); 
		if(ret_point_on_curve == 0)
			return 0x43;
		else 
			return ERR_get_error(); 
	}

	
	// create empty shared key BN
	//
	private_key = EC_KEY_new();
	if (private_key == NULL) {
		EC_POINT_free(given_point);
                BN_clear_free(b_x); 
                BN_clear_free(b_y);
		return 0x42;//ret = SGX_ERROR_OUT_OF_MEMORY;
	}

	// init private key group (set curve)
	//
	if (EC_KEY_set_group (private_key, ec_group) != 1) {
		EC_POINT_free(given_point); 
		EC_KEY_free(private_key); 
                BN_clear_free(b_x); 
                BN_clear_free(b_y);
		return ERR_get_error(); //break;
	}

	// init private key with BN value
	//
	BIGNUM *BN_dh_privB = NULL; 
	BN_dh_privB = BN_bin2bn(priv_key, 32, NULL);
	if(BN_dh_privB == NULL) {
		EC_POINT_free(given_point); 
                EC_KEY_free(private_key); 
                BN_clear_free(b_x); 
                BN_clear_free(b_y);
		return 0x44;
	}

	if (EC_KEY_set_private_key(private_key, BN_dh_privB) != 1) {
		EC_POINT_free(given_point);
		EC_KEY_free(private_key); 
        	BN_clear_free(b_x); 
	        BN_clear_free(b_y);
		BN_clear_free(BN_dh_privB); 
		return ERR_get_error(); 
	}

	// calculate shared dh key
	//
	int shared_key_len = 32; //sizeof(sgx_ec256_dh_shared_t);
	shared_key_len = ECDH_compute_key(shared_key, 32, given_point, private_key, NULL);
	if (shared_key_len <= 0) {
		long_ret = ERR_get_error(); 
	}
	else 
		long_ret = generate_sha256_hash(shared_key, 32, derived_key); 

	BN_clear_free(BN_dh_privB);
	BN_clear_free(b_x); 
	BN_clear_free(b_y);
	EC_POINT_free(given_point);
	EC_KEY_free(private_key); 
	return long_ret; 
}

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)
	{
		//printf("EVP_MD_CTX_create returned NULL - could not create context\n"); fflush(stdout); 
		return 0x1; 
	}

	if(EVP_DigestInit_ex(mdctx, EVP_sha256(), NULL) != 1)
	{
                //printf("EVP_DigestInit_ex returned 0 - could not initialize hash with SHA256\n"); fflush(stdout); 
		return 0x2; 
	}

        if(EVP_DigestUpdate(mdctx, message, message_len) != 1)        
        {
                //printf("EVP_DigestUpdate returned 0 - could not compute SHA256 hash\n"); fflush(stdout); 
		return 0x3; 
        }

	if(1 != EVP_DigestFinal_ex(mdctx, digest, &digest_len))     
        {
                //printf("EVP_DigestFinal_ex returned 0 - could not finalize SHA256 hash\n"); fflush(stdout); 
		return 0x4; 
        }

	if(digest_len != 32) 
	{
//		printf("EVP_DigestFinal_ex returned a digest length of 0x%x instead of 0x20\n", digest_len); fflush(stdout); 
	return 0x5; 
	}
	EVP_MD_CTX_destroy(mdctx);
	return 0; 
}
EVP_CIPHER_CTX* ctx; 

// Code adapted from here: https://wiki.openssl.org/index.php/EVP_Authenticated_Encryption_and_Decryption 
int aes_gcm(int enc, unsigned char *key, unsigned char *iv, unsigned char* plaintext, int plaintext_len, unsigned char *ciphertext,  int* 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_256_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 0x3;
	}
	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 0x4;
	}
	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;
		}
	}

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

	*op_ciphertext_len=ciphertext_len; 
		//EVP_CIPHER_CTX_free(ctx); // TODO: memory leaks here - need to free this for erroneous cases too.
	return 0;
}