/**
 * @file tees_secure_object.c
 * @brief Blowfish-compatible Secure object API implementation for qsee
 * @author Iaroslav Makarchuk (i.makarchuk@samsung.com)
 * @date Created Oct 3, 2016
 * @par In Samsung Ukraine R&D Center (SURC) under a contract between
 * @par LLC "Samsung Electronics Ukraine Company" (Kiev, Ukraine) and
 * @par "Samsung Elecrtronics Co", Ltd (Seoul, Republic of Korea)
 * @par Copyright: (c) Samsung Electronics Co, Ltd 2015. All rights reserved.
 *
 * This software is proprietary of Samsung Electronics.
 * No part of this software, either material or conceptual may be copied
 * or distributed, transmitted, transcribed, stored in a retrieval system
 * or translated into any human or computer language in any form by any means,
 * electronic, mechanical, manual or otherwise, or disclosed to third parties
 * without the express written permission of Samsung Electronics.
 */

#include <tees_secure_object.h>

#include <tee_prop_names.h>
#include <tees_log.h>
#include <qsee_kdf.h>
#include <qsee_message.h>
#include <qsee_prng.h>

#include <uuid_utils.h>

#define AES_GCM_TAG_LEN         16
#define AES_GCM_IV_LEN          16
#define AES_KEY_LEN             32
#define QSEE_SO_MAGIC_LEN       4
#define QSEE_OWN_SO_HEADER_LEN  (QSEE_SO_MAGIC_LEN + \
                                 AES_GCM_TAG_LEN +   \
                                 AES_GCM_IV_LEN)
#define QSEE_OWN_SO_MAGIC       0x005E1F50

#define MAX_UUID_LEN            16
#define MAX_QSEE_APP_NAME       128

/* Needed for older than msm8996 nhlos builds */
#ifdef _UINT32_DEFINED
    #define qc_uint32_t  uint32
#else
    #define qc_uint32_t  uint32_t
#endif

/*
 * 144 is taken from qsee_message.h. It is not defined,
 * just mentioned in function description.
 */
#define QSEE_ENCAPSULATED_HEADER_LEN  (144)

/*
 * TODO: Analyze input parameters and offesets for SCL usage.
 */

static TEE_Result QseeGetOwnKey(uint8_t *key, uint32_t key_len) {
  TEE_Result result = TEE_SUCCESS;
  const char key_label[] = "SecureObjectLabel";
  const char salt[] = "SecureObjectSalt";

  if (qsee_kdf(NULL, key_len, (void *)&key_label, sizeof(key_label) - 1,
               (void *)salt, sizeof(salt) - 1, (void *)key, key_len)) {
    result = TEE_ERROR_GENERIC;
  }

  return result;
}

static TEE_Result QseeGetIv(uint8_t *iv, uint32_t iv_len) {
  TEE_Result result = TEE_SUCCESS;
  int32_t res;
  uint8_t *out_ptr = iv;

  while (iv_len) {
    res = qsee_prng_getdata(out_ptr, (iv_len < QSEE_MAX_PRNG) ?
                            iv_len : QSEE_MAX_PRNG);
    if (res < 1) {
      result = TEE_ERROR_GENERIC;
      goto exit;
    }

    out_ptr += res;
    iv_len -= res;
  }

exit:
  return result;
}

static TEE_Result QseeSoEncrypt(const uint8_t *in, uint32_t in_len,
                                uint8_t *out, uint32_t *out_len,
                                const uint8_t *key, uint32_t key_len) {
  TEE_Result result = TEE_ERROR_BAD_PARAMETERS;
  uint8_t *iv = out;
  uint8_t *tag = out + AES_GCM_IV_LEN;
  uint8_t aad[MAX_UUID_LEN] = { 0 };
  uint32_t tag_len = AES_GCM_TAG_LEN;
  uint32_t encrypted_len = *out_len - AES_GCM_TAG_LEN - AES_GCM_IV_LEN;

  TEE_UUID uuid;

  if (*out_len < (AES_GCM_TAG_LEN + AES_GCM_IV_LEN)) {
    result = TEE_ERROR_SHORT_BUFFER;
    goto exit;
  }

  /* Generate iv with qsee_prng */
  result = QseeGetIv(iv, AES_GCM_IV_LEN);

  if (TEE_SUCCESS != result) {
    goto exit;
  }

  /* Use UUID as AAD */
  result = TEE_GetPropertyAsUUID(TEE_PROPSET_CURRENT_TA,
                                 TA_APPID_PROP_NAME, &uuid);
  if (TEE_SUCCESS != result) {
    goto exit;
  }

  result = TeeUuidToUuid(&uuid, aad, sizeof(aad));

  if (TEE_SUCCESS != result) {
    goto exit;
  }

  /* Set up encryption */
  TEE_Attribute attrs[1];
  TEE_OperationHandle encrypt_operation_handle;
  TEE_ObjectHandle key_handle;

  result = TEE_AllocateOperation(&encrypt_operation_handle,
                                 TEE_ALG_AES_GCM, TEE_MODE_ENCRYPT,
                                 key_len * 8);

  if (TEE_SUCCESS != result) {
    goto free_operation;
  }

  result = TEE_AllocateTransientObject(TEE_TYPE_AES, key_len * 8, &key_handle);

  if (TEE_SUCCESS != result) {
    goto free_operation;
  }

  TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, key, key_len);

  result = TEE_PopulateTransientObject(key_handle, attrs, 1);
  if (TEE_SUCCESS != result) {
      goto close_object;
  }

  result = TEE_SetOperationKey(encrypt_operation_handle, key_handle);

  if (TEE_SUCCESS != result) {
    goto close_object;
  }

  result = TEE_AEInit(encrypt_operation_handle, iv, AES_GCM_IV_LEN,
                      AES_GCM_TAG_LEN * 8, 0, 0);

  if (TEE_SUCCESS != result) {
    goto close_object;
  }

  TEE_AEUpdateAAD(encrypt_operation_handle, aad, sizeof(aad));
  /* Encrypt */
  result = TEE_AEEncryptFinal(encrypt_operation_handle, in, in_len,
                              tag + AES_GCM_TAG_LEN, &encrypted_len,
                              tag, &tag_len);

  if (TEE_SUCCESS != result) {
    goto close_object;
  }

  *out_len = encrypted_len + AES_GCM_TAG_LEN + AES_GCM_IV_LEN;
  result = TEE_SUCCESS;

close_object:
  TEE_CloseObject(key_handle);
free_operation:
  TEE_FreeOperation(encrypt_operation_handle);
exit:
  if (TEE_SUCCESS != result) {
    TEE_MemFill(out, 0x00, *out_len);
  }

  return result;
}

static TEE_Result QseeSoDecrypt(const uint8_t *in, uint32_t in_len,
                                uint8_t *out, uint32_t *out_len,
                                const uint8_t *key, uint32_t key_len) {
  TEE_Result result = TEE_ERROR_BAD_PARAMETERS;
  uint8_t *iv = (uint8_t *)in;
  uint8_t *tag = (uint8_t *)in + AES_GCM_IV_LEN;
  uint8_t aad[MAX_UUID_LEN] = { 0 };
  uint32_t encrypted_len = in_len - AES_GCM_TAG_LEN - AES_GCM_IV_LEN;
  TEE_UUID uuid;

  if (in_len < (AES_GCM_TAG_LEN + AES_GCM_IV_LEN)) {
    result = TEE_ERROR_SHORT_BUFFER;
    goto exit;
  }

  /* Use UUID as AAD */
  result = TEE_GetPropertyAsUUID(TEE_PROPSET_CURRENT_TA,
                                 TA_APPID_PROP_NAME, &uuid);
  if (TEE_SUCCESS != result) {
    goto exit;
  }

  result = TeeUuidToUuid(&uuid, aad, sizeof(aad));

  if (TEE_SUCCESS != result) {
    goto exit;
  }

  /* Set up decryption */
  TEE_Attribute attrs[1];
  TEE_OperationHandle decrypt_operation_handle;
  TEE_ObjectHandle key_handle;

  result = TEE_AllocateOperation(&decrypt_operation_handle,
                                 TEE_ALG_AES_GCM, TEE_MODE_DECRYPT,
                                 key_len * 8);

  if (TEE_SUCCESS != result) {
    goto free_operation;
  }

  result = TEE_AllocateTransientObject(TEE_TYPE_AES, key_len * 8, &key_handle);

  if (TEE_SUCCESS != result) {
    goto free_operation;
  }

  TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, key, key_len);

  result = TEE_PopulateTransientObject(key_handle, attrs, 1);
  if (TEE_SUCCESS != result) {
      goto close_object;
  }

  result = TEE_SetOperationKey(decrypt_operation_handle, key_handle);

  if (TEE_SUCCESS != result) {
    goto close_object;
  }

  result = TEE_AEInit(decrypt_operation_handle, iv, AES_GCM_IV_LEN,
                      AES_GCM_TAG_LEN * 8, 0, 0);

  if (TEE_SUCCESS != result) {
    goto close_object;
  }

  TEE_AEUpdateAAD(decrypt_operation_handle, aad, sizeof(aad));
  /* Decrypt and check tag */

  result = TEE_AEDecryptFinal(decrypt_operation_handle, tag + AES_GCM_TAG_LEN,
                              encrypted_len, out, out_len, tag,
                              AES_GCM_TAG_LEN);

close_object:
  TEE_CloseObject(key_handle);
free_operation:
  TEE_FreeOperation(decrypt_operation_handle);
exit:
  if (TEE_SUCCESS != result) {
    TEE_MemFill(out, 0x00, *out_len);
  }

  return result;
}

static TEE_Result QseeSelfWrap(const unsigned char *in, uint32_t in_len,
                               unsigned char *out, uint32_t *out_len) {
  TEE_Result result = TEE_ERROR_BAD_PARAMETERS;
  uint8_t aes_key[AES_KEY_LEN] = { 0 };
  uint32_t magic = QSEE_OWN_SO_MAGIC;
  uint32_t wrapped_len = 0;

  if (!in || !in_len || !out_len) {
    goto exit;
  }
  if (*out_len < in_len + QSEE_OWN_SO_HEADER_LEN) {
    *out_len = in_len + QSEE_OWN_SO_HEADER_LEN;
    result = TEE_ERROR_SHORT_BUFFER;
    goto exit;
  }
  if (!out) {
    goto exit;
  }

  /* Generate key with qsee_kdf */
  result = QseeGetOwnKey(aes_key, sizeof(aes_key));

  if (TEE_SUCCESS != result) {
    goto exit;
  }

  /* Encrypt data with AES GCM */
  wrapped_len = *out_len - QSEE_SO_MAGIC_LEN;
  result = QseeSoEncrypt(in, in_len, out + QSEE_SO_MAGIC_LEN, &wrapped_len,
                         aes_key, sizeof(aes_key));

  if (TEE_SUCCESS != result) {
    goto exit;
  }

  TEE_MemMove(out, &magic, QSEE_SO_MAGIC_LEN);
  *out_len = wrapped_len + QSEE_SO_MAGIC_LEN;

exit:

  TEE_MemFill(aes_key, 0x00, sizeof(aes_key));
  return result;
}

static TEE_Result QseeSelfUnwrap(const unsigned char *in, uint32_t in_len,
                                 unsigned char *out, uint32_t *out_len) {
  TEE_Result result = TEE_ERROR_BAD_PARAMETERS;
  uint8_t aes_key[AES_KEY_LEN] = { 0 };

  if (!in || !in_len || !out_len) {
    goto exit;
  }

  if (*out_len < in_len - QSEE_OWN_SO_HEADER_LEN) {
    *out_len = in_len - QSEE_OWN_SO_HEADER_LEN;
    result = TEE_ERROR_SHORT_BUFFER;
    goto exit;
  }

  if (!out) {
    goto exit;
  }
  /* Generate key with qsee_kdf */
  result = QseeGetOwnKey(aes_key, sizeof(aes_key));

  if (TEE_SUCCESS != result) {
    goto exit;
  }

  /* Decrypt data with AES GCM */
  result = QseeSoDecrypt(in + QSEE_SO_MAGIC_LEN, in_len - QSEE_SO_MAGIC_LEN,
                         out, out_len, aes_key, sizeof(aes_key));
exit:

  TEE_MemFill(aes_key, 0x00, sizeof(aes_key));
  return result;
}

static TEE_Result QseeDelegatedWrap(const unsigned char *in, uint32_t in_len,
                                    unsigned char *out, uint32_t *out_len,
                                    const unsigned char *target) {
  TEE_Result result = TEE_SUCCESS;

  if (!in || !in_len || !out_len || !target) {
    result = TEE_ERROR_BAD_PARAMETERS;
    goto exit;
  }

  if (*out_len < QSEE_ENCAPSULATED_HEADER_LEN + in_len) {
    *out_len = QSEE_ENCAPSULATED_HEADER_LEN + in_len;
    result = TEE_ERROR_SHORT_BUFFER;
    goto exit;
  }

  if (!out) {
    goto exit;
  }

  if (qsee_encapsulate_inter_app_message((char *)target, (uint8_t *)in,
                                         in_len, out,
                                         (qc_uint32_t *)out_len)) {
    result = TEE_ERROR_GENERIC;
    goto exit;
  }

exit:
  return result;
}

static TEE_Result QseeDelegatedUnwrap(const unsigned char *in, uint32_t in_len,
                                      unsigned char *out, uint32_t *out_len) {
  TEE_Result result = TEE_SUCCESS;
  unsigned char src_app_name[MAX_QSEE_APP_NAME] = { 0 };

  if (!in || !in_len || !out_len ||
      *out_len < in_len - QSEE_ENCAPSULATED_HEADER_LEN) {
    result = TEE_ERROR_BAD_PARAMETERS;
    goto exit;
  }

  if (*out_len < in_len - QSEE_ENCAPSULATED_HEADER_LEN) {
    *out_len = in_len - QSEE_ENCAPSULATED_HEADER_LEN;
    result = TEE_ERROR_SHORT_BUFFER;
    goto exit;
  }

  if (!out) {
    goto exit;
  }

  if (qsee_decapsulate_inter_app_message((char *)src_app_name,
                                         (uint8_t *)in,
                                         in_len, out,
                                         (qc_uint32_t *)out_len)) {
    result = TEE_ERROR_GENERIC;
    goto exit;
  }

exit:
  return result;
}

TEE_Result TEES_WrapSecureObject(const unsigned char *in, uint32_t in_len,
                                 unsigned char *out, uint32_t *out_len,
                                 SO_AccessControlInfoType *ac) {

  TEE_Result result = TEE_ERROR_BAD_PARAMETERS;

  /* On QSEE we cannot wrap for any TA,
   * so ac != NULL and ac->access_flags != 0
   */
  if (!ac || !ac->access_flags || !in || !in_len || !out_len) {
    goto exit;
  }
  /* Currently we don't support wrappting for the whole authority */
  if (ac->access_flags & DELEGATED_AUTH_ID_AC ||
      ac->access_flags & AUTH_ID_AC) {
    result = TEE_ERROR_NOT_SUPPORTED;
    goto exit;
  }

  /* For self-usage use encryption on qsee_kdf-derived key */
  if (ac->access_flags & TA_ID_AC) {
    result = QseeSelfWrap(in, in_len, out, out_len);
  } else { /* Assume delegated wrap */
    result = QseeDelegatedWrap(in, in_len, out, out_len,
                               (const unsigned char *)ac->ta_id.uuid);
  }

exit:
  return result;
}

TEE_Result TEES_UnwrapSecureObject(const unsigned char *in, uint32_t in_len,
                                   unsigned char *out, uint32_t *out_len) {
  TEE_Result result = TEE_ERROR_BAD_PARAMETERS;
  uint32_t magic = 0;

  if (!in || !in_len || !out_len || in_len < QSEE_SO_MAGIC_LEN) {
    goto exit;
  }

  TEE_MemMove(&magic, in, QSEE_SO_MAGIC_LEN);

  switch (magic) {
    case QSEE_OWN_SO_MAGIC: {
      result = QseeSelfUnwrap(in, in_len, out, out_len);
      break;
    }
    default: { /* If no own header, user QSEEs encapsulate/decapsulate */
      result = QseeDelegatedUnwrap(in, in_len, out, out_len);
      break;
    }
  }

exit:
  return result;
}