// Copyright 2023 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//
// Reference implementation utilities of OEMCrypto APIs
//
#include "signed_csr_payload_validator.h"

#include <vector>

#include "string_conversions.h"

namespace wvoec {
namespace util {
namespace {
enum CoseKeyAlgorithm : int {
  AES_GCM_256 = 3,
  HMAC_256 = 5,
  ES256 = -7,  // ECDSA with SHA-256
  EDDSA = -8,
  ECDH_ES_HKDF_256 = -25,
  ES384 = -35,  // ECDSA with SHA-384
};

void AddMessages(std::stringstream& ss,
                 const std::vector<std::string>& fmt_msgs, int indent) {
  const std::string spaces = std::string(indent * 2, ' ');
  for (auto& msg : fmt_msgs) {
    ss << spaces << msg << "\n";
  }
}
}  // namespace

// clang-format off
// Signed CSR payload CBOR data to be verified:
// SignedData<CsrPayload> = [
//     protected: bstr .cbor { 1 : AlgorithmEdDSA / AlgorithmES256 /
//                                 AlgorithmES384 },
//     unprotected: {},
//     payload: bstr .cbor DataToBeSigned / nil,
//     signature: bstr      ; PureEd25519(priv_key, Sig_structure<DataToBeSigned>) /
//                          ; ECDSA(priv_key, Sig_structure<DataToBeSigned>)
// ]
//
// DataToBeSigned = [
//     challenge: bstr .size (0..64),
//     bstr .cbor CsrPayload,
// ]
//
// CsrPayload = [         ; CBOR Array defining the payload for CSR.
//     version: 3,        ; The CsrPayload CDDL Schema version.
//     CertificateType: "widevine"  ; The type of certificate being requested.
//     DeviceInfo,        ; Defined in Android DeviceInfo.aidl
//     KeysToSign: []     ; Empty list
// ]
//
// Sig_structure<DataToBeSigned> = [
//    context: "Signature1",
//    protected: bstr .cbor { 1 : AlgorithmEdDSA / AlgorithmES256 / AlgorithmES384 },
//    external_aad: bstr .size 0,
//    payload: bstr .cbor DataToBeSigned / nil,
// ]
// clang-format on

// Caller ensures the pointer is not NULL.
CborMessageStatus SignedCsrPayloadValidator::ValidateProtectedParams(
    const cppbor::Bstr* protected_params) {
  auto parse_result = cppbor::parse(protected_params);
  std::unique_ptr<cppbor::Item> parsed_protected_params =
      std::move(std::get<0>(parse_result));
  const std::string error_message = std::move(std::get<2>(parse_result));
  if (!parsed_protected_params || !error_message.empty()) {
    AddValidationMessage(kCborValidateFatal,
                         "Unable to parse protectedParams: " + error_message);
    return message_status_;
  }
  // |parsed_protected_params| is a CBOR map of 1 entry
  const cppbor::Map* protected_params_map = parsed_protected_params->asMap();
  if (!protected_params_map) {
    AddValidationMessage(
        kCborValidateFatal,
        "ProtectedParams must be a CBOR map. Actual type: " +
            CppborMajorTypeToString(parsed_protected_params->type()));
    return message_status_;
  }
  if (protected_params_map->size() == 0) {
    AddValidationMessage(kCborValidateFatal, "ProtectedParams is empty.");
    return message_status_;
  }
  if (protected_params_map->size() > 1) {
    AddValidationMessage(kCborValidateWarning,
                         "ProtectedParams expects 1 entry, but got " +
                             std::to_string(protected_params_map->size()));
  }
  // TODO(b/314141962): Replace this with the map lookup function in cppbor
  // library
  bool algo_found = false;
  for (auto const& entry : *protected_params_map) {
    if (!entry.first->asInt()) {
      AddValidationMessage(kCborValidateWarning,
                           "Unsupported key type: " +
                               CppborMajorTypeToString(entry.first->type()));
    } else if (entry.first->asInt()->value() != 1) {
      AddValidationMessage(kCborValidateWarning,
                           "Unsupported key value: " +
                               std::to_string(entry.first->asInt()->value()));
    } else {
      auto& algorithm = entry.second;
      if (!algorithm) {
        AddValidationMessage(kCborValidateFatal, "Algorithm value is missing.");
        return message_status_;
      }
      if (!algorithm->asInt()) {
        AddValidationMessage(kCborValidateFatal,
                             "Unsupported signature algorithm value type: " +
                                 CppborMajorTypeToString(algorithm->type()));
        return message_status_;
      }
      std::string kv = "1: ";
      const int64_t algo = algorithm->asInt()->value();
      if (algo == EDDSA)
        kv += "EDDSA";
      else if (algo == ES256)
        kv += "ES256";
      else if (algo == ES384)
        kv += "ES384";
      else {
        kv += std::to_string(algo);
        AddValidationMessage(
            kCborValidateFatal,
            "Unsupported signature algorithm value: " + std::to_string(algo));
        return message_status_;
      }
      AddMessages(msg_ss_, {std::move(kv)}, 1);
      algo_found = true;
    }
  }
  if (!algo_found) {
    AddValidationMessage(kCborValidateFatal,
                         "ProtectedParams has no signature algorithm.");
  }
  return message_status_;
}

// Caller ensures the pointer is not NULL.
CborMessageStatus SignedCsrPayloadValidator::ValidateDataToBeSigned(
    const cppbor::Bstr* data) {
  if (data->value().empty()) {
    AddValidationMessage(kCborValidateFatal, "Payload to be signed is empty.");
    return message_status_;
  }
  auto parse_result = cppbor::parse(data);
  std::unique_ptr<cppbor::Item> parsed_payload_to_be_signed =
      std::move(std::get<0>(parse_result));
  std::string error_message = std::move(std::get<2>(parse_result));
  if (!parsed_payload_to_be_signed || !error_message.empty()) {
    AddValidationMessage(
        kCborValidateFatal,
        "Unable to parse the payload to be signed: " + error_message);
    return message_status_;
  }
  // Verify that |parsed_payload_to_be_signed| is a valid array.
  const cppbor::Array* payload_to_be_signed_array =
      parsed_payload_to_be_signed->asArray();
  if (!payload_to_be_signed_array) {
    AddValidationMessage(
        kCborValidateFatal,
        "Payload to be signed must be a CBOR array. Actual type: " +
            CppborMajorTypeToString(parsed_payload_to_be_signed->type()));
    return message_status_;
  }
  if (payload_to_be_signed_array->size() != 2U) {
    AddValidationMessage(
        kCborValidateFatal,
        "Payload to be signed must contain the challenge and "
        "payload. Actual size: " +
            std::to_string(payload_to_be_signed_array->size()));
    return message_status_;
  }

  // Element 0: challenge.
  std::string msg = "- Challenge: ";
  const cppbor::Bstr* challenge = payload_to_be_signed_array->get(0)->asBstr();
  if (!challenge) {
    AddValidationMessage(kCborValidateFatal,
                         "Challenge must be Bstr. Actual type: " +
                             CppborMajorTypeToString(
                                 payload_to_be_signed_array->get(0)->type()));
    return message_status_;
  }
  if (challenge->value().size() > 64) {
    AddValidationMessage(
        kCborValidateError,
        "Challenge size must be between 0 and 64 bytes inclusive. "
        "However, challenge is " +
            std::to_string(challenge->value().size()) + " bytes long.");
  }
  msg += wvutil::b2a_hex(challenge->value());
  AddMessages(msg_ss_, {std::move(msg)}, 1);

  // Element 1: CsrPayload.
  AddMessages(msg_ss_, {"- CsrPayload:"}, 1);
  const cppbor::Bstr* csr_payload =
      payload_to_be_signed_array->get(1)->asBstr();
  if (!csr_payload) {
    AddValidationMessage(kCborValidateFatal, "CSR payload is missing.");
    return message_status_;
  }
  parse_result = cppbor::parse(csr_payload);
  std::unique_ptr<cppbor::Item> parsed_csr_payload =
      std::move(std::get<0>(parse_result));
  error_message = std::move(std::get<2>(parse_result));
  if (!parsed_csr_payload) {
    AddValidationMessage(kCborValidateFatal,
                         "Unable to parse CSR payload: " + error_message);
    return message_status_;
  }
  if (!parsed_csr_payload->asArray()) {
    AddValidationMessage(
        kCborValidateFatal,
        "CSR payload must be an array. Actual: " +
            CppborMajorTypeToString(parsed_csr_payload->type()));
    return message_status_;
  }
  if (parsed_csr_payload->asArray()->size() != 4U) {
    AddValidationMessage(
        kCborValidateFatal,
        "CSR payload must contain 4 elements. Actual size: " +
            std::to_string(parsed_csr_payload->asArray()->size()));
    return message_status_;
  }
  // |parsed_csr_payload| is an array of 4 elements.
  const cppbor::Uint* version = parsed_csr_payload->asArray()->get(0)->asUint();
  if (!version) {
    AddValidationMessage(kCborValidateFatal,
                         "CSR payload version must be an unsigned integer.");
    return message_status_;
  }
  AddMessages(msg_ss_, {"- version: " + std::to_string(version->value())}, 2);
  if (version->value() != 3U) {
    AddValidationMessage(kCborValidateError,
                         "CSR payload version must be must be equal to 3.");
  }

  const cppbor::Tstr* certificate_type =
      parsed_csr_payload->asArray()->get(1)->asTstr();
  if (!certificate_type) {
    // Certificate type is allowed to be extended by vendor, i.e. we can't
    // enforce its value.
    AddValidationMessage(
        kCborValidateFatal,
        "Certificate type must be Tstr. Actual type: " +
            CppborMajorTypeToString(
                parsed_csr_payload->asArray()->get(1)->type()));
    return message_status_;
  }
  AddMessages(msg_ss_, {"- certificate_type: " + certificate_type->value()}, 2);

  const cppbor::Map* device_info =
      parsed_csr_payload->asArray()->get(2)->asMap();
  if (!device_info) {
    AddValidationMessage(
        kCborValidateFatal,
        "Device info must be a CBOR map. Actual type: " +
            CppborMajorTypeToString(
                parsed_csr_payload->asArray()->get(2)->type()));
    return message_status_;
  }
  AddMessages(msg_ss_, {"- device_info: " + cppbor::prettyPrint(device_info)},
              2);

  const cppbor::Array* keys = parsed_csr_payload->asArray()->get(3)->asArray();
  if (!keys) {
    AddValidationMessage(
        kCborValidateFatal,
        "Keys must be a CBOR array. Actual type: " +
            CppborMajorTypeToString(
                parsed_csr_payload->asArray()->get(3)->type()));
    return message_status_;
  }
  AddMessages(msg_ss_, {"- keys_to_sign: " + cppbor::prettyPrint(keys)}, 2);
  return message_status_;
}

CborMessageStatus SignedCsrPayloadValidator::Validate() {
  if (message_status_ != kCborParseOk) return message_status_;
  const cppbor::Item* parsed_signed_csr_payload =
      std::get<0>(parse_result()).get();
  if (!parsed_signed_csr_payload) {
    AddValidationMessage(kCborValidateFatal, "Signed CSR payload is empty.");
    return message_status_;
  }

  // Verify that |parsed_signed_csr_payload| is a valid array.
  if (!parsed_signed_csr_payload->asArray()) {
    AddValidationMessage(
        kCborValidateFatal,
        "Signed CSR payload must be a CBOR array. Actual type: " +
            CppborMajorTypeToString(parsed_signed_csr_payload->type()));
    return message_status_;
  }
  const cppbor::Array* signed_csr_payload_array =
      parsed_signed_csr_payload->asArray();
  if (signed_csr_payload_array->size() != 4U) {
    AddValidationMessage(kCborValidateFatal,
                         "Invalid CoseSign1 size. Actual: " +
                             std::to_string(signed_csr_payload_array->size()));
    return message_status_;
  }

  // Writes YAML-formatted output to |msg_ss_| during validation.
  msg_ss_.str(std::string());
  msg_ss_ << "---"
          << "\n";

  // Element 0: protected params.
  msg_ss_ << "- Protected params:\n";
  const cppbor::Bstr* protected_params =
      signed_csr_payload_array->get(0)->asBstr();
  CborMessageStatus status = ValidateProtectedParams(protected_params);
  if (status == kCborValidateFatal) return kCborValidateFatal;

  // Element 1: unprotected params map.
  std::string msg = "- Unprotected params: ";
  const cppbor::Map* unprotected_params =
      signed_csr_payload_array->get(1)->asMap();
  if (!unprotected_params) {
    AddValidationMessage(
        kCborValidateFatal,
        "UnprotectedParams must be a CBOR map. Actual type: " +
            CppborMajorTypeToString(signed_csr_payload_array->get(1)->type()));
    return message_status_;
  }
  if (unprotected_params->size() == 0) {
    msg += "<null>";
  } else {
    msg += "non-null map";
    AddValidationMessage(kCborValidateWarning,
                         "UnprotectedParams is expected to be empty.");
  }
  msg_ss_ << msg << "\n";

  // Element 2: payload (DataToBeSigned).
  msg_ss_ << "- Payload:\n";
  const cppbor::Bstr* payload_to_be_signed =
      signed_csr_payload_array->get(2)->asBstr();
  if (!payload_to_be_signed) {
    AddValidationMessage(
        kCborValidateFatal,
        "Payload to be signed must be Bstr. Actual type: " +
            CppborMajorTypeToString(signed_csr_payload_array->get(2)->type()));
    return message_status_;
  }
  status = ValidateDataToBeSigned(payload_to_be_signed);
  if (status == kCborValidateFatal) return kCborValidateFatal;

  // Element 3: signature.
  std::string sig_msg = "- Signature: ";
  const cppbor::Bstr* signature = signed_csr_payload_array->get(3)->asBstr();
  if (!signature) {
    AddValidationMessage(
        kCborValidateFatal,
        "CoseSign1 signature must be Bstr. Actual type: " +
            CppborMajorTypeToString(signed_csr_payload_array->get(3)->type()));
    return message_status_;
  }
  // Skip CoseSign1 signature verification since the validator doesn't have
  // verifying keys.
  if (signature->value().empty()) {
    sig_msg += "<null>";
    AddValidationMessage(kCborValidateError, "CoseSign1 signature is missing.");
  } else {
    sig_msg += wvutil::b2a_hex(signature->value());
  }
  msg_ss_ << sig_msg << "\n";
  msg_ss_ << "...\n";
  if (message_status_ == kCborParseOk) message_status_ = kCborValidateOk;
  return message_status_;
}

std::string SignedCsrPayloadValidator::GetFormattedMessage() const {
  if (message_status_ == kCborUninitialized ||
      message_status_ == kCborParseError) {
    return std::string();
  }
  const cppbor::Item* parsed_item = std::get<0>(parse_result()).get();
  if (parsed_item == nullptr) {
    return "<null>";
  }
  return msg_ss_.str();
}
}  // namespace util
}  // namespace wvoec