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

#include "intermediary.h"
#include "util.h"

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

#define HASH_BYTES 32
#define HASH_CHUNK_BYTES 32
#define KEY_BYTES 32

struct IntermediaryArgs {
    char* key_path;
    char* firmware_path;
};


/*
static int generate_key(EVP_PKEY** key) {
    OSSL_PARAM key_params[2];
    EVP_PKEY_CTX* gctx;

    gctx = EVP_PKEY_CTX_new_from_name(NULL, "EC", NULL);
    if (gctx == NULL)
        return (1);

    if (EVP_PKEY_keygen_init(gctx) != 1)
        return (2);

    key_params[0] = OSSL_PARAM_construct_utf8_string(OSSL_PKEY_PARAM_GROUP_NAME, "prime256v1", 0);
    key_params[1] = OSSL_PARAM_construct_end();

    if (EVP_PKEY_CTX_set_params(gctx, key_params) != 1)
        return (3);

    if(EVP_PKEY_generate(gctx, key) != 1)
        return (4);

    EVP_PKEY_CTX_free(gctx);
    return (0);
}
*/


char* intermediary_syntax(void) {
    return
        "intermediary       mock up implementation of the employee binary\n"
        "                   outputs signature on stdout\n"
        "                   WARNING: output is in binary format, may mess up terminal\n"
        "    -ekey <path>   file path of the PEM encoded private key of the employee\n"
        "    -firm <path>   path of the firmware\n";
}

int handle_intermediary(int argc, char** argv) {
    struct IntermediaryArgs args = {
        NULL,
        NULL
    };
    FILE* key_file = NULL;
    FILE* firmware_file = NULL;
    uint8_t firmware_chunk[HASH_CHUNK_BYTES];
    EVP_PKEY* key = NULL;
    EVP_MD_CTX *mdctx = NULL;
    size_t sig_len;
    unsigned char* sig = NULL;
    int status = EXIT_FAILURE;

    /*
     * Parse Input
     */

    int i = 0;
    while(i < argc) {
        if(strcmp(argv[i], "-ekey")==0 && argc-i >=2){
            args.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.key_path == NULL)
        syntax_exit();


    /*
     * Load Signing Key
     */

    key_file = fopen(args.key_path, "rb");
    if(key_file == NULL){
        perror("Error opening key file");
        status = EXIT_FAILURE;
        goto cleanup;
    }
    key = PEM_read_PrivateKey(key_file, &key, NULL, NULL);
    if(key == NULL) {
        fprintf(stderr, "failed to read key");
        fclose(key_file);
        status = EXIT_FAILURE;
        goto cleanup;
    }

    fclose(key_file);

    /*
     * Sign Firmware
     */

    firmware_file = fopen(args.firmware_path, "rb");
    if(firmware_file == NULL){
        perror("Error opening firmware file");
        status = EXIT_FAILURE;
        goto cleanup;
    }

    mdctx = EVP_MD_CTX_new();
    if (EVP_DigestSignInit(mdctx, NULL, EVP_sha256(), NULL, key) != 1) {
        fprintf(stderr, "Message digest initialization failed.\n");
        fclose(firmware_file);
        status = EXIT_FAILURE;
        goto cleanup;
    }


    size_t chunk_len = HASH_CHUNK_BYTES;
    while(chunk_len==HASH_CHUNK_BYTES) {
        chunk_len = fread(&firmware_chunk, 1, HASH_CHUNK_BYTES, firmware_file);
        if(chunk_len!=HASH_CHUNK_BYTES&&ferror(firmware_file)!=0){
            perror("Failed to read firmware file");
            exit(EXIT_FAILURE);
        }

        if (EVP_DigestSignUpdate(mdctx, firmware_chunk, chunk_len) != 1) {
            printf("Message digest update failed.\n");
            exit(EXIT_FAILURE);
        }
    }

    fclose(firmware_file);

    // call with empty sig to get length
    if (EVP_DigestSignFinal(mdctx, NULL, &sig_len) != 1) {
        printf("Message digest finalization failed.\n");
        status = EXIT_FAILURE;
        goto cleanup;
    }

    // allocate signature buffer
    sig = malloc(sizeof(unsigned char) * sig_len);
    if(sig == NULL){
        perror("could not initialize digest buffer");
        status = EXIT_FAILURE;
        goto cleanup;
    }

    // load signature into buffer
    if (EVP_DigestSignFinal(mdctx, sig, &sig_len) != 1) {
        printf("Message digest finalization failed.\n");
        EVP_MD_CTX_free(mdctx);
        status = EXIT_FAILURE;
        goto cleanup;
    }

    fwrite(sig, sig_len, 1, stdout);
    if (ferror(stdout) != 0) {
        fprintf(stdout, "failed to write signature to stdout\n");
        status = EXIT_FAILURE;
        goto cleanup;
    }

    fflush(stdout);

    status = EXIT_SUCCESS;

    // free all allocated resources
cleanup:
    if(sig != NULL)
        free(sig);
    if (mdctx != NULL)
        EVP_MD_CTX_free(mdctx);
    if (key != NULL)
        EVP_PKEY_free(key);
    exit(status);
}
