// 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 <gmock/gmock.h>
#include <gtest/gtest.h>

#include "device_info_validator.h"

using ::testing::AllOf;
using ::testing::Ge;
using ::testing::HasSubstr;
using ::testing::Le;

namespace wvoec {
namespace util {
namespace {
constexpr int kDeviceVersion1 = 1;
constexpr int kDeviceVersion2 = 2;
constexpr int kDeviceVersion3 = 3;

cppbor::Map BuildDeviceInfoMap(int version) {
  cppbor::Map device_info =
      cppbor::Map()
          .add("brand", "brand")
          .add("manufacturer", "manufacturer")
          .add("product", "product")
          .add("model", "model")
          .add("vb_state", "green")
          .add("bootloader_state", "unlocked")
          .add("vbmeta_digest", cppbor::Bstr(std::vector<uint8_t>()))
          .add("os_version", "os_version")
          .add("system_patch_level", 202312)
          .add("boot_patch_level", 20231201)
          .add("vendor_patch_level", 20231201)
          .add("security_level", "tee");
  switch (version) {
    case kDeviceVersion1:
      device_info.add("board", "board");
      device_info.add("version", 1);
      device_info.add("att_id_state", "open");
      break;
    case kDeviceVersion2:
      device_info.add("device", "device");
      device_info.add("version", 2);
      device_info.add("fused", 0);
      break;
    case kDeviceVersion3:
      device_info.add("device", "device");
      device_info.add("fused", 0);
      break;
  }
  return device_info;
}

std::vector<uint8_t> BuildDeviceInfo(int version) {
  auto map = BuildDeviceInfoMap(version);
  return map.canonicalize().encode();
}
}  // namespace

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoParseError) {
  const std::vector<uint8_t> device_info = BuildDeviceInfo(kDeviceVersion3);
  const std::vector<uint8_t> device_info_bad(device_info.begin(),
                                             device_info.end() - 1);
  DeviceInfoValidator validator(kDeviceVersion3);
  CborMessageStatus result = validator.Parse(device_info_bad);
  EXPECT_EQ(kCborParseError, result);
  result = validator.Validate();
  EXPECT_EQ(kCborParseError, result);
  EXPECT_EQ("", validator.GetRawMessage());
  EXPECT_EQ("", validator.GetFormattedMessage());
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoNotMap) {
  cppbor::Array array = cppbor::Array().add("make").add(123).add("model");
  const std::vector<uint8_t> device_info_bad = array.encode();
  DeviceInfoValidator validator(kDeviceVersion3);
  CborMessageStatus result = validator.Parse(device_info_bad);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_THAT(result, kCborValidateFatal);
  const std::vector<std::pair<CborMessageStatus, std::string>> msgs =
      validator.GetValidateMessages();
  EXPECT_EQ(1u, msgs.size());
  EXPECT_EQ(kCborValidateFatal, msgs[0].first);
  EXPECT_THAT(msgs[0].second, HasSubstr("Device info is not a CBOR map"));
}

TEST(OEMCryptoDeviceInfoValidatorTest,
     DeviceInfoV3WrongKeyValueTypeAndMissingField) {
  const std::vector<uint8_t> device_info_bad =
      cppbor::Map()
          .add("brand", "brand")
          .add("manufacturer", "manufacturer")
          .add(123, 456)                            // Non-Tstr key type
          .add("system_patch_level", "not a uint")  // Non-uint value type
          .canonicalize()
          .encode();
  DeviceInfoValidator validator(kDeviceVersion3);
  CborMessageStatus result = validator.Parse(device_info_bad);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_EQ(kCborValidateError, result);
  const std::vector<std::pair<CborMessageStatus, std::string>> msgs =
      validator.GetValidateMessages();
  const bool unexpected_key_type_found = std::any_of(
      msgs.begin(), msgs.end(),
      [](const std::pair<CborMessageStatus, std::string>& p) {
        return p.second.find("Unexpected entry key type") != std::string::npos;
      });
  EXPECT_EQ(true, unexpected_key_type_found);
  const bool unexpected_value_type_found = std::any_of(
      msgs.begin(), msgs.end(),
      [](const std::pair<CborMessageStatus, std::string>& p) {
        return p.second.find("system_patch_level has the wrong type") !=
               std::string::npos;
      });
  EXPECT_EQ(true, unexpected_value_type_found);
  const bool missing_model_found = std::any_of(
      msgs.begin(), msgs.end(),
      [](const std::pair<CborMessageStatus, std::string>& p) {
        return p.second.find("model is missing") != std::string::npos;
      });
  EXPECT_EQ(true, missing_model_found);
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV3NonCanonical) {
  const cppbor::Map map = BuildDeviceInfoMap(kDeviceVersion3);
  const std::vector<uint8_t> device_info = map.encode();
  DeviceInfoValidator validator(kDeviceVersion3);
  CborMessageStatus result = validator.Parse(device_info);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_THAT(result, kCborValidateError);
  const std::vector<std::pair<CborMessageStatus, std::string>> msgs =
      validator.GetValidateMessages();
  EXPECT_EQ(1u, msgs.size());
  EXPECT_EQ(kCborValidateError, msgs[0].first);
  EXPECT_THAT(msgs[0].second,
              HasSubstr("Device info ordering is non-canonical"));
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV3) {
  const std::vector<uint8_t> device_info = BuildDeviceInfo(kDeviceVersion3);
  DeviceInfoValidator validator(kDeviceVersion3);
  CborMessageStatus result = validator.Parse(device_info);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_THAT(result, AllOf(Ge(kCborValidateOk), Le(kCborValidateWarning)));
  const std::string out = validator.GetFormattedMessage();
  EXPECT_THAT(out, HasSubstr("manufacturer: manufacturer"));
  EXPECT_THAT(out, HasSubstr("model: model"));
  EXPECT_THAT(out, HasSubstr("fused: 0"));
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV2) {
  const std::vector<uint8_t> device_info = BuildDeviceInfo(kDeviceVersion2);
  DeviceInfoValidator validator(kDeviceVersion2);
  CborMessageStatus result = validator.Parse(device_info);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_THAT(result, AllOf(Ge(kCborValidateOk), Le(kCborValidateWarning)));
  const std::string out = validator.GetFormattedMessage();
  EXPECT_THAT(out, HasSubstr("manufacturer: manufacturer"));
  EXPECT_THAT(out, HasSubstr("model: model"));
  EXPECT_THAT(out, HasSubstr("fused: 0"));
  EXPECT_THAT(out, HasSubstr("version: 2"));
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV1MissingField) {
  const std::vector<uint8_t> device_info = cppbor::Map()
                                               .add("brand", "brand")
                                               .add("security_level", "tee")
                                               .add("version", 1)
                                               .canonicalize()
                                               .encode();
  DeviceInfoValidator validator(kDeviceVersion1);
  CborMessageStatus result = validator.Parse(device_info);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_THAT(result, kCborValidateError);
  const std::vector<std::pair<CborMessageStatus, std::string>> msgs =
      validator.GetValidateMessages();
  EXPECT_EQ(1u, msgs.size());
  EXPECT_EQ(kCborValidateError, msgs[0].first);
  EXPECT_THAT(msgs[0].second, HasSubstr("att_id_state is missing"));
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV1) {
  DeviceInfoValidator validator(kDeviceVersion1);
  const std::vector<uint8_t> device_info = BuildDeviceInfo(kDeviceVersion1);
  CborMessageStatus result = validator.Parse(device_info);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_THAT(result, AllOf(Ge(kCborValidateOk), Le(kCborValidateWarning)));
  const std::string out = validator.GetFormattedMessage();
  EXPECT_THAT(out, HasSubstr("board: board"));
  EXPECT_THAT(out, HasSubstr("version: 1"));
}
}  // namespace util
}  // namespace wvoec