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

#include <stddef.h>

#include <cstdint>
#include <iomanip>
#include <string>
#include <unordered_set>
#include <vector>

#include <openssl/evp.h>

#include "oemcrypto_ecc_key.h"
#include "string_conversions.h"

namespace wvoec {
namespace util {
namespace {
// The BCC is encoded using RFC 8949- Concise Binary Object Representation
// (CBOR).

// The full definition of the following enums can be found here:
// go/remote-provisioning-hal#bcc.

// The device key is encoded in a cbor map. The key values are a mix of
// positive and negative integer values.
enum {
  MAP_KEY_DEVICE_KEY_TYPE = 1,
  MAP_KEY_DEVICE_KEY_ALGORITHM = 3,
  MAP_KEY_DEVICE_KEY_OPS = 4,
  MAP_KEY_DEVICE_KEY_CURVE = -1,
  MAP_KEY_DEVICE_KEY_BYTES_0 = -2,
  MAP_KEY_DEVICE_KEY_BYTES_1 = -3,
};

// The device key may be encoded in the BCC as either X,Y elliptic curve
// coordinates, or as raw bytes. The value is identified using
// MAP_KEY_DEVICE_KEY_TYPE.
enum {
  DEVICE_KEY_ENCODING_UNKNOWN = 0,
  DEVICE_KEY_BYTE_STRING = 1,
  DEVICE_KEY_OCTET_PAIR = 2,
};

// Android/Widevine Dice Attestation allows two signing models. This is
// identified using MAP_KEY_DEVICE_KEY_ALGORITHM.
enum {
  DEVICE_KEY_ALGORITHM_ES256 = -7,   // EC key with SHA-256
  DEVICE_KEY_ALGORITHM_EDDSA = -8,   // Pure ED25519.
  DEVICE_KEY_ALGORITHM_ES384 = -35,  // EC key with SHA-384
};

// The curve used to generate the device public key is identified using the
// MAP_KEY_DEVICE_KEY_CURVE.
enum {
  DEVICE_KEY_CURVE_P256 = 1,
  DEVICE_KEY_CURVE_P384 = 2,
  DEVICE_KEY_CURVE_ED25519 = 6,
};

// Sized to hold a component of a P256 public key uncompressed point compatible
// with X9.62. The key is formatted in an Z/X/Y format in which Z == 0x04 and X
// and Y are the public key coordinates. X and Y are each 32 bytes.
constexpr int kP256KeyComponentSize = 256 / 8;
// Sized to hold a P384 public key uncompressed point compatible with X9.62.
// The key is formatted in an Z/X/Y format in which Z == 0x04 and X and Y are
// the public key coordinates. X and Y are each 48 bytes.
constexpr int kP384KeyComponentSize = 384 / 8;
constexpr int kMarshaledP256KeySize = kP256KeyComponentSize * 2 + 1;
constexpr int kMarshaledP384KeySize = kP384KeyComponentSize * 2 + 1;
constexpr char kMarshaledECKeyZValue = 0x04;
constexpr int kED25519KeyDataItemSize = 32;
// The Issuer field key in BccEntryPayload.
constexpr int64_t kIssuer = 1;
// The Subject field key in BccEntryPayload.
constexpr int64_t kSubject = 2;
// The SubjectPublicKey field key in BccEntryPayload.
constexpr int64_t kSubjectPublicKey = -4670552;
// This signature context is defined by COSE SIGN1.
constexpr char kSignatureContextString[] = "Signature1";

struct IssuerSubject {
  std::string issuer;
  std::string subject;
  bool IsValid() const { return !issuer.empty() && !subject.empty(); }
  void PrintTo(std::vector<std::string>& fmt_msgs) const {
    fmt_msgs.push_back("Issuer: ");
    fmt_msgs.back().append(issuer.empty() ? "<missing>" : issuer);
    fmt_msgs.push_back("Subject: ");
    fmt_msgs.back().append(subject.empty() ? "<missing>" : subject);
  }
};

IssuerSubject GetIssuerSubjectFromBccEntryPayload(
    const cppbor::Map* bcc_entry_payload) {
  IssuerSubject ret;
  for (size_t i = 0; i < bcc_entry_payload->size(); ++i) {
    const auto& entry = (*bcc_entry_payload)[i];
    if (entry.first == nullptr || entry.first->asInt() == nullptr ||
        entry.second == nullptr || entry.second->asTstr() == nullptr) {
      continue;
    }
    const auto& value = entry.second->asTstr()->value();
    if (entry.first->asInt()->value() == kIssuer) {
      ret.issuer = value.empty() ? "<empty>" : value;
    } else if (entry.first->asInt()->value() == kSubject) {
      ret.subject = value.empty() ? "<empty>" : value;
    }
  }
  return ret;
}

const cppbor::Bstr* GetSubjectPublicKeyFromBccEntryPayload(
    const cppbor::Map* bcc_entry_payload) {
  for (size_t i = 0; i < bcc_entry_payload->size(); ++i) {
    const auto& entry = (*bcc_entry_payload)[i];
    if (entry.first == nullptr || entry.first->asInt() == nullptr ||
        entry.second == nullptr) {
      continue;
    }
    if (entry.first->asInt()->value() == kSubjectPublicKey) {
      return entry.second->asBstr();
    }
  }
  return nullptr;
}

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

bool BccValidator::VerifySignature(const BccPublicKeyInfo& signing_key,
                                   const std::vector<uint8_t>& message,
                                   const std::vector<uint8_t>& signature) {
  if (signing_key.signature_algorithm == kBccEdDsa) {
    constexpr size_t kEd25519SignatureLength = 64;
    // ED25519 incorporates SHA512 into the signing algorithm.
    if (signature.size() != kEd25519SignatureLength) {
      AddValidationMessage(
          kCborValidateError,
          "Signature has unexpected size: " + std::to_string(signature.size()));
      return false;
    }
    EVP_PKEY* pkey = nullptr;
    if ((pkey = EVP_PKEY_new_raw_public_key(
             EVP_PKEY_ED25519, nullptr,
             reinterpret_cast<const uint8_t*>(signing_key.key_bytes.data()),
             signing_key.key_bytes.size())) == nullptr) {
      AddValidationMessage(
          kCborValidateError,
          "Can not create EVP_PKEY_ED25519 from the public key info.");
      return false;
    }
    EVP_MD_CTX* md_ctx = EVP_MD_CTX_new();
    const bool res =
        EVP_DigestVerifyInit(md_ctx, nullptr, nullptr, nullptr, pkey) &&
        EVP_DigestVerify(md_ctx, signature.data(), signature.size(),
                         message.data(), message.size()) == 1;
    EVP_MD_CTX_free(md_ctx);
    EVP_PKEY_free(pkey);
    return res;
  }
  if (signing_key.signature_algorithm == kBccEcdsaSha256 ||
      signing_key.signature_algorithm == kBccEcdsaSha384) {
    const EccCurve curve = (signing_key.signature_algorithm == kBccEcdsaSha256)
                               ? EccCurve::kEccSecp256r1
                               : EccCurve::kEccSecp384r1;
    std::unique_ptr<EccPublicKey> key = EccPublicKey::LoadKeyPoint(
        curve, reinterpret_cast<const uint8_t*>(signing_key.key_bytes.data()),
        signing_key.key_bytes.size());
    if (!key) {
      AddValidationMessage(kCborValidateError,
                           "Can not create ECPublicKey from raw EC KeyPoint.");
      return false;
    }
    const OEMCryptoResult res = key->VerifyRawSignature(
        message.data(), message.size(), signature.data(), signature.size());
    return (res == OEMCrypto_SUCCESS);
  }
  AddValidationMessage(kCborValidateError,
                       "Unknown signature algorithm: " +
                           std::to_string(signing_key.signature_algorithm));
  return false;
}

CborMessageStatus BccValidator::Validate() {
  if (message_status_ != kCborParseOk) return message_status_;
  const cppbor::Item* parsed_bcc = std::get<0>(parse_result()).get();
  if (parsed_bcc == nullptr) {
    AddValidationMessage(kCborValidateFatal, "BCC is empty.");
    return message_status_;
  }
  if (parsed_bcc->asArray() == nullptr) {
    AddValidationMessage(kCborValidateFatal,
                         "BCC is not a CBOR array. Actual type: " +
                             CppborMajorTypeToString(parsed_bcc->type()));
    return message_status_;
  }
  const cppbor::Array* bcc_array = parsed_bcc->asArray();
  if (bcc_array->size() < 2) {
    AddValidationMessage(kCborValidateFatal,
                         "BCC should contain at least two elements. Actual: " +
                             std::to_string(bcc_array->size()));
    return message_status_;
  }

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

  // The first element in the array contains the root device public key
  // definition.
  const cppbor::Map* device_public_key_map = (*bcc_array)[0]->asMap();
  if (device_public_key_map == nullptr) {
    AddValidationMessage(
        kCborValidateFatal,
        "Device public key info is not a CBOR map. Actual type: " +
            CppborMajorTypeToString((*bcc_array)[0]->type()));
    return message_status_;
  }
  BccPublicKeyInfo root_pub_key;
  std::vector<std::string> key_value_texts;  // for pretty print
  CborMessageStatus status = ProcessSubjectPublicKeyInfo(
      *device_public_key_map, key_value_texts, &root_pub_key);
  AddMessages(msg_ss_, key_value_texts, 1);
  if (status == kCborValidateFatal) return status;

  BccPublicKeyInfo leaf_pub_key = root_pub_key;
  msg_ss_ << "BCC ENTRY:\n";
  // Parse and verify each certificate in the chain. The structure of thr
  // entries are COSE_Sign1 (untagged). leaf_pub_key is being updated while we
  // process the chain.
  for (size_t i = 1; i < bcc_array->size(); ++i) {
    msg_ss_ << "- CDI PUBLIC KEY INDEX: " << i << "\n";
    const cppbor::Array* bcc_entry = (*bcc_array)[i]->asArray();
    if (bcc_entry == nullptr) {
      AddValidationMessage(kCborValidateFatal,
                           "BCC entry is empty at index " + std::to_string(i));
      return message_status_;
    }
    if (bcc_entry->size() != 4) {
      AddValidationMessage(kCborValidateFatal,
                           "BCC entry should contain 4 items. Actual: " +
                               std::to_string(bcc_entry->size()));
      return message_status_;
    }
    // Skip CoseSign1 signature verification here, only extract pub keys
    if ((*bcc_entry)[0]->type() != cppbor::BSTR ||
        (*bcc_entry)[1]->type() != cppbor::MAP ||
        (*bcc_entry)[2]->type() != cppbor::BSTR ||
        (*bcc_entry)[3]->type() != cppbor::BSTR) {
      AddValidationMessage(kCborValidateFatal, "Invalid BCC entry type.");
      return message_status_;
    }

    // Signature verification Step 1: construct and encode signature input
    const std::vector<uint8_t>& protected_bytes =
        (*bcc_entry)[0]->asBstr()->value();
    // Index 1 is unprotected parameters, which is ignored.
    const std::vector<uint8_t>& payload = (*bcc_entry)[2]->asBstr()->value();
    const std::vector<uint8_t>& actual_signature =
        (*bcc_entry)[3]->asBstr()->value();

    const std::vector<uint8_t> signature_input =
        cppbor::Array()
            .add(kSignatureContextString)
            .add(protected_bytes)
            .add(/* AAD */ std::vector<uint8_t>())
            .add(payload)
            .encode();

    // Signature verification Step 2: verify
    if (!VerifySignature(leaf_pub_key, signature_input, actual_signature)) {
      AddValidationMessage(
          kCborValidateError,
          "Failed to verify the signature for BCC entry index: " +
              std::to_string(i));
    }

    key_value_texts.clear();
    BccPublicKeyInfo entry_pub_key;
    status =
        ProcessDiceChainEntryPayload(payload, key_value_texts, &entry_pub_key);
    AddMessages(msg_ss_, key_value_texts, 1);
    if (status == kCborValidateFatal) return status;
    leaf_pub_key = std::move(entry_pub_key);
  }
  // If the size of the BCC array (including device pub key) is 2, then it
  // must be a de-generated BCC, which means the second element in the array
  // is a self-signed entry. The entry's public key should be identical to the
  // device's public key.
  if (bcc_array->size() == 2) {
    // self-signed BCC entry
    if (leaf_pub_key.key_bytes != root_pub_key.key_bytes) {
      AddValidationMessage(kCborValidateError,
                           "The public key of a self-signed entry should be "
                           "identical to its device public key.");
    }
  }
  msg_ss_ << "...\n";
  if (message_status_ == kCborParseOk) message_status_ = kCborValidateOk;
  return message_status_;
}

CborMessageStatus BccValidator::ProcessSubjectPublicKeyInfo(
    const cppbor::Map& public_key_info_map, std::vector<std::string>& fmt_msgs,
    BccPublicKeyInfo* public_key_info) {
  int key_encoding_format = DEVICE_KEY_ENCODING_UNKNOWN;
  std::vector<uint8_t> device_key_bytes_0;
  std::vector<uint8_t> device_key_bytes_1;
  std::unordered_set<int64_t> key_set;
  for (size_t index = 0; index < public_key_info_map.size(); ++index) {
    std::pair<const std::unique_ptr<cppbor::Item>&,
              const std::unique_ptr<cppbor::Item>&>
        entry = public_key_info_map[index];
    if (entry.first->type() != cppbor::NINT &&
        entry.first->type() != cppbor::UINT) {
      AddValidationMessage(kCborValidateFatal,
                           "Invalid key type in public key info map: " +
                               CppborMajorTypeToString(entry.first->type()));
      return kCborValidateFatal;
    }
    const int64_t map_key = entry.first->asInt()->value();
    switch (map_key) {
      case MAP_KEY_DEVICE_KEY_TYPE: {
        if (entry.second->type() != cppbor::UINT) {
          AddValidationMessage(
              kCborValidateFatal,
              "Invalid value type in public key info map for "
              "key MAP_KEY_DEVICE_KEY_TYPE: " +
                  CppborMajorTypeToString(entry.second->type()));
          return kCborValidateFatal;
        }
        std::string kv = "key encoding format: ";
        const int64_t value = entry.second->asUint()->value();
        if (value == DEVICE_KEY_OCTET_PAIR) {
          key_encoding_format = DEVICE_KEY_OCTET_PAIR;
          kv += "DEVICE_KEY_OCTET_PAIR";
        } else if (value == DEVICE_KEY_BYTE_STRING) {
          key_encoding_format = DEVICE_KEY_BYTE_STRING;
          kv += "DEVICE_KEY_BYTE_STRING";
        } else {
          AddValidationMessage(kCborValidateError,
                               "Invalid value in public key info map for key "
                               "MAP_KEY_DEVICE_KEY_TYPE: " +
                                   std::to_string(value));
        }
        fmt_msgs.push_back(kv);
      } break;
      case MAP_KEY_DEVICE_KEY_ALGORITHM: {
        if (entry.second->type() != cppbor::NINT) {
          AddValidationMessage(
              kCborValidateFatal,
              "Invalid value type in public key info map for "
              "key MAP_KEY_DEVICE_KEY_ALGORITHM: " +
                  CppborMajorTypeToString(entry.second->type()));
          return kCborValidateFatal;
        }
        std::string kv = "key algorithm type: ";
        const int64_t value = entry.second->asNint()->value();
        if (value == DEVICE_KEY_ALGORITHM_ES256) {
          kv += "ECDSA_SHA256";
          public_key_info->signature_algorithm = kBccEcdsaSha256;
        } else if (value == DEVICE_KEY_ALGORITHM_ES384) {
          kv += "ECDSA_SHA384";
          public_key_info->signature_algorithm = kBccEcdsaSha384;
        } else if (value == DEVICE_KEY_ALGORITHM_EDDSA) {
          kv += "EDDSA";
          public_key_info->signature_algorithm = kBccEdDsa;
        } else {
          AddValidationMessage(kCborValidateError,
                               "Invalid value in public key info map for key "
                               "MAP_KEY_DEVICE_KEY_ALGORITHM: " +
                                   std::to_string(value));
        }
        fmt_msgs.push_back(kv);
      } break;
      case MAP_KEY_DEVICE_KEY_OPS:
        // The OPS is an array. Ignored for now.
        break;
      case MAP_KEY_DEVICE_KEY_CURVE: {
        if (entry.second->type() != cppbor::UINT) {
          AddValidationMessage(
              kCborValidateFatal,
              "Invalid value type in public key info map for "
              "key MAP_KEY_DEVICE_KEY_CURVE: " +
                  CppborMajorTypeToString(entry.second->type()));
          return kCborValidateFatal;
        }
        std::string kv = "curve: ";
        const int64_t value = entry.second->asUint()->value();
        if (value == DEVICE_KEY_CURVE_P256) {
          public_key_info->curve = kBccP256;
          kv += "P256";
        } else if (value == DEVICE_KEY_CURVE_P384) {
          public_key_info->curve = kBccP384;
          kv += "P384";
        } else if (value == DEVICE_KEY_CURVE_ED25519) {
          public_key_info->curve = kBccEd25519;
          kv += "ED25519";
        } else {
          AddValidationMessage(kCborValidateError,
                               "Invalid value in public key info map for key "
                               "MAP_KEY_DEVICE_KEY_CURVE: " +
                                   std::to_string(value));
        }
        fmt_msgs.push_back(kv);
      } break;
      case MAP_KEY_DEVICE_KEY_BYTES_0:
      case MAP_KEY_DEVICE_KEY_BYTES_1:
        // BCC encodes keys as either two X, Y octet strings or a single
        // octet string. The format used depends on the key type.
        if (entry.second->type() != cppbor::BSTR) {
          AddValidationMessage(
              kCborValidateFatal,
              "Invalid value type in public key info map for "
              "key MAP_KEY_DEVICE_KEY_BYTES_0/1: " +
                  CppborMajorTypeToString(entry.second->type()));
          return kCborValidateFatal;
        }
        const std::vector<uint8_t>& key_bytes = entry.second->asBstr()->value();
        // Key byte length depends upon the key type.
        if (key_bytes.size() != kED25519KeyDataItemSize &&
            key_bytes.size() != kP256KeyComponentSize &&
            key_bytes.size() != kP384KeyComponentSize) {
          AddValidationMessage(kCborValidateFatal,
                               "Malformed public key data size of: " +
                                   std::to_string(key_bytes.size()));
          return kCborValidateFatal;
        }
        if (map_key == MAP_KEY_DEVICE_KEY_BYTES_0) {
          device_key_bytes_0 = key_bytes;
        } else {
          device_key_bytes_1 = key_bytes;
        }
    }
    key_set.insert(map_key);
  }
  if (key_set.find(MAP_KEY_DEVICE_KEY_TYPE) == key_set.end()) {
    AddValidationMessage(kCborValidateError,
                         "Missing MAP_KEY_DEVICE_KEY_TYPE.");
  }
  if (key_set.find(MAP_KEY_DEVICE_KEY_ALGORITHM) == key_set.end()) {
    AddValidationMessage(kCborValidateError,
                         "Missing MAP_KEY_DEVICE_KEY_ALGORITHM.");
  }
  if (key_set.find(MAP_KEY_DEVICE_KEY_CURVE) == key_set.end()) {
    AddValidationMessage(kCborValidateError,
                         "Missing MAP_KEY_DEVICE_KEY_CURVE.");
  }
  if (device_key_bytes_0.empty() ||
      (key_encoding_format == DEVICE_KEY_OCTET_PAIR &&
       device_key_bytes_1.empty())) {
    AddValidationMessage(
        kCborValidateFatal,
        "Malformed public key definition. Missing device public key bytes.");
    return kCborValidateFatal;
  }
  std::vector<uint8_t> device_key_bytes;
  if (key_encoding_format == DEVICE_KEY_OCTET_PAIR) {
    // Key is an ECDSA elliptic key. We need to return the ANSI X9.62
    // marshaled public key. Generate the marshaled key if needed. The
    // marshaled key is needed to create an ECPublicKey object.
    device_key_bytes.push_back(kMarshaledECKeyZValue);
    device_key_bytes.insert(device_key_bytes.end(), device_key_bytes_0.begin(),
                            device_key_bytes_0.end());
    device_key_bytes.insert(device_key_bytes.end(), device_key_bytes_1.begin(),
                            device_key_bytes_1.end());
    if (device_key_bytes.size() != kMarshaledP384KeySize &&
        device_key_bytes.size() != kMarshaledP256KeySize) {
      AddValidationMessage(kCborValidateFatal,
                           "Invalid ECDSA public key size: " +
                               std::to_string(device_key_bytes.size()));
      return kCborValidateFatal;
    }
  } else {
    device_key_bytes = std::move(device_key_bytes_0);
  }
  fmt_msgs.push_back("public key bytes: " + wvutil::b2a_hex(device_key_bytes));
  public_key_info->key_bytes = std::move(device_key_bytes);
  return message_status_;
}

CborMessageStatus BccValidator::ProcessDiceChainEntryPayload(
    const std::vector<uint8_t>& payload, std::vector<std::string>& fmt_msgs,
    BccPublicKeyInfo* entry_public_key_info) {
  if (payload.empty()) {
    AddValidationMessage(kCborValidateFatal, "Empty bcc entry payload.");
    return kCborValidateFatal;
  }
  auto parse_result = cppbor::parse(payload);
  std::unique_ptr<cppbor::Item> item = std::move(std::get<0>(parse_result));
  std::string error_message = std::move(std::get<2>(parse_result));
  if (item == nullptr || !error_message.empty()) {
    AddValidationMessage(kCborValidateFatal,
                         "Unable to parse bcc entry payload: " + error_message);
    return kCborValidateFatal;
  }
  if (item->type() != cppbor::MAP) {
    AddValidationMessage(kCborValidateFatal,
                         "Unexpected bcc entry payload type: " +
                             CppborMajorTypeToString(item->type()));
    return kCborValidateFatal;
  }
  const IssuerSubject issuer_subject =
      GetIssuerSubjectFromBccEntryPayload(item->asMap());
  if (!issuer_subject.IsValid()) {
    AddValidationMessage(kCborValidateError, "Missing Issuer or Subject.");
  }
  issuer_subject.PrintTo(fmt_msgs);
  const cppbor::Bstr* subject_public_key =
      GetSubjectPublicKeyFromBccEntryPayload(item->asMap());
  if (subject_public_key == nullptr) {
    AddValidationMessage(kCborValidateFatal,
                         "Bcc entry payload has no subject public key.");
    return kCborValidateFatal;
  }

  // Now parse the serialized subject public key.
  parse_result = cppbor::parse(subject_public_key->value());
  item = std::move(std::get<0>(parse_result));
  error_message = std::move(std::get<2>(parse_result));
  if (item == nullptr || !error_message.empty()) {
    AddValidationMessage(
        kCborValidateFatal,
        "Unable to parse serialized subject public key: " + error_message);
    return kCborValidateFatal;
  }
  const cppbor::Map* subject_public_key_info = item->asMap();
  if (subject_public_key_info == nullptr) {
    AddValidationMessage(kCborValidateFatal,
                         "Invalid subject public key type. Expected Map.");
    return kCborValidateFatal;
  }
  return ProcessSubjectPublicKeyInfo(*subject_public_key_info, fmt_msgs,
                                     entry_public_key_info);
}

std::string BccValidator::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