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

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

#include "absl/log/absl_log.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/escaping.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "cbor/cbor.h"
#include "cbor/parse.h"
#include "cppbor/cppbor.h"
#include "crypto/algorithms.h"

namespace credentio {
namespace {
// 'alg' in https://datatracker.ietf.org/doc/html/rfc8152#section-3.1
constexpr int kSigningAlgorithmKey = 1;
// See
// https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html#_obtaining_the_time_stamp
// for more details.
constexpr absl::string_view kTimestampV2HeaderKey = "sigTst2";
// DEPRECATED. DO NOT GENERATE. Contents of 'sigTst' CBOR map.
constexpr absl::string_view kTimestampV1HeaderKey = "sigTst";
// See
// https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html#_create_a_temporary_claim_and_signature
// for more details about padding.
constexpr absl::string_view kPadHeaderKey = "pad";
// See
// https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html#_going_back_and_filling_in
// about using `pad2`.
constexpr absl::string_view kPad2HeaderKey = "pad2";
// https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html#x509_certificates
constexpr int kCertificateChainKey = 33;
constexpr absl::string_view kCertificateChainDeprecatedKey = "x5chain";
constexpr absl::string_view kOcspValsKey = "ocspVals";
constexpr absl::string_view kRvalsKey = "rVals";

// https://www.iana.org/assignments/cose/cose.xhtml#algorithms
absl::StatusOr<int64_t> GetAlgorithm(SigningAlgorithm e) {
  switch (e) {
    case SigningAlgorithm::kEs256:
      return -7;
    case SigningAlgorithm::kEdDsa:
      return -8;
    case SigningAlgorithm::kEs384:
      return -35;
    case SigningAlgorithm::kEs512:
      return -36;
    case SigningAlgorithm::kPs256:
      return -37;
    case SigningAlgorithm::kPs384:
      return -38;
    case SigningAlgorithm::kPs512:
      return -39;
  }
  return absl::InvalidArgumentError("Unsupported algorithm");
}

absl::StatusOr<SigningAlgorithm> GetSigningAlgorithm(cbor::MapView headers) {
  auto alg = headers.GetInt64(kSigningAlgorithmKey);
  if (!alg.ok()) {
    return absl::Status(
        alg.status().code(),
        absl::StrCat("no COSE algorithm specified in protected headers: ",
                     alg.status().message()));
  }
  // https://www.iana.org/assignments/cose/cose.xhtml#algorithms
  ABSL_DVLOG(1) << ".cose_alg = " << *alg << ",";
  switch (*alg) {
    case -7:
      return SigningAlgorithm::kEs256;
    case -35:
      return SigningAlgorithm::kEs384;
    case -36:
      return SigningAlgorithm::kEs512;
    case -37:
      return SigningAlgorithm::kPs256;
    case -38:
      return SigningAlgorithm::kPs384;
    case -39:
      return SigningAlgorithm::kPs512;
    case -8:
      return SigningAlgorithm::kEdDsa;
  }
  return absl::UnimplementedError(
      absl::StrCat("unsupported COSE algorithm: ", *alg));
}

absl::StatusOr<std::vector<std::string>> GetCertificateChain(
    cbor::MapView headers) {
  // Check if there is a single certificate in the chain. In this case, label 33
  // of the header map (or "x5chain" key) is a bstr.
  // Claim generators should only use 33 (integer) key but the validators should
  // both accept 33 and "x5chain" keys. If both 33 (integer) and "x5chain"
  // (string) keys are present, the label 33 should be used.
  auto certificate_chain = headers.GetByteString(kCertificateChainKey);
  if (!certificate_chain.ok()) {
    certificate_chain = headers.GetByteString(kCertificateChainDeprecatedKey);
  }
  if (certificate_chain.ok()) {
    return std::vector<std::string>({std::string(*certificate_chain)});
  }
  // Check if there is a list of certificates in the chain. In this case, label
  // 33 of the header map (or "x5chain" key) is an array of bstrs.
  auto certificate_chain_array = headers.GetArray(kCertificateChainKey);
  if (!certificate_chain_array.ok()) {
    certificate_chain_array = headers.GetArray(kCertificateChainDeprecatedKey);
  }
  if (certificate_chain_array.ok()) {
    std::vector<std::string> certificate_chain;
    for (int i = 0; i < certificate_chain_array->size(); ++i) {
      auto certificate_chain_element =
          certificate_chain_array->GetByteString(i);
      if (!certificate_chain_element.ok()) {
        return absl::InvalidArgumentError(
            absl::StrCat("Failed to parse the certificate chain element: ",
                         certificate_chain_element.status()));
      }
      ABSL_DVLOG(1) << ".certs_b64[" << i << "] = R\"("
                    << absl::Base64Escape(*certificate_chain_element) << ")\",";
      certificate_chain.push_back(std::string(*certificate_chain_element));
    }
    return certificate_chain;
  }
  return std::vector<std::string>();
}

std::vector<std::string> DecodeOcspResponses(cbor::MapView map) {
  auto rvals = map.GetOptionalMap(kRvalsKey);
  if (!rvals.has_value()) {
    return {};
  }
  auto ocsp_responses = rvals->GetArray(kOcspValsKey);
  if (!ocsp_responses.ok()) {
    return {};
  }
  std::vector<std::string> v;
  v.reserve(ocsp_responses->size());
  for (int i = 0; i < ocsp_responses->size(); ++i) {
    auto ocsp_response = ocsp_responses->GetByteString(i);
    if (ocsp_response.ok()) {
      v.push_back(std::string(*ocsp_response));
    }
  }
  return v;
}

}  // namespace

absl::StatusOr<std::string> EncodeProtectedHeader(
    const ProtectedHeader& protected_header) {
  cppbor::Map m;
  ABSL_ASSIGN_OR_RETURN(auto alg, GetAlgorithm(protected_header.alg));
  m.add(cppbor::Uint(kSigningAlgorithmKey), cppbor::Nint(alg));
  if (protected_header.certificate_chain.size() == 1) {
    m.add(cppbor::Uint(kCertificateChainKey),
          cppbor::Bstr(protected_header.certificate_chain[0]));
  } else {
    cppbor::Array chain;
    for (const auto& cert : protected_header.certificate_chain) {
      chain.add(cppbor::Bstr(cert));
    }
    m.add(cppbor::Uint(kCertificateChainKey), std::move(chain));
  }
  return m.toString();
}

absl::StatusOr<ProtectedHeader> DecodeProtectedHeader(absl::string_view cbor) {
  ProtectedHeader protected_header;
  ABSL_ASSIGN_OR_RETURN(auto res, cbor::Parse(cbor));
  ABSL_ASSIGN_OR_RETURN(auto map, res->AsMap());
  ABSL_ASSIGN_OR_RETURN(auto alg, GetSigningAlgorithm(map));
  protected_header.alg = alg;

  ABSL_ASSIGN_OR_RETURN(auto certificate_chain, GetCertificateChain(map));
  protected_header.certificate_chain = std::move(certificate_chain);
  return protected_header;
}

void EncodeTstToken(const TstToken& tst_token, cppbor::Map& m) {
  m.add(cppbor::Tstr("val"), cppbor::Bstr(tst_token.val));
}

absl::StatusOr<TstToken> DecodeTstToken(cbor::MapView map) {
  ABSL_ASSIGN_OR_RETURN(auto val, map.GetByteString("val"));
  return TstToken{.val = std::string(val)};
}

void EncodeTstContainer(const TstContainer& tst_container, cppbor::Map& m) {
  cppbor::Array tstTokens;
  for (const auto& tstToken : tst_container.tst_tokens) {
    cppbor::Map tstTokenMap;
    EncodeTstToken(tstToken, tstTokenMap);
    tstTokens.add(std::move(tstTokenMap));
  }
  m.add(cppbor::Tstr("tstTokens"), std::move(tstTokens));
}

absl::StatusOr<TstContainer> DecodeTstContainer(cbor::MapView map) {
  ABSL_ASSIGN_OR_RETURN(auto tst_tokens, map.GetArray("tstTokens"));
  TstContainer tst_container;
  for (int i = 0; i < tst_tokens.size(); ++i) {
    ABSL_ASSIGN_OR_RETURN(auto tst_token, tst_tokens.GetMap(i));
    ABSL_ASSIGN_OR_RETURN(auto tst_token_decoded, DecodeTstToken(tst_token));
    tst_container.tst_tokens.push_back(std::move(tst_token_decoded));
  }
  return tst_container;
}

void EncodeUnprotectedHeader(const UnprotectedHeader& unprotected_header,
                             cppbor::Map& m) {
  if (unprotected_header.sig_tst2.has_value()) {
    cppbor::Map tstContainer;
    EncodeTstContainer(*unprotected_header.sig_tst2, tstContainer);
    m.add(cppbor::Tstr(kTimestampV2HeaderKey), std::move(tstContainer));
  }
  if (!unprotected_header.ocsp_responses.empty()) {
    cppbor::Array ocsp_responses;
    for (const auto& ocsp_response : unprotected_header.ocsp_responses) {
      ocsp_responses.add(cppbor::Bstr(ocsp_response));
    }
    cppbor::Map rvals;
    rvals.add(cppbor::Tstr(kOcspValsKey), std::move(ocsp_responses));
    m.add(cppbor::Tstr(kRvalsKey), std::move(rvals));
  }
  if (unprotected_header.pad.has_value()) {
    m.add(cppbor::Tstr(kPadHeaderKey), cppbor::Bstr(*unprotected_header.pad));
  }
  if (unprotected_header.pad2.has_value()) {
    m.add(cppbor::Tstr(kPad2HeaderKey), cppbor::Bstr(*unprotected_header.pad2));
  }
}

absl::StatusOr<UnprotectedHeader> DecodeUnprotectedHeader(cbor::MapView map) {
  UnprotectedHeader unprotected_header;
  auto sig_tst2 = map.GetOptionalMap(kTimestampV2HeaderKey);
  if (sig_tst2.has_value()) {
    ABSL_ASSIGN_OR_RETURN(unprotected_header.sig_tst2,
                          DecodeTstContainer(*sig_tst2));
  }
  auto sig_tst = map.GetOptionalMap(kTimestampV1HeaderKey);
  if (sig_tst.has_value()) {
    ABSL_ASSIGN_OR_RETURN(unprotected_header.sig_tst,
                          DecodeTstContainer(*sig_tst));
  }
  unprotected_header.ocsp_responses = DecodeOcspResponses(map);
  auto pad = map.GetOptionalByteString(kPadHeaderKey);
  if (pad.has_value()) {
    unprotected_header.pad = std::string(*pad);
  }
  auto pad2 = map.GetOptionalByteString(kPad2HeaderKey);
  if (pad2.has_value()) {
    unprotected_header.pad2 = std::string(*pad2);
  }
  auto certificate_chain = GetCertificateChain(map);
  if (certificate_chain.ok()) {
    unprotected_header.certificate_chain = *std::move(certificate_chain);
  }
  return unprotected_header;
}

std::string EncodeSig1Structure(const Sig1Structure& sig_structure) {
  cppbor::Array s;
  s.add(cppbor::Tstr(sig_structure.context));
  s.add(cppbor::Bstr(sig_structure.body_protected));
  s.add(cppbor::Bstr(sig_structure.external_aad));
  s.add(cppbor::Bstr(sig_structure.payload));
  return s.toString();
}

std::string EncodeCoseSign1TaggedStructure(const CoseSign1TaggedStructure& s) {
  cppbor::Array array;
  // Add protected_header.
  array.add(cppbor::Bstr(s.protected_header));
  // Add unprotected_header.
  cppbor::Map unprotected_header_map;
  EncodeUnprotectedHeader(s.unprotected_header, unprotected_header_map);
  array.add(std::move(unprotected_header_map));
  // Add payload.
  if (s.payload.has_value()) {
    array.add(cppbor::Bstr(*s.payload));
  } else {
    array.add(cppbor::Null());
  }
  // Add signature.
  array.add(cppbor::Bstr(s.signature));
  // tag 18 see https://datatracker.ietf.org/doc/html/rfc8152#section-4.2
  cppbor::SemanticTag tagged_cbor(18, std::move(array));
  return tagged_cbor.toString();
}

absl::StatusOr<CoseSign1TaggedStructure> DecodeCoseSign1TaggedStructure(
    absl::string_view cbor) {
  // COSE_Sign1 structure is a CBOR array of size 4.
  // CDDL schema for the case of a single signer is:
  //
  // COSE_Sign1 = [
  //     protected: bstr,
  //     unprotected: map,
  //     payload: bstr / nil,
  //     signature: bstr
  // ]
  //
  // See https://datatracker.ietf.org/doc/html/rfc9052 for more
  // details.
  ABSL_ASSIGN_OR_RETURN(auto res, cbor::Parse(cbor));
  ABSL_ASSIGN_OR_RETURN(auto array, res->AsArray());
  CoseSign1TaggedStructure s;
  ABSL_ASSIGN_OR_RETURN(s.protected_header, array.GetByteString(0));
  ABSL_ASSIGN_OR_RETURN(auto unprotect_header_cbor, array.GetMap(1));
  ABSL_ASSIGN_OR_RETURN(s.unprotected_header,
                        DecodeUnprotectedHeader(unprotect_header_cbor));
  ABSL_ASSIGN_OR_RETURN(auto is_payload_null, array.IsNull(2));
  if (!is_payload_null) {
    ABSL_ASSIGN_OR_RETURN(auto payload, array.GetByteString(2));
    s.payload.emplace(payload);
  }

  // In COSE standard, there can also be an array of signatures instead of a
  // single signature, if there are multiple signers. However, in C2PA we will
  // always have a single signer.
  ABSL_ASSIGN_OR_RETURN(auto signature, array.GetByteString(3));
  s.signature = std::string(signature);
  return s;
}

std::string CoseSign1TaggedStructure::signature_bstr() const {
  return cppbor::Bstr(signature).toString();
}

}  // namespace credentio