#include <errno.h>
#include <sgx_urts.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <openssl/core_names.h>
#include <openssl/evp.h>
#include <openssl/param_build.h>
#include <openssl/pem.h>

#include <sgx_tcrypto.h>

#include "enclave_u.h"
#include "proxysetup.h"
#include "util.h"

struct ProxysetupArgs {
    char* sealed_key_file_path;
    char* sgx_token_path;
};

char* proxysetup_syntax(void) {
    return
        "proxysetup          implementation of the enclave-powered SignatureProxy\n"
        "                    outputs public key on stdout\n"
        "    -pkey <path>    file path of the sealed proxy key\n"
        "    -token <path>   (optional) file path of the sgx token\n";
}

/*
 * This function is a modified version of the `get_pub_key_from_coords` function in the [Intel SGX crypto library](https://github.com/intel/linux-sgx/blob/c1ceb4fe146e0feb1097dee81c7e89925443e43c/sdk/tlibcrypto/sgxssl/sgx_ecc256.cpp).
 * The specified License applies.
 */
static EVP_PKEY *sgx_public_to_EVP_PKEY(const sgx_ec256_public_t *p_public)
{
	EVP_PKEY *evp_key = NULL;
	EVP_PKEY_CTX *pkey_ctx = NULL;
	BIGNUM *bn_pub_x = NULL;
	BIGNUM *bn_pub_y = NULL;
	EC_POINT *point = NULL;
    EC_GROUP* group = NULL;
	OSSL_PARAM_BLD *params_build = NULL;
	OSSL_PARAM *params = NULL;
	const char *curvename = NULL;
	int nid = 0;
	size_t key_len;
	unsigned char pub_key[SGX_ECP256_KEY_SIZE+4];

    group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
	if (group == NULL)
        return NULL;

	do {
		// converts the x value of public key, represented as positive integer in little-endian into a BIGNUM
		bn_pub_x = BN_lebin2bn((unsigned char*)p_public->gx, sizeof(p_public->gx), bn_pub_x);
		if (NULL == bn_pub_x) {
			break;
		}
		// converts the y value of public key, represented as positive integer in little-endian into a BIGNUM
		bn_pub_y = BN_lebin2bn((unsigned char*)p_public->gy, sizeof(p_public->gy), bn_pub_y);
		if (NULL == bn_pub_y) {
			break;
		}
		// creates new point and assigned the group object that the point relates to
		point = EC_POINT_new(group);
		if (NULL == point) {
			break;
		}

		// sets point based on public key's x,y coordinates
		if (1 != EC_POINT_set_affine_coordinates(group, point, bn_pub_x, bn_pub_y, NULL)) {
			break;
		}

		// check point if the point is on curve
		if (1 != EC_POINT_is_on_curve(group, point, NULL)) {
			break;
		}

		// convert point to octet string
		key_len = EC_POINT_point2oct(group, point, POINT_CONVERSION_COMPRESSED, pub_key, sizeof(pub_key), NULL);
		if (key_len == 0) {
			break;
		}

		// build OSSL_PARAM
		params_build = OSSL_PARAM_BLD_new();
		if (NULL == params_build) {
			break;
		}
		nid = EC_GROUP_get_curve_name(group);
		if (nid == NID_undef) {
			break;
		}
		curvename = OBJ_nid2sn(nid);
		if (curvename == NULL) {
			break;
		}
		if (1 != OSSL_PARAM_BLD_push_utf8_string(params_build, "group", curvename, 0)) {
			break;
		}
		if (1 != OSSL_PARAM_BLD_push_octet_string(params_build, OSSL_PKEY_PARAM_PUB_KEY, pub_key, key_len)) {
			break;
		}
		params = OSSL_PARAM_BLD_to_param(params_build);
		if (NULL == params) {
			break;
		}

		// get pkey from params
		pkey_ctx = EVP_PKEY_CTX_new_from_name(NULL, "EC", NULL);
		if (NULL == pkey_ctx) {
			break;
		}
		if (1 != EVP_PKEY_fromdata_init(pkey_ctx)) {
			break;
		}
		if (1 != EVP_PKEY_fromdata(pkey_ctx, &evp_key, EVP_PKEY_PUBLIC_KEY, params)) {
			EVP_PKEY_free(evp_key);
			evp_key = NULL;
		}
	} while(0);

	BN_clear_free(bn_pub_x);
	BN_clear_free(bn_pub_y);
	EC_POINT_clear_free(point);
	OSSL_PARAM_free(params);
	OSSL_PARAM_BLD_free(params_build);
	EVP_PKEY_CTX_free(pkey_ctx);
    EC_GROUP_free(group);

	return evp_key;
}

int handle_proxysetup(int argc, char** argv) {
    struct ProxysetupArgs args = {
        NULL,
        NULL
    };
    FILE* sealed_file;
    int sealed_size;
    uint8_t* sealed;
    sgx_ec256_public_t sgx_public_key;
    EVP_PKEY* public_key;
    sgx_status_t sgx_ret;

    /*
     * Parse Input
     */

    int i = 0;
    while(i < argc) {
        if(strcmp(argv[i], "-pkey")==0 && argc-i >=2){
            args.sealed_key_file_path = argv[i+1];
            i += 2;
        }else if(strcmp(argv[i], "-token")==0 && argc-i >=2){
            args.sgx_token_path = argv[i+1];
            i += 2;
        }else
            syntax_exit();
    }

    if(args.sealed_key_file_path == NULL)
        syntax_exit();

    /*
     * Initialize SGX Enclave
     */

    if (initialize_enclave(args.sgx_token_path) != 0)
        exit (EXIT_FAILURE);

    /*
     * Setup Sealed Keypair
     */

    sealed_file = fopen(args.sealed_key_file_path, "wb");
    if(sealed_file == NULL){
        perror("Error opening sealed_key_file file");
        exit (EXIT_FAILURE);
    }

    sgx_ret = get_sealed_size(get_global_eid(), &sealed_size);
    if (sgx_ret != SGX_SUCCESS) {
        sgx_print_error_message(sgx_ret);
        exit (EXIT_FAILURE);
    }
    sealed = malloc(sizeof(uint8_t)*sealed_size);
    if (sealed == NULL) {
        fprintf(stderr, "failed to allocate for sealed key");
        exit (EXIT_FAILURE);
    }

    /*
     * Use Enclave To Generate Keypair
     */

    generate_key_pair(get_global_eid(), &sgx_ret, sealed, sealed_size);
    if (sgx_ret != SGX_SUCCESS) {
        sgx_print_error_message(sgx_ret);
        exit (EXIT_FAILURE);
    }

    /*
     * Store Sealed Keypair
     */

    if (fwrite(sealed, sealed_size, 1, sealed_file) != 1 || ferror(sealed_file) != 0) {
        fprintf(stderr, "failed to write sealed key");
        exit (EXIT_FAILURE);
    }
    fflush(sealed_file);
    fclose(sealed_file);

    /*
     * Fetch Public Key From Enclave And Print
     */

    get_public_key(get_global_eid(), &sgx_ret, sealed, sealed_size, (uint8_t*)&sgx_public_key);
    if (sgx_ret != SGX_SUCCESS) {
        sgx_print_error_message(sgx_ret);
        exit (EXIT_FAILURE);
    }

    public_key = sgx_public_to_EVP_PKEY(&sgx_public_key);
    if (PEM_write_PUBKEY(stdout, public_key) != 1) {
        fprintf(stderr, "could not write publickey\n");
        exit (EXIT_FAILURE);
    }
    fflush(stdout);

    EVP_PKEY_free(public_key);
    free(sealed);

    exit (EXIT_SUCCESS);
}