// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

#include <vector>

#include "api/golden_data.h"
#include "api/license_whitebox_test_base.h"
#include "api/result.h"
#include "api/test_license_builder.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace widevine {
namespace {
using RemoteAttestation = TestLicenseBuilder::RemoteAttestation;
using VerificationStatus = TestLicenseBuilder::VerificationStatus;

// We can't use the actual keys with TEST_P, so define an enum that we can use
// to communicate which key to use.
enum class Key {
  kCrypto,
  kDecode,
  kHardware,
};

}  // namespace

class LicenseWhiteboxChromeOSTest
    : public LicenseWhiteboxTestBase,
      public testing::WithParamInterface<
          std::tuple<Key, RemoteAttestation, VerificationStatus>> {
 protected:
  void SetUp() override {
    LicenseWhiteboxTestBase::SetUp();

    std::tie(key_, remote_attestation_, verification_status_) = GetParam();

    LoadLicense();

    // We know that the plaintext will be smaller than the ciphertext, so we
    // use that to ensure the buffer is large enough.
    plaintext_size_ = golden_data_.CBCContent().ciphertext.size();
    plaintext_.resize(plaintext_size_);

    // We make the assumption that the secret string can never be longer than
    // the ciphertext.
    secret_string_size_ = golden_data_.CBCContent().ciphertext.size();
    secret_string_.resize(secret_string_size_);
  }

  // This requires that the remote attestation and verification status to be
  // set before being called.
  void LoadLicense() {
    TestLicenseBuilder builder;

    // Only use CBC keys so that we can always use the CBC content.
    builder.AddContentKey(golden_data_.CBCCryptoKey().level,
                          golden_data_.CBCCryptoKey().id,
                          golden_data_.CBCCryptoKey().content->key);

    builder.AddContentKey(golden_data_.CBCDecodeKey().level,
                          golden_data_.CBCDecodeKey().id,
                          golden_data_.CBCDecodeKey().content->key);

    builder.AddContentKey(golden_data_.CBCHardwareKey().level,
                          golden_data_.CBCHardwareKey().id,
                          golden_data_.CBCHardwareKey().content->key);

    builder.SetRemoteAttestation(remote_attestation_);
    builder.SetVerificationStatus(verification_status_);

    License license;
    builder.Build(*public_key_, &license);

    ASSERT_EQ(WB_License_ProcessLicenseResponse(
                  whitebox_, license.message.data(), license.message.size(),
                  license.signature.data(), license.signature.size(),
                  license.session_key.data(), license.session_key.size(),
                  license.request.data(), license.request.size()),
              WB_RESULT_OK);
  }

  // This is the strictest level of enforcement. It will be the last one
  // checked. Any basic overrides, for example a crypto key in decrypt,
  // will short-circuit it.
  WB_Result GetExpectedPlatformVerificationResult() {
    if (remote_attestation_ == RemoteAttestation::kUnverified) {
      return WB_RESULT_INSUFFICIENT_SECURITY_LEVEL;
    }

    if (verification_status_ == VerificationStatus::kOther) {
      return WB_RESULT_INSUFFICIENT_SECURITY_LEVEL;
    }

    if (remote_attestation_ == RemoteAttestation::kVerified) {
      return WB_RESULT_OK;
    }

    if (verification_status_ == VerificationStatus::kHardwareVerified) {
      return WB_RESULT_OK;
    }

    return WB_RESULT_INSUFFICIENT_SECURITY_LEVEL;
  }

  WB_Result GetExpectedDecryptResult() {
    if (key_ == Key::kCrypto) {
      return WB_RESULT_OK;
    }

    return GetExpectedPlatformVerificationResult();
  }

  // Since MaskedDecrypt() and GetSecretString() go hand-in-hand. This function
  // will be used for both.
  WB_Result GetExpectedMaskedDecryptResult() {
    if (key_ == Key::kCrypto || key_ == Key::kDecode) {
      return WB_RESULT_OK;
    }

    return GetExpectedPlatformVerificationResult();
  }

  const GoldenData::Key* GetContentKey() const {
    switch (key_) {
      case Key::kCrypto:
        return &golden_data_.CBCCryptoKey();
      case Key::kDecode:
        return &golden_data_.CBCDecodeKey();
      case Key::kHardware:
        return &golden_data_.CBCHardwareKey();
    }
    return nullptr;
  }

  Key key_;
  RemoteAttestation remote_attestation_;
  VerificationStatus verification_status_;

  // This is the buffer used to store the output of each decrypt and unmask
  // call.
  size_t plaintext_size_;
  std::vector<uint8_t> plaintext_;

  // This is the buffer used to store the secret string used to demask the
  // result of WB_License_MaskedDecrypt().
  size_t secret_string_size_;
  std::vector<uint8_t> secret_string_;
};

TEST_P(LicenseWhiteboxChromeOSTest, Decrypt) {
  const WB_Result expected_result = GetExpectedDecryptResult();

  const GoldenData::Key* key = GetContentKey();
  ASSERT_NE(key, nullptr);

  const auto* content = key->content;
  ASSERT_EQ(WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC, key->id.data(),
                               key->id.size(), content->ciphertext.data(),
                               content->plaintext.size(), content->iv.data(),
                               content->iv.size(), plaintext_.data(),
                               &plaintext_size_),
            expected_result);
}

TEST_P(LicenseWhiteboxChromeOSTest, MaskedDecrypt) {
  const WB_Result expected_result = GetExpectedMaskedDecryptResult();

  const GoldenData::Key* key = GetContentKey();
  ASSERT_NE(key, nullptr);

  const auto* content = key->content;
  ASSERT_EQ(WB_License_MaskedDecrypt(
                whitebox_, WB_CIPHER_MODE_CBC, key->id.data(), key->id.size(),
                content->ciphertext.data(), content->plaintext.size(),
                content->iv.data(), content->iv.size(), plaintext_.data(),
                &plaintext_size_),
            expected_result);
}

TEST_P(LicenseWhiteboxChromeOSTest, GetSecretString) {
  const WB_Result expected_result = GetExpectedMaskedDecryptResult();

  const GoldenData::Key* key = GetContentKey();
  ASSERT_NE(key, nullptr);

  ASSERT_EQ(WB_License_GetSecretString(
                whitebox_, WB_CIPHER_MODE_CBC, key->id.data(), key->id.size(),
                secret_string_.data(), &secret_string_size_),
            expected_result);
}

INSTANTIATE_TEST_SUITE_P(
    AllCombinations,
    LicenseWhiteboxChromeOSTest,
    ::testing::Combine(
        ::testing::Values(Key::kCrypto, Key::kDecode, Key::kHardware),
        ::testing::Values(RemoteAttestation::kUnavailable,
                          RemoteAttestation::kUnverified,
                          RemoteAttestation::kVerified),
        ::testing::Values(VerificationStatus::kUnavailable,
                          VerificationStatus::kOther,
                          VerificationStatus::kHardwareVerified)));

}  // namespace widevine