#include "device_compromised.h"
#include "api.h"
#include "config.h"
#include "pa_tz_api.h"
#include "driver_ipc.h"
#include "driver_log.h"
#include "serialize.h"
#include "PaDriverCommand.h"
#include "PaDriverCommandResponse.h"

#include <tee_internal_api.h>

/**
 * @brief Checks whether pointer size is correct
 * @note Pointer size can be counterfeited by third party
 * @param [in] addr_size Pointer size from NWD
 * @result 1 if address size is correct, 0 0therwise
 */
static int IsCorrectPointerSize(size_t addr_size);

/**
 * @brief Checks authentication of process in SWD
 * @details  The function checks authentication of process in SWD
 * @param [in] memory Pointer of command structure
 * @param [out] response Pointer to response structure
 * @result ::PA_TZ_SUCCESS if command is processed successfully,
 *         ::PA_TZ_GENERAL_ERROR if command can not be processed
 */
static PaTzResult HandleAuthenticateProcess(
    const PaDriverCommandAuthenticate_t *command,
    PaDriverCommandAuthenticateResponse_t *response);

/**
 * @brief Read data from Nwd process
 * @param [in] request Request structure
 * @param [out] response Result of operation
 * @return ::PA_TZ_SUCCESS if command is processed successfully,
 *         ::PA_TZ_GENERAL_ERROR if command can not be processed
 */
static PaTzResult HandleReadFromNwdTask(
    const PaDriverCommandReadFromNwdTask_t *request,
    PaDriverCommandReadResponse_t *response);

/**
 * @brief Write data to Nwd process
 * @param [in] request Request structure with data
 * @param [out] response Result of operation
 * @return ::PA_TZ_SUCCESS if command is processed successfully,
 *         ::PA_TZ_GENERAL_ERROR if command can not be processed
 */
static PaTzResult HandleWriteToNwdTask(
    const PaDriverCommandWriteToNwdTask_t *request,
    PaDriverCommandResult_t *response);

/**
 * @brief Create new certificate
 * @param [in] request Request structure with data
 * @param [out] response Result of operation
 * @return ::PA_TZ_SUCCESS if command is processed successfully,
 *         ::PA_TZ_GENERAL_ERROR if command can not be processed
 */
static PaTzResult HandleNewCertificate(
    const PaDriverCommandNewCertificate_t *request,
    PaDriverCommandProvisioningResponse_t *response);


static int IsCorrectPointerSize(size_t addr_size) {
  const uint8_t kArch32PointerSize = 4;
  const uint8_t kArch64PointerSize = 8;

  if (addr_size > sizeof(ProcessAddress)) {
    return 0;
  }

  if (!(addr_size == kArch32PointerSize
      || addr_size == kArch64PointerSize)) {
    return 0;
  }

  return 1;
}

static PaDriverAuthenticateResult_t ConvertAuthenticateResult(PaTzResult result) {
  switch (result) {
    case PA_TZ_SUCCESS: return PaDriverAuthenticateResult_paAuthenticated;
    case PA_TZ_AF_TASK_IS_NOT_FOUND: return PaDriverAuthenticateResult_paTaskIsNotFound;
    case PA_TZ_AF_INTEGRITY_IS_NONE: return PaDriverAuthenticateResult_paIntegrityIsNone;
    case PA_TZ_AF_INTEGRITY_IS_NOT_READY: return PaDriverAuthenticateResult_paIntegrityIsNotready;
    case PA_TZ_AF_CERTIFICATE_IS_ABSENT: return PaDriverAuthenticateResult_paCertificateIsAbsent;
    case PA_TZ_AF_CERTIFICATE_IS_INCORRECT: return PaDriverAuthenticateResult_paCertificateIsIncorrect;
    case PA_TZ_AF_CERTIFICATE_IS_NOT_MATCH: return PaDriverAuthenticateResult_paCertificateIsNotMatch;
    case PA_TZ_AF_APPNAME_IS_INCORRECT: return PaDriverAuthenticateResult_paAppNameIsIncorrect;
    case PA_TZ_AF_TASK_HAS_DUPLICATE: return PaDriverAuthenticateResult_paAppNameIsNotUnique;
    case PA_TZ_PROCA_NOT_SUPPORTED: return PaDriverAuthenticateResult_paNotSupported;
    default: return PaDriverAuthenticateResult_paNotAuthenticated;
  }
}

static inline PaDriverCommandResult_t ConvertReadWriteResult(PaTzResult result) {
  switch (result) {
    case PA_TZ_SUCCESS: return PaDriverCommandResult_paSuccess;
    case PA_TZ_VIRT_TO_PHYS_ERROR: return PaDriverCommandResult_paVirtToPhysError;
    case PA_TZ_MEMORY_PERM_ERROR: return PaDriverCommandResult_paMemPermissionError;
    case PA_TZ_GENERAL_ERROR: return PaDriverCommandResult_paGeneralError;
    default: return PaDriverCommandResult_paGeneralError;
  }
}

static PaTzResult HandleAuthenticateProcess(
    const PaDriverCommandAuthenticate_t *command,
    PaDriverCommandAuthenticateResponse_t *response) {
  if (command == NULL || response == NULL) {
    LOG_E("Pointer of memory is NULL\n");
    return PA_TZ_GENERAL_ERROR;
  }

  LOG_V("Process name: %s, size: %d\n", command->processName.buf,
      command->processName.size);
  LOG_V("Memory range: address: 0x%x, size:%d type:%d\n",
      command->memoryStartAddress, command->memorySize, command->memoryType);

  PaHandler *handler = (PaHandler *)command->handler.buf;

  LOG_I("Start authenticate user process.\n");
  LOG_V("User data pointer is\n");
  LOGM_V(&handler->data, sizeof(handler->data));

  const char *names = NULL;
  size_t names_size = 0;
  PaTzMemoryRange *memory = NULL;
  PaTzMemoryRange memory_range = {0};

  if (command->processName.buf && command->processName.size != 0) {
    names = (const char *)command->processName.buf;
    names_size = command->processName.size;
  }

  if (command->memoryStartAddress && command->memorySize != 0) {
    memory_range.type = (PaTzMemoryType)command->memoryType;
    memory_range.addr = command->memoryStartAddress;
    memory_range.size = command->memorySize;
    memory = &memory_range;
  }

  ProcessInfo process_info = {0};
  PaTzResult result = Authenticate(*handler, names, names_size, memory, &process_info);
  if (result == PA_TZ_SUCCESS) {
    LOG_I("Authentication finished.\n");
    OCTET_STRING_fromBuf(&response->processInfo, (const char *)&process_info,
                         sizeof(process_info));
  } else if (result == PA_TZ_PROCA_NOT_SUPPORTED) {
    LOG_I("This device does not support PROCA authentication.\n");
  } else {
    LOG_E("Authenticate return error.\n");
    LOG_D("Received result: 0x%x\n", result);
  }

  response->result = ConvertAuthenticateResult(result);

  return PA_TZ_SUCCESS;
}

static PaTzResult HandleReadFromNwdTask(
    const PaDriverCommandReadFromNwdTask_t *request,
    PaDriverCommandReadResponse_t *response) {
  if (request == NULL || response == NULL) {
    LOG_E("Pointer of memory is NULL\n");
    return PA_TZ_GENERAL_ERROR;
  }

  response->result = PaDriverCommandResult_paGeneralError;
  PaHandler *handler = (PaHandler *)request->handler.buf;
  LOG_V("User process PID from handler is %d.\n", handler->pid);

  ProcessAddress user_addr = 0;
  void *trustlet_addr = NULL;

  if (!request->inputAddress.buf
      || request->inputAddress.size < sizeof(user_addr)) {
    LOG_E("Input address is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_MemMove(&user_addr, request->inputAddress.buf, sizeof(user_addr));

  LOGADDR_V("User process buffer is : ", user_addr);

  if (!request->outputAddress.buf
      || !IsCorrectPointerSize(request->outputAddress.size)) {
    LOG_E("Output address is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_MemMove(&trustlet_addr, request->outputAddress.buf, sizeof(trustlet_addr));

  PaTzResult result = TaskReadToTrustlet(*handler, user_addr, request->inputSize, trustlet_addr);
  if (result == PA_TZ_PROCA_NOT_SUPPORTED) {
    LOG_I("This device does not support PROCA.\n");
  } else if (result != PA_TZ_SUCCESS) {
    LOG_E("TaskReadMemoryForTrustlet return error.\n");
    response->result = ConvertReadWriteResult(result);
  } else {
    response->result = PaDriverCommandResult_paSuccess;
  }

  return PA_TZ_SUCCESS;
}

static PaTzResult HandleWriteToNwdTask(
    const PaDriverCommandWriteToNwdTask_t *request,
    PaDriverCommandResult_t *response) {
  if (request == NULL || response == NULL) {
    LOG_E("Pointer of memory is NULL\n");
    return PA_TZ_GENERAL_ERROR;
  }

  *response = PaDriverCommandResult_paGeneralError;
  PaHandler *handler = (PaHandler *)request->handler.buf;

  LOG_V("User process PID from handler is %d.\n", handler->pid);

  ProcessAddress user_addr = 0;
  void *trustlet_addr = NULL;

  if (!request->inputAddress.buf
      || !IsCorrectPointerSize(request->inputAddress.size)) {
    LOG_E("Input address is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_MemMove(&trustlet_addr, request->inputAddress.buf, sizeof(trustlet_addr));

  if (!request->outputAddress.buf
      || request->outputAddress.size < sizeof(user_addr)) {
    LOG_E("Output address is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_MemMove(&user_addr, request->outputAddress.buf, sizeof(user_addr));

  LOGADDR_V("User process buffer is: ", user_addr);

  PaTzResult result = TaskWriteFromTrustlet(*handler, trustlet_addr, request->inputSize, user_addr);

  if (result == PA_TZ_PROCA_NOT_SUPPORTED) {
    LOG_I("This device does not support PROCA.\n");
  } else if (result != PA_TZ_SUCCESS) {
    LOG_E("TaskWriteMemory return error.\n");
    *response = ConvertReadWriteResult(result);
  } else {
    *response = PaDriverCommandResult_paSuccess;
  }

  return PA_TZ_SUCCESS;
}

static PaTzResult HandleNewCertificate(
    const PaDriverCommandNewCertificate_t *request,
    PaDriverCommandProvisioningResponse_t *response) {
  if (request == NULL || response == NULL) {
    LOG_E("Pointer of memory is NULL\n");
    return PA_TZ_GENERAL_ERROR;
  }

  response->result = PaDriverCommandResult_paGeneralError;

  PaHandler *handler = (PaHandler *)request->handler.buf;
  LOG_V("User process PID from handler is %d.\n", handler->pid);

  ProcessAddress mapped_apk = 0;
  if (!request->apkAddress.buf
      || !IsCorrectPointerSize(request->apkAddress.size)) {
    LOG_E("APK address is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_MemMove(&mapped_apk, request->apkAddress.buf, request->apkAddress.size);

  const char *package_name = NULL;
  if (!request->packageName.buf || request->packageName.size == 0) {
    LOG_E("APK package name is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  package_name = (char *)request->packageName.buf;

  const uint8_t *rsa = NULL;
  if (!request->rsaPublicKey.buf || request->rsaPublicKey.size == 0) {
    LOG_E("APK rsa is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  rsa = request->rsaPublicKey.buf;
  const size_t rsa_size = request->rsaPublicKey.size;

  uint8_t output_xattr[kSerializedDataMaxSize];
  uint32_t output_xattr_size = sizeof(output_xattr);

  PaTzResult result = IssueNewCertificate(*handler, mapped_apk, package_name,
                                          rsa, rsa_size, output_xattr,
                                          &output_xattr_size);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("IssueNewCertificate return error.\n");
    response->result = PaDriverCommandResult_paGeneralError;
  } else {
    LOG_D("Returned certificate with size: %d\n", output_xattr_size);
    LOGM_D(output_xattr, output_xattr_size);
    OCTET_STRING_fromBuf(&response->signedCertificate, (const char *)output_xattr, output_xattr_size);
    response->result = PaDriverCommandResult_paSuccess;
  }

  return PA_TZ_SUCCESS;
}

TEE_Result TEES_HandleDriverCommand(uint32_t cmd,
    const void *in, uint32_t in_size, void *out, uint32_t *out_size) {
  int checker = PA_TZ_SUCCESS;

  if (!in || !out || !out_size) {
    LOG_E("Pointer of memory is NULL\n");
    return TEE_ERROR_BAD_PARAMETERS;
  }

  PaDriverCommand_t *request = NULL;
  PaDriverCommandResponse_t response = {PaDriverCommandResponse_PR_NOTHING};

  if (!ConfigIsInited()) {
    LOG_E("Config is not initialized in SWD\n");
    return TEE_ERROR_BAD_STATE;
  }

  // Step 1: Decode incoming command buffer
  checker = PaDecodeDriverCommand(in, in_size, &request);
  if (checker != PA_TZ_SUCCESS || !request) {
    LOG_E("Failed decoding.\n");
    LOG_D("Received result: 0x%x\n", checker);
    return TEE_ERROR_BAD_PARAMETERS;
  }

  // Step 2: Run handler for present command
  switch (request->present) {
    case PaDriverCommand_PR_authenticateCommand: {
      response.present = PaDriverCommandResponse_PR_authenticateResponse;
      if (PlatformIsDeviceCompromised()) {
        // Diagnostic messages are printed by PlatformIsDeviceCompromised 
        response.choice.result = PaDriverAuthenticateResult_paDeviceIsCompromised;
        break;
      }
      if (HandleAuthenticateProcess(&request->choice.authenticateCommand,
                                    &response.choice.authenticateResponse)
          != PA_TZ_SUCCESS) {
        LOG_E("HalAuthenticateProcess can not processes command.\n");
      }
      break;
    }
    case PaDriverCommand_PR_readFromNwdTask: {
      response.present = PaDriverCommandResponse_PR_readResponse;
      if (PlatformIsDeviceCompromised()) {
        // Diagnostic messages are printed by PlatformIsDeviceCompromised 
        response.choice.result = PaDriverAuthenticateResult_paDeviceIsCompromised;
        break;
      }
      if (HandleReadFromNwdTask(&request->choice.readFromNwdTask,
                                &response.choice.readResponse) != PA_TZ_SUCCESS) {
        LOG_E("HalReadFromNwdTask can not processes command.\n");
      }
      break;
    }
    case PaDriverCommand_PR_writeToNwdTask: {
      response.present = PaDriverCommandResponse_PR_result;
      if (PlatformIsDeviceCompromised()) {
        // Diagnostic messages are printed by PlatformIsDeviceCompromised 
        response.choice.result = PaDriverAuthenticateResult_paDeviceIsCompromised;
        break;
      }
      if (HandleWriteToNwdTask(&request->choice.writeToNwdTask,
                               &response.choice.result) != PA_TZ_SUCCESS) {
        LOG_E("HalWriteToNwdTask can not processes command.\n");
      }
      break;
    }
    case PaDriverCommand_PR_newCertificate: {
      response.present = PaDriverCommandResponse_PR_provisioningResponse;
      if (PlatformIsDeviceCompromised()) {
        // Diagnostic messages are printed by PlatformIsDeviceCompromised 
        response.choice.result = PaDriverAuthenticateResult_paDeviceIsCompromised;
        break;
      }
      if (HandleNewCertificate(&request->choice.newCertificate,
                               &response.choice.provisioningResponse) != PA_TZ_SUCCESS) {
        LOG_E("HandleNewCertificate can not processes command.\n");
      }
      break;
    }
    case PaDriverCommand_PR_ping: {
      response.present = PaDriverCommandResponse_PR_result;
      response.choice.result = PaDriverCommandResult_paSuccess;
      break;
    }
    default: {
      response.present = PaDriverCommandResponse_PR_result;
      response.choice.result = PaDriverCommandResult_paGeneralError;
      LOG_E("Failed unknown message has been received.\n");
      break;
    }
  }

  // Step 3: Encode response to buffer
  checker = PaEncodeDriverCommandResponse(&response, out, out_size);
  if (checker != PA_TZ_SUCCESS) {
    LOG_E("Failed encoding response.\n");
    LOG_D("Received result: 0x%x\n", checker);
  }

  ASN_STRUCT_FREE(asn_DEF_PaDriverCommand, request);
  ASN_STRUCT_FREE_CONTENTS_ONLY(asn_DEF_PaDriverCommandResponse, &response);

  return ((checker == 0) ? TEE_SUCCESS : TEE_ERROR_GENERIC);
}