// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox_benchmark.h"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <vector>

#include "api/license_whitebox.h"
#include "api/test_license_builder.h"
#include "api/test_license_whitebox_keys.h"
#include "base/check_op.h"
#include "benchmarking/data_source.h"
#include "crypto_utils/crypto_util.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace widevine {
namespace {
constexpr size_t kBlockSize = 16;  // This must align with the AES block size.
}  // namespace

void LicenseWhiteboxBenchmark::SetUp() {
  // We use size=4 for all our test key ids.
  key_data_.id = data_source_.Get<4>();
  key_data_.key = data_source_.Get<kBlockSize>();
  // Use secure crypto as it will work with both Decrypt() and MaskedDecrypt().
  key_data_.level = SecurityLevel::kSoftwareSecureCrypto;

  iv_ = data_source_.Get<kBlockSize>();

  const auto public_key_data = GetEncryptionKey().public_key;
  encryption_public_key_.reset(RsaPublicKey::Create(
      std::string(public_key_data.begin(), public_key_data.end())));
  ASSERT_TRUE(encryption_public_key_);
}

License LicenseWhiteboxBenchmark::CreateLicense(WB_LicenseKeyMode key_mode,
                                                int number_of_content_keys,
                                                SecurityLevel security_level,
                                                size_t provider_key_id) {
  TestLicenseBuilder::Settings settings;
  settings.provider_key_id = provider_key_id;

  TestLicenseBuilder license_builder;
  license_builder.SetSettings(settings);
  license_builder.AddSigningKey(TestLicenseBuilder::DefaultSigningKey());

  // Update the first key to match the security level.
  key_data_.level = security_level;
  license_builder.AddContentKey(key_data_);

  CHECK_GE(number_of_content_keys, 1);
  for (int i = 1; i < number_of_content_keys; ++i) {
    ContentKeyData key;
    key.id = data_source_.Get<4>();
    key.key = data_source_.Get<kBlockSize>();
    key.level = security_level;
    license_builder.AddContentKey(key);
  }

  std::unique_ptr<TestServer> server;

  switch (key_mode) {
    case WB_LICENSE_KEY_MODE_SINGLE_KEY:
      server = TestServer::CreateSingleKey();
      break;
    case WB_LICENSE_KEY_MODE_DUAL_KEY:
      server = TestServer::CreateDualKey();
      break;
    default:
      CHECK(false) << "Could not create test server, Unknown key mode.";
      break;
  }

  License license;
  license_builder.Build(*server, &license);

  return license;
}

std::vector<uint8_t> LicenseWhiteboxBenchmark::SignAsServer(
    const std::vector<uint8_t>& message) const {
  // The server key is the first half of the signing key.
  const auto key = TestLicenseBuilder::DefaultSigningKey();
  const std::string server_key(key.begin(),
                               key.begin() + crypto_util::kSigningKeySizeBytes);

  // crypto util uses strings, so we will need to convert the result back to a
  // vector before we return it.
  const auto signature = crypto_util::CreateSignatureHmacSha256(
      server_key, std::string(message.begin(), message.end()));
  return std::vector<uint8_t>(signature.begin(), signature.end());
}

}  // namespace widevine