#include <stdarg.h>
#include <math.h>
#include <stdio.h>

#include "dk_crypto.h"
#include "dk_utils.h"

static byte KDF_COUNTER_separator[1] = {0x00};

DK_Result sha256(byte* in, size_t in_len, byte* out)
{
    TEE_OperationHandle hndl;
	DK_Result rc;

    size_t out_len = SHA256_HASH_BITS / 8;

    rc = TEE_AllocateOperation(&hndl, TEE_ALG_SHA256, TEE_MODE_DIGEST, 0);
	if (rc)
		return rc;

	TEE_DigestUpdate(hndl, in, in_len);
	rc = TEE_DigestDoFinal(hndl, NULL, 0, out, &out_len);
	if (rc)
		goto catch_error;

    rc = DK_SUCCESS;

catch_error:
    TEE_FreeOperation(hndl);

    return rc;
}

DK_Result calculate_cmac(byte* s_key, size_t s_key_len, byte* out, size_t out_len, size_t n, ...)
{
    TEE_OperationHandle hndl;
	TEE_ObjectHandle obj_hndl;
	TEE_Attribute attrs[1];
	DK_Result rc;

    byte iv[32] = {0};
    size_t iv_len = sizeof(iv);
    size_t i;
    byte* data;
    size_t data_len;

    va_list valist;

    rc = TEE_AllocateOperation(&hndl, TEE_ALG_AES_CMAC, TEE_MODE_MAC, s_key_len * 8);
    if (rc) {
        return rc;
    }
    
    rc = TEE_AllocateTransientObject(TEE_TYPE_AES, s_key_len * 8, &obj_hndl);
    if (rc) {
        goto catch_error;
    }

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, (void*)s_key, s_key_len);
    TEE_PopulateTransientObject(obj_hndl, attrs, 1);
    rc = TEE_SetOperationKey(hndl, obj_hndl);
    if (rc) {
        goto catch_error;
    }
    
    TEE_MACInit(hndl, iv, iv_len);

    va_start(valist, n);
    for(i = 0; i < n; i++) {
        data = va_arg(valist, uint8_t*);
        data_len = va_arg(valist, size_t);
        TEE_MACUpdate(hndl, data, data_len);
    }
    va_end(valist);

    rc = TEE_MACComputeFinal(hndl, NULL, 0, out, &out_len);
    if (rc) {
        goto catch_error;
    }

    rc = DK_SUCCESS;
    
catch_error:
    TEE_CloseObject(obj_hndl);
    TEE_FreeOperation(hndl);

    return rc;
}

DK_Result kdf_counter(byte *ki, size_t ki_len, byte *label, size_t label_len, byte *context, size_t context_len, uint32_t L, byte **out)
{
    TEE_OperationHandle hndl;
	TEE_ObjectHandle obj_hndl;
	TEE_Attribute attrs[1];
	DK_Result rc;

    char iv[32] = {0};
    size_t iv_len = sizeof(iv);
    int32_t n, i;

    byte derived_key[512] = {0};
    size_t derived_key_len;
    byte* result;
    size_t result_len;
    byte* tmp;
    size_t tmp_len;

    byte* data;
    size_t data_len;

    byte counter[1];
    byte array_L[2];

    if (ki == NULL || label == NULL || context == NULL) {
        return DK_ERROR_BAD_PARAM;
    }

    if (L != KDF_COUNTER_L_64  && L != KDF_COUNTER_L_128 && 
        L != KDF_COUNTER_L_192 && L != KDF_COUNTER_L_256) {
        return DK_ERROR_NOT_SUPPORTED;
    }

    rc = TEE_AllocateOperation(&hndl, TEE_ALG_AES_CMAC, TEE_MODE_MAC, ki_len * 8);
    if (rc)
        return rc;
    
    rc = TEE_AllocateTransientObject(TEE_TYPE_AES, ki_len * 8, &obj_hndl);
    if (rc)
        goto catch_error;

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, ki, ki_len);
    TEE_PopulateTransientObject(obj_hndl, attrs, 1);
    rc = TEE_SetOperationKey(hndl, obj_hndl);
    if (rc)
        goto catch_error;
    
    // TEE_MACInit(hndl, iv, iv_len);

    n = (int)ceilf((float)L / (float)KDF_COUNTER_h);
    if (n > (int)powf(2, KDF_COUNTER_r_8) - 1) {
        rc = DK_ERROR_NOT_SUPPORTED;
        goto catch_error;
    }

    derived_key_len = L;
    // derived_key = dk_malloc(derived_key_len, 0);
    result_len = 0;
    result = dk_malloc(result_len);
    tmp_len = 0;
    // tmp = dk_malloc(tmp_len);

    array_L[1] = (byte)(L >> 0u);
    array_L[0] = (byte)(L >> 8u);

    for (i = 1; i <= n; i++) {
        counter[0] = i;
        data = ARRAY_CONCATN(byte, 5, label, label_len, 
                                      KDF_COUNTER_separator, 1, 
                                      array_L, 2, counter, 1, 
                                      context, context_len);
        
        if (data == NULL) {
            rc = DK_ERROR_OUT_OF_MEMORY;
            goto catch_error;
        }

        data_len = label_len + 1 + 2 + 1 + context_len;

        TEE_MACInit(hndl, iv, iv_len);
        TEE_MACUpdate(hndl, data, data_len);
        rc = TEE_MACComputeFinal(hndl, NULL, 0, derived_key, &derived_key_len);
        if (rc) 
            goto catch_error;

        tmp = ARRAY_CONCAT(byte, result, result_len, derived_key, L);
        
        dk_free(result);
        dk_free(data);

        result = tmp;
        result_len = result_len + derived_key_len;
    }
    
    (*out) = dk_malloc((L / 8) * sizeof(byte));
    if (out == NULL) {
        rc = DK_ERROR_OUT_OF_MEMORY;
        goto catch_error;
    }

    dk_memcpy(*out, result, L / 8);

    rc = DK_SUCCESS;

catch_error:
    TEE_CloseObject(obj_hndl);
    TEE_FreeOperation(hndl);

    dk_free(result);
    // dk_free(derived_key);

    return rc;
}

DK_Result aes_encrypt(byte* plaintext, size_t plaintext_len, byte* key, size_t key_len, byte* iv, size_t iv_len,  byte* cyphertext, size_t* cyphertext_len)
{
    DK_Result rc;
    TEE_Attribute attrs[1];
	TEE_OperationHandle hndl;
	TEE_ObjectHandle obj_hndl;

    size_t key_bit_len;
    key_bit_len = key_len * 8;
    if (key_bit_len != 128 && key_bit_len != 192 && key_bit_len != 256) {
        return DK_ERROR_NOT_SUPPORTED;
    }

    rc = TEE_AllocateOperation(&hndl, TEE_ALG_AES_CBC_NOPAD, TEE_MODE_ENCRYPT, key_bit_len);
    if (rc) {
        return rc;
    }

    rc = TEE_AllocateTransientObject(TEE_TYPE_AES, key_bit_len, &obj_hndl);
    if (rc) {
        goto catch_error;
    }

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, key, key_len);
    TEE_PopulateTransientObject(obj_hndl, attrs, 1);
    rc = TEE_SetOperationKey(hndl, obj_hndl);
    if (rc) {
        goto catch_error;
    }
    
    TEE_CipherInit(hndl, iv, iv_len);
    
    rc = TEE_CipherDoFinal(hndl, plaintext, plaintext_len, cyphertext, cyphertext_len);
    if (rc) {
        goto catch_error;
    }

    rc = DK_SUCCESS;

catch_error:
    TEE_CloseObject(obj_hndl);
 	TEE_FreeOperation(hndl);
    return rc;
}

DK_Result aes_decrypt(byte* cyphertext, size_t cyphertext_len, byte* key, size_t key_len, byte* iv, size_t iv_len, byte* plaintext, size_t* plaintext_len)
{
    DK_Result rc;
    TEE_Attribute attrs[1];
	TEE_OperationHandle hndl;
	TEE_ObjectHandle obj_hndl;
    size_t key_bit_len;

    key_bit_len = key_len * 8;
    if (key_bit_len != 128 && key_bit_len != 192 && key_bit_len != 256) {
        return DK_ERROR_NOT_SUPPORTED;
    }

    rc = TEE_AllocateOperation(&hndl, TEE_ALG_AES_CBC_NOPAD, TEE_MODE_DECRYPT, key_bit_len);
    if (rc) {
        return rc;
    }

    rc = TEE_AllocateTransientObject(TEE_TYPE_AES, key_bit_len, &obj_hndl);
    if (rc)
        goto catch_error;

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, key, key_len);
    TEE_PopulateTransientObject(obj_hndl, attrs, 1);
    rc = TEE_SetOperationKey(hndl, obj_hndl);
    if (rc) 
        goto catch_error;
    
    TEE_CipherInit(hndl, iv, iv_len);

    rc = TEE_CipherDoFinal(hndl, cyphertext, cyphertext_len, plaintext, plaintext_len);
    if (rc)
        goto catch_error;

    rc = DK_SUCCESS;

catch_error:
    TEE_CloseObject(obj_hndl);
 	TEE_FreeOperation(hndl);
    return rc;
}