#include <stdlib.h>
#include <string.h>
#include "ra_ecp.h"

#include "crypto/crypto.h"


#define MAC_KEY_SIZE       16

errno_t memcpy_s(
    void *dest,
    size_t numberOfElements,
    const void *src,
    size_t count) {
    if(numberOfElements<count)
        return -1;
    memcpy(dest, src, count);
    return 0;
}

bool verify_cmac128(
    sgx_ec_key_128bit_t mac_key,
    const uint8_t *p_data_buf,
    uint32_t buf_size,
    const uint8_t *p_mac_buf) {
    sgx_cmac_128bit_tag_t data_mac;

    if(cmac128((sgx_cmac_128bit_key_t *)&mac_key, p_data_buf, buf_size, &data_mac))
        return false;
    // In real implementation, should use a time safe version of memcmp here,
    // in order to avoid side channel attack.
    if(memcmp(p_mac_buf, data_mac, SGX_CMAC_MAC_SIZE))
        return false;
    return true;
}

#define EC_DERIVATION_BUFFER_SIZE(label_length) ((label_length) +4)

const char str_SMK[] = "SMK";
const char str_SK[] = "SK";
const char str_MK[] = "MK";
const char str_VK[] = "VK";

// Derive key from shared key and key id.
// key id should be sample_derive_key_type_t.
bool derive_key(
    const sgx_ec256_dh_shared_t *p_shared_key,
    uint8_t key_id,
    sgx_ec_key_128bit_t* derived_key
    ) {
    int ret = 0;
    uint8_t cmac_key[MAC_KEY_SIZE];
    sgx_ec_key_128bit_t key_derive_key;
    
    memset(&cmac_key, 0, MAC_KEY_SIZE);

    ret = cmac128(
        (const sgx_cmac_128bit_key_t *)&cmac_key,
        (const uint8_t*)p_shared_key,
        sizeof(sgx_ec256_dh_shared_t),
        (sgx_cmac_128bit_tag_t *)&key_derive_key);
    if (ret) {
        // memset here can be optimized away by compiler, so please use memset_s on
        // windows for production code and similar functions on other OSes.
        memset(&key_derive_key, 0, sizeof(key_derive_key));
        return false;
    }

    const char *label = NULL;
    uint32_t label_length = 0;
    switch (key_id) {
    case DERIVE_KEY_SMK:
        label = str_SMK;
        label_length = sizeof(str_SMK) -1;
        break;
    case DERIVE_KEY_SK:
        label = str_SK;
        label_length = sizeof(str_SK) -1;
        break;
    case DERIVE_KEY_MK:
        label = str_MK;
        label_length = sizeof(str_MK) -1;
        break;
    case DERIVE_KEY_VK:
        label = str_VK;
        label_length = sizeof(str_VK) -1;
        break;
    default:
        // memset here can be optimized away by compiler, so please use memset_s on
        // windows for production code and similar functions on other OSes.
        memset(&key_derive_key, 0, sizeof(key_derive_key));
        return false;
        break;
    }
    /* derivation_buffer = counter(0x01) || label || 0x00 || output_key_len(0x0080) */
    uint32_t derivation_buffer_length = EC_DERIVATION_BUFFER_SIZE(label_length);
    uint8_t *p_derivation_buffer = (uint8_t *)malloc(derivation_buffer_length);
    if (p_derivation_buffer == NULL) {
        // memset here can be optimized away by compiler, so please use memset_s on
        // windows for production code and similar functions on other OSes.
        memset(&key_derive_key, 0, sizeof(key_derive_key));
        return false;
    }
    memset(p_derivation_buffer, 0, derivation_buffer_length);

    /*counter = 0x01 */
    p_derivation_buffer[0] = 0x01;
    /*label*/
    memcpy_s(&p_derivation_buffer[1], derivation_buffer_length - 1, label, label_length);
    /*output_key_len=0x0080*/
    uint16_t *key_len = (uint16_t *)(&(p_derivation_buffer[derivation_buffer_length - 2]));
    *key_len = 0x0080;

    ret = cmac128(
        (sgx_cmac_128bit_key_t *)&key_derive_key,
        p_derivation_buffer,
        derivation_buffer_length,
        (sgx_cmac_128bit_tag_t *)derived_key);
    free(p_derivation_buffer);
    // memset here can be optimized away by compiler, so please use memset_s on
    // windows for production code and similar functions on other OSes.
    memset(&key_derive_key, 0, sizeof(key_derive_key));
    if (ret) {
        return false;
    }
    return true;
}