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

#ifndef QSEE_KDF_SUCCESS
  #define QSEE_KDF_SUCCESS 0
#endif /* !QSEE_KDF_SUCCESS */

#define AES_256_KEY_SIZE 32

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 = NULL;
    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";

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

    tr = &object->tr;

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

    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 = qsee_kdf(0, AES_256_KEY_SIZE, lable, lableLen, context, contextLen, (void *) tr->attr.buffer, outputKeyLen);
            if (result != QSEE_KDF_SUCCESS) {
                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:
    return result;
}