// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox_benchmark.h"

#include <stddef.h>
#include <stdint.h>

#include <memory>
#include <vector>

#include "api/license_whitebox.h"
#include "api/test_data.h"
#include "api/test_license_builder.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() {
  key_id_ = data_source_.Get(8);  // The id size is not meaningful.
  key_ = data_source_.Get(kBlockSize);
  iv_ = data_source_.Get(kBlockSize);

  const auto public_key_data = GetMatchingLicensePublicKey();
  public_key_.reset(widevine::RsaPublicKey::Create(
      std::string(public_key_data.begin(), public_key_data.end())));
  ASSERT_TRUE(public_key_);
}

License LicenseWhiteboxBenchmark::CreateLicense() const {
  widevine::TestLicenseBuilder license_builder;
  license_builder.AddSigningKey(TestLicenseBuilder::DefaultSigningKey());
  // Use secure crypto as it will work with both Decrypt() and
  // MaskedDecrypt().
  license_builder.AddContentKey(
      video_widevine::License_KeyContainer_SecurityLevel_SW_SECURE_CRYPTO,
      key_id_, key_);

  widevine::License license;
  license_builder.Build(*public_key_, &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