#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 "proxy.h"
#include "util.h"


struct ProxyArgs {
    char* sealed_key_file_path;
    char* sgx_token_path;
    char* employee_public_key_path;
    char* firmware_path;
};

char* proxy_syntax(void) {
    return
        "proxy               implementation of the enclave-powered SignatureProxy\n"
        "                    expects intermediary signature on stdin\n"
        "                    outputs proxied signature on stdout\n"
        "                    WARNING: output is binary format, may mess up terminal\n"
        "    -pkey <path>    file path of the sealed proxy key\n"
        "    -epub <path>    path of the PEM encoded employee public key\n"
        "    -firm <path>    path of the firmware\n"
        "    -token <path>   (optional) file path of the sgx token\n";
}

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;
}

static int EVP_PKEY_to_sgx_public(EVP_PKEY* ecdsa_key, sgx_ec256_public_t* sgx_public) {
    EC_GROUP* group = NULL;
	EC_POINT *point = NULL;
    BIGNUM* pub_x = NULL;
	BIGNUM* pub_y = NULL;
    size_t ec_key_buf_len = 0;
    unsigned char ec_key_buf[1024];
    int ret;
    int retval;

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

    point = EC_POINT_new(group);
    if (point == NULL)
        return 2;

    ret = EVP_PKEY_get_octet_string_param(ecdsa_key, OSSL_PKEY_PARAM_PUB_KEY, ec_key_buf, 1024, &ec_key_buf_len);
    if (ret != 1)
        return 3;

    ret = EC_POINT_oct2point(group, point, ec_key_buf, ec_key_buf_len, NULL);
    if (ret != 1){
        retval = 4;
        goto cleanup;
    }

    pub_x = BN_new();
    pub_y = BN_new();

    ret = EC_POINT_get_affine_coordinates(group, point, pub_x, pub_y, NULL);
    if (ret != 1){
        retval = 5;
        goto cleanup;
    }

    ret = BN_bn2lebinpad(pub_x, sgx_public->gx, SGX_ECP256_KEY_SIZE);
    if (ret == -1){
        retval = 6;
        goto cleanup;
    }

    ret = BN_bn2lebinpad(pub_y, sgx_public->gy, SGX_ECP256_KEY_SIZE);
    if (ret == -1){
        retval = 7;
        goto cleanup;
    }

    retval = 0;

cleanup:
    if (pub_x != NULL)
        BN_clear_free(pub_x);
    if (pub_y != NULL)
        BN_clear_free(pub_y);
    if (point != NULL)
        EC_POINT_clear_free(point);
    if (group != NULL)
        EC_GROUP_free(group);

    return (retval);
}

static int ECDSA_SIG_to_sgx_signature(ECDSA_SIG* ecdsa_sig, sgx_ec256_signature_t* sgx_signature) {
    const BIGNUM* r = NULL;
	const BIGNUM* s = NULL;
    int ret;

    r = ECDSA_SIG_get0_r(ecdsa_sig);
    s = ECDSA_SIG_get0_s(ecdsa_sig);

    ret = BN_bn2lebinpad(r, (unsigned char*)sgx_signature->x, SGX_ECP256_KEY_SIZE);
    if (ret == -1)
        return (1);

    ret = BN_bn2lebinpad(s, (unsigned char*)sgx_signature->y, SGX_ECP256_KEY_SIZE);
    if (ret == -1)
        return (2);

    return (0);
}

/*
 * This function is a modified version of the `sgx_ecdsa_verify_hash` function in the [Intel SGX crypto library](https://github.com/intel/linux-sgx/blob/main/sdk/tlibcrypto/sgxssl/sgx_ecc256_ecdsa.cpp).
 * The specified License applies.
 */
static int sgx_signature_to_ECDSA_SIG(sgx_ec256_signature_t* sgx_signature, ECDSA_SIG** ecdsa_signature) {
	BIGNUM *bn_r = NULL;
	BIGNUM *bn_s = NULL;

	// converts the x value of the signature, represented as positive integer in little-endian into a BIGNUM
	//
	bn_r = BN_lebin2bn((unsigned char*)sgx_signature->x, sizeof(sgx_signature->x), 0);
	if (bn_r == NULL)
        return (1);


	// converts the y value of the signature, represented as positive integer in little-endian into a BIGNUM
	//
	bn_s = BN_lebin2bn((unsigned char*)sgx_signature->y, sizeof(sgx_signature->y), 0);
	if (NULL == bn_s) {
        if (bn_r != NULL)
		    BN_clear_free(bn_r);

        return (2);
	}


	// allocates a new ECDSA_SIG structure (note: this function also allocates the BIGNUMs) and initialize it
	//
	*ecdsa_signature = ECDSA_SIG_new();
	if (NULL == *ecdsa_signature) {
        if (bn_r != NULL)
		    BN_clear_free(bn_r);
        if (bn_s != NULL)
		    BN_clear_free(bn_s);

        return (3);
	}

	// setes the r and s values of ecdsa_sig
	// calling this function transfers the memory management of the values to the ECDSA_SIG object,
	// and therefore the values that have been passed in should not be freed directly after this function has been called
	//
	if (1 != ECDSA_SIG_set0(*ecdsa_signature, bn_r, bn_s)) {
		ECDSA_SIG_free(*ecdsa_signature);
		*ecdsa_signature = NULL;

        if (bn_r != NULL)
		    BN_clear_free(bn_r);

        if (bn_s != NULL)
		    BN_clear_free(bn_s);

        return (4);
	}

    return (0);
}

int handle_proxy(int argc, char** argv) {
    struct ProxyArgs args = {
        NULL,
        NULL,
        NULL
    };
    uint8_t* sealed;
    size_t sealed_len;
    uint8_t* firmware;
    size_t firmware_len;
	unsigned char* ecdsa_signature_data;
	size_t ecdsa_signature_size = 0;
    uint8_t signature[70];
    sgx_status_t ret;
    ECDSA_SIG* ecdsa_signature;
    sgx_ec256_signature_t sgx_signature;

    /*
     * 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 if(strcmp(argv[i], "-epub")==0 && argc-i >=2){
            args.employee_public_key_path = argv[i+1];
            i += 2;
        }else if(strcmp(argv[i], "-firm")==0 && argc-i >=2){
            args.firmware_path = argv[i+1];
            i += 2;
        }else
            syntax_exit();
    }

    if(args.sealed_key_file_path == NULL || args.employee_public_key_path == NULL || args.firmware_path == NULL)
        syntax_exit();

    /*
     * Read And Parse Signature Input
     */

    ecdsa_signature_data = malloc(1024);
    if (ecdsa_signature_data == NULL) {
        perror("failed to allocate signature");
        exit (EXIT_FAILURE);
    }

    ecdsa_signature_size = fread(ecdsa_signature_data, 1, 1024, stdin);

    if (ferror(stdin) != 0) {
        fprintf(stderr, "failed to read signature from stdin\n");
        exit (EXIT_FAILURE);
    }

    ecdsa_signature = ECDSA_SIG_new();
    ecdsa_signature = d2i_ECDSA_SIG(&ecdsa_signature, (const unsigned char**)&ecdsa_signature_data, ecdsa_signature_size);
    ecdsa_signature_data = NULL;
    if (ecdsa_signature == NULL) {
        fprintf(stderr, "failed to read signature\n");
        exit (EXIT_FAILURE);
    }

    if (ECDSA_SIG_to_sgx_signature(ecdsa_signature, &sgx_signature) != 0) {
        fprintf(stderr, "failed to transform signature\n");
        exit (EXIT_FAILURE);
    }
    ECDSA_SIG_free(ecdsa_signature);
    ecdsa_signature = NULL;


    // Initialize SGX Enclave
    if (initialize_enclave(args.sgx_token_path) != 0)
        exit(1);

    /*
     * Read Employee Public Key
     */

    EVP_PKEY* key = NULL;
    FILE* key_file = fopen(args.employee_public_key_path, "rb");
    if(key_file == NULL){
        perror("Error opening employee public key file");
        exit(1);
    }
    key = PEM_read_PUBKEY(key_file, &key, NULL, NULL);
    if(key == NULL) {
        fprintf(stderr, "failed to read employee public key\n");
        exit (EXIT_FAILURE);
    }
    fclose(key_file);

    sgx_ec256_public_t sgx_public;
    if (EVP_PKEY_to_sgx_public(key, &sgx_public) != 0) {
        fprintf(stderr, "failed transform employee public key\n");
        exit (EXIT_FAILURE);
    }

    //Read Sealed Proxy Keypair
    if (load_file(args.sealed_key_file_path, &sealed, &sealed_len)!=0){
        fprintf(stderr, "failed to read sealed key\n");
        exit (EXIT_FAILURE);
    }

    // Read Firmware
    if (load_file(args.firmware_path, &firmware, &firmware_len)!=0){
        fprintf(stderr, "failed to read firmware\n");
        exit (EXIT_FAILURE);
    }

    /*
     * Use Enclave To Resign the Firmware
     */

    sign_firmware(get_global_eid(), &ret, firmware, firmware_len, sealed, sealed_len, (uint8_t*)&sgx_public, (uint8_t*)&sgx_signature);
    if (ret != SGX_SUCCESS) {
        sgx_print_error_message(ret);
        exit (1);
    }

    free(sealed);
    free(firmware);

    /*
     * Output Signature
     */
    if (sgx_signature_to_ECDSA_SIG(&sgx_signature, &ecdsa_signature) != 0) {
        fprintf(stderr, "could not convert signature\n");
        exit (EXIT_FAILURE);
    }

    ecdsa_signature_data = NULL;
	ecdsa_signature_size = i2d_ECDSA_SIG(ecdsa_signature, &ecdsa_signature_data);
	if (ecdsa_signature_size <= 0) {
        fprintf(stderr, "could not convert signature\n");
        exit (EXIT_FAILURE);
	}

    fwrite(ecdsa_signature_data, ecdsa_signature_size, 1, stdout);

    if (ferror(stdout) != 0) {
        fprintf(stderr, "could not write signature to stdout\n");
        exit (EXIT_FAILURE);
    }

    fflush(stdout);

    free(ecdsa_signature_data);
    ECDSA_SIG_free(ecdsa_signature);

    exit (EXIT_SUCCESS);
}