/*
 * 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 DG_SIZE
#define DG_SIZE SK_SIZE
#endif

#ifndef SI_SIZE
#define SI_SIZE 2*SK_SIZE
#endif


const sgx_ec256_public_t authorized[2] = {
    {
        0,
        0
    },
    {
        0,
        0
    }
};


int get_digest_size() {
    return DG_SIZE;
}

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

static sgx_status_t seal_key_pair(const sgx_ec256_private_t *private, const sgx_ec256_public_t *public, uint8_t **sealed) {
    // allocate temporary buffers on stack
    uint8_t pk[PK_SIZE] = {0};
    uint8_t sk[SK_SIZE] = {0};

    // copy key pair into buffers
    memcpy(pk, public->gx, PK_SIZE);
    memcpy(sk, private->r, SK_SIZE);

    // seal keypair
    return sgx_seal_data(PK_SIZE, (const uint8_t *)pk, SK_SIZE, (const uint8_t *)sk, get_sealed_size(), (sgx_sealed_data_t *) *sealed);
}

static sgx_status_t unseal_key_pair(const uint8_t *sealed, sgx_ec256_private_t *private, sgx_ec256_public_t *public) {
    // invalid parameter handling
    if(sealed == NULL) {
        return SGX_ERROR_INVALID_PARAMETER;
    }

    // allocate temporary buffers on stack
    uint8_t pk[PK_SIZE] = {0};
    uint8_t sk[SK_SIZE] = {0};
    
    // 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 != PK_SIZE) || (sk_size != SK_SIZE)) {
        return SGX_ERROR_UNEXPECTED;
    }

    // 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) {
        return status;
    }

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

    // return success
    return status;
}

sgx_status_t generate_key_pair(uint8_t *sealed, uint32_t sealed_size) {
    // invalid parameter handling
    if((sealed == NULL) || (sealed_size != get_sealed_size())) {
        return SGX_ERROR_INVALID_PARAMETER;   
    }

    // declare needed structs
    sgx_ecc_state_handle_t ecc_handle;
    sgx_ec256_private_t private;
    sgx_ec256_public_t public;
    sgx_status_t status;
    
    // open ecc handle
    if((status = sgx_ecc256_open_context(&ecc_handle)) != SGX_SUCCESS) {
        return status;
    }

    // create ecc keypair
    if((status = sgx_ecc256_create_key_pair(&private, &public, ecc_handle)) != SGX_SUCCESS) {
        sgx_ecc256_close_context(ecc_handle);
        return status;
    }

    // return status of sealing
    return seal_key_pair(&private, &public, &sealed);
}

sgx_status_t get_public_key(const uint8_t *sealed, uint32_t sealed_size, uint8_t *public) {
    // invalid parameter handling
    if((sealed == NULL) || (sealed_size != get_sealed_size())) {
        return SGX_ERROR_INVALID_PARAMETER;
    } else if(public == NULL) {
        return SGX_ERROR_INVALID_PARAMETER;
    }

    // unseal public key
    sgx_status_t status;
    if((status = unseal_key_pair(sealed, NULL, (sgx_ec256_public_t *)public)) != SGX_SUCCESS) {
        return status;
    }

    // return success
    return status;
}

static sgx_status_t verify_signature(const uint8_t *data, uint32_t data_size, const sgx_ec256_public_t *public, const sgx_ec256_signature_t* ecc_signature) {
    // invalid parameter handling
    if((data == NULL) || (data_size == 0) || (public == NULL) || (ecc_signature == NULL)) {
        return SGX_ERROR_INVALID_PARAMETER;
    }

    // declare needed structure
    sgx_ecc_state_handle_t ecc_handle;

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

    // verify signature
    uint8_t result;
    sgx_status_t verification_status = sgx_ecdsa_verify(data, data_size, public, ecc_signature, &result, ecc_handle);

    // handle failed verification process
    if(verification_status != SGX_SUCCESS) {
        result = verification_status;
    }

    // close context and return valid signature
    sgx_ecc256_close_context(ecc_handle);
    return result;
}

sgx_status_t sign_firmware(const uint8_t *data, uint32_t data_size, const uint8_t *sealed, uint32_t sealed_size, uint8_t *public_key, uint8_t *signature) {
    // invalid parameter handling
    if((data == NULL) || (data_size == 0) || (public_key == NULL) || (signature == NULL)) {
        return SGX_ERROR_INVALID_PARAMETER;
    } else if((sealed == NULL) || (sealed_size != get_sealed_size())) {
        return SGX_ERROR_INVALID_PARAMETER;
    }

    // verify public key
    for(size_t i = 0; i < sizeof(authorized)/sizeof(authorized[0]); i++) {
        if(memcmp(public_key, authorized[i].gx, PK_SIZE) == 0) {
            continue;
        }
        goto sign;
    }
    return SGX_ERROR_UNEXPECTED;
    
    sign: ;

    // declare need structures
    sgx_ec256_signature_t ecc_signature;
    sgx_ecc_state_handle_t ecc_handle;
    sgx_ec256_private_t private;

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

    // verify request
    /*
    if((status = verify_signature(data, data_size, (const sgx_ec256_public_t *)public_key, (const sgx_ec256_signature_t *)signature)) != SGX_EC_VALID) {
        sgx_ecc256_close_context(ecc_handle);
        return status;
    }*/

    // try unseal keypair
    sgx_status_t seal_status;
    if(seal_status = unseal_key_pair(sealed, &private, (sgx_ec256_public_t *)public_key) != SGX_SUCCESS) {
        sgx_ecc256_close_context(ecc_handle);
        return seal_status;
    }

    // create signature
    if((status = sgx_ecdsa_sign(data, DG_SIZE, &private, &ecc_signature, ecc_handle)) != SGX_SUCCESS) {
        sgx_ecc256_close_context(ecc_handle);
        return status;
    }

    // copy signature to return buffer
    if(signature == NULL) {
        sgx_ecc256_close_context(ecc_handle);
        return SGX_ERROR_INVALID_PARAMETER;
    }

    memcpy(signature, ecc_signature.x, SI_SIZE);

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

sgx_status_t verify_firmware(const uint8_t *data, uint32_t data_size, const uint8_t *sealed, uint32_t sealed_size, const uint8_t *public_key, const uint8_t *signature) {
    // invalid parameter handling
    if((data == NULL) || (data_size == 0) || (signature == NULL)) {
        return SGX_ERROR_INVALID_PARAMETER;
    } else if((sealed == NULL) && (public_key == NULL)) {
        return SGX_ERROR_INVALID_PARAMETER;
    } else if((sealed != NULL) && (public_key != NULL)) {
        return SGX_ERROR_INVALID_PARAMETER;
    }

    // declare needed structures
    sgx_ec256_public_t public;
    sgx_status_t status;

    // verify signature from staff or enclave
    if(public_key != NULL) {
        // verification only with authorized public keys
        for(size_t i = 0; i < sizeof(authorized)/sizeof(authorized[0]); i++) {

        }
        
        // copy public key into struct
        memcpy(public.gx, public_key, PK_SIZE);
    } else {
        // unseal public key
        if((status = unseal_key_pair(sealed, NULL, &public)) != SGX_SUCCESS) {
            return status;
        }
    }

    // verify signature and return result
    return verify_signature(data, data_size, &public, (const sgx_ec256_signature_t *)signature);
}