#include "pa_tz_api.h"
#include "driver_ipc.h"
#include "serialize.h"
#include "pa_tz_internal.h"

#include "platform_driver_api.h"
#include "driver_log.h"

#include <tee_internal_api.h>

#include <scl/arithmetic.h>
#include <scl/string.h>

#if defined(VERSION_NAME) && defined(VERSION_SUFFIX)
#define VERSION_STRING (VERSION_NAME VERSION_SUFFIX)
#else
#error "Please, define VERSION_NAME and VERSION_SUFFIX"
#endif

//Maximal size of name rule (for example: path of process, socket)
typedef enum {
  kMaximalSizeNameRule = 1024
} UtilsLimits;

static char process_names_buff[kMaximalNameNumbers * kMaximalSizeNameRule];

typedef struct {
  char *process_names;
  PaTzMemoryRange memory_range;
  uint32_t names_number;
  uint32_t names_size;
} PaRulesInternal;

static PaRulesInternal singleton;

static const char kVersionString[] __attribute__((used)) = VERSION_STRING;

/**
 * @brief Always print tz lib version to log
 */
static void PrintVersion(void);

PaTzRules PaTzRulesCreate(void) {
  TEE_MemFill(&singleton, 0, sizeof(singleton));
  return (PaTzRules)&singleton;
}

PaTzResult PaTzRulesAddProcessName(PaTzRules rules, const char *process_name) {
  PaRulesInternal *rules_internal = (PaRulesInternal *)rules;
  if (!rules_internal) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if (rules_internal->names_number >= kMaximalNameNumbers) {
    LOG_E("So many process names!\n");
    return PA_TZ_AF_APPNAME_IS_INCORRECT;
  }

  if (!rules_internal->process_names) {
    rules_internal->process_names = process_names_buff;
  }

  scl_size_t current_name_len = 0;
  if (!scl_strlen(process_name, kMaximalSizeNameRule, &current_name_len)) {
    // Failed to obtain string length
    return PA_TZ_AF_APPNAME_IS_INCORRECT;
  }

  if ((current_name_len > 0) &&
      !scl_strcpy(rules_internal->process_names + rules_internal->names_size,
                  kMaximalSizeNameRule, process_name)) {
    // Failed to string copy
    return PA_TZ_AF_APPNAME_IS_INCORRECT;
  }

  if (!SCL_ADD(&rules_internal->names_size,
               rules_internal->names_size, current_name_len + 1)) {
    // Failed to add two numbers
    return PA_TZ_AF_APPNAME_IS_INCORRECT;
  }
  rules_internal->names_number++;

  return PA_TZ_SUCCESS;
}

PaTzResult PaTzRulesAddMemoryRange(PaTzRules rules,
                                   const PaTzMemoryRange *memory) {
  PaRulesInternal *rules_internal = (PaRulesInternal *)rules;
  if (rules_internal) {
    TEE_MemMove(&rules_internal->memory_range, memory,
                sizeof(rules_internal->memory_range));
  }

  return PA_TZ_SUCCESS;
}

void PaTzRulesDestroy(PaTzRules rules) {
  PaRulesInternal *rules_internal = (PaRulesInternal *)rules;
  TEE_MemFill(rules_internal, 0, sizeof(PaRulesInternal));
}

PaTzResult PaTzAuthenticate(const PaHandler handler) {
  return PaTzAuthenticateWithRules(handler, NULL, NULL);
}

PaTzResult PaTzAuthenticateWithCommandBuffer(const PaHandler handler,
                                             const void *buffer, size_t bufferSize,
                                             ProcessInfo *process_info) {
  PaTzResult res;
  PaTzRules rules = PaTzRulesCreate();
  PaTzMemoryRange memory = {(uintptr_t)buffer, bufferSize, PA_MEMORY_SWD_VA};
  res = PaTzRulesAddMemoryRange(rules, &memory);
  if (PA_TZ_SUCCESS != res) {
    PaTzRulesDestroy(rules);
    return res;
  }

  res = PaTzAuthenticateWithRules(handler, rules, process_info);
  PaTzRulesDestroy(rules);

  return res;
}

PaTzResult PaTzAuthenticateWithRules(const PaHandler handler,
                                     const PaTzRules rules,
                                     ProcessInfo *process_info) {
  PrintVersion();

  static uint8_t data_encode_buffer[kSerializedDataMaxSize];
  PaTzResult result = PA_TZ_SUCCESS;

  PaDriverCommand_t driver_command = {PaDriverCommand_PR_NOTHING};

  driver_command.present = PaDriverCommand_PR_authenticateCommand;
  PaDriverCommandAuthenticate_t *command_authenticate = &driver_command.choice.authenticateCommand;

  TEE_MemFill(data_encode_buffer, 0, sizeof(data_encode_buffer));

  command_authenticate->handler.buf = (void *)&handler;
  command_authenticate->handler.size = sizeof(handler);

  PaRulesInternal *rules_internal = (PaRulesInternal *)rules;
  if (rules_internal != NULL) {
    command_authenticate->processName.buf = (uint8_t *)rules_internal->process_names;
    command_authenticate->processName.size = rules_internal->names_size;

    command_authenticate->memoryStartAddress = rules_internal->memory_range.addr;
    command_authenticate->memorySize = rules_internal->memory_range.size;
    command_authenticate->memoryType = rules_internal->memory_range.type;

    if (rules_internal->memory_range.type == PA_MEMORY_SWD_VA) {
      PlatformRegisterTrustletInputBuffer(
          (void *)(uintptr_t)rules_internal->memory_range.addr,
          rules_internal->memory_range.size);
    }
  }

  if (handler.flags == kPaHandlerTidUnusedMark && handler.data) {
    if (command_authenticate->processName.buf) {
      LOG_E("AppName is set in handler and in Rules.\n");
      return PA_TZ_INCOMPATIBLE_RULES;
    }

    char *name = (char *)(uintptr_t)handler.data;
    scl_size_t name_len = 0;
    if (!scl_strlen(name, kMaximalSizeNameRule, &name_len)) {
      // Failed to obtain string length
      return PA_TZ_AF_APPNAME_IS_INCORRECT;
    }

    command_authenticate->processName.buf = (uint8_t *)name;
    command_authenticate->processName.size = name_len;
  }

  uint32_t buffer_size = sizeof(data_encode_buffer);
  // Encode the command type as DER(BER)
  result = PaEncodeDriverCommand(
      &driver_command, data_encode_buffer, &buffer_size);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Failed encoding.\n");
    return PA_TZ_ENCODE_AUTHENTICATION_FAILED;
  }

  // Send encoding command to driver
  result = PlatformCallDriver(DRV_PA_TZ_ID, kFidDrvAuthenticate,
      data_encode_buffer, buffer_size);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Communication with PA is failed.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  // Decoding the answer command
  PaDriverCommandResponse_t *command_response = NULL;
  do {
    result = PaDecodeDriverCommandResponse(
        data_encode_buffer, kSerializedDataMaxSize, &command_response);
    if (result != PA_TZ_SUCCESS || !command_response) {
      LOG_E("Failed decoding.\n");
      LOG_D("Received result: %x.\n", result);
      result = PA_TZ_DECODE_AUTHENTICATION_FAILED;
      break;
    }

    if (command_response->present != PaDriverCommandResponse_PR_authenticateResponse) {
      LOG_E("Failed decoding: response is not valid\n");
      result = PA_TZ_DECODE_AUTHENTICATION_FAILED;
      break;
    }

    // Finally to check answer
    result = ConvertAuthenticateResult(&command_response->choice.authenticateResponse);
    if (result == PA_TZ_PROCA_NOT_SUPPORTED) {
      LOG_I("This device does not support PROCA authentication.\n");
      break;
    } else if (result != PA_TZ_SUCCESS) {
      LOG_E("Failed authenticate.\n");
      LOG_D("Received result: 0x%08x.\n", result);
      break;
    }

    if (process_info) {
      const OCTET_STRING_t *response_process_info = &command_response->choice.authenticateResponse.processInfo;

      if (response_process_info->buf
          && response_process_info->size == sizeof(ProcessInfo)) {
        TEE_MemMove(process_info, response_process_info->buf,
                    response_process_info->size);
      } else {
        LOG_E("Failed get process info.\n");
        LOG_D("Received result: %d.\n", result);
        break;
      }
    }
  } while (0);

  if (command_response) {
    ASN_STRUCT_FREE(asn_DEF_PaDriverCommandResponse, command_response);
  }

  return result;
}

PaTzResult PaTzHandlerCreateFromPid(int pid, PaHandler *handler)
{
  if (!handler) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_MemFill(handler, 0, sizeof(*handler));

  handler->flags = kPaHandlerTidUnusedMark;
  handler->pid = pid;

  return PA_TZ_SUCCESS;
}

PaTzResult PaTzHandlerCreateFromProcessName(const char *process_name, PaHandler *handler)
{
  if (!process_name || !handler) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  scl_size_t current_name_len = 0;
  if (!scl_strlen(process_name, kMaximalSizeNameRule, &current_name_len)) {
    // Failed to obtain string length
    return PA_TZ_AF_APPNAME_IS_INCORRECT;
  }

  TEE_MemFill(handler, 0, sizeof(*handler));

  handler->flags = kPaHandlerTidUnusedMark;
  handler->data = (uintptr_t)process_name;

  return PA_TZ_SUCCESS;
}

void PaTzHandlerDestroy(PaHandler *handler)
{
  TEE_MemFill(handler, 0, sizeof(*handler));
}

static void PrintVersion(void) {
  TEES_Log(TEES_LOG_LEVEL_INFO, __FUNCTION__, __LINE__, PLATFORM_LOG_TAG, "Version: %s\n", kVersionString);
}
