// 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 <string>

#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/escaping.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "cbor/cbor.h"
#include "cbor/parse.h"
#include "crypto/algorithms.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "testing/cbor_utils.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;

absl::StatusOr<Sig1Structure> DecodeSig1Structure(absl::string_view cbor) {
  ABSL_ASSIGN_OR_RETURN(auto res, cbor::Parse(cbor));
  ABSL_ASSIGN_OR_RETURN(auto array, res->AsArray());
  Sig1Structure s;

  auto context = array.GetString(0);
  if (!context.ok()) {
    return absl::Status(
        context.status().code(),
        absl::StrCat("could not get `context`: ", context.status().message()));
  }
  s.context = std::string(*context);

  auto body_protected = array.GetByteString(1);
  if (!body_protected.ok()) {
    return absl::Status(body_protected.status().code(),
                        absl::StrCat("could not get `body_protected`: ",
                                     body_protected.status().message()));
  }
  s.body_protected = std::string(*body_protected);

  auto external_aad = array.GetByteString(2);
  if (!external_aad.ok()) {
    return absl::Status(external_aad.status().code(),
                        absl::StrCat("could not get `external_aad`: ",
                                     external_aad.status().message()));
  }
  s.external_aad = std::string(*external_aad);

  auto payload = array.GetByteString(3);
  if (!payload.ok()) {
    return absl::Status(
        payload.status().code(),
        absl::StrCat("could not get `payload`: ", payload.status().message()));
  }
  s.payload = std::string(*payload);

  return s;
}

std::string Base64Unescape(absl::string_view base64) {
  std::string bytes;
  CHECK(absl::Base64Unescape(base64, &bytes));
  return bytes;
}

// Asserts the value is a byte string, and equals to the `hex` argument in hex
// representation.
MATCHER_P(IsBytesEqHexStr, hex, "") {
  auto hex_arg = absl::BytesToHexString(arg);
  *result_listener << absl::StrCat("Hex strings ", hex, " and ", hex_arg,
                                   " are not equal");
  return hex == hex_arg;
}

TEST(SigStructureTest, Sig1StructureTestVectorSuccess) {
  const Sig1Structure payload = {
      .context = "Signature1",
      .body_protected = "body_protected",
      .external_aad = "external_aad",
      .payload = "payload",
  };

  const std::string cbor = EncodeSig1Structure(payload);
  EXPECT_THAT(
      cbor, IsBytesEqHexStr(
                "846a5369676e6174757265314e626f64795f70726f7465637465644c657874"
                "65726e616c5f616164477061796c6f6164"));

  EXPECT_THAT(DecodeSig1Structure(cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderEdDsa) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kEdDsa,
      .certificate_chain = {"abcd"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderEs256) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kEs256,
      .certificate_chain = {"abcd"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}
TEST(SigStructureTest, ProtectedHeaderEs384) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kEs384,
      .certificate_chain = {"abcd"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderEs512) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kEs512,
      .certificate_chain = {"abcd"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderPs256) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kPs256,
      .certificate_chain = {"abcd"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderPs384) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kPs384,
      .certificate_chain = {"abcd"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderPs512) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kPs512,
      .certificate_chain = {"abcd"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderMultipleCerts) {
  const ProtectedHeader payload = {
      .alg = SigningAlgorithm::kPs512,
      .certificate_chain = {"cert1", "cert2"},
  };

  auto cbor = EncodeProtectedHeader(payload);
  ASSERT_TRUE(cbor.ok());
  EXPECT_THAT(DecodeProtectedHeader(*cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, ProtectedHeaderUnsupportedAlgorithm) {
  // hex for {1: -100, 33: "abcd"}
  auto cbor = absl::HexStringToBytes("a2013863214461626364");
  EXPECT_THAT(DecodeProtectedHeader(cbor),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       "unsupported COSE algorithm: -100"));
}

TEST(SigStructureTest, CoseSign1TaggedStructureEncodeNullPayload) {
  const CoseSign1TaggedStructure payload = {
      .protected_header = "abc",
      .signature = "def",
  };
  const std::string cbor = EncodeCoseSign1TaggedStructure(payload);
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(cbor), IsOkAndHolds(payload));
}
TEST(SigStructureTest, CoseSign1TaggedStructureEncodeNonNullPayload) {
  const CoseSign1TaggedStructure payload = {
      .protected_header = "abc",
      .payload = "xyz",
      .signature = "def",
  };
  const std::string cbor = EncodeCoseSign1TaggedStructure(payload);
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(cbor), IsOkAndHolds(payload));
}
TEST(SigStructureTest, CoseSign1TaggedStructureEncodeTstTokens) {
  const CoseSign1TaggedStructure payload = {
      .protected_header = "abc",
      .unprotected_header = {.sig_tst2 =
                                 TstContainer{.tst_tokens = {{.val = "val1"},
                                                             {.val = "val2"}}}},
      .signature = "def",
  };
  const std::string cbor = EncodeCoseSign1TaggedStructure(payload);
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, CoseSign1TaggedStructureEncodeOcspResponses) {
  const CoseSign1TaggedStructure payload = {
      .protected_header = "abc",
      .unprotected_header = {.ocsp_responses = {"ocsp1", "ocsp2"}},
      .signature = "def",
  };
  const std::string cbor = EncodeCoseSign1TaggedStructure(payload);
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, CoseSign1TaggedStructureEncodePad) {
  const CoseSign1TaggedStructure payload = {
      .protected_header = "abc",
      .unprotected_header = {.pad = "padding"},
      .signature = "def",
  };
  const std::string cbor = EncodeCoseSign1TaggedStructure(payload);
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(cbor), IsOkAndHolds(payload));
}

TEST(SigStructureTest, CoseSign1TaggedStructureDecodeSignature) {
  auto protected_headers_b64 =
      R"(cHJvdGVjdGVkX2hlYWRlcg==)";           // "protected_header"
  auto timestamp_cms_b64 = R"(dGltZXN0YW1w)";  // "timestamp"
  auto signature_b64 = R"(c2lnbmF0dXJl)";      // "signature"
  auto signature = cbor::FromJson(absl::Substitute(
      R"json(["b64'$0'", {"sigTst2": {"tstTokens": [{"val": "b64'$1'"}]}}, null, "b64'$2'"])json",
      protected_headers_b64, timestamp_cms_b64, signature_b64));
  // Prepend the COSE_Sign1 tag (Major type 6 addl value 18).
  signature = "\xD2" + signature;
  CoseSign1TaggedStructure expected{
      .protected_header = Base64Unescape(protected_headers_b64),
      .unprotected_header =
          {
              .sig_tst2 =
                  TstContainer{
                      .tst_tokens = {{.val =
                                          Base64Unescape(timestamp_cms_b64)}},
                  },
          },
      .signature = Base64Unescape(signature_b64),
  };
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(signature),
              IsOkAndHolds(expected));
  EXPECT_EQ(EncodeCoseSign1TaggedStructure(expected), signature);
};

TEST(SigStructureTest, DecodeUnprotectedHeaderContainingSigTst) {
  auto protected_headers_b64 =
      R"(cHJvdGVjdGVkX2hlYWRlcg==)";           // "protected_header"
  auto timestamp_cms_b64 = R"(dGltZXN0YW1w)";  // "timestamp"
  auto signature_b64 = R"(c2lnbmF0dXJl)";      // "signature"
  auto signature = cbor::FromJson(absl::Substitute(
      R"json(["b64'$0'", {"sigTst": {"tstTokens": [{"val": "b64'$1'"}]}}, null, "b64'$2'"])json",
      protected_headers_b64, timestamp_cms_b64, signature_b64));
  // Prepend the COSE_Sign1 tag (Major type 6 addl value 18).
  signature = "\xD2" + signature;
  CoseSign1TaggedStructure expected{
      .protected_header = Base64Unescape(protected_headers_b64),
      .unprotected_header =
          {
              .sig_tst =
                  TstContainer{
                      .tst_tokens = {{.val =
                                          Base64Unescape(timestamp_cms_b64)}},
                  },
          },
      .signature = Base64Unescape(signature_b64),
  };
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(signature),
              IsOkAndHolds(expected));
}

TEST(SigStructureTest, DecodeUnprotectedHeaderContainingCertificateChain) {
  auto protected_headers_b64 =
      R"(cHJvdGVjdGVkX2hlYWRlcg==)";            // "protected_header"
  auto timestamp_cms_b64 = R"(dGltZXN0YW1w)";   // "timestamp"
  auto signature_b64 = R"(c2lnbmF0dXJl)";       // "signature"
  auto certificate_64 = R"(Y2VydGlmaWNhdGU=)";  // "certificate"
  auto signature = cbor::FromJson(absl::Substitute(
      R"json(["b64'$0'", {"sigTst": {"tstTokens": [{"val": "b64'$1'"}]}, "x5chain": ["b64'$2'"]}, null, "b64'$3'"])json",
      protected_headers_b64, timestamp_cms_b64, certificate_64, signature_b64));
  // Prepend the COSE_Sign1 tag (Major type 6 addl value 18).
  signature = "\xD2" + signature;
  CoseSign1TaggedStructure expected{
      .protected_header = Base64Unescape(protected_headers_b64),
      .unprotected_header =
          {
              .sig_tst =
                  TstContainer{
                      .tst_tokens = {{.val =
                                          Base64Unescape(timestamp_cms_b64)}},
                  },
              .certificate_chain = {Base64Unescape(certificate_64)},
          },
      .signature = Base64Unescape(signature_b64),
  };
  EXPECT_THAT(DecodeCoseSign1TaggedStructure(signature),
              IsOkAndHolds(expected));
}

}  // namespace
}  // namespace credentio