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

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

#include "absl/functional/function_ref.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/default/hasher.h"
#include "crypto/hash.h"
#include "openssl/asn1.h"
#include "openssl/base.h"
#include "openssl/bytestring.h"
#include "openssl/mem.h"
#include "openssl/obj.h"
#include "openssl/x509.h"

namespace credentio_testing {

namespace {

constexpr absl::string_view kOidMessageDigest = "1.2.840.113549.1.9.4";
constexpr absl::string_view kOidContentType = "1.2.840.113549.1.9.3";
constexpr absl::string_view kOidTstInfo = "1.2.840.113549.1.9.16.1.4";
constexpr absl::string_view kOidData = "1.2.840.113549.1.7.1";
constexpr absl::string_view kOidSigningTime = "1.2.840.113549.1.9.5";
constexpr absl::string_view kOidSmimeCaps = "1.2.840.113549.1.9.15";
constexpr absl::string_view kOidAes256Cbc = "2.16.840.1.101.3.4.1.42";
constexpr absl::string_view kOidAes192Cbc = "2.16.840.1.101.3.4.1.22";
constexpr absl::string_view kOidAes128Cbc = "2.16.840.1.101.3.4.1.2";
constexpr absl::string_view kOidDesEde3Cbc = "1.2.840.113549.3.7";
constexpr absl::string_view kOidRsaesOaep = "1.2.840.113549.1.1.7";
constexpr absl::string_view kOidSha256 = "2.16.840.1.101.3.4.2.1";
constexpr absl::string_view kOidSignedData = "1.2.840.113549.1.7.2";
constexpr absl::string_view kOidRsaEncryption = "1.2.840.113549.1.1.1";

bool SerializeCertificate(const X509& certificate, std::vector<uint8_t>* data) {
  int len = i2d_X509(&certificate, nullptr);
  if (len <= 0) {
    return false;
  }
  data->resize(len);
  uint8_t* dataptr = data->data();
  return i2d_X509(&certificate, &dataptr) == len;
}

bool SerializeCertificateIssuerName(const X509& certificate,
                                    std::vector<uint8_t>* data) {
  int len = i2d_X509_NAME(X509_get_issuer_name(&certificate), nullptr);
  if (len <= 0) {
    return false;
  }
  data->resize(len);
  uint8_t* dataptr = data->data();
  return i2d_X509_NAME(X509_get_issuer_name(&certificate), &dataptr) == len;
}

bool SerializeCertificateSerialNumber(const X509& certificate,
                                      std::vector<uint8_t>* data) {
  const ASN1_INTEGER* serial_asn1 = X509_get0_serialNumber(&certificate);
  if (serial_asn1 == nullptr) {
    return false;
  }
  int len = i2d_ASN1_INTEGER(serial_asn1, nullptr);
  if (len <= 0) {
    return false;
  }
  data->resize(len);
  uint8_t* dataptr = data->data();
  return i2d_ASN1_INTEGER(serial_asn1, &dataptr) == len;
}

bool SerializeSigningTime(time_t signing_time, std::vector<uint8_t>* data) {
  bssl::UniquePtr<ASN1_TIME> asn1_time(ASN1_TIME_set(nullptr, signing_time));
  if (asn1_time == nullptr) {
    return false;
  }
  int len = i2d_ASN1_TIME(asn1_time.get(), nullptr);
  if (len <= 0) {
    return false;
  }
  data->resize(len);
  uint8_t* dataptr = data->data();
  return i2d_ASN1_TIME(asn1_time.get(), &dataptr) == len;
}

bool AddOid(CBB* cbb, absl::string_view oid) {
  return CBB_add_asn1_oid_from_text(cbb, oid.data(), oid.length());
}

}  // namespace

absl::StatusOr<std::string> CreateSignedCms(
    absl::Span<const absl::string_view> contents,
    const X509& signer_certificate,
    absl::FunctionRef<absl::StatusOr<std::vector<uint8_t>>(absl::string_view)>
        sign_callback,
    const SignedCmsOptions& options) {
  ABSL_ASSIGN_OR_RETURN(auto hasher,
                        CreateHasher(credentio::HashAlgorithm::kSha256));
  for (const auto& part : contents) {
    hasher->Update(part);
  }
  std::string hash = hasher->Digest();

  std::vector<uint8_t> issuer_name, serial_number;
  if (!SerializeCertificateIssuerName(signer_certificate, &issuer_name) ||
      !SerializeCertificateSerialNumber(signer_certificate, &serial_number)) {
    return absl::InternalError("Failed to serialize certificate components.");
  }

  bssl::ScopedCBB attributes;
  CBB att_seq, oid, value_set, value;
  if (CBB_init(attributes.get(), 1024) != 1) {
    return absl::InternalError("Memory allocation failed for attributes.");
  }

  // Message Digest
  if (CBB_add_asn1(attributes.get(), &att_seq, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&att_seq, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidMessageDigest) ||
      CBB_add_asn1(&att_seq, &value_set, CBS_ASN1_SET) != 1 ||
      CBB_add_asn1(&value_set, &value, CBS_ASN1_OCTETSTRING) != 1 ||
      CBB_add_bytes(&value, reinterpret_cast<const uint8_t*>(hash.data()),
                    hash.size()) != 1) {
    return absl::InternalError("Failed to add message digest attribute.");
  }

  // Content Type
  if (CBB_add_asn1(attributes.get(), &att_seq, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&att_seq, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidContentType) ||
      CBB_add_asn1(&att_seq, &value_set, CBS_ASN1_SET) != 1 ||
      CBB_add_asn1(&value_set, &value, CBS_ASN1_OBJECT) != 1) {
    return absl::InternalError("Failed to add content type attribute.");
  }
  if (options.is_time_stamp_token) {
    if (!AddOid(&value, kOidTstInfo)) {
      return absl::InternalError("Failed to add TSTInfo OID.");
    }
  } else {
    if (!AddOid(&value, kOidData)) {
      return absl::InternalError("Failed to add data OID.");
    }
  }

  // Signing Time
  if (options.signing_time.has_value()) {
    std::vector<uint8_t> signing_time_der;
    if (!SerializeSigningTime(absl::ToUnixSeconds(*options.signing_time),
                              &signing_time_der) ||
        CBB_add_asn1(attributes.get(), &att_seq, CBS_ASN1_SEQUENCE) != 1 ||
        CBB_add_asn1(&att_seq, &oid, CBS_ASN1_OBJECT) != 1 ||
        !AddOid(&oid, kOidSigningTime) ||
        CBB_add_asn1(&att_seq, &value_set, CBS_ASN1_SET) != 1 ||
        CBB_add_bytes(&value_set, signing_time_der.data(),
                      signing_time_der.size()) != 1) {
      return absl::InternalError("Failed to add signing time attribute.");
    }
  }

  // SMIME-CAPS
  CBB cap_seq, cap_item;
  if (CBB_add_asn1(attributes.get(), &att_seq, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&att_seq, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidSmimeCaps) ||
      CBB_add_asn1(&att_seq, &value_set, CBS_ASN1_SET) != 1 ||
      CBB_add_asn1(&value_set, &cap_seq, CBS_ASN1_SEQUENCE) != 1) {
    return absl::InternalError("Failed to add SMIME-CAPS attribute.");
  }

  // AES-256-CBC
  if (CBB_add_asn1(&cap_seq, &cap_item, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&cap_item, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidAes256Cbc)) {
    return absl::InternalError("Failed to add AES-256-CBC capability.");
  }
  // AES-192-CBC
  if (CBB_add_asn1(&cap_seq, &cap_item, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&cap_item, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidAes192Cbc)) {
    return absl::InternalError("Failed to add AES-192-CBC capability.");
  }
  // AES-128-CBC
  if (CBB_add_asn1(&cap_seq, &cap_item, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&cap_item, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidAes128Cbc)) {
    return absl::InternalError("Failed to add AES-128-CBC capability.");
  }
  // DES-EDE3-CBC
  if (CBB_add_asn1(&cap_seq, &cap_item, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&cap_item, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidDesEde3Cbc)) {
    return absl::InternalError("Failed to add DES-EDE3-CBC capability.");
  }
  // RSAES-OAEP
  if (CBB_add_asn1(&cap_seq, &cap_item, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&cap_item, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidRsaesOaep)) {
    return absl::InternalError("Failed to add RSAES-OAEP capability.");
  }
  // SHA256
  if (CBB_add_asn1(&cap_seq, &cap_item, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&cap_item, &oid, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&oid, kOidSha256)) {
    return absl::InternalError("Failed to add SHA256 capability.");
  }

  // Extra Signed Attributes
  if (!options.extra_signed_attributes.empty()) {
    if (CBB_add_bytes(attributes.get(),
                      reinterpret_cast<const uint8_t*>(
                          options.extra_signed_attributes.data()),
                      options.extra_signed_attributes.length()) != 1) {
      return absl::InternalError("Failed to add extra signed attributes.");
    }
  }

  // Sort signed attributes into canonical DER SET OF order (RFC 5652 §5.4).
  uint8_t* raw_attr_data;
  size_t raw_attr_len;
  if (CBB_flush_asn1_set_of(attributes.get()) != 1 ||
      CBB_finish(attributes.get(), &raw_attr_data, &raw_attr_len) != 1) {
    return absl::InternalError("Failed to finish attributes CBB.");
  }
  bssl::UniquePtr<uint8_t> raw_attr_data_uniq(raw_attr_data);

  // Sign the attributes. The signature is computed over the SET of attributes.
  bssl::ScopedCBB pk_signed_attributes;
  CBB pk_attributes_set;
  if (CBB_init(pk_signed_attributes.get(), raw_attr_len + 5) != 1 ||
      CBB_add_asn1(pk_signed_attributes.get(), &pk_attributes_set,
                   CBS_ASN1_SET) != 1 ||
      CBB_add_bytes(&pk_attributes_set, raw_attr_data, raw_attr_len) != 1 ||
      CBB_flush(pk_signed_attributes.get()) != 1) {
    return absl::InternalError("Failed to prepare attributes for signing.");
  }

  ABSL_ASSIGN_OR_RETURN(
      std::vector<uint8_t> signature,
      sign_callback(absl::string_view(
          reinterpret_cast<const char*>(CBB_data(pk_signed_attributes.get())),
          CBB_len(pk_signed_attributes.get()))));

  // Assemble CMS
  bssl::ScopedCBB cms;
  CBB content_type, signed_data, digest_algs, digest_alg_seq, content_seq,
      certs, algo_params, temp, signed_seq;
  if (CBB_init(cms.get(), 2048) != 1) {
    return absl::InternalError("Memory allocation for CMS failed.");
  }

  uint64_t version = options.is_time_stamp_token ? 3 : 1;

  if (CBB_add_asn1(cms.get(), &content_type, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&content_type, &temp, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&temp, kOidSignedData) ||
      CBB_add_asn1(&content_type, &signed_data,
                   CBS_ASN1_CONSTRUCTED | CBS_ASN1_CONTEXT_SPECIFIC) != 1 ||
      CBB_add_asn1(&signed_data, &signed_seq, CBS_ASN1_SEQUENCE) != 1 ||
      // Version
      CBB_add_asn1_uint64(&signed_seq, version) != 1 ||
      // Digest algorithms
      CBB_add_asn1(&signed_seq, &digest_algs, CBS_ASN1_SET) != 1 ||
      CBB_add_asn1(&digest_algs, &digest_alg_seq, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&digest_alg_seq, &temp, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&temp, kOidSha256) ||
      CBB_add_asn1(&digest_alg_seq, &algo_params, CBS_ASN1_NULL) != 1 ||
      // Content
      CBB_add_asn1(&signed_seq, &content_seq, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&content_seq, &temp, CBS_ASN1_OBJECT) != 1) {
    return absl::InternalError("Failed to build CMS header.");
  }

  if (options.is_time_stamp_token) {
    if (!AddOid(&temp, kOidTstInfo)) {
      return absl::InternalError("Failed to add TSTInfo OID to header.");
    }
  } else {
    if (!AddOid(&temp, kOidData)) {
      return absl::InternalError("Failed to add data OID to header.");
    }
  }

  // Opaque Content
  CBB octet_stream;
  if (CBB_add_asn1(&content_seq, &octet_stream,
                   CBS_ASN1_CONSTRUCTED | CBS_ASN1_CONTEXT_SPECIFIC) != 1 ||
      CBB_add_asn1(&octet_stream, &temp, CBS_ASN1_OCTETSTRING) != 1) {
    return absl::InternalError("Failed to add opaque content wrapper.");
  }
  for (const auto& part : contents) {
    if (CBB_add_bytes(&temp, reinterpret_cast<const uint8_t*>(part.data()),
                      part.size()) != 1) {
      return absl::InternalError("Failed to add opaque content bytes.");
    }
  }

  // Certificates
  if (options.include_user_certificate || !options.ca_certificates.empty()) {
    if (CBB_add_asn1(&signed_seq, &certs,
                     CBS_ASN1_CONSTRUCTED | CBS_ASN1_CONTEXT_SPECIFIC) != 1) {
      return absl::InternalError("Failed to add certificates wrapper.");
    }
    std::vector<uint8_t> der;
    for (const X509* cert : options.ca_certificates) {
      if (!SerializeCertificate(*cert, &der)) {
        return absl::InternalError("Failed to serialize CA certificate.");
      }
      if (CBB_add_bytes(&certs, der.data(), der.size()) != 1) {
        return absl::InternalError("Failed to add CA certificate.");
      }
    }
    if (options.include_user_certificate) {
      if (!SerializeCertificate(signer_certificate, &der)) {
        return absl::InternalError("Failed to serialize user certificate.");
      }
      if (CBB_add_bytes(&certs, der.data(), der.size()) != 1) {
        return absl::InternalError("Failed to add user certificate.");
      }
    }
  }

  // Signer Info
  CBB signer_info_set, signer_info, issuer_and_sn;
  if (CBB_add_asn1(&signed_seq, &signer_info_set, CBS_ASN1_SET) != 1 ||
      CBB_add_asn1(&signer_info_set, &signer_info, CBS_ASN1_SEQUENCE) != 1 ||
      // Version
      CBB_add_asn1_uint64(&signer_info, 1) != 1 ||
      CBB_add_asn1(&signer_info, &issuer_and_sn, CBS_ASN1_SEQUENCE) != 1 ||
      // Signer info: Signer identifier: Issuer and serial number
      CBB_add_bytes(&issuer_and_sn, issuer_name.data(), issuer_name.size()) !=
          1 ||
      CBB_add_bytes(&issuer_and_sn, serial_number.data(),
                    serial_number.size()) != 1 ||
      // Signer info: Digest algorithm
      CBB_add_asn1(&signer_info, &digest_alg_seq, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&digest_alg_seq, &temp, CBS_ASN1_OBJECT) != 1 ||
      !AddOid(&temp, kOidSha256) ||
      CBB_add_asn1(&digest_alg_seq, &algo_params, CBS_ASN1_NULL) != 1 ||
      // Signer info: Signed attributes
      CBB_add_asn1(&signer_info, &temp,
                   CBS_ASN1_CONSTRUCTED | CBS_ASN1_CONTEXT_SPECIFIC) != 1 ||
      CBB_add_bytes(&temp, raw_attr_data, raw_attr_len) != 1 ||
      // Signer info: Signature algorithm
      CBB_add_asn1(&signer_info, &digest_alg_seq, CBS_ASN1_SEQUENCE) != 1 ||
      CBB_add_asn1(&digest_alg_seq, &temp, CBS_ASN1_OBJECT) != 1) {
    return absl::InternalError("Failed to build SignerInfo.");
  }

  // Hardcode RsaEncryption for now to match simple_cms behavior in tests, as
  // TimestampVerifier doesn't strictly check this.
  if (!AddOid(&temp, kOidRsaEncryption) ||
      CBB_add_asn1(&digest_alg_seq, &algo_params, CBS_ASN1_NULL) != 1 ||
      // Signer info: Signature
      CBB_add_asn1(&signer_info, &temp, CBS_ASN1_OCTETSTRING) != 1 ||
      CBB_add_bytes(&temp, signature.data(), signature.size()) != 1) {
    return absl::InternalError("Failed to add signature to SignerInfo.");
  }

  uint8_t* cms_data;
  size_t cms_len;
  if (CBB_finish(cms.get(), &cms_data, &cms_len) != 1) {
    return absl::InternalError("Failed to finish CMS CBB.");
  }
  bssl::UniquePtr<uint8_t> cms_data_uniq(cms_data);

  return std::string(reinterpret_cast<char*>(cms_data), cms_len);
}

}  // namespace credentio_testing