/*
 * Copyright (C) 2011-2018 Intel Corporation. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in
 *     the documentation and/or other materials provided with the
 *     distribution.
 *   * Neither the name of Intel Corporation nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */


#include <stdarg.h>
#include <stdio.h>      /* vsnprintf */
#include <string.h>
#include <stdlib.h>

#include "Enclave.h"
#include "Enclave_t.h"
#include <sgx_tseal.h>
#include <sgx_error.h>
#include <sgx_tcrypto.h>

#ifndef SK_SIZE
#define SK_SIZE SGX_ECP256_KEY_SIZE
#endif

#ifndef PK_SIZE
#define PK_SIZE 2*SK_SIZE
#endif

#ifndef SI_SIZE
#define SI_SIZE 2*SK_SIZE
#endif

int get_sealed_size() {
    return sgx_calc_sealed_data_size(PK_SIZE, SK_SIZE);
}

int get_signature_size() {
    return SI_SIZE;
}

int get_public_key_size() {
    return PK_SIZE;
}

int get_private_key_size() {
    return SK_SIZE;
}

sgx_status_t seal_key_pair(sgx_ec256_private_t *private, sgx_ec256_public_t *public, uint8_t *sealed, uint32_t *sealed_size) {
    // handle missing keypair
    if((private == NULL) || (public == NULL))
        return SGX_ERROR_UNEXPECTED;

    // allocate temporary buffers on stack
    uint8_t pk[2*SGX_ECP256_KEY_SIZE] = {0};
    uint8_t sk[SGX_ECP256_KEY_SIZE] = {0};

    // copy keypair into buffers
    memcpy(pk, public->gx, SGX_ECP256_KEY_SIZE);
    memcpy(pk + SGX_ECP256_KEY_SIZE, public->gy, SGX_ECP256_KEY_SIZE);
    memcpy(sk, private->r, SGX_ECP256_KEY_SIZE);

    // calculate needed size
    *sealed_size = sgx_calc_sealed_data_size((uint32_t)(2*SGX_ECP256_KEY_SIZE), (uint32_t)SGX_ECP256_KEY_SIZE);
    if(*sealed_size == UINT32_MAX)
        return SGX_ERROR_UNEXPECTED;

    // allocate buffer on heap
    sealed = (uint8_t *)malloc(*sealed_size);
    if(sealed == NULL) {
        free(sealed);
        return SGX_ERROR_OUT_OF_MEMORY;
    }

    // seal keypair
    sgx_status_t status = sgx_seal_data((uint32_t)2*SGX_ECP256_KEY_SIZE, (const uint8_t *)pk, (uint32_t)SGX_ECP256_KEY_SIZE, (uint8_t *)sk, *sealed_size, (sgx_sealed_data_t *) sealed);
    if(status != SGX_SUCCESS) {
        free(sealed);
        return SGX_ERROR_UNEXPECTED;
    }

    return SGX_SUCCESS;
}

sgx_status_t unseal_key_pair(uint8_t *sealed, uint32_t *sealed_size, sgx_ec256_private_t *private, sgx_ec256_public_t *public) {
    // handle missing sealed data
    if((sealed == NULL) || (sealed_size == 0))
        return SGX_ERROR_UNEXPECTED;

    // calculate public_key size and return error for unexpected results
    uint32_t pk_size = sgx_get_add_mac_txt_len((const sgx_sealed_data_t *)sealed);
    uint32_t sk_size = sgx_get_encrypt_txt_len((const sgx_sealed_data_t *)sealed);
    if ((pk_size != 2*SGX_ECP256_KEY_SIZE) || (sk_size != SGX_ECP256_KEY_SIZE))
        return SGX_ERROR_UNEXPECTED;

    // allocate memory for public and secret key
    uint8_t *pk =(uint8_t *)malloc(pk_size);
    uint8_t *sk =(uint8_t *)malloc(pk_size);
    if((pk == NULL) || (sk == NULL)) {
        free(pk);
        free(sk);
        return SGX_ERROR_OUT_OF_MEMORY;
    }

    // unseal ecc key pair
    sgx_status_t status = sgx_unseal_data((const sgx_sealed_data_t *)sealed, pk, &pk_size, sk, &sk_size);
    if (status != SGX_SUCCESS) {
        free(pk);
        free(sk);
        return status;
    }

    // copy buffers into key structs
    if(public != NULL) {
        memcpy(public->gx, pk, SGX_ECP256_KEY_SIZE);
        memcpy(public->gy, pk + SGX_ECP256_KEY_SIZE, SGX_ECP256_KEY_SIZE);
    }
    if (private != NULL) {
        memcpy(private->r, sk, SGX_ECP256_KEY_SIZE);
    }

    // free temporary buffers
    free(pk);
    free(sk);

    // return success
    return SGX_SUCCESS;
}

sgx_status_t public_key(const uint8_t *sealed, const size_t sealed_size, uint8_t *gx, uint8_t *gy) {
    // return if no sealed data provided
    if(sealed == NULL)
        return SGX_ERROR_UNEXPECTED;

    // unseal public key
    sgx_ec256_public_t public;
    if(unseal_key_pair(sealed, sealed_size, NULL, &public) != SGX_SUCCESS) {
        return SGX_ERROR_UNEXPECTED;
    }

    // copy public key into return buffers
    if((gx != NULL) && (gy != NULL)) {
        memcpy(gx, public.gx, SGX_ECP256_KEY_SIZE);
        memcpy(gy, public.gy, SGX_ECP256_KEY_SIZE);
    }

    return SGX_SUCCESS;
}

sgx_status_t sign_firmware(const uint8_t *sealed, size_t sealed_size, uint8_t *data, size_t data_size, uint8_t *signature, size_t signature_size) {
    // handle missing sealed buffer
    if((sealed == NULL) || (sealed_size == 0)) {
        return SGX_ERROR_UNEXPECTED;
    }
    
    // handle missing firmware buffer
    if((data == NULL) || (data_size == 0)) {
        return SGX_ERROR_UNEXPECTED;
    }

    // declare need structures
    sgx_ecc_state_handle_t ecc_handle;
    sgx_ec256_private_t private;
    sgx_ec256_public_t public;

    // open ecc handle
    sgx_status_t status;
    if((status = sgx_ecc256_open_context(&ecc_handle)) != SGX_SUCCESS) {
        return status;
    }

    // try unseal keypair
    sgx_status_t seal_status;
    if(seal_status = unseal_key_pair(sealed, &sealed_size, &private, NULL) != SGX_SUCCESS) {
        if((status = sgx_ecc256_create_key_pair(&private, &public, ecc_handle)) != SGX_SUCCESS) {
            sgx_ecc256_close_context(ecc_handle);
            return status;
        }
    }
    
    // create signature
    sgx_ec256_signature_t ecc_signature;
    if((status = sgx_ecdsa_sign(data, data_size, &private, &ecc_signature, ecc_handle)) != SGX_SUCCESS) {
        sgx_ecc256_close_context(ecc_handle);
        return status;
    }
    
    // copy signature to return buffer
    // TODO: endian swap
    if((signature == NULL) || (signature_size == 0)) {
        sgx_ecc256_close_context(ecc_handle);
        return SGX_ERROR_UNEXPECTED;
    }
    memcpy(signature, ecc_signature.x, SGX_ECP256_KEY_SIZE);
    memcpy(signature + SGX_ECP256_KEY_SIZE, ecc_signature.y, SGX_ECP256_KEY_SIZE);
    
    if(seal_status != SGX_SUCCESS) {
        seal_status = seal_key_pair(&private, &public, sealed, &sealed_size);
        // TODO: return sealed keypair
    }

    // close ecc handle and return success
    sgx_ecc256_close_context(ecc_handle);
    return seal_status;
}