// 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"
#include "log.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>(32, 0xCC)))
          .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

static void DumpValidatorOutput(const DeviceInfoValidator& validator) {
  const std::string out = validator.GetFormattedMessage();
  LOGI("%s", out.c_str());
  LOGE("Validation results\n%s", validator.PrintValidateMessage().c_str());
}

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"));
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

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
          .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 missing_model_found =
      std::any_of(msgs.begin(), msgs.end(),
                  [](const std::pair<CborMessageStatus, std::string>& p) {
                    return p.second.find("missing important field model") !=
                           std::string::npos;
                  });
  EXPECT_EQ(true, missing_model_found);
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV3NonCanonical) {
  const cppbor::Map map = BuildDeviceInfoMap(kDeviceVersion3);
  const std::vector<uint8_t> device_info = map.encode();
  DeviceInfoValidator validator(kDeviceVersion3, true /* is_gms */);
  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"));
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV3) {
  const std::vector<uint8_t> device_info = BuildDeviceInfo(kDeviceVersion3);
  DeviceInfoValidator validator(kDeviceVersion3, true /* is_gms */);
  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"));
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV3InvalidFields) {
  cppbor::Map device_info_map =
      cppbor::Map()
          .add("brand", "brand")
          .add("manufacturer", "manufacturer")
          .add("product", "product")
          .add("model", "model")
          .add("vb_state", "invalid_green")             // invalid value
          .add("bootloader_state", "invalid_unlocked")  // invalid value
          .add("vbmeta_digest", cppbor::Bstr(std::vector<uint8_t>(32, 0xCC)))
          .add("os_version", "os_version")
          .add("system_patch_level", 100)  // invalid value, expects "YYYYMM"
          .add("boot_patch_level",
               12345678)  // invalid value, expectes "YYYYMMDD"
          .add("vendor_patch_level",
               "20231201")  // invalid value, expects YYYYMMDD in int
          .add("security_level", "tee")
          .add("device", "device")
          .add("fused", 9);  // invalid value, expects 0 or 1
  auto device_info = device_info_map.canonicalize().encode();
  DeviceInfoValidator validator(kDeviceVersion3, true /* is_gms */);
  CborMessageStatus result = validator.Parse(device_info);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_THAT(result, Ge(kCborValidateError));
  const std::vector<std::pair<CborMessageStatus, std::string>> msgs =
      validator.GetValidateMessages();
  std::string out = "";
  for (auto& msg : msgs) {
    out += (msg.second + "\n");
  }
  EXPECT_THAT(out, HasSubstr("invalid value for system_patch_level"));
  EXPECT_THAT(out, HasSubstr("invalid value for boot_patch_level"));
  EXPECT_THAT(out, HasSubstr("missing required field vendor_patch_level"));
  EXPECT_THAT(out, HasSubstr("unexpected value for vb_state"));
  EXPECT_THAT(out, HasSubstr("unexpected value for bootloader_state"));
  EXPECT_THAT(out, HasSubstr("unexpected value for fused"));
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV2) {
  const std::vector<uint8_t> device_info = BuildDeviceInfo(kDeviceVersion2);
  DeviceInfoValidator validator(kDeviceVersion2, true /* is_gms */);
  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"));
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV1MissingField) {
  const std::vector<uint8_t> device_info =
      cppbor::Map()
          .add("manufacturer", "manufacturer")
          .add("model", "model")
          .add("brand", "brand")
          .add("security_level", "tee")
          .add("version", 1)
          .canonicalize()
          .encode();
  DeviceInfoValidator validator(kDeviceVersion1, true /* is_gms */);
  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("missing required field att_id_state"));
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoV1) {
  DeviceInfoValidator validator(kDeviceVersion1, true /* is_gms */);
  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"));
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}

TEST(OEMCryptoDeviceInfoValidatorTest, DeviceInfoMissingFused) {
  const std::vector<uint8_t> device_info_bad =
      cppbor::Map()
          .add("model", "model")
          .add("manufacturer", "manufacturer")
          .canonicalize()
          .encode();
  DeviceInfoValidator validator(kDeviceVersion1, false /* is_gms */);
  CborMessageStatus result = validator.Parse(device_info_bad);
  EXPECT_EQ(kCborParseOk, result);
  result = validator.Validate();
  EXPECT_EQ(kCborValidateWarning, result);
  const std::vector<std::pair<CborMessageStatus, std::string>> msgs =
      validator.GetValidateMessages();
  const bool missing_fused_found = std::any_of(
      msgs.begin(), msgs.end(),
      [](const std::pair<CborMessageStatus, std::string>& p) {
        return p.second.find("missing field fused") != std::string::npos;
      });
  EXPECT_EQ(true, missing_fused_found);
  if (result >= kCborValidateWarning) {
    DumpValidatorOutput(validator);
  }
}
}  // namespace util
}  // namespace wvoec