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

#include <tees_kdf.h>

#include <openssl/evp.h>
#include <wb_enc.h>
#include <storage.h>

#define KDF_ITERATIONS 3000
#define WB_IV_LEN 16

static TEE_Result FillWithPattern(void *output, uint32_t output_len,
                                  const void *pattern, uint32_t pattern_len) {
  TEE_Result result = TEE_ERROR_BAD_PARAMETERS;

  if (!output || !pattern || !output_len || !pattern_len) {
    goto exit;
  }
  uint32_t filled = 0;

  while (filled < output_len) {
    uint32_t to_copy = (output_len - filled) > pattern_len ? pattern_len :
                                                             (output_len - filled);

    TEE_MemMove(output, pattern, to_copy);
    output += to_copy;
    filled += to_copy;
  }

  result = TEE_SUCCESS;

exit:
  return result;
}

static TEE_Result GetKeyFromWB(
  const void* lable, uint32_t lableLen,
  const void* context, uint32_t contextLen,
  uint32_t outputKeyLen, void *out_buffer) {
  TEE_Result result = TEE_SUCCESS;
  unsigned char *wb_input = TEE_Malloc(outputKeyLen, HINT_DONT_FILL_WITH_ZEROS);
  unsigned char wb_iv[WB_IV_LEN] = {0};
  unsigned char *wb_output = NULL;

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

  result = FillWithPattern(wb_input, outputKeyLen, context, contextLen);
  if (TEE_SUCCESS != result) {
    goto exit;
  }

  result = FillWithPattern(wb_iv, WB_IV_LEN, lable, lableLen);
  if (TEE_SUCCESS != result) {
    goto exit;
  }

  long enc_result = cbc_wb_encrypt(&wb_output, &wb_input, (long)outputKeyLen, wb_iv);

  if ((long)outputKeyLen > enc_result) {
    MB_LOGE("wb_encryption failed(%ld instead of %ld)!\n", enc_result, (long)outputKeyLen);
    result = TEE_ERROR_GENERIC;
    goto exit;
  }

  TEE_MemMove(out_buffer, wb_output, outputKeyLen);
exit:
  if (wb_output) {
    cleanup(&wb_output);
  }

  if (wb_input) {
    TEE_Free(wb_input);
  }
  return result;
}

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:
      if (TEE_SUCCESS != GetKeyFromWB(lable, lableLen, context,
                                      contextLen, outputKeyLen,
                                      (void *) tr->attr.buffer)) {
        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;
}