// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "tsp/test_helpers.h"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <variant>
#include <vector>

#include "absl/base/nullability.h"
#include "absl/log/absl_check.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "crypto/algorithms.h"
#include "crypto/cbs_utils.h"
#include "crypto/default/hasher.h"
#include "crypto/hash.h"
#include "openssl/base.h"
#include "openssl/bio.h"
#include "openssl/bytestring.h"
#include "openssl/digest.h"
#include "openssl/err.h"
#include "openssl/evp.h"
#include "openssl/mem.h"
#include "openssl/obj.h"
#include "openssl/pem.h"  // IWYU pragma: keep
#include "openssl/x509.h"
#include "proto/common.pb.h"
#include "testing/boringssl_utils.h"
#include "testing/cms.h"
#include "tink/cleartext_keyset_handle.h"
#include "tink/keyset_handle.h"
#include "tink/public_key_sign.h"
#include "tink/signature/config_v0.h"
#include "tink/signature/signature_pem_keyset_reader.h"
#include "tsp/constants.h"
#include "tsp/status_codes.h"

namespace credentio {
namespace {

using ::crypto::tink::KeysetHandle;
using ::crypto::tink::PublicKeySign;
using ::crypto::tink::SignaturePemKeysetReaderBuilder;

absl::StatusOr<std::unique_ptr<crypto::tink::KeysetHandle>>
LoadEcdsaPrivateKeyFromPem(absl::string_view pem_data) {
  crypto::tink::SignaturePemKeysetReaderBuilder builder(
      SignaturePemKeysetReaderBuilder::PemReaderType::PUBLIC_KEY_SIGN);
  builder.Add({.serialized_key = std::string(pem_data),
               .parameters = {
                   .key_type = crypto::tink::PemKeyType::PEM_EC,
                   .algorithm = crypto::tink::PemAlgorithm::ECDSA_DER,
                   .key_size_in_bits = 256,
                   .hash_type = google::crypto::tink::HashType::SHA256,
               }});
  ABSL_ASSIGN_OR_RETURN(auto reader, builder.Build());
  return crypto::tink::CleartextKeysetHandle::Read(std::move(reader));
}

absl::StatusOr<std::vector<uint8_t>> SignWithBoringSsl(
    const EVP_PKEY& private_key, absl::string_view data) {
  if (EVP_PKEY_id(&private_key) != EVP_PKEY_RSA) {
    return absl::UnimplementedError("Only RSA keys are supported.");
  }
  bssl::ScopedEVP_MD_CTX md_context;
  EVP_PKEY_CTX* pkey_ctx = nullptr;
  if (EVP_DigestSignInit(md_context.get(), &pkey_ctx, EVP_sha256(), nullptr,
                         const_cast<EVP_PKEY*>(&private_key)) != 1 ||
      EVP_DigestSignUpdate(md_context.get(), data.data(), data.length()) != 1) {
    return absl::InternalError("EVP_DigestSignInit/Update failed.");
  }
  size_t sig_len = 0;
  if (EVP_DigestSignFinal(md_context.get(), nullptr, &sig_len) != 1) {
    return absl::InternalError("EVP_DigestSignFinal failed.");
  }
  std::vector<uint8_t> signature(sig_len, 0);
  if (EVP_DigestSignFinal(md_context.get(), signature.data(), &sig_len) != 1) {
    return absl::InternalError("EVP_DigestSignFinal failed.");
  }
  return signature;
}

}  // namespace

std::string WrapDerSequence(absl::string_view payload) {
  CBB cbb;
  ABSL_CHECK(CBB_init(&cbb, payload.size() + 5));
  CBB seq_cbb;
  ABSL_CHECK(CBB_add_asn1(&cbb, &seq_cbb, CBS_ASN1_SEQUENCE));
  ABSL_CHECK(CBB_add_bytes(&seq_cbb,
                           reinterpret_cast<const uint8_t*>(payload.data()),
                           payload.size()));
  uint8_t* data;
  size_t len;
  ABSL_CHECK(CBB_finish(&cbb, &data, &len));
  bssl::UniquePtr<uint8_t> data_uniq(data);
  return std::string(reinterpret_cast<char*>(data), len);
}

std::string EncodeDerInteger(uint64_t n) {
  CBB cbb;
  ABSL_CHECK(CBB_init(&cbb, 10));
  ABSL_CHECK(CBB_add_asn1_int64(&cbb, n));
  uint8_t* cbb_data;
  size_t cbb_len;
  ABSL_CHECK(CBB_finish(&cbb, &cbb_data, &cbb_len));
  bssl::UniquePtr<uint8_t> cbb_data_uniq(cbb_data);
  return std::string(reinterpret_cast<char*>(cbb_data), cbb_len);
}

std::string CreateTimeStampResp(absl::string_view ts_token) {
  CBB cbb;
  ABSL_CHECK(CBB_init(&cbb, 40));
  CBB ts_resp_cbb;
  ABSL_CHECK(CBB_add_asn1(&cbb, &ts_resp_cbb, CBS_ASN1_SEQUENCE));

  CBB pki_status_cbb;
  ABSL_CHECK(CBB_add_asn1(&ts_resp_cbb, &pki_status_cbb, CBS_ASN1_SEQUENCE));
  ABSL_CHECK(CBB_add_asn1_int64(&pki_status_cbb,
                                static_cast<int64_t>(TspPkiStatus::kGranted)));

  ABSL_CHECK(CBB_add_bytes(&ts_resp_cbb,
                           reinterpret_cast<const uint8_t*>(ts_token.data()),
                           ts_token.size()));

  uint8_t* cbb_data_ptr;
  size_t cbb_len;
  ABSL_CHECK(CBB_finish(&cbb, &cbb_data_ptr, &cbb_len));
  bssl::UniquePtr<uint8_t> cbb_uniq(cbb_data_ptr);

  std::string cbb_str(cbb_data_ptr, cbb_data_ptr + cbb_len);
  return cbb_str;
}

std::string CreateFailedTimeStampResp(
    TspPkiStatus status, absl::Span<const absl::string_view> status_string,
    absl::Span<const int32_t> failure_info) {
  CBB cbb;
  ABSL_CHECK(CBB_init(&cbb, 512));
  CBB resp;
  ABSL_CHECK(CBB_add_asn1(&cbb, &resp, CBS_ASN1_SEQUENCE));
  CBB pki_status_info;
  ABSL_CHECK(CBB_add_asn1(&resp, &pki_status_info, CBS_ASN1_SEQUENCE));
  ABSL_CHECK(
      CBB_add_asn1_int64(&pki_status_info, static_cast<int64_t>(status)));
  if (!status_string.empty()) {
    CBB status_string_cbb;
    ABSL_CHECK(
        CBB_add_asn1(&pki_status_info, &status_string_cbb, CBS_ASN1_SEQUENCE));
    for (absl::string_view entry : status_string) {
      CBB utf8_string;
      ABSL_CHECK(
          CBB_add_asn1(&status_string_cbb, &utf8_string, CBS_ASN1_UTF8STRING));
      ABSL_CHECK(CBB_add_bytes(&utf8_string,
                               reinterpret_cast<const uint8_t*>(entry.data()),
                               entry.length()));
      ABSL_CHECK(CBB_flush(&status_string_cbb));
    }
    ABSL_CHECK(CBB_flush(&pki_status_info));
  }
  if (!failure_info.empty()) {
    CBB failure_info_cbb;
    ABSL_CHECK(
        CBB_add_asn1(&pki_status_info, &failure_info_cbb, CBS_ASN1_BITSTRING));
    ABSL_CHECK(CBB_add_u8(&failure_info_cbb, 0 /* no padding bits */));
    int32_t max_bit = 0;
    for (int32_t entry : failure_info) {
      if (entry > max_bit) {
        max_bit = entry;
      }
    }
    int32_t bytes = (max_bit / 8) + 1;
    std::string data(bytes, 0);
    for (int32_t entry : failure_info) {
      data[entry >> 3] |= 1 << (7 - (entry & 7));
    }
    ABSL_CHECK(CBB_add_bytes(&failure_info_cbb,
                             reinterpret_cast<const uint8_t*>(data.data()),
                             data.size()));
    ABSL_CHECK(CBB_flush(&pki_status_info));
  }

  uint8_t* cbb_data;
  size_t cbb_len;
  ABSL_CHECK(CBB_finish(&cbb, &cbb_data, &cbb_len));
  bssl::UniquePtr<uint8_t> cbb_data_uniq(cbb_data);

  return std::string(reinterpret_cast<char*>(cbb_data), cbb_len);
}

std::string CreateTimeStampTokenRsa(
    absl::Span<const absl::string_view> der_data, absl::string_view key_pem,
    absl::string_view cert_pem, bool embed_cert,
    std::variant<absl::string_view, bool> signing_certificate_attribute) {
  bssl::UniquePtr<EVP_PKEY> key(PemToKey(key_pem));
  ABSL_CHECK(key != nullptr);
  bssl::UniquePtr<X509> cert(PemToCertificate(cert_pem));
  ABSL_CHECK(cert != nullptr);
  std::string extra_attrs;
  if (std::holds_alternative<absl::string_view>(
          signing_certificate_attribute)) {
    extra_attrs = std::get<absl::string_view>(signing_certificate_attribute);
  } else {
    if (std::get<bool>(signing_certificate_attribute)) {
      X509* certs[] = {cert.get()};
      extra_attrs =
          CreateSigningCertificateV2Attribute(certs, HashAlgorithm::kSha256);
    }
  }

  auto cms_or = credentio_testing::CreateSignedCms(
      der_data, *cert,
      [&key](absl::string_view data) { return SignWithBoringSsl(*key, data); },
      {.signing_time = absl::FromUnixSeconds(1427738483),
       .include_user_certificate = embed_cert,
       .extra_signed_attributes = extra_attrs,
       .is_time_stamp_token = true});
  ABSL_CHECK_OK(cms_or.status());
  return *cms_or;
}

std::string CreateTimeStampTokenEcdsa(
    absl::Span<const absl::string_view> der_data,
    const crypto::tink::PublicKeySign& signer, absl::string_view cert_pem,
    bool embed_cert,
    std::variant<absl::string_view, bool> signing_certificate_attribute) {
  bssl::UniquePtr<X509> cert(PemToCertificate(cert_pem));
  ABSL_CHECK(cert != nullptr);
  std::string extra_attrs;
  if (std::holds_alternative<absl::string_view>(
          signing_certificate_attribute)) {
    extra_attrs = std::get<absl::string_view>(signing_certificate_attribute);
  } else {
    if (std::get<bool>(signing_certificate_attribute)) {
      X509* certs[] = {cert.get()};
      extra_attrs =
          CreateSigningCertificateV2Attribute(certs, HashAlgorithm::kSha256);
    }
  }

  auto cms_or = credentio_testing::CreateSignedCms(
      der_data, *cert,
      [&signer](
          absl::string_view data) -> absl::StatusOr<std::vector<uint8_t>> {
        ABSL_ASSIGN_OR_RETURN(auto signature, signer.Sign(data));
        return std::vector<uint8_t>(signature.begin(), signature.end());
      },
      {.signing_time = absl::FromUnixSeconds(1427738483),
       .include_user_certificate = embed_cert,
       .extra_signed_attributes = extra_attrs,
       .is_time_stamp_token = true});
  ABSL_CHECK_OK(cms_or.status());
  return *cms_or;
}

std::string CreateTimeStampTokenEcdsa(
    absl::Span<const absl::string_view> der_data, absl::string_view key_pem,
    absl::string_view cert_pem, bool embed_cert,
    std::optional<std::vector<const X509*>> extra_certs,
    std::variant<absl::string_view, bool> signing_certificate_attribute,
    bool use_wrong_e_content_type) {
  auto keyset_handle = LoadEcdsaPrivateKeyFromPem(key_pem);
  ABSL_CHECK_OK(keyset_handle);
  auto signer = (*keyset_handle)
                    ->GetPrimitive<crypto::tink::PublicKeySign>(
                        crypto::tink::ConfigSignatureV0());
  ABSL_CHECK_OK(signer);
  std::unique_ptr<crypto::tink::PublicKeySign> signer_ptr = *std::move(signer);

  bssl::UniquePtr<EVP_PKEY> key(PemToKey(key_pem));
  ABSL_CHECK(key != nullptr);
  bssl::UniquePtr<X509> cert(PemToCertificate(cert_pem));
  ABSL_CHECK(cert != nullptr);
  std::string extra_attrs;
  if (std::holds_alternative<absl::string_view>(
          signing_certificate_attribute)) {
    extra_attrs = std::get<absl::string_view>(signing_certificate_attribute);
  } else {
    if (std::get<bool>(signing_certificate_attribute)) {
      std::vector<X509*> certs;
      certs.push_back(cert.get());
      if (extra_certs.has_value()) {
        for (const X509* extra_cert : *extra_certs) {
          certs.push_back(const_cast<X509*>(extra_cert));
        }
      }
      extra_attrs = CreateSigningCertificateV2Attribute(absl::MakeSpan(certs),
                                                        HashAlgorithm::kSha256);
    }
  }

  std::vector<const X509*> ca_certs;
  if (extra_certs.has_value()) {
    ca_certs = *extra_certs;
  }

  auto cms_or = credentio_testing::CreateSignedCms(
      der_data, *cert,
      [&signer_ptr](
          absl::string_view data) -> absl::StatusOr<std::vector<uint8_t>> {
        ABSL_ASSIGN_OR_RETURN(auto signature, signer_ptr->Sign(data));
        return std::vector<uint8_t>(signature.begin(), signature.end());
      },
      {.signing_time = absl::FromUnixSeconds(1427738483),
       .include_user_certificate = embed_cert,
       .ca_certificates = ca_certs,
       .extra_signed_attributes = extra_attrs,
       .is_time_stamp_token = !use_wrong_e_content_type});
  ABSL_CHECK_OK(cms_or.status());
  return *cms_or;
}

std::string CreateTstInfo(absl::string_view gen_time,
                          absl::string_view message_imprint_hash,
                          absl::string_view message_imprint_algorithm_txt,
                          std::optional<absl::string_view> nonce) {
  CBB cbb;
  ABSL_CHECK(CBB_init(&cbb, 30));

  CBB sequence_cbb;
  ABSL_CHECK(CBB_add_asn1(&cbb, &sequence_cbb, CBS_ASN1_SEQUENCE));
  // version
  ABSL_CHECK(CBB_add_asn1_int64(&sequence_cbb, 1));

  // policy
  CBB policy_cbb;
  ABSL_CHECK(CBB_add_asn1(&sequence_cbb, &policy_cbb, CBS_ASN1_OBJECT));
  ABSL_CHECK(
      CBB_add_asn1_oid_from_text(&policy_cbb, "1.2.840.113554.4.1.72585", 24));
  ABSL_CHECK(CBB_flush(&sequence_cbb));

  // messageImprint
  CBB mi_cbb;
  ABSL_CHECK(CBB_add_asn1(&sequence_cbb, &mi_cbb, CBS_ASN1_SEQUENCE));
  // messageImprint.hashAlgorithm
  CBB algorithm_identifier_cbb;
  ABSL_CHECK(
      CBB_add_asn1(&mi_cbb, &algorithm_identifier_cbb, CBS_ASN1_SEQUENCE));
  // messageImprint.hashAlgorithm.algorithm
  CBB algorithm_cbb;
  ABSL_CHECK(
      CBB_add_asn1(&algorithm_identifier_cbb, &algorithm_cbb, CBS_ASN1_OBJECT));
  ABSL_CHECK(CBB_add_asn1_oid_from_text(
      &algorithm_cbb, message_imprint_algorithm_txt.data(),
      message_imprint_algorithm_txt.length()));
  ABSL_CHECK(CBB_flush(&mi_cbb));

  // messageImprint.hashedMessage
  ABSL_CHECK(CBB_add_asn1_octet_string(
      &mi_cbb, reinterpret_cast<const uint8_t*>(message_imprint_hash.data()),
      message_imprint_hash.length()));
  ABSL_CHECK(CBB_flush(&sequence_cbb));

  // serial_number
  ABSL_CHECK(CBB_add_asn1_int64(&sequence_cbb, 1));

  // genTime
  CBB gen_time_cbb;
  ABSL_CHECK(
      CBB_add_asn1(&sequence_cbb, &gen_time_cbb, CBS_ASN1_GENERALIZEDTIME));
  ABSL_CHECK(CBB_add_bytes(&gen_time_cbb,
                           reinterpret_cast<const uint8_t*>(gen_time.data()),
                           gen_time.length()));
  ABSL_CHECK(CBB_flush(&sequence_cbb));

  // accuracy (this is an empty SEQUENCE because currently code under test will
  // never actually read it)
  CBB accuracy_cbb;
  ABSL_CHECK(CBB_add_asn1(&sequence_cbb, &accuracy_cbb, CBS_ASN1_SEQUENCE));

  // ordering
  ABSL_CHECK(CBB_add_asn1_bool(&sequence_cbb, true));

  // nonce
  if (nonce.has_value()) {
    ABSL_CHECK(CBB_add_bytes(&sequence_cbb,
                             reinterpret_cast<const uint8_t*>(nonce->data()),
                             nonce->length()));
  }

  uint8_t* der_data;
  size_t der_len;
  ABSL_CHECK(CBB_finish(&cbb, &der_data, &der_len));

  std::string der(der_data, der_data + der_len);

  OPENSSL_free(der_data);
  return der;
}

std::string CreateSigningCertificateV2Attribute(
    absl::Span<X509* absl_nonnull> certs, HashAlgorithm hash_algorithm) {
  CBB cbb;
  ABSL_CHECK(CBB_init(&cbb, 1000));
  CBB attr_seq_cbb;
  ABSL_CHECK(CBB_add_asn1(&cbb, &attr_seq_cbb, CBS_ASN1_SEQUENCE));
  CBB oid_cbb;
  ABSL_CHECK(CBB_add_asn1(&attr_seq_cbb, &oid_cbb, CBS_ASN1_OBJECT));
  ABSL_CHECK(CBB_add_bytes(&oid_cbb, OBJ_get0_data(IdAaSigningCertificateV2()),
                           OBJ_length(IdAaSigningCertificateV2())));
  ABSL_CHECK(CBB_flush(&attr_seq_cbb));

  CBB value_set_cbb;
  ABSL_CHECK(CBB_add_asn1(&attr_seq_cbb, &value_set_cbb, CBS_ASN1_SET));
  CBB value_seq_cbb;
  ABSL_CHECK(CBB_add_asn1(&value_set_cbb, &value_seq_cbb, CBS_ASN1_SEQUENCE));
  CBB certs_cbb;
  ABSL_CHECK(CBB_add_asn1(&value_seq_cbb, &certs_cbb, CBS_ASN1_SEQUENCE));
  for (X509* x509 : certs) {
    uint8_t* der_ptr = nullptr;
    size_t der_len = i2d_X509(x509, &der_ptr);
    bssl::UniquePtr<uint8_t> der(der_ptr);

    CBS der_cbs;
    CBS_init(&der_cbs, der_ptr, der_len);
    auto hasher = CreateHasher(hash_algorithm);
    ABSL_CHECK_OK(hasher);
    (*hasher)->Update(ToStringView(der_cbs));
    std::string hash = (*hasher)->Digest();
    CBS hash_cbs = FromStringView(hash);

    CBB certs_entry_cbb;
    ABSL_CHECK(CBB_add_asn1(&certs_cbb, &certs_entry_cbb, CBS_ASN1_SEQUENCE));
    if (hash_algorithm != HashAlgorithm::kSha256) {
      ABSL_CHECK(
          OBJ_nid2cbb(&certs_entry_cbb, static_cast<int>(hash_algorithm)));
    }
    ABSL_CHECK(CBB_add_asn1_octet_string(&certs_entry_cbb, CBS_data(&hash_cbs),
                                         CBS_len(&hash_cbs)));
    ABSL_CHECK(CBB_flush(&certs_cbb));
  }

  auto result = FinishToString(&cbb);
  ABSL_CHECK_OK(result);
  return *result;
}

}  // namespace credentio