#include <tlStd.h>
#include <TlApi/TlApi.h>

#define PLATFORM_TA_LOG_TAG "GPAPI_TEST_KDF"
#include <gpapi_log.h>
#include <tees_kdf.h>
#include <storage.h>

TEE_Result TEES_DeriveKeyKDF(
    const void* lable, uint32_t lableLen,
    const void* context, uint32_t contextLen,
    uint32_t outputKeyLen, TEE_ObjectHandle object)
{
    struct TransientObject* tr = &object->tr;
    TEE_Result result = TEE_ERROR_BAD_PARAMETERS;
    const uint32_t keysize_bits = 8 * outputKeyLen;
    const unsigned char dummy_context[] = "dummycontext";
    const unsigned char dummy_label[] = "dummylabel";
    uint8_t *salt = NULL;
    uint32_t salt_len = 0;

    if (!object || !outputKeyLen) {
        goto exit;
    }

#ifdef STORAGE_HANDLES_VALIDATION
    if (!in_objects_list(object)) {
        MB_LOGE("Panic Reason: can't find object in list\n");
        TEE_Panic(ID_TEES_DeriveKeyKDF);
        result = TEE_ERROR_GENERIC;
        goto exit;
    }
#endif

    if (tr->info.handleFlags & TEE_HANDLE_FLAG_INITIALIZED) {
        MB_LOGE("Panic Reason: object already initialized\n");
        TEE_Panic(ID_TEES_DeriveKeyKDF);
        result = TEE_ERROR_GENERIC;
        goto exit;
    }

    if (keysize_bits > tr->info.maxKeySize) {
        MB_LOGE("Panic Reason: key size exceeds max key size\n");
        TEE_Panic(ID_TEES_DeriveKeyKDF);
        result = TEE_ERROR_GENERIC;
        goto exit;
    }

    if (!lable || !lableLen) {
        lable = dummy_label;
        lableLen = sizeof(dummy_label);
    }

    if (!context || !contextLen) {
        context = dummy_context;
        contextLen = sizeof(dummy_context);
    }

    salt_len = lableLen + contextLen;
    salt = TEE_Malloc(salt_len, HINT_DONT_FILL_WITH_ZEROS); /* always zero here !!! HINT_DONT_FILL_WITH_ZEROS is not working */

    if (!salt) {
        result = TEE_ERROR_OUT_OF_MEMORY;
        goto exit;
    }

    memcpy(salt, lable, lableLen);
    memcpy(salt + lableLen, context, contextLen);

    switch (tr->info.objectType) {
        case TEE_TYPE_AES:
        case TEE_TYPE_DES:
        case TEE_TYPE_DES3:
        case TEE_TYPE_HMAC_MD5:
        case TEE_TYPE_HMAC_SHA1:
        case TEE_TYPE_HMAC_SHA224:
        case TEE_TYPE_HMAC_SHA256:
        case TEE_TYPE_HMAC_SHA384:
        case TEE_TYPE_HMAC_SHA512:
        case TEE_TYPE_GENERIC_SECRET:
            result = tlApiDeriveKey(salt, salt_len, (void *) tr->attr.buffer, outputKeyLen, MC_SO_CONTEXT_TLT, MC_SO_LIFETIME_PERMANENT);
            if (result != TLAPI_OK) {
                result = TEE_ERROR_GENERIC;
                goto exit;
            }
            tr->attr.attr_array[0].content.ref.buffer = tr->attr.buffer;
            tr->attr.attr_array[0].content.ref.length = outputKeyLen;
            tr->attr.attr_array[0].attributeID = TEE_ATTR_SECRET_VALUE;
            tr->attr.attr_number = 1;
            tr->info.keySize = keysize_bits;
            tr->info.handleFlags |= TEE_HANDLE_FLAG_INITIALIZED;
            result = TEE_SUCCESS;
            break;
        default:
            MB_LOGE("Object type is not supported by KDF\n");
            result = TEE_ERROR_NOT_SUPPORTED;
            break;
    }

exit:
    if (salt) {
        TEE_Free(salt);
    }

    return result;
}