#include "SCP03_kdf.h"
#include "dk_log.h"
#include <TZ_Vendor_tl.h>

#define LOGE DK_LOG_ERR

void kdf(uint8_t* input_key, uint8_t* label, uint8_t* context, int contextLength, uint8_t* output_length_in_bit, uint8_t* output_key) {
    int n;
    int L;
    uint8_t i;
    int prf_input_data_size = 0;
    uint8_t* prf_input_data = NULL;

    L = (output_length_in_bit[0] << 8) | output_length_in_bit[1]; // calculate output key length in bit
    n = L / PRF_CMAC_OUTLENGTH_IN_BIT; // calculate max counter in counter mode
    if (n == 0 && L != 0) {
        n = 1;
    }

    prf_input_data_size = KDF_COUNTER_SIZE + KDF_LABEL_SIZE + KDF_SEPERATION_INDICATOR_SIZE + contextLength + KDF_L_SIZE;
    prf_input_data = (uint8_t*) TZ_malloc(prf_input_data_size);
    if ( NULL == prf_input_data )
    {
        LOGE("memory is not allocated");
        return;
    }

    memcpy(prf_input_data, label, KDF_LABEL_SIZE); // 1st 12 byte for label
    prf_input_data[12] = 0x00; // 2nd 1 byte for seperation indicator
    memcpy(prf_input_data + KDF_LABEL_SIZE + KDF_SEPERATION_INDICATOR_SIZE , output_length_in_bit, KDF_L_SIZE); // 3rd 2 byte for L
    memcpy(prf_input_data + KDF_LABEL_SIZE + KDF_SEPERATION_INDICATOR_SIZE + KDF_L_SIZE + KDF_COUNTER_SIZE, context, contextLength); // 5th context (length can be changed)
    for (i = 1 ; i <= n ; i++) {
        *(prf_input_data+(KDF_LABEL_SIZE + KDF_SEPERATION_INDICATOR_SIZE +KDF_L_SIZE)) = i; // 4th 1 byte for counter
        SCP03_AES_CMAC (input_key, prf_input_data, prf_input_data_size, output_key + (PRF_CMAC_OUTLENGTH_IN_BIT/8) * (i-1)); // PRF (CMAC)
    }

    TZ_free(prf_input_data);
}

void genSessionKey(uint8_t* input_key, uint8_t* output_key, uint8_t* hostChallenge, uint8_t* cardChallenge, int keyType) {
    uint8_t label[12] = {0,};
    uint8_t context[16] = {0,};
    uint8_t L[2];

    if (keyType == SMAC) {
        label[11] = 0x06; // first byte of label should be 0x06 for SMAC
    } else if (keyType == SENC) {
        label[11] = 0x04; // first byte of label should be 0x06 for SENC
    } else if (keyType == SRMAC) {
        label[11] = 0x07; // first byte of label should be 0x06 for SRMAC
    }
    memcpy(context, hostChallenge, 8);
    memcpy(context + 8, cardChallenge, 8); // context should be set to the concatenation of the hostChallenge and the cardChallenge

    if (KEY_LEN_BIT == 256) {
        L[0] = 0x01; // should be set as 256 because SMAC key size is 256 bit
        L[1] = 0x00;
    } else if (KEY_LEN_BIT == 128) {
        L[0] = 0x00; // should be set as 256 because SMAC key size is 256 bit
        L[1] = 0x80;
    }

    kdf(input_key, label, context, 16, L, output_key);
}

void genCardChallenge(uint8_t* KEY_ENC, uint8_t* sequenceCounter, uint8_t* AID, int AIDLength, uint8_t* cCardChallenge) {
    uint8_t label[12] = {0,};
    uint8_t* context = NULL;
    uint8_t L[2];
    int contextLength = KDF_SEQUENCE_COUNTER_SIZE + AIDLength;

    //label[0] = 0x02; // first byte of label should be 0x02 for card challenge generation
    label[11] = 0x02;

    context = (uint8_t*) TZ_malloc (contextLength);
    if ( NULL == context )
    {
        LOGE("memory is not allocated");
        return;
    }

    L[0] = 0x00;
    L[1] = 0x40;

    memcpy(context, sequenceCounter, KDF_SEQUENCE_COUNTER_SIZE);
    memcpy(context + KDF_SEQUENCE_COUNTER_SIZE, AID, AIDLength);

    kdf(KEY_ENC, label, context, contextLength, L, cCardChallenge);

    TZ_free(context);
}

void genCardCryptogram(uint8_t* S_MAC, uint8_t* hostChallenge, uint8_t* cardChallenge, uint8_t* cCardCryptogram) {
    uint8_t label[12] = {0,};
    uint8_t context[16] = {0,};
    uint8_t L[2];

    label[11] = 0x00; // first byte of label should be 0x00 for card cryptogram

    L[0] = 0x00;
    L[1] = 0x40;

    memcpy(context, hostChallenge, 8);
    memcpy(context + 8, cardChallenge, 8); // context should be set to the concatenation of the hostChallenge and the cardChallenge

    kdf(S_MAC, label, context, 16, L, cCardCryptogram);
}

void genHostCryptogram(uint8_t* S_MAC, uint8_t* hostCryptogram, uint8_t* hostChallenge, uint8_t* cardChallenge) {
    uint8_t label[12] = {0,};
    uint8_t context[16] = {0,};
    uint8_t cHostCryptogram[16] = {0,};
    uint8_t L[2];

    label[11] = 0x01; // first byte of label should be 0x01 for host cryptogram

    L[0] = 0x00;
    L[1] = 0x40;

    memcpy(context, hostChallenge, 8);
    memcpy(context + 8, cardChallenge, 8); // context should be set to the concatenation of the hostChallenge and the cardChallenge

    kdf(S_MAC, label, context, 16, L, cHostCryptogram);

    memcpy(hostCryptogram, cHostCryptogram, 8);
}