// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

#include <string>
#include <vector>

#include "api/golden_data.h"
#include "api/license_whitebox_test_base.h"
#include "api/test_data.h"
#include "api/test_license_builder.h"
#include "crypto_utils/rsa_key.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace widevine {

class LicenseWhiteboxGetSecretStringTest : public LicenseWhiteboxTestBase {
 protected:
  void SetUp() override {
    LicenseWhiteboxTestBase::SetUp();

    // The secret string size is implementation specific, so this number just
    // needs to be reasonably large to accommodate different implementations.
    secret_string_size_ = 256;
    secret_string_.resize(secret_string_size_);

    golden_data_.MakeKeyIdDifferent(&non_content_key_id_);
    golden_data_.MakeKeyIdDifferent(&missing_key_id_);
  }

  void LoadLicense() {
    TestLicenseBuilder builder;

    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.AddOperatorSessionKey(non_content_key_id_);

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

  // We need two special keys for this test, one that will be used for a
  // non-content key and one that will never be in the license.
  std::vector<uint8_t> non_content_key_id_ = {0, 0, 0};
  std::vector<uint8_t> missing_key_id_ = {1, 0, 0};

  size_t secret_string_size_;
  std::vector<uint8_t> secret_string_;
};

TEST_F(LicenseWhiteboxGetSecretStringTest, SuccessForCBCWithCryptoKey) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  ASSERT_GT(secret_string_size_, 0u);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, SuccessForCTRWithCryptoKey) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  ASSERT_GT(secret_string_size_, 0u);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, SuccessForCBCWithDecodeKey) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCDecodeKey().id.data(),
                                 golden_data_.CBCDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  ASSERT_GT(secret_string_size_, 0u);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, SuccessForCTRWithDecodeKey) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CBCDecodeKey().id.data(),
                                 golden_data_.CBCDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  ASSERT_GT(secret_string_size_, 0u);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, InvalidParameterForNullWhitebox) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(nullptr, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxGetSecretStringTest,
       InvalidParameterForInvalidCipherMode) {
  LoadLicense();

  // In order to trick the compiler into letting us pass an invalid enum value
  // to WB__License_Decrypt(), we need to cast it. If we don't do this, the
  // compiler tries to save us.
  const WB_CipherMode invalid_mode = static_cast<WB_CipherMode>(0xFF);

  ASSERT_EQ(WB_License_GetSecretString(
                whitebox_, invalid_mode, golden_data_.CBCCryptoKey().id.data(),
                golden_data_.CBCCryptoKey().id.size(), secret_string_.data(),
                &secret_string_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, InvalidParameterForNullKeyId) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC, nullptr,
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, InvalidParameterForZeroKeyIdSize) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCCryptoKey().id.data(), 0,
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxGetSecretStringTest,
       InvalidParameterForNullSecretString) {
  LoadLicense();

  ASSERT_EQ(WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                       golden_data_.CBCCryptoKey().id.data(),
                                       golden_data_.CBCCryptoKey().id.size(),
                                       nullptr, &secret_string_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxGetSecretStringTest,
       InvalidParameterForNullSecretStringSize) {
  LoadLicense();

  ASSERT_EQ(WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                       golden_data_.CBCCryptoKey().id.data(),
                                       golden_data_.CBCCryptoKey().id.size(),
                                       secret_string_.data(), nullptr),
            WB_RESULT_INVALID_PARAMETER);
}

// For this test, "missing key id" specifically means a key id that was never
// in the license to start with. This is different than "non content key"
// and "dropped content key", as those keys were in the license but ignored.
TEST_F(LicenseWhiteboxGetSecretStringTest, KeyUnavailableForMissingKeyId) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 missing_key_id_.data(), missing_key_id_.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_KEY_UNAVAILABLE);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, KeyUnavailableForNonContentKey) {
  LoadLicense();

  ASSERT_EQ(WB_License_GetSecretString(
                whitebox_, WB_CIPHER_MODE_CBC, non_content_key_id_.data(),
                non_content_key_id_.size(), secret_string_.data(),
                &secret_string_size_),
            WB_RESULT_KEY_UNAVAILABLE);
}

// Under normal circumstances, a hardware key should be dropped. The exception
// to this rule is on ChromeOS with a special license.
TEST_F(LicenseWhiteboxGetSecretStringTest,
       InsufficientSecurityLevelForHardwareContentKey) {
  LoadLicense();

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCHardwareKey().id.data(),
                                 golden_data_.CBCHardwareKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_INSUFFICIENT_SECURITY_LEVEL);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, BufferTooSmall) {
  LoadLicense();

  // Since the secret string is implementation specific, we don't want to make
  // any big assumptions about what would be too small. Using 1 is fairly safe
  // as it would not introduce enough variation to be effective. We avoid using
  // zero here so that we can verify that we are not just checking for zero.
  secret_string_size_ = 1;
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_BUFFER_TOO_SMALL);

  // Make sure that the output included the required size. We don't know what
  // it is, so we rely on checking that it is just bigger than the "too small"
  // size.
  ASSERT_GT(secret_string_size_, 1u);
}

TEST_F(LicenseWhiteboxGetSecretStringTest, InvalidState) {
  // Purposely do not load a license so that we won't have any keys, causing
  // use to be in an invalid state.

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_INVALID_STATE);
}

}  // namespace widevine