#include "api.h"

#include "access_control.h"
#include "authentication.h"
#include "config.h"
#include "driver_log.h"
#include "pa_certificate.h"
#include "provisioning.h"
#include "serialize.h"
#include "task.h"
#include "task_access.h"
#include "task_parser.h"
#include "PaFlagBits.h"

#include <tee_internal_api.h>

/**
 * @brief Check whether PaTzHandler has valid flags field
 * @param [in] handler
 * @return ::PA_TZ_SUCCESS, ::PA_TZ_GENERAL_ERROR
 */
static int IsPaTzHandlerValid(const PaHandler *handler) {
  return handler->flags == handler->pid ||
         handler->flags == kPaHandlerTidUnusedMark;
}

PaTzResult Authenticate(PaHandler handler,
                        const char *process_names, const size_t process_names_size,
                        const PaTzMemoryRange *memory, ProcessInfo *info) {
  TaskInfo task = {0};
  PaTzResult result = PA_TZ_AUTHENTICATION_FAILED;

  if (!IsPaTzHandlerValid(&handler)) {
    LOG_E("Incoming PaHandler is invalid.\n");
    return PA_TZ_INVALID_HANDLER;
  }

  if (!IsFiveEnabled()) {
    LOG_I("This device does not support FIVE.\n");
    return PA_TZ_PROCA_NOT_SUPPORTED;
  }

  do {
    if (!AccessControlIsAllowedOperation(kAuthentication)) {
      LOG_E("Current caller can not call Authenticate.\n");
      result = PA_TZ_CALLER_IS_FORBIDEN;
      break;
    }

    if (handler.pid) {
      result = TaskFindByPid(handler.pid, &task);
    } else if (process_names) {
      result = TaskFindByAppName(process_names, process_names_size, &task);
    } else {
      LOG_E("Got invalid PaTzHandler, both pid and application name are NULL.\n");
      result = PA_TZ_INVALID_HANDLER;
      break;
    }

    if (PA_TZ_SUCCESS != result) {
      LOG_E("Task was not found.\n");
      LOG_D("PID: %d.\n", handler.pid);
      break;
    }

    result = ProcessAuthentication(&task, process_names, process_names_size, memory);
    if (PA_TZ_SUCCESS != result) {
      LOG_E("Task with PID is not authenticated.\n");
      LOG_D("PID: %d.\n", handler.pid);
      break;
    }

    // Pass additional information about process to client
    if (info) {
      TEE_MemMove(info->id, task.certificate->paData.paId.buf,
                  task.certificate->paData.paId.size);

      // Set bit fields of process information.
      info->is_android_app =
          (((uint64_t)task.certificate->paData.paFlags
              & (1 << PaFlagBits_bitAndroid)) ? 1 : 0);
      info->is_thirdparty_app =
          (((uint64_t)task.certificate->paData.paFlags
              & (1 << PaFlagBits_bitThirdParty)) ? 1 : 0);
      info->is_weak_integrity = ((PA_TZ_SUCCESS == CheckIntegrityWeak(&task)) ? 1 : 0);
    }

    result = PA_TZ_SUCCESS;
  } while (0);

  FreeTaskInfo(&task);

  return result;
}

PaTzResult TaskReadToTrustlet(PaHandler handler, ProcessAddress user_address,
                              size_t size, void *out_trustlet) {
  if (!out_trustlet) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if (!IsPaTzHandlerValid(&handler)) {
    LOG_E("Incoming PaHandler is invalid.\n");
    return PA_TZ_INVALID_HANDLER;
  }

  if (!IsFiveEnabled()) {
    LOG_I("This device does not support FIVE.\n");
    return PA_TZ_PROCA_NOT_SUPPORTED;
  }

  if (!AccessControlIsAllowedOperation(kReadToTrustlet)) {
    LOG_E("Current caller can not call TaskReadToTrustlet.\n");
    return PA_TZ_CALLER_IS_FORBIDEN;
  } else {
    LOG_V("Current caller can call TaskReadToTrustlet.\n");
  }

  TaskInfo task = {0};

  void *virt = NULL;

  PaTzResult result = PA_TZ_GENERAL_ERROR;

  do {
    result = TaskFindByPid(handler.pid, &task);
    if (PA_TZ_SUCCESS != result) {
      LOG_E("Task by PID is not find.\n");
      LOG_D("PID: %d.\n", handler.pid);
      break;
    }

    result = PA_TZ_SUCCESS;

  } while (0);

  if (result != PA_TZ_SUCCESS) {
    FreeTaskInfo(&task);

    return result;
  }

  do {
    // Map output trustlet address for Write
    result = PlatformSysMapTrustlet(out_trustlet, size, kMemoryAccessWrite,
                                    &virt);
    if (result != PA_TZ_SUCCESS || !virt) {
      LOG_E("PlatformSysMapTrustlet return error.\n");
      break;
    }

    result = TaskAccessGetBytes(&task, user_address, size, virt);
    if (result != PA_TZ_SUCCESS) {
      LOG_E("TaskAccessGetBytes return error.\n");
    }

    if (PlatformSysUnmapTrustlet(virt, size) != PA_TZ_SUCCESS) {
      LOG_E("PlatformSysUnmapTrustlet return error.\n");
    }

  } while (0);

  FreeTaskInfo(&task);

  return result;
}

PaTzResult TaskWriteFromTrustlet(PaHandler handler, const void *in_trustlet,
                                 size_t size, ProcessAddress user_address) {
  if (!in_trustlet) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if (!IsPaTzHandlerValid(&handler)) {
    LOG_E("Incoming PaHandler is invalid.\n");
    return PA_TZ_INVALID_HANDLER;
  }

  if (!IsFiveEnabled()) {
    LOG_I("This device does not support FIVE.\n");
    return PA_TZ_PROCA_NOT_SUPPORTED;
  }

  if (!AccessControlIsAllowedOperation(kWriteFromTrustlet)) {
    LOG_E("Current caller can not call TaskWriteFromTrustlet.\n");
    return PA_TZ_CALLER_IS_FORBIDEN;
  } else {
    LOG_V("Current caller can call TaskWriteFromTrustlet.\n");
  }

  TaskInfo task = {0};

  void *virt = NULL;

  PaTzResult result = PA_TZ_GENERAL_ERROR;

  do {
    result = TaskFindByPid(handler.pid, &task);
    if (PA_TZ_SUCCESS != result) {
      LOG_E("Task by PID is not find.\n");
      LOG_D("PID: %d.\n", handler.pid);
      break;
    }

    result = PA_TZ_SUCCESS;

  } while (0);

  if (result != PA_TZ_SUCCESS) {
    FreeTaskInfo(&task);

    return result;
  }

  do {
    // Map input trustlet address for Read
    result = PlatformSysMapTrustlet(in_trustlet, size, kMemoryAccessRead,
                                    &virt);
    if (result != PA_TZ_SUCCESS || !virt) {
      LOG_E("PlatformSysMapTrustlet return error.\n");
      break;
    }

    result = TaskAccessPutBytes(&task, virt, size, user_address);
    if (result != PA_TZ_SUCCESS) {
      LOG_E("TaskAccessPutBytes return error.\n");
    }

    if (PlatformSysUnmapTrustlet(virt, size) != PA_TZ_SUCCESS) {
      LOG_E("PlatformSysUnmapTrustlet return error.\n");
    }

  } while (0);

  FreeTaskInfo(&task);

  return result;
}

PaTzResult IssueNewCertificate(PaHandler handler,
                               ProcessAddress mapped_apk,
                               const char *package_name,
                               const uint8_t *rsa, size_t rsa_size,
                               uint8_t *out_xattr, uint32_t *out_xattr_size) {
  TaskInfo task = {0};
  PaCertificate_t *new_certificate = NULL;
  PaTzResult result = PA_TZ_GENERAL_ERROR;

  if (!IsPaTzHandlerValid(&handler)) {
    LOG_E("Incoming PaHandler is invalid.\n");
    return PA_TZ_INVALID_HANDLER;
  }

  if (!IsFiveEnabled()) {
    LOG_I("This device does not support FIVE.\n");
    return PA_TZ_PROCA_NOT_SUPPORTED;
  }

  result = TaskFindByPid(handler.pid, &task);
  if (PA_TZ_SUCCESS != result) {
    LOG_E("Task by PID is not find.\n");
    LOG_D("PID: %d.\n", handler.pid);
    return result;
  }

  do {
    result = CreateNewCertificate(&task, mapped_apk, package_name,
                                  rsa, rsa_size, &new_certificate);
    if (result != PA_TZ_SUCCESS) {
      LOG_E("Can not create new certificate. Package_name:\n");
      LOG_E(package_name);
      break;
    }

    result = PaEncodeCertificate(new_certificate, (void *)out_xattr, out_xattr_size);
    if (result != PA_TZ_SUCCESS) {
      LOG_E("Can not encode PA certificate\n");
      break;
    }
  } while (0);

  FreeTaskInfo(&task);
  PaCertificateDestroy(new_certificate);

  return result;
}