// Copyright 2023 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.

#include "BccParser.h"

#include <cppbor.h>
#include <cppbor_parse.h>
#include <stddef.h>

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

#include "DiceCborConstants.h"

namespace widevine {
namespace {

// Sized to hold 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 kP256KeyCoordinateSizeBytes = 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 kP384KeyCoordinateSizeBytes = 384 / 8;
constexpr int kMarshaledP384KeySizeBytes = kP384KeyCoordinateSizeBytes * 2 + 1;
constexpr int kMaxMarshaledECKeySizeBytes = kMarshaledP384KeySizeBytes;
constexpr char kMarshaledECKeyZValue = 0x04;
constexpr int kED25519KeyDataItemSizeBytes = 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;

std::string TypeNameFromType(cppbor::MajorType type) {
  switch (type) {
    case cppbor::UINT:
      return "UINT";
    case cppbor::NINT:
      return "NINT";
    case cppbor::BSTR:
      return "BSTR";
    case cppbor::TSTR:
      return "TSTR";
    case cppbor::ARRAY:
      return "ARRAY";
    case cppbor::MAP:
      return "MAP";
    case cppbor::SEMANTIC:
      return "SEMANTIC";
    case cppbor::SIMPLE:
      return "SIMPLE";
    default:
      return "undefined type";
  }
}

std::string GetIssuerSubjectFromBccEntryPayload(
    const cppbor::Map* bcc_entry_payload) {
  std::string issuer = "Issuer: ";
  std::string subject = "Subject: ";
  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() == kIssuer) {
      issuer += (entry.second->asTstr()->value() + "\n");
    } else if (entry.first->asInt()->value() == kSubject) {
      subject += (entry.second->asTstr()->value() + "\n");
    }
  }
  return issuer + subject;
}

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;
}

}  // namespace

// BCC/DiceCertChain definition:
// https://source.corp.google.com/android-internal/hardware/interfaces/security/rkp/aidl/android/hardware/security/keymint/IRemotelyProvisionedComponent.aidl
std::string BccParser::Parse(const std::vector<uint8_t>& bcc) {
  std::stringstream ss;
  auto [parsed_bcc, _, err_msg] = cppbor::parse(bcc);
  if (parsed_bcc == nullptr) {
    ss << "Failed to parse input BCC: " << err_msg.c_str() << "\n";
    return ss.str();
  }
  if (parsed_bcc->asArray() == nullptr) {
    ss << "Input BCC is not a CBOR array: "
       << TypeNameFromType(parsed_bcc->type()) << "\n";
    return ss.str();
  }
  const cppbor::Array* bcc_array = parsed_bcc->asArray();
  if (bcc_array->size() < 2) {
    ss << "Input BCC should contain at least two elements, actual: "
       << bcc_array->size() << "\n";
    return ss.str();
  }
  // The first element in the array contains the root device public key
  // definition.
  ss << "ROOT DEVICE PUBLIC KEY: \n";
  const cppbor::Map* device_public_key_info = (*bcc_array)[0]->asMap();
  if (device_public_key_info == nullptr) {
    ss << "Invalid dice cert chain type for the first element"
       << "\n";
    return ss.str();
  }
  if (!ProcessDevicePublicKeyInfo(ss, *device_public_key_info)) {
    return ss.str();
  }

  // Parse each certificate in the chain. The structure of thr entries are
  // COSE_Sign1 (untagged).
  for (size_t i = 1; i < bcc_array->size(); ++i) {
    ss << "\nCDI PUBLIC KEY " << i << ": \n";
    const cppbor::Array* bcc_entry = (*bcc_array)[i]->asArray();
    if (bcc_entry == nullptr) {
      ss << "Invalid dice cert chain type"
         << "\n";
      return ss.str();
    }

    // Skip CoseSign1 signature verification here, only extract pub keys
    if (bcc_entry->size() != 4 || (*bcc_entry)[0]->type() != cppbor::BSTR ||
        (*bcc_entry)[1]->type() != cppbor::MAP ||
        (*bcc_entry)[2]->type() != cppbor::BSTR ||
        (*bcc_entry)[3]->type() != cppbor::BSTR) {
      ss << "Invalid signature array"
         << "\n";
      return ss.str();
    }

    const std::vector<uint8_t>& key_payload =
        (*bcc_entry)[2]->asBstr()->value();
    std::string payload(key_payload.begin(), key_payload.end());
    if (!ProcessDiceChainEntryPayload(ss, payload)) {
      return ss.str();
    }
  }
  return ss.str();
}

bool BccParser::ProcessDevicePublicKeyInfo(
    std::stringstream& ss, const cppbor::Map& public_key_info_map) {
  int key_encoding_format = DEVICE_KEY_ENCODING_UNKNOWN;
  std::string device_key_bytes_0;
  std::string device_key_bytes_1;
  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) {
      ss << "Invalid map key type " << TypeNameFromType(entry.first->type())
         << " in device key info"
         << "\n";
      return false;
    }
    int map_key = entry.first->type() == cppbor::NINT
                      ? entry.first->asNint()->value()
                      : entry.first->asInt()->value();
    switch (map_key) {
      case MAP_KEY_DEVICE_KEY_TYPE: {
        if (entry.second->type() != cppbor::UINT) {
          ss << "Invalid map value type "
             << TypeNameFromType(entry.second->type()) << " for device key"
             << "\n";
          return false;
        }
        ss << "key encoding format: ";
        int64_t value = entry.second->asUint()->value();
        switch (value) {
          case DEVICE_KEY_OCTET_PAIR:
            key_encoding_format = DEVICE_KEY_OCTET_PAIR;
            ss << "DEVICE_KEY_OCTET_PAIR"
               << "\n";
            break;
          case DEVICE_KEY_BYTE_STRING:
            key_encoding_format = DEVICE_KEY_BYTE_STRING;
            ss << "DEVICE_KEY_BYTE_STRING"
               << "\n";
            break;
          default:
            ss << "Unhandled cbor value for device key format: " << value
               << "\n";
            return false;
        }
      } break;
      case MAP_KEY_DEVICE_KEY_ALGORITHM: {
        if (entry.second->type() != cppbor::NINT) {
          ss << "Invalid map value type "
             << TypeNameFromType(entry.second->type())
             << " for device key algorithm"
             << "\n";
          return false;
        }
        ss << "key algorithm type: ";
        int64_t value = entry.second->asNint()->value();
        switch (value) {
          case DEVICE_KEY_ALGORITHM_ES256:
            ss << "ECDSA_SHA256";
            break;
          case DEVICE_KEY_ALGORITHM_ES384:
            ss << "ECDSA_SHA384";
            break;
          case DEVICE_KEY_ALGORITHM_EDDSA:
            ss << "EDDSA";
            break;
        }
        ss << "\n";
      } 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) {
          ss << "Invalid map value type "
             << TypeNameFromType(entry.second->type())
             << " for device key curve"
             << "\n";
          return false;
        }
        ss << "curve: ";
        int64_t value = entry.second->asUint()->value();
        switch (value) {
          case DEVICE_KEY_CURVE_ED25519:
            ss << "ED25519";
            break;
          case DEVICE_KEY_CURVE_P256:
            ss << "P256";
            break;
          case DEVICE_KEY_CURVE_P384:
            ss << "P384";
            break;
          default:
            ss << "Invalid map value " << value << " for device key curve"
               << "\n";
            return false;
        }
        ss << "\n";
      } 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) {
          ss << "Unexpected cbor type for device key bytes"
             << "\n";
          return false;
        }
        const std::vector<uint8_t>& key_bytes = entry.second->asBstr()->value();
        // Key byte length depends upon the key type.
        if (key_bytes.size() != kED25519KeyDataItemSizeBytes &&
            key_bytes.size() != kP256KeyCoordinateSizeBytes &&
            key_bytes.size() != kP384KeyCoordinateSizeBytes) {
          ss << "Malformed device public key data size of" << key_bytes.size()
             << "\n";
          return false;
        }
        std::string& key_bytes_str = map_key == MAP_KEY_DEVICE_KEY_BYTES_0
                                         ? device_key_bytes_0
                                         : device_key_bytes_1;
        key_bytes_str.assign(key_bytes.begin(), key_bytes.end());
    }
  }
  if (device_key_bytes_0.empty() ||
      (key_encoding_format == DEVICE_KEY_OCTET_PAIR &&
       device_key_bytes_1.empty())) {
    ss << "Malformed device public key definition. Missing device public key "
          "bytes"
       << "\n";
    return false;
  }
  std::string 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.
    // std::string* device_key_bytes = public_key_info.mutable_key_bytes();
    device_key_bytes.reserve(kMaxMarshaledECKeySizeBytes);
    device_key_bytes.resize(1);
    device_key_bytes[0] = kMarshaledECKeyZValue;
    device_key_bytes.append(device_key_bytes_0);
    device_key_bytes.append(device_key_bytes_1);
  } else {
    device_key_bytes = device_key_bytes_0;
  }
  std::ostringstream ss_hex;
  ss_hex << std::setfill('0');
  for (size_t i = 0; i < device_key_bytes.size(); i++) {
    ss_hex << std::setw(2) << std::hex << static_cast<int>(device_key_bytes[i]);
  }
  ss << "public key bytes: " << ss_hex.str() << "\n";
  return true;
}

bool BccParser::ProcessDiceChainEntryPayload(std::stringstream& ss,
                                             std::string& payload) {
  if (payload.empty()) {
    ss << "Empty bcc entry payload"
       << "\n";
    return false;
  }

  auto parse_result = cppbor::parse(
      reinterpret_cast<const uint8_t*>(payload.data()),
      reinterpret_cast<const uint8_t*>(payload.data() + payload.size()));
  std::unique_ptr<cppbor::Item> item = std::move(std::get<0>(parse_result));
  std::string error_message = std::get<2>(parse_result);
  if (item == nullptr || !error_message.empty()) {
    ss << "Unable to parse bcc entry payload: " << error_message << "\n";
    return false;
  }
  if (item->type() != cppbor::MAP) {
    ss << "Unexpected bcc entry payload type"
       << "\n";
    return false;
  }

  const std::string issuer_subject =
      GetIssuerSubjectFromBccEntryPayload(item->asMap());
  ss << issuer_subject;

  const cppbor::Bstr* subject_public_key =
      GetSubjectPublicKeyFromBccEntryPayload(item->asMap());
  if (subject_public_key == nullptr) {
    ss << "Bcc entry payload has no subject public key"
       << "\n";
    return false;
  }

  // Now parse the serialized subject public key.
  parse_result = cppbor::parse(
      subject_public_key->value().data(),
      subject_public_key->value().data() + subject_public_key->value().size());
  item = std::move(std::get<0>(parse_result));
  error_message = std::get<2>(parse_result);
  if (item == nullptr || !error_message.empty()) {
    ss << "Unable to parse serialized subject public key: " << error_message
       << "\n";
    return false;
  }
  const cppbor::Map* subject_public_key_info = item->asMap();
  if (subject_public_key_info == nullptr) {
    ss << "Invalid subject public key type"
       << "\n";
    return false;
  }
  return ProcessDevicePublicKeyInfo(ss, *subject_public_key_info);
}

}  // namespace widevine