// 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 "testing/local_key_crypto_write_handler.h"

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

#include "absl/base/no_destructor.h"
#include "absl/base/nullability.h"
#include "absl/functional/any_invocable.h"
#include "absl/log/absl_check.h"
#include "absl/log/die_if_null.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/clock_interface.h"
#include "absl/time/time.h"
#include "async/async_context.h"
#include "crypto/algorithms.h"
#include "crypto/crypto_write_handler.h"
#include "crypto/timestamp_write_handler.h"
#include "tink/insecure_secret_key_access.h"
#include "tink/key_status.h"
#include "tink/keyset_handle.h"
#include "tink/partial_key_access.h"
#include "tink/pem/signature_key_parser.h"
#include "tink/public_key_sign.h"
#include "tink/signature/config_v0.h"
#include "tink/signature/ecdsa_parameters.h"
#include "tink/signature/ecdsa_private_key.h"
#include "tink/signature/signature_config.h"
#include "tsp/test_helpers.h"

namespace credentio {
namespace {

using ::crypto::tink::EcdsaParameters;
using ::crypto::tink::EcdsaPrivateKey;
using ::crypto::tink::GetPartialKeyAccess;
using ::crypto::tink::InsecureSecretKeyAccess;
using ::crypto::tink::KeysetHandle;
using ::crypto::tink::KeysetHandleBuilder;
using ::crypto::tink::KeyStatus;
using ::crypto::tink::PublicKeySign;
using ::crypto::tink::SignatureConfig;

// Extra bytes reserved on top of a measured signature or timestamp size, since
// DER-encoded ECDSA signatures vary in length.
constexpr uint64_t kSizeMargin = 10;

// OCSP responses are dropped once this much time has passed since the signer
// was created.
constexpr absl::Duration kOcspResponseLifetime = absl::Hours(24 * 7);

// Returns the OID of `hash_algorithm` (see RFC 5754).
absl::StatusOr<absl::string_view> HashAlgorithmOid(
    HashAlgorithm hash_algorithm) {
  switch (hash_algorithm) {
    case HashAlgorithm::kSha256:
      return "2.16.840.1.101.3.4.2.1";
    case HashAlgorithm::kSha384:
      return "2.16.840.1.101.3.4.2.2";
    case HashAlgorithm::kSha512:
      return "2.16.840.1.101.3.4.2.3";
    default:
      return absl::InvalidArgumentError(
          absl::StrCat("Unsupported hash algorithm: ", hash_algorithm));
  }
}

// Returns an ES256 signer for the PEM-encoded EC P-256 private key `key_pem`.
absl::StatusOr<std::unique_ptr<PublicKeySign>> CreateEs256Signer(
    absl::string_view key_pem,
    EcdsaParameters::SignatureEncoding signature_encoding) {
  // Populates the global registry that the config-less
  // `KeysetHandleBuilder::Build()` below relies on.
  static const absl::NoDestructor<absl::Status> registration_status(
      SignatureConfig::Register());
  ABSL_RETURN_IF_ERROR(*registration_status);

  ABSL_ASSIGN_OR_RETURN(EcdsaParameters parameters,
                        EcdsaParameters::Builder()
                            .SetCurveType(EcdsaParameters::CurveType::kNistP256)
                            .SetHashType(EcdsaParameters::HashType::kSha256)
                            .SetSignatureEncoding(signature_encoding)
                            .SetVariant(EcdsaParameters::Variant::kNoPrefix)
                            .Build());
  ABSL_ASSIGN_OR_RETURN(EcdsaPrivateKey private_key,
                        tink_pem::PemToEcdsaPrivateKey(
                            key_pem, parameters, InsecureSecretKeyAccess::Get(),
                            GetPartialKeyAccess()));
  ABSL_ASSIGN_OR_RETURN(
      KeysetHandle keyset_handle,
      KeysetHandleBuilder()
          .AddEntry(KeysetHandleBuilder::Entry::CreateFromCopyableKey(
              std::move(private_key), KeyStatus::kEnabled,
              /*is_primary=*/true))
          .Build());
  return keyset_handle.GetPrimitive<PublicKeySign>(
      crypto::tink::ConfigSignatureV0());
}

// Returns a DER-encoded RFC 3161 TimeStampToken over `hashed_message`, signed
// by `tsa_signer` at the current time on `clock`.
absl::StatusOr<std::string> CreateTimeStampToken(
    HashAlgorithm hash_algorithm, absl::string_view hashed_message,
    absl::Clock& clock, const PublicKeySign& tsa_signer,
    absl::string_view tsa_cert_pem,
    std::optional<absl::Time> tsa_signing_time) {
  ABSL_ASSIGN_OR_RETURN(absl::string_view oid,
                        HashAlgorithmOid(hash_algorithm));
  const std::string gen_time =
      absl::FormatTime("%Y%m%d%H%M%SZ", clock.TimeNow(), absl::UTCTimeZone());
  const std::string tst_info = CreateTstInfo(gen_time, hashed_message, oid);
  return CreateTimeStampTokenEcdsa({tst_info}, tsa_signer, tsa_cert_pem,
                                   /*embed_cert=*/true,
                                   /*signing_certificate_attribute=*/true,
                                   tsa_signing_time);
}

}  // namespace

LocalKeySigner::LocalKeySigner(
    std::unique_ptr<PublicKeySign> absl_nonnull signer,
    std::vector<std::string> certificate_chain,
    std::vector<std::string> ocsp_responses, absl::Clock* absl_nonnull clock,
    SigningAlgorithm algorithm)
    : signer_(std::move(ABSL_DIE_IF_NULL(signer))),
      certificate_chain_(std::move(certificate_chain)),
      ocsp_responses_(std::move(ocsp_responses)),
      clock_(ABSL_DIE_IF_NULL(clock)),
      create_time_(clock_->TimeNow()),
      algorithm_(algorithm) {}

const std::vector<std::string>& LocalKeySigner::ocsp_responses() const {
  if (clock_->TimeNow() > create_time_ + kOcspResponseLifetime) {
    static const absl::NoDestructor<std::vector<std::string>> kNoResponses;
    return *kNoResponses;
  }
  return ocsp_responses_;
}

uint64_t LocalKeySigner::max_signature_byte_size() const {
  absl::MutexLock lock(mutex_);
  if (!max_signature_byte_size_.has_value()) {
    // The interface cannot report errors, and a signer that cannot sign is
    // unusable, so crash.
    absl::StatusOr<SignResponse> response = Sign("measure signature size");
    ABSL_CHECK_OK(response);
    max_signature_byte_size_ = response->signature.size() + kSizeMargin;
  }
  return *max_signature_byte_size_;
}

absl::StatusOr<SignResponse> LocalKeySigner::Sign(
    absl::string_view to_be_signed) const {
  ABSL_ASSIGN_OR_RETURN(std::string signature, signer_->Sign(to_be_signed));
  return SignResponse{.signature = std::move(signature)};
}

LocalKeyCryptoWriteHandler::LocalKeyCryptoWriteHandler(
    absl::Clock* absl_nonnull clock,
    std::unique_ptr<PublicKeySign> absl_nonnull signer,
    std::unique_ptr<PublicKeySign> absl_nullable tsa_signer,
    bool is_timestamping_disabled, std::vector<std::string> certificate_chain,
    std::vector<std::string> ocsp_responses, SigningAlgorithm algorithm,
    std::string tsa_cert_pem, std::optional<absl::Time> tsa_signing_time)
    : clock_(ABSL_DIE_IF_NULL(clock)),
      signer_(std::move(signer), std::move(certificate_chain),
              std::move(ocsp_responses), clock, algorithm),
      tsa_signer_(std::move(tsa_signer)),
      is_timestamping_disabled_(is_timestamping_disabled),
      tsa_cert_pem_(std::move(tsa_cert_pem)),
      tsa_signing_time_(tsa_signing_time) {}

void LocalKeyCryptoWriteHandler::CreateTimestamp(
    HashAlgorithm hash_algorithm, std::string hashed_message,
    TimestampWriteHandler::Options options,
    absl::AnyInvocable<void(absl::StatusOr<std::string>) &&> cb) const {
  AsyncResolver<std::string> resolver(options.async_context, std::move(cb));
  if (is_timestamping_disabled_ || tsa_signer_ == nullptr) {
    resolver.Return(absl::UnavailableError("TSA is disabled."));
    return;
  }
  resolver.Return(CreateTimeStampToken(hash_algorithm, hashed_message, *clock_,
                                       *tsa_signer_, tsa_cert_pem_,
                                       tsa_signing_time_));
}

uint64_t LocalKeyCryptoWriteHandler::max_timestamp_byte_size(
    HashAlgorithm hash_algorithm) const {
  if (tsa_signer_ == nullptr) {
    return 0;
  }
  absl::MutexLock lock(mutex_);
  auto [it, inserted] =
      max_timestamp_byte_size_map_.try_emplace(hash_algorithm);
  if (inserted) {
    // The interface cannot report errors, and a TSA that cannot sign is
    // unusable, so crash.
    absl::StatusOr<std::string> token = CreateTimeStampToken(
        hash_algorithm, std::string(HashSizeBytes(hash_algorithm), '\0'),
        *clock_, *tsa_signer_, tsa_cert_pem_, tsa_signing_time_);
    ABSL_CHECK_OK(token);
    it->second = token->size() + kSizeMargin;
  }
  return it->second;
}

absl::StatusOr<std::unique_ptr<CryptoWriteHandler>>
CreateLocalKeyCryptoWriteHandler(LocalKeyCryptoWriteHandlerOptions options) {
  if (options.claim_signer.key_pem.empty()) {
    return absl::InvalidArgumentError("Claim signer key is missing.");
  }
  if (options.claim_signer.certificate_chain.empty()) {
    return absl::InvalidArgumentError(
        "Claim signer certificate chain is missing.");
  }
  if (!options.disable_timestamping) {
    if (options.tsa.key_pem.empty()) {
      return absl::InvalidArgumentError("TSA key is missing.");
    }
    if (options.tsa.cert_pem.empty()) {
      return absl::InvalidArgumentError("TSA certificate is missing.");
    }
  } else if (options.tsa.key_pem.empty() != options.tsa.cert_pem.empty()) {
    return absl::InvalidArgumentError(
        "TSA key and certificate must either both be provided or both be "
        "omitted.");
  }

  // COSE requires raw (IEEE P1363) ECDSA signatures, while CMS in RFC 3161
  // timestamps requires DER.
  ABSL_ASSIGN_OR_RETURN(
      std::unique_ptr<PublicKeySign> signer,
      CreateEs256Signer(options.claim_signer.key_pem,
                        EcdsaParameters::SignatureEncoding::kIeeeP1363));
  std::unique_ptr<PublicKeySign> tsa_signer;
  if (!options.tsa.key_pem.empty()) {
    ABSL_ASSIGN_OR_RETURN(
        tsa_signer,
        CreateEs256Signer(options.tsa.key_pem,
                          EcdsaParameters::SignatureEncoding::kDer));
  }
  absl::Clock* clock =
      options.clock != nullptr ? options.clock : &absl::Clock::GetRealClock();
  return std::make_unique<LocalKeyCryptoWriteHandler>(
      clock, std::move(signer), std::move(tsa_signer),
      options.disable_timestamping,
      std::move(options.claim_signer.certificate_chain),
      std::move(options.claim_signer.ocsp_responses), SigningAlgorithm::kEs256,
      std::move(options.tsa.cert_pem));
}

}  // namespace credentio