#include "provisioning.h"

#include "authentication.h"
#include "crypto.h"
#include "driver_log.h"
#include "pa_certificate.h"
#include "scl/string.h"
#include "serialize.h"
#include "task.h"
#include "task_access.h"
#include "task_parser.h"

#include "PaFlagBits.h"
#include "PaCertificate.h"

#include <tee_internal_api.h>

enum {
  kSecureStorageKeyMaterialSize = 128,
  kHmacSignatureMaxSize = 256
};


/**
 * @brief Generate PA ID for APK
 * @param [in] package_name Package name of APK
 * @param [in] rsa,rsa_size Public key of APK
 * @param [out] id,id_size Output PA ID
 * @return ::PA_TZ_SUCCESS in case of success, ::PA_TZ_GENERAL_ERROR
 */
static PaTzResult GeneratePaIdForApk(const char *package_name,
                                     const uint8_t *rsa, size_t rsa_size,
                                     void *id, size_t *id_size);

/**
 * @brief Return FIVE signature of mapped file
 * @param [in] task Task structure
 * @param [in] mapped_file Process address of mapped file
 * @param [out] five_signature,five_signature_length FIVE signature
 * @return ::PA_TZ_SUCCESS in case of success, ::PA_TZ_GENERAL_ERROR
 */
static PaTzResult GetFiveSignatureOfMappedFile(
    const TaskInfo *task, ProcessAddress mapped_file,
    uint8_t *five_signature, size_t *five_signature_length);

PaTzResult CreateNewCertificate(const TaskInfo *issuer,
                                ProcessAddress mapped_apk,
                                const char *package_name,
                                const uint8_t *rsa, size_t rsa_size,
                                PaCertificate_t **new_certificate) {
  if (!issuer || !package_name || mapped_apk == 0 ||
      !rsa || rsa_size == 0 ||
      !new_certificate) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  PaTzResult result = CheckIntegritySigningRights(issuer);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Requested process has no rights for signing.\n");
    return result;
  }

  uint8_t five_signature[kMaxPaSignatureLength];
  size_t five_signature_length = sizeof(five_signature);

  result = GetFiveSignatureOfMappedFile(issuer, mapped_apk, five_signature, &five_signature_length);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not obtain FIVE signature for APK.\n");
    return result;
  }

  uint8_t id[kPaIdLength] = {0};
  size_t id_size = sizeof(id);
  result = GeneratePaIdForApk(package_name, rsa, rsa_size, id, &id_size);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not generate id for APK.\n");
    return result;
  }

  int flags = (1 << PaFlagBits_bitAndroid | 1 << PaFlagBits_bitThirdParty);

  size_t package_name_len = 0;
  if (!scl_strlen(package_name, kPaAppNameMaxLength, (scl_size_t *)&package_name_len)) {
    return PA_TZ_GENERAL_ERROR;
  }
  result = PaCertificateCreate(id, id_size, five_signature, five_signature_length,
                               (const uint8_t *)package_name, package_name_len,
                               flags, new_certificate);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not create PA certificate.\n");
    return result;
  }

  return result;
}

static PaTzResult GeneratePaIdForApk(const char *package_name,
                                     const uint8_t *rsa, size_t rsa_size,
                                     void *id, size_t *id_size) {
  if (!package_name || !rsa || !id || !id_size) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if (rsa_size < kSecureStorageKeyMaterialSize) {
    LOG_E("RSA modulus is too small!\n");
    LOG_D("RSA modulus: %d.\n", rsa_size);
    return PA_TZ_GENERAL_ERROR;
  }

  size_t package_name_len = 0;
  if (!scl_strlen(
      package_name, kPaAppNameMaxLength, (scl_size_t *)&package_name_len)) {
    LOG_E("Package name is too long.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_OperationHandle sha256operation;
  TEE_Result result = TEE_AllocateOperation(
      &sha256operation, TEE_ALG_SHA256, TEE_MODE_DIGEST, 0);
  if (result != TEE_SUCCESS) {
    LOG_E("Can not allocate hash operation.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  TEE_DigestUpdate(sha256operation, package_name, package_name_len);

  result = TEE_DigestDoFinal(
      sha256operation, rsa, rsa_size, id, (uint32_t *)id_size);

  TEE_FreeOperation(sha256operation);

  return ((result == TEE_SUCCESS) ? PA_TZ_SUCCESS : PA_TZ_GENERAL_ERROR);
}

static PaTzResult GetFiveSignatureOfMappedFile(
    const TaskInfo *task, ProcessAddress mapped_file,
    uint8_t *five_signature, size_t *five_signature_length) {
  if (!task || !five_signature || !five_signature_length) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  VmaInfo file = {0};
  PaTzResult result = TaskFindVmaWithProcessAddress(task, mapped_file, &file);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Provided address of mapped file is incorrect.\n");
    return result;
  }

  result = TaskParseFiveSignature(file.vm_file, five_signature, five_signature_length);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not obtain FIVE signature for mapped file.\n");
    return result;
  }

  return result;
}