#include <cstdio>

#include "sgx_urts.h"
#include "sgx_tcrypto.h"
#include "sgx_tseal.h"
#include "Untrusted.hpp"

static void dump(const char *label, void *p, size_t len)
{
    unsigned char *pc = (unsigned char *)p;
    if (label) {
        printf("%s: ", label);
    }
    for (size_t i=0; i<len; ++i) {
        printf("%02x", pc[i]);
    }
    printf("\n");
}

static void genkey()
{
    size_t sealedprivsize =
        sizeof(sgx_sealed_data_t) + sizeof(sgx_ec256_private_t) + 19;

    sgx_ec256_public_t pubkey;
    sgx_sealed_data_t *sealedprivkey =
        (sgx_sealed_data_t *)malloc(sealedprivsize);

    ecall_identity_key_new(&pubkey, sealedprivkey);

    dump("Pubkey", &pubkey, sizeof(pubkey));
    printf("Saving sealed private key\n");

    sgx_destroy_enclave(global_eid);

    FILE *sprivf = fopen("privkey.seal", "wb");
    fwrite(sealedprivkey, sealedprivsize, 1, sprivf);
    fclose(sprivf);

    free(sealedprivkey);
}

static void loadkey(FILE *sprivf)
{
    size_t sealedprivsize =
        sizeof(sgx_sealed_data_t) + sizeof(sgx_ec256_private_t) + 19;

    sgx_ec256_public_t pubkey;
    sgx_sealed_data_t *sealedprivkey =
        (sgx_sealed_data_t *)malloc(sealedprivsize);

    if (fread(sealedprivkey, sealedprivsize, 1, sprivf) != 1) {
        fprintf(stderr, "Could not read privkey.seal file\n");
        exit(1);
    }

    bool res = ecall_identity_key_load(&pubkey, sealedprivkey);

    if (!res) {
        fprintf(stderr, "Key load failed\n");
        exit(1);
    }

    printf("Loaded sealed private key\n");
    dump("Pubkey", &pubkey, sizeof(pubkey));

    free(sealedprivkey);
}

int main(int argc, char **argv)
{
    if (initialize_enclave() < 0) {
        return -1;
    }

    FILE *sprivf = fopen("privkey.seal", "rb");
    if (sprivf) {
        loadkey(sprivf);
        fclose(sprivf);
    } else {
        genkey();
    }

    return 0;
}