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

#include <set>

#include "string_conversions.h"

namespace wvoec {
namespace util {
namespace {
// Number of required device info properties returned from TEE for DeviceInfo
// version v3.
constexpr uint32_t kNumTeeDeviceInfoEntriesV3 = 14;
const std::vector<std::string> kDeviceInfoKeysV3 = {"brand",
                                                    "manufacturer",
                                                    "product",
                                                    "model",
                                                    "device",
                                                    "vb_state",
                                                    "bootloader_state",
                                                    "vbmeta_digest",
                                                    "os_version",
                                                    "system_patch_level",
                                                    "boot_patch_level",
                                                    "vendor_patch_level",
                                                    "security_level",
                                                    "fused"};

struct AttestationIdEntry {
  const char* id;
  bool alwaysValidate;
};

// Attestation Id and whether it is required.
constexpr AttestationIdEntry kAttestationIdEntrySet[] = {{"brand", false},
                                                         {"manufacturer", true},
                                                         {"product", true},
                                                         {"model", true},
                                                         {"device", false}};
}  // namespace

CborMessageStatus DeviceInfoValidator::Parse(
    const std::vector<uint8_t>& device_info) {
  message_status_ = CborValidator::Parse(device_info);
  device_info_bytes_ = device_info;
  return message_status_;
}

CborMessageStatus DeviceInfoValidator::Validate() {
  if (message_status_ != kCborParseOk) return message_status_;
  std::unique_ptr<cppbor::Item> parsed_device_info =
      std::get<0>(parse_result())->clone();
  if (!parsed_device_info) {
    AddValidationMessage(kCborValidateFatal, "Device info is empty.");
    return message_status_;
  }
  cppbor::Map* device_info_map = parsed_device_info->asMap();
  if (!device_info_map) {
    AddValidationMessage(
        kCborValidateFatal,
        "Device info is not a CBOR map. Actual type: " +
            CppborMajorTypeToString(parsed_device_info->type()));
    return message_status_;
  }
  if (device_info_map->canonicalize().encode() != device_info_bytes_) {
    AddValidationMessage(kCborValidateError,
                         "Device info ordering is non-canonical.");
  }
  const cppbor::Item* security_level =
      GetMapEntry(*device_info_map, "security_level");
  const bool is_tee_device_info = security_level && security_level->asTstr() &&
                                  security_level->asTstr()->value() == "tee";
  std::set<std::string> previous_keys;
  switch (version_number_) {
    case 3:
      if (is_tee_device_info &&
          device_info_map->size() != kNumTeeDeviceInfoEntriesV3) {
        AddValidationMessage(
            kCborValidateError,
            "Incorrect number of TEE device info entries. Expected " +
                std::to_string(kNumTeeDeviceInfoEntriesV3) + " but got " +
                std::to_string(device_info_map->size()));
      }
      // TEE IRPC instances require all entries to be present in device info.
      // Non-TEE instances may omit `os_version`.
      if (!is_tee_device_info &&
          (device_info_map->size() != kNumTeeDeviceInfoEntriesV3 &&
           device_info_map->size() != kNumTeeDeviceInfoEntriesV3 - 1)) {
        AddValidationMessage(
            kCborValidateError,
            "Incorrect number of non-TEE device info entries. Expected " +
                std::to_string(kNumTeeDeviceInfoEntriesV3 - 1) + " but got " +
                std::to_string(device_info_map->size()));
      }
      for (auto const& entry : *device_info_map) {
        if (!entry.first->asTstr()) {
          AddValidationMessage(
              kCborValidateError,
              "Unexpected entry key type. Expected TSTR, but got " +
                  CppborMajorTypeToString(entry.first->type()));
          continue;
        }
        const std::string& key = entry.first->asTstr()->value();
        if (previous_keys.find(key) != previous_keys.end()) {
          AddValidationMessage(kCborValidateError,
                               "Duplicate device info entry: " + key);
        }
        previous_keys.insert(key);
        if (std::find(kDeviceInfoKeysV3.begin(), kDeviceInfoKeysV3.end(),
                      key) == kDeviceInfoKeysV3.end()) {
          AddValidationMessage(kCborValidateError,
                               "Unrecognized device info entry: " + key);
        }
      }
      // Checks for the required fields that only apply to v3.
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::UINT,
                              "system_patch_level");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::UINT,
                              "boot_patch_level");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::UINT,
                              "vendor_patch_level");
      // Fall through
      CORE_UTIL_FALLTHROUGH;
    case 2:
      for (const auto& entry : kAttestationIdEntrySet) {
        if (entry.alwaysValidate) {
          CheckDeviceInfoMapEntry(*device_info_map, cppbor::TSTR, entry.id);
        }
      }
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::TSTR, "vb_state");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::TSTR,
                              "bootloader_state");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::BSTR, "vbmeta_digest");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::UINT,
                              "system_patch_level");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::UINT,
                              "boot_patch_level");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::UINT,
                              "vendor_patch_level");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::UINT, "fused");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::TSTR, "security_level");
      if (is_tee_device_info) {
        CheckDeviceInfoMapEntry(*device_info_map, cppbor::TSTR, "os_version");
      }
      break;
    case 1:
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::TSTR, "security_level");
      CheckDeviceInfoMapEntry(*device_info_map, cppbor::TSTR, "att_id_state");
      break;
    default:
      AddValidationMessage(
          kCborValidateFatal,
          "Unrecognized version: " + std::to_string(version_number_));
      return message_status_;
  }

  if (message_status_ == kCborParseOk) message_status_ = kCborValidateOk;
  return message_status_;
}

void DeviceInfoValidator::CheckDeviceInfoMapEntry(
    const cppbor::Map& device_info, cppbor::MajorType major_type,
    const std::string& entry_name) {
  const std::string error = CheckMapEntry(device_info, major_type, entry_name);
  if (!error.empty()) {
    AddValidationMessage(kCborValidateError, error);
  }
}

std::string DeviceInfoValidator::GetFormattedMessage() const {
  if (message_status_ == kCborUninitialized ||
      message_status_ == kCborParseError ||
      message_status_ == kCborValidateFatal) {
    return std::string();
  }
  const cppbor::Item* parsed_item = std::get<0>(parse_result()).get();
  if (parsed_item == nullptr) {
    return "<null>";
  }
  // Writes YAML-formatted output to |msg_ss_|.
  std::stringstream msg_ss;
  msg_ss << "---\n";
  msg_ss << "DEVICE INFO MAP:\n";

  for (auto const& entry : *(parsed_item->asMap())) {
    auto const& entry_value = entry.second;
    // Device info map only allows TSTR key type.
    if (!entry.first->asTstr()) continue;
    const std::string& name = entry.first->asTstr()->value();
    msg_ss << "  " << name << ": ";
    switch (entry_value->type()) {
      case cppbor::TSTR: {
        const std::string val = entry_value->asTstr()->value().empty()
                                    ? "<null>"
                                    : entry_value->asTstr()->value();
        msg_ss << val << "\n";
        break;
      }
      case cppbor::UINT:
        msg_ss << std::to_string(entry_value->asUint()->value()) << "\n";
        break;
      case cppbor::NINT:
        msg_ss << std::to_string(entry_value->asNint()->value()) << "\n";
        break;
      case cppbor::BSTR: {
        const std::vector<uint8_t>& bytes = entry_value->asBstr()->value();
        const std::string val =
            bytes.empty() ? "<null>" : wvutil::b2a_hex(bytes);
        msg_ss << val << "\n";
        break;
      }
      default:
        msg_ss << "Unsupported type ("
               << CppborMajorTypeToString(entry_value->type()) << ")\n";
        break;
    }
  }
  msg_ss << "...\n";
  return msg_ss.str();
}
}  // namespace util
}  // namespace wvoec