/**
 * @file       activation_gp_handler.c
 * @brief      Activation trustlet (SWd) processes requests from Activation library (NWd).
 *
 * Implementation is based on GP API and Multibuild SDK.
 *
 * @author     Oleksandr Fadieiev (o.fadieiev@samsung.com)
 * @author     Andrii Kravchenko (a.kravchenko@samsung.com)
 * @version    1.0
 * @date       April 21, 2017
 * @copyright  In Samsung Ukraine R&D Center (SURC) under a contract between
 * @par        LLC "Samsung Electronics Ukraine Company" (Kiev, Ukraine) and
 * @par        "Samsung Electronics Co", Ltd (Seoul, Republic of Korea)
 * @par        Copyright: (c) Samsung Electronics Co, Ltd 2017. All rights reserved.
 */
#include <tee_internal_api.h>
#include <tees_secure_object.h>
#include <activation.h>
#include <activation_tee_command.h>
#include <drk_parser.h>

#define PLATFORM_LOG_TAG "Activation TA"
#include <tees_log.h>

#define RSA_PRIVATE_KEY_STORAGE_NAME "RsaPrivateKeyStorage"
#define AES_STORAGE_NAME "AesStorage"

#define kAesMaxSize (256)
#define kAesBits (128)
#define kAesADataSize (20)

static TEE_Result HandleActivationStorePrivateKey(TEE_Param params[4]);
static TEE_Result HandleActivationStoreServerKey(TEE_Param params[4]);
static TEE_Result HandleActivationGenerateCredentials(TEE_Param params[4]);

// Version can be found by "strings <file> | grep Version"
static const volatile char *g_version =
    PLATFORM_LOG_TAG " Version: " CONFIG_VERSION ", " CONFIG_BUILD_TYPE;

TEE_Result TA_CreateEntryPoint(void) {
  return TEE_SUCCESS;
}

void TA_DestroyEntryPoint(void) {
}

TEE_Result TA_OpenSessionEntryPoint(uint32_t paramTypes, TEE_Param params[4],
                                    void **sessionContext) {
  (void) paramTypes;
  (void) params;
  (void) sessionContext;

  return TEE_SUCCESS;
}

void TA_CloseSessionEntryPoint(void *sessionContext) {
  (void) sessionContext;
}

TEE_Result TA_InvokeCommandEntryPoint(void *sessionContext,
                                      uint32_t commandID,
                                      uint32_t paramTypes,
                                      TEE_Param params[4]) {
  (void) sessionContext;
  (void) paramTypes;
  (void) params;

  TEE_Result res = TEE_ERROR_GENERIC;

  switch (commandID) {
    case kActivationStorePrivateKey: {
      if (paramTypes != TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
                                        TEE_PARAM_TYPE_MEMREF_OUTPUT,
                                        TEE_PARAM_TYPE_MEMREF_OUTPUT,
                                        TEE_PARAM_TYPE_NONE)) {
        MB_LOGE("Type of input parameters have other type\n");
        return TEE_ERROR_BAD_PARAMETERS;
      }
      res = HandleActivationStorePrivateKey(params);
      break;
    }

    case kActivationStoreServerKey: {
      if (paramTypes != TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
                                        TEE_PARAM_TYPE_NONE,
                                        TEE_PARAM_TYPE_NONE,
                                        TEE_PARAM_TYPE_NONE)) {
        MB_LOGE("Type of input parameters have other type\n");
        return TEE_ERROR_BAD_PARAMETERS;
      }
      res = HandleActivationStoreServerKey(params);
      break;
    }

    case kActivationGenerateCredentials: {
      if (paramTypes != TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
                                        TEE_PARAM_TYPE_MEMREF_OUTPUT,
                                        TEE_PARAM_TYPE_MEMREF_OUTPUT,
                                        TEE_PARAM_TYPE_MEMREF_OUTPUT)) {
        MB_LOGE("Type of input parameters have other type\n");
        return TEE_ERROR_BAD_PARAMETERS;
      }
      res = HandleActivationGenerateCredentials(params);
      break;
    }

    default: {
      MB_LOGE("Incorrect command\n");
      break;
    }
  }

  return res;
}

static TEE_Result HandleActivationStorePrivateKey(TEE_Param params[4]) {
  MB_LOGD("Store DRK (RSA) private key command\n");

  const uint8_t *wrapped = params[0].memref.buffer;
  const uint32_t wrapped_size = params[0].memref.size;

  uint8_t unwrapped[kActivationMaxBufferSize];
  uint32_t unwrapped_size = sizeof(unwrapped);

  TEE_Result res = TEES_UnwrapSecureObject(wrapped, wrapped_size,
                                           unwrapped, &unwrapped_size);
  if (res != TEE_SUCCESS) {
    MB_LOGE("Cannot unwrap key\n");
    return TEE_ERROR_BAD_FORMAT;
  }

  const void *drk_cert, *rsa_cert, *rsa_priv;
  size_t drk_cert_size, rsa_cert_size, rsa_priv_size;

  DrkResult drk_res = DrkParserGetCertificates(
                            unwrapped, unwrapped_size,
                            &drk_cert, &drk_cert_size,
                            &rsa_cert, &rsa_cert_size,
                            &rsa_priv, &rsa_priv_size);
  if (drk_res != kDrkResultSuccess) {
    return TEE_ERROR_GENERIC;
  }

  if (drk_cert_size > params[1].memref.size ||
      rsa_cert_size > params[2].memref.size) {
    MB_LOGE("Output buffer is too short\n");
    return TEE_ERROR_SHORT_BUFFER;
  }

  TEE_ObjectHandle storage = TEE_HANDLE_NULL;

  // Save RSA private key.
  res = TEE_CreatePersistentObject(
      TEE_STORAGE_PRIVATE, RSA_PRIVATE_KEY_STORAGE_NAME,
      sizeof(RSA_PRIVATE_KEY_STORAGE_NAME),
      TEE_DATA_FLAG_ACCESS_WRITE | TEE_DATA_FLAG_OVERWRITE, TEE_HANDLE_NULL,
      rsa_priv, rsa_priv_size, &storage);
  if (res != TEE_SUCCESS) {
    MB_LOGE("Cannot store RSA private key\n");
    return res;
  }

  TEE_CloseObject(storage);

  // Return RSA x.509 and DRK certificates.
  memcpy(params[1].memref.buffer, drk_cert, drk_cert_size);
  params[1].memref.size = drk_cert_size;
  memcpy(params[2].memref.buffer, rsa_cert, rsa_cert_size);
  params[2].memref.size = rsa_cert_size;

  return TEE_SUCCESS;
}


static TEE_Result HandleActivationStoreServerKey(TEE_Param params[4]) {
  MB_LOGD("Store AES server key command\n");

  TEE_Result res;

  uint8_t priv_cert[kActivationMaxBufferSize];
  uint32_t priv_cert_size = sizeof(priv_cert);

  TEE_ObjectHandle storage = TEE_HANDLE_NULL;

  // Open previously stored RSA private key.
  res = TEE_OpenPersistentObject(TEE_STORAGE_PRIVATE,
                                 RSA_PRIVATE_KEY_STORAGE_NAME,
                                 sizeof(RSA_PRIVATE_KEY_STORAGE_NAME),
                                 TEE_DATA_FLAG_ACCESS_READ, &storage);
  if (res != TEE_SUCCESS) {
    MB_LOGE("Cannot open RSA private key object\n");
    return res;
  }

  // Read RSA private key.
  TEE_SeekObjectData(storage, 0, TEE_DATA_SEEK_SET);
  res = TEE_ReadObjectData(storage, priv_cert, priv_cert_size, &priv_cert_size);
  TEE_CloseObject(storage);

  if (res != TEE_SUCCESS) {
    MB_LOGE("RSA private key is too long\n");
    return res;
  }

  const void *n, *e, *d, *prime1, *prime2, *exp1, *exp2, *coef;
  size_t n_size, e_size, d_size, prime1_size, prime2_size, exp1_size, exp2_size,
         coef_size;

  // Retrieve RSA parameters.
  DrkResult drk_res = DrkParserGetRsaParameters(
                          priv_cert, priv_cert_size,
                          &n, &n_size, &e, &e_size, &d, &d_size,
                          &prime1, &prime1_size, &prime2, &prime2_size,
                          &exp1, &exp1_size, &exp2, &exp2_size,
                          &coef, &coef_size);
  if (res != kDrkResultSuccess) {
    MB_LOGE("DRK private key certificate is invalid\n");
    return TEE_ERROR_GENERIC;
  }

  TEE_Attribute attrs[8];

  // Set RSA parameters previously read from certificate.
  TEE_InitRefAttribute(&attrs[0], TEE_ATTR_RSA_MODULUS, n, n_size);
  TEE_InitRefAttribute(&attrs[1], TEE_ATTR_RSA_PUBLIC_EXPONENT, e, e_size);
  TEE_InitRefAttribute(&attrs[2], TEE_ATTR_RSA_PRIVATE_EXPONENT, d, d_size);
  TEE_InitRefAttribute(&attrs[3], TEE_ATTR_RSA_PRIME1, prime1, prime1_size);
  TEE_InitRefAttribute(&attrs[4], TEE_ATTR_RSA_PRIME2, prime2, prime2_size);
  TEE_InitRefAttribute(&attrs[5], TEE_ATTR_RSA_EXPONENT1, exp1, exp1_size);
  TEE_InitRefAttribute(&attrs[6], TEE_ATTR_RSA_EXPONENT2, exp2, exp2_size);
  TEE_InitRefAttribute(&attrs[7], TEE_ATTR_RSA_COEFFICIENT, coef, coef_size);

  TEE_ObjectHandle rsa_key_object = TEE_HANDLE_NULL;

  // Allocate RSA key object.
  res = TEE_AllocateTransientObject(TEE_TYPE_RSA_KEYPAIR, n_size * 8,
                                    &rsa_key_object);
  if (res != TEE_SUCCESS) {
    MB_LOGE("Cannot allocate RSA transient object\n");
    return res;
  }

  // Set RSA parameters to RSA key object.
  TEE_PopulateTransientObject(rsa_key_object, attrs, 8);

  TEE_OperationHandle operation = TEE_HANDLE_NULL;

  // Allocate RSA decrypt operation.
  res = TEE_AllocateOperation(&operation, TEE_ALG_RSAES_PKCS1_V1_5,
                              TEE_MODE_DECRYPT, n_size * 8);
  if (res != TEE_SUCCESS) {
    MB_LOGE("Cannot allocate RSA operation\n");
    TEE_CloseObject(rsa_key_object);
    return res;
  }

  // Set RSA key for operation.
  res = TEE_SetOperationKey(operation, rsa_key_object);
  if (res != TEE_SUCCESS) {
    MB_LOGD("Cannot set RSA key to operation\n");
    TEE_FreeOperation(operation);
    TEE_CloseObject(rsa_key_object);
    return res;
  }

  // Get encrypted AES key from parameters.
  void *encrypted_aes_key = params[0].memref.buffer;
  uint32_t encrypted_aes_size = params[0].memref.size;

  // Decrypted AES key.
  uint8_t aes_key[kAesBits / 8];
  uint32_t aes_size = sizeof(aes_key);

  // Decrypt AES key.
  res = TEE_AsymmetricDecrypt(operation, TEE_HANDLE_NULL, 0, encrypted_aes_key,
                              encrypted_aes_size, aes_key, &aes_size);

  TEE_FreeOperation(operation);
  TEE_CloseObject(rsa_key_object);

  if (res != TEE_SUCCESS) {
    MB_LOGE("Fail to decrypt AES key\n");
    return res;
  }

  if (aes_size * 8 != kAesBits) {
    MB_LOGE("Corrupted AES key\n");
    return TEE_ERROR_GENERIC;
  }

  // Create AES persistent storage.
  res = TEE_CreatePersistentObject(
      TEE_STORAGE_PRIVATE, AES_STORAGE_NAME, sizeof(AES_STORAGE_NAME),
      TEE_DATA_FLAG_ACCESS_WRITE | TEE_DATA_FLAG_OVERWRITE,
      TEE_HANDLE_NULL, aes_key, aes_size, &storage);

  TEE_CloseObject(storage);

  if (res != TEE_SUCCESS) {
    MB_LOGE("Cannot store object\n");
    return res;
  }

  return TEE_SUCCESS;
}

static const uint8_t kAad[kAesADataSize] = {
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
};

static TEE_Result HandleActivationGenerateCredentials(TEE_Param params[4]) {
  MB_LOGD("Generate credentials command\n");

  if (params[0].memref.size > kActivationMaxBufferSize) {
    MB_LOGD("Input buffer size > kActivationMaxBufferSize\n");
    return TEE_ERROR_BAD_PARAMETERS;
  }

  TEE_ObjectHandle aes_storage;
  TEE_Result tee_api_res = TEE_OpenPersistentObject(TEE_STORAGE_PRIVATE,
      AES_STORAGE_NAME, sizeof(AES_STORAGE_NAME), TEE_DATA_FLAG_ACCESS_READ,
      &aes_storage);
  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot open AES object\n");
    return tee_api_res;
  }

  uint8_t key[kAesBits / 8];
  uint32_t key_size = sizeof(key);
  TEE_SeekObjectData(aes_storage, 0, TEE_DATA_SEEK_SET);
  tee_api_res = TEE_ReadObjectData(aes_storage, key, key_size, &key_size);

  // One way or another we need to close object. Check state is below.
  TEE_CloseObject(aes_storage);

  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot read AES object\n");
    return tee_api_res;
  }

  TEE_ObjectHandle ase_key_object = NULL;
  tee_api_res = TEE_AllocateTransientObject(TEE_TYPE_AES, kAesBits,
      &ase_key_object);
  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot allocate transient object\n");
    return tee_api_res;
  }

  TEE_Attribute attrs[1];
  TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, key, key_size);

  tee_api_res = TEE_PopulateTransientObject(ase_key_object, attrs, 1);
  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot populate transient object\n");
    TEE_CloseObject(ase_key_object);
    return tee_api_res;
  }

  TEE_OperationHandle aes_encrypt_operation = NULL;
  tee_api_res = TEE_AllocateOperation(&aes_encrypt_operation, TEE_ALG_AES_GCM,
      TEE_MODE_ENCRYPT, kAesBits);
  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot allocate operation\n");
    TEE_CloseObject(ase_key_object);
    return tee_api_res;
  }

  tee_api_res = TEE_SetOperationKey(aes_encrypt_operation, ase_key_object);

  // Once more. Whatever result is we need to close object. The result of
  // TEE_SetOperationKey will be checked further
  TEE_CloseObject(ase_key_object);

  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot set operation key\n");
    TEE_FreeOperation(aes_encrypt_operation);
    return tee_api_res;
  }

  uint8_t *iv = params[2].memref.buffer;
  uint32_t iv_size = params[2].memref.size;
  // _TEE_Panic occurs when TEE_GenerateRandom is called in blowfish
  // Use zero iv instead.
#if defined (USE_BLOWFISH)
  if (iv == NULL)
    MB_LOGD("iv is NULL, iv_size is %d\n", iv_size);
  else if(iv_size == 0)
    MB_LOGD("iv_size is 0\n");
  iv_size = 12;
  memset(iv, 0x00, iv_size);
#else
  TEE_GenerateRandom((void *)iv, iv_size);
#endif

  tee_api_res = TEE_AEInit(aes_encrypt_operation, iv, iv_size, kAesBits, 0, 0);
  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot init AE\n");
    TEE_FreeOperation(aes_encrypt_operation);
    return tee_api_res;
  }

  TEE_AEUpdateAAD(aes_encrypt_operation, kAad, kAesADataSize);

  uint8_t *input_data = (uint8_t *)params[0].memref.buffer;
  uint32_t input_data_size = params[0].memref.size;
  uint8_t *encrypted_data = (uint8_t *)params[1].memref.buffer;
  uint32_t encrypted_data_size = params[1].memref.size;
  uint8_t *tag = (uint8_t *)params[3].memref.buffer;
  uint32_t tag_size = params[3].memref.size;
  tee_api_res = TEE_AEEncryptFinal(aes_encrypt_operation, input_data,
      input_data_size, encrypted_data, &encrypted_data_size, tag, &tag_size);

  // Free and then check TEE_AEEncryptFinal operation status
  TEE_FreeOperation(aes_encrypt_operation);

  if (tee_api_res != TEE_SUCCESS) {
    MB_LOGD("Cannot encrypt\n");
    return tee_api_res;
  }

  params[1].memref.size = encrypted_data_size;
  params[2].memref.size = iv_size;
  params[3].memref.size = tag_size;

  return TEE_SUCCESS;
}