/**
 * @file       activation_lib.c
 * @brief      Activation TEE API static library.
 *
 * Implementation is based on GP API and Multibuild SDK as a back-end.
 * Activation API is defined in activation.h as a main header for clients.
 *
 * @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 "activation_tee_command.h"
#include <activation.h>

#include <stdio.h>
#include <stdlib.h>
#include <tee_client_api.h>
#include <libdk_native_client.h>

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

static const TEEC_UUID kUuid = CONFIG_TA_UUID_STRUCT;

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

#define UNWRAP_RETRIES 3

ActivationResult ActivationGenerateSessionCertificate(
    void *drk_cert, size_t *drk_cert_size,
    void *rsa_cert, size_t *rsa_cert_size) {
  if (!rsa_cert || !rsa_cert_size || !drk_cert || !drk_cert_size) {
    MB_LOGD("Parameters cannot be NULL\n");
    return kActivationErrorBadParameters;
  }

  // Check DRK existance.
  int ret = isExistDeviceRootKey(KEY_TYPE_RSA);
  if (DRK_IS_EXIST != ret) {
    MB_LOGE("DRK service is not launched: %d\n", ret);
    return kActivationErrorDrk;
  }

  // Start the connection to TEE.
  // Initialize context, open session.
  // Need to load Activation TA first for generating DRK service key.

  TEEC_Context context;
  TEEC_Session session;
  TEEC_Operation operation;
  uint32_t return_origin;

  TEEC_Result res = TEEC_InitializeContext(NULL, &context);
  if (res != TEEC_SUCCESS) {
    MB_LOGD("InitializeContext failed: %x\n", res);
    return kActivationErrorTee;
  }

  ActivationResult result = kActivationErrorGeneric;

  res = TEEC_OpenSession(&context, &session, &kUuid, 0, NULL, NULL, &return_origin);
  if (res != TEEC_SUCCESS || return_origin != TEEC_ORIGIN_TRUSTED_APP) {
    MB_LOGD("TEEC_OpenSession returned 0x%x from 0x%x\n", res, return_origin);
    result = kActivationErrorTee;
    goto finalize_context;
  }

  uint8_t wrapped_cert[kActivationMaxBufferSize];
  int wrapped_len;

  int i = 0;

  res = TEEC_ERROR_BAD_FORMAT;

  // Workaround: retries are needed to fix the issue when Activation TA
  // cannon unwrap (decapsulate) DRK wrapped object.
  // qsee_decapsulate_inter_app_message() returns QSEE_MESSAGE_ERROR_BAD_SOURCE_NAME
  // (0xFF000FF3) sometimes. The reason is unclear: wrong encapsulation by DRK,
  // or wrong decapsulation by QSEE. Anyway, retries help to workaround the issue.
  for (i = 0; res == TEEC_ERROR_BAD_FORMAT && i < UNWRAP_RETRIES; i++) {
    // Generate RSA key pair. Encrypted object is returned.
    // TEE must decapsulate the wrapped blob and return public part
    // containing public DRK certificates.
    wrapped_len = createServiceKeySession("CHNACTIV", KEY_TYPE_RSA, 0,
                                          (char *)wrapped_cert,
                                          sizeof(wrapped_cert));
    if (wrapped_len <= NO_ERROR) {
      MB_LOGE("DRK service internal error: %d", wrapped_len);
      result = kActivationErrorDrk;
      goto close_session;
    }

    MB_LOGD("DRK certificate is here. Length is %d.\n", wrapped_len);

    // Send commands, close session.
    // Output size of unwrapped certificates is checked inside.
    operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT,
                                            TEEC_MEMREF_TEMP_OUTPUT,
                                            TEEC_MEMREF_TEMP_OUTPUT,
                                            TEEC_NONE);
    operation.params[0].tmpref.buffer = wrapped_cert;
    operation.params[0].tmpref.size = wrapped_len;
    operation.params[1].tmpref.buffer = drk_cert;
    operation.params[1].tmpref.size = *drk_cert_size;
    operation.params[2].tmpref.buffer = rsa_cert;
    operation.params[2].tmpref.size = *rsa_cert_size;

    res = TEEC_InvokeCommand(&session, kActivationStorePrivateKey,
                             &operation, &return_origin);
    if (res != TEEC_SUCCESS || return_origin != TEEC_ORIGIN_TRUSTED_APP) {
      MB_LOGE("TEEC_InvokeCommand returned 0x%x from 0x%x\n", res, return_origin);

      // Return "short buffer" for usability purporses.
      if (res == TEEC_ERROR_SHORT_BUFFER) {
        result = kActivationErrorShortBuffer;
      } else {
        result = kActivationErrorTee;
      }

      releaseServiceKeySession();
      continue;
    }

    // Update output result sizes. And unwrapped certificates are already in output buffer.
    *drk_cert_size = operation.params[1].tmpref.size;
    *rsa_cert_size = operation.params[2].tmpref.size;

    result = kActivationSuccess;

    releaseServiceKeySession();
  }

close_session:
  TEEC_CloseSession(&session);

finalize_context:
  TEEC_FinalizeContext(&context);

  return result;
}

ActivationResult ActivationStoreServerKey(
    const void *key,
    size_t key_size) {
  ActivationResult result = kActivationErrorGeneric;
  TEEC_Result res;
  TEEC_Context context;
  TEEC_Session session;
  TEEC_Operation operation;
  uint32_t return_origin;

  res = TEEC_InitializeContext(NULL, &context);
  if (res != TEEC_SUCCESS) {
    MB_LOGD("InitializeContext failed: 0x%x\n", res);
    return kActivationErrorTee;
  }

  res = TEEC_OpenSession(&context, &session, &kUuid, 0, NULL, NULL, &return_origin);
  if (res != TEEC_SUCCESS || return_origin != TEEC_ORIGIN_TRUSTED_APP) {
    MB_LOGD("TEEC_OpenSession returned 0x%x from 0x%x\n", res, return_origin);
    result = kActivationErrorTee;
    goto finalize_context;
  }

  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT,
                                          TEEC_NONE,
                                          TEEC_NONE,
                                          TEEC_NONE);
  operation.params[0].tmpref.buffer = (void *)key;
  operation.params[0].tmpref.size = key_size;

  res = TEEC_InvokeCommand(&session, kActivationStoreServerKey,
                           &operation, &return_origin);
  if (res != TEEC_SUCCESS || return_origin != TEEC_ORIGIN_TRUSTED_APP) {
    MB_LOGD("TEEC_InvokeCommand returned 0x%x from 0x%x\n", res, return_origin);
    result = kActivationErrorTee;
    goto close_session;
  }

  result = kActivationSuccess;

close_session:
  TEEC_CloseSession(&session);
finalize_context:
  TEEC_FinalizeContext(&context);
  return result;
}

ActivationResult ActivationGenerateCredentials(
    const uint8_t *buffer, size_t buffer_size,
    uint8_t *enc, size_t *enc_size,
    uint8_t *iv, size_t *iv_size,
    uint8_t *tag, size_t *tag_size) {
  if (!buffer || !buffer_size || !enc || !enc_size ||
      *enc_size < buffer_size || !tag || !tag_size ||
      !iv || !iv_size) {
    return kActivationErrorBadParameters;
  }

  ActivationResult activation_result = kActivationErrorGeneric;

  TEEC_Context context;
  TEEC_Result res = TEEC_InitializeContext(NULL, &context);
  if (res != TEEC_SUCCESS) {
    MB_LOGD("InitializeContext failed: %x\n", res);
    activation_result = kActivationErrorTee;
    goto exit;
  }

  TEEC_Session session;
  uint32_t return_origin;
  res = TEEC_OpenSession(&context, &session, &kUuid, 0, NULL, NULL, &return_origin);
  if (res != TEEC_SUCCESS || return_origin != TEEC_ORIGIN_TRUSTED_APP) {
    MB_LOGD("TEEC_OpenSession returned %x from %x\n", res, return_origin);
    activation_result = kActivationErrorTee;
    goto finalize_context;
  }

  TEEC_Operation operation;
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT,
                                          TEEC_MEMREF_TEMP_OUTPUT,
                                          TEEC_MEMREF_TEMP_OUTPUT,
                                          TEEC_MEMREF_TEMP_OUTPUT);

  operation.params[0].tmpref.buffer = (void *)buffer;
  operation.params[0].tmpref.size = buffer_size;
  operation.params[1].tmpref.buffer = (void *)enc;
  operation.params[1].tmpref.size = *enc_size;
  operation.params[2].tmpref.buffer = (void *)iv;
  operation.params[2].tmpref.size = *iv_size;
  operation.params[3].tmpref.buffer = (void *)tag;
  operation.params[3].tmpref.size = *tag_size;

  res = TEEC_InvokeCommand(&session, kActivationGenerateCredentials,
                           &operation, &return_origin);
  if (res != TEEC_SUCCESS || return_origin != TEEC_ORIGIN_TRUSTED_APP) {
    MB_LOGD("TEEC_InvokeCommand returned %x from %x\n", res, return_origin);
    if (res == TEEC_ERROR_SHORT_BUFFER) {
      activation_result = kActivationErrorShortBuffer;
    } else if(res == TEEC_ERROR_BAD_PARAMETERS) {
      activation_result = kActivationErrorBadParameters;
    } else {
       activation_result = kActivationErrorTee;
    }
    goto close_session;
  }

  *enc_size = operation.params[1].tmpref.size;
  *iv_size = operation.params[2].tmpref.size;
  *tag_size = operation.params[3].tmpref.size;
  activation_result = kActivationSuccess;

close_session:
  TEEC_CloseSession(&session);

finalize_context:
  TEEC_FinalizeContext(&context);

exit:
  return activation_result;
}