#include "pa_certificate.h"
#include "config.h"
#include "driver_log.h"
#include "serialize.h"
#include "validation_key.h"

#include <tee_internal_api.h>

/**
 * @brief Sign PaData of certificate using HMAC
 * @param [in] data PaData of certificate
 * @param [out] signature,signature_size Signature
 * @return ::PA_TZ_SUCCESS in case of success, ::PA_TZ_GENERAL_ERROR
 */
static PaTzResult SignCertificate(const PaData_t *data,
                                  void *signature, size_t *signature_size);


static const int kCertificateVersion = 1;

enum {
  kHmacSignatureMaxSize = 256
};

PaTzResult PaCertificateCreate(const uint8_t *id, size_t id_size,
                               const uint8_t *five_signature, size_t five_signature_size,
                               const uint8_t *pa_app_name, size_t pa_app_name_size,
                               int flags,
                               PaCertificate_t **out_certificate) {
  if (!id || !five_signature || !pa_app_name || !out_certificate) {
    LOG_E("Empty invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if ((id_size != kPaIdLength) || (pa_app_name_size > kPaAppNameMaxLength)) {
    LOG_E("Incorrect arguments size.\n");
    LOG_D("id_size %d, five_signature_size %d, pa_app_name_size %d.\n",
          id_size, five_signature_size, pa_app_name_size);
    return PA_TZ_GENERAL_ERROR;
  }

  PaCertificate_t *certificate = TEE_Malloc(sizeof(PaCertificate_t),
                                            HINT_FILL_WITH_ZEROS);
  if (!certificate) {
    LOG_E("Can not allocate memory for PA certificate.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  certificate->paData.paVersion = kCertificateVersion;

  OCTET_STRING_fromBuf(&certificate->paData.paId, (char *)id, id_size);
  OCTET_STRING_fromBuf(&certificate->paData.paAppName, (char *)pa_app_name, pa_app_name_size);

  uint8_t five_signature_hash[kSha256Size];
  PaTzResult result = CryptoSha256(five_signature, five_signature_size, five_signature_hash);
  if (PA_TZ_SUCCESS != result) {
    LOG_E("Cannot calculate hash of FIVE signature.\n");
    LOG_D("Received result: 0x%x\n", result);
    return result;
  }

  OCTET_STRING_fromBuf(&certificate->paData.fiveSignatureHash,
      (char *)five_signature_hash, sizeof(five_signature_hash));

  certificate->paData.paFlags = flags | (1 << PaFlagBits_bitHmac);

  uint8_t hmac_signature[kHmacSignatureMaxSize] = {0};
  size_t hmac_signature_size = sizeof(hmac_signature);

  result = SignCertificate(&certificate->paData, hmac_signature, &hmac_signature_size);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not sign certificate.\n");
    PaCertificateDestroy(certificate);
    return result;
  }

  OCTET_STRING_fromBuf(&certificate->signature, (char *)hmac_signature, hmac_signature_size);

  *out_certificate = certificate;

  return PA_TZ_SUCCESS;
}

void PaCertificateDestroy(PaCertificate_t *certificate) {
  if (!certificate) {
    return;
  }

  ASN_STRUCT_FREE(asn_DEF_PaCertificate, certificate);
}

PaTzResult PaCertificateValidate(const PaCertificate_t *certificate) {
  if (!certificate) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  PaTzResult result = PA_TZ_GENERAL_ERROR;

  do {
    uint8_t memory[kSerializedDataMaxSize];
    uint32_t size_encoded = sizeof(memory);

    int checker = PaEncodePaData(&certificate->paData, memory,
                                 &size_encoded);
    if (checker == -1) {
      LOG_E("Failed encoding certificate.\n");
      LOG_D("Received result: %d\n", checker);
      break;
    }

    if (certificate->paData.paVersion != kCertificateVersion) {
      LOG_E("Didn't match current and supported certificate versions.\n");
      LOG_D("Current certificate version = %d, supported certificate version = %d\n",
            certificate->paData.paVersion, kCertificateVersion);
      return PA_TZ_GENERAL_ERROR;
    }

    if ((certificate->paData.paFlags & (1 << PaFlagBits_bitHmac)) == 0) {
      RsaPublicKey validation_key;
      result = ValidationKeyGetByType((BuildType)certificate->paData.paKeyId, &validation_key);
      if (result != PA_TZ_SUCCESS) {
        LOG_E("Can not get return key value.\n");
        LOG_D("Received result: %d\n", result);
        return PA_TZ_GENERAL_ERROR;
      }

      LOG_D("RSA signature verification.\n");
      result = RsaSignatureVerification(memory, size_encoded,
          certificate->signature.buf,
          certificate->signature.size,
          &validation_key, kSha256);
    } else {
      LOG_D("HMAC signature verification.\n");
      result = CryptoHmacSignatureVerification(memory, size_encoded,
          kSha256,
          certificate->signature.buf,
          certificate->signature.size);
    }

    if (result != PA_TZ_SUCCESS) {
      LOG_E("Signature validation error.\n");
      LOG_D("Received result: 0x%x\n", result);
      break;
    }

    result = PA_TZ_GENERAL_ERROR;

    if (!certificate->paData.paId.buf || certificate->paData.paId.size != kPaIdLength) {
      LOG_E("Can't get PA ID.\n");
      LOG_D("PA ID has paId.size: %d.\n", certificate->paData.paId.size);
      break;
    }

    if (!certificate->paData.paAppName.buf ||
        certificate->paData.paAppName.size > kPaAppNameMaxLength) {
      LOG_E("Can't get PA app name.\n");
      LOG_D("PA app name size: %d.\n", certificate->paData.paAppName.size);
      break;
    }

    if (certificate->paData.fiveSignatureHash.size != kSha256Size) {
      LOG_E("FIVE signature size is incorrect! Should be Sha256(32).\n");
      LOG_D("FIVE signature size: %d.\n",
          certificate->paData.fiveSignatureHash.size);
      break;
    }

    result = PA_TZ_SUCCESS;

    LOG_D("PA certificate signature has been verified\n");
  } while (0);

  return result;
}

static PaTzResult SignCertificate(const PaData_t *data, void *signature, size_t *signature_size) {
  uint8_t buffer[kSerializedDataMaxSize];
  uint32_t buffer_size = sizeof(buffer);

  int ret = PaEncodePaData(data, buffer, &buffer_size);
  if (ret) {
    LOG_E("Cannot encode PA data.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  PaTzResult result = CryptoHmacSignatureGenerate(buffer, buffer_size,
                                                  kSha256,
                                                  signature, signature_size);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("The HMAC signing of certificate is failed.\n");
  } else {
    LOG_I("The HMAC certificate is signed successfully.\n");
  }

  return result;
}