// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/test_license_builder.h"

#include <ctime>

#include "base/logging.h"
#include "cdm/keys/certs.h"
#include "crypto_utils/aes_cbc_encryptor.h"
#include "crypto_utils/crypto_util.h"
#include "crypto_utils/sha_util.h"

namespace widevine {

namespace {

void InitializeRequest(video_widevine::LicenseRequest* request) {
  request->set_request_time(std::time(nullptr));  // Use time=now.

  auto* client_id = request->mutable_client_id();
  auto* capabilities = client_id->mutable_client_capabilities();
  capabilities->set_video_resolution_constraints(true);
  capabilities->set_client_token(true);
  capabilities->set_session_token(false);
  capabilities->set_max_hdcp_version(
      video_widevine::ClientIdentification::ClientCapabilities::HDCP_V1);

  auto* client_info = client_id->add_client_info();
  client_info->set_name("architecture_name");
  client_info->set_value("x86-64");
  client_info = client_id->add_client_info();
  client_info->set_name("company_name");
  client_info->set_value("Google");
  client_info = client_id->add_client_info();
  client_info->set_name("model_name");
  client_info->set_value("ChromeCDM");
  client_info = client_id->add_client_info();
  client_info->set_name("platform_name");
  client_info->set_value("Windows");
  client_info = client_id->add_client_info();
  client_info->set_name("widevine_cdm_version");
  client_info->set_value("4.10.1686.29");

  client_id->set_type(
      video_widevine::ClientIdentification_TokenType_DRM_DEVICE_CERTIFICATE);
  client_id->set_token(wvcdm::kRsaDrmCertificate,
                       wvcdm::kRsaDrmCertificateSize);

  auto* content_id = request->mutable_content_id();
  auto* webm_key_id = content_id->mutable_webm_key_id();
  webm_key_id->set_license_type(video_widevine::STREAMING);
  webm_key_id->set_request_id("REQUEST_ID");
  webm_key_id->set_header("01234567890123456");

  request->set_protocol_version(video_widevine::VERSION_2_1);
  request->set_type(video_widevine::LicenseRequest::NEW);
}

void InitializeResponse(const video_widevine::LicenseRequest& request,
                        video_widevine::License* response) {
  auto* id = response->mutable_id();
  id->set_request_id("REQUEST_ID");
  id->set_session_id("SESSION_ID");
  id->set_type(video_widevine::STREAMING);
  id->set_version(0);

  auto* policy = response->mutable_policy();
  policy->set_can_play(true);
  policy->set_can_persist(false);
  policy->set_can_renew(true);
  policy->set_license_duration_seconds(600);
  policy->set_renewal_delay_seconds(30);
  policy->set_renewal_retry_interval_seconds(10);
  policy->set_renew_with_usage(false);

  response->set_license_start_time(request.request_time());
  response->set_remote_attestation_verified(false);
  response->set_platform_verification_status(
      video_widevine::PlatformVerificationStatus::PLATFORM_UNVERIFIED);
}

std::string EncryptKey(const std::string& key,
                       const std::string& iv,
                       const std::vector<uint8_t>& plaintext) {
  AesCbcEncryptor encryptor;
  encryptor.SetKey(reinterpret_cast<const uint8_t*>(key.data()), key.size());

  std::vector<uint8_t> ciphertext(plaintext.size());
  CHECK(encryptor.Encrypt(reinterpret_cast<const uint8_t*>(iv.data()),
                          iv.size(), plaintext.data(), plaintext.size(),
                          ciphertext.data()));

  return std::string(ciphertext.begin(), ciphertext.end());
}

std::string DeriveIV(const std::vector<uint8_t>& context) {
  const std::string context_str(context.begin(), context.end());
  return crypto_util::DeriveIv(context_str);
}

}  // namespace

// static
std::vector<uint8_t> TestLicenseBuilder::NoPadding() {
  return {};
}

// static
std::vector<uint8_t> TestLicenseBuilder::PKSC8Padding() {
  return {
      0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10,
      0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10,
  };
}

// static
std::vector<uint8_t> TestLicenseBuilder::DefaultSigningKey() {
  return {
      0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15,
      16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
      32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
      48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
  };
}

TestLicenseBuilder::TestLicenseBuilder() {
  DCHECK_EQ(session_key_.size(), 16u);

  // Initialize the request and the response with the static fields that will
  // be common across all licences.
  InitializeRequest(&request_);
  InitializeResponse(request_, &response_);

  serialized_request_ = request_.SerializeAsString();
  container_key_ = crypto_util::DeriveKey(
      session_key_, crypto_util::kWrappingKeyLabel, serialized_request_,
      crypto_util::kWrappingKeySizeBits);
}

void TestLicenseBuilder::AddSigningKey(const std::vector<uint8_t>& key,
                                       const std::vector<uint8_t>& padding) {
  DCHECK_EQ(key.size(), 64u);  // 512 bits

  auto* container = response_.add_key();
  container->set_type(video_widevine::License_KeyContainer_KeyType_SIGNING);

  // To avoid having to define a key iv for each key, derive a key iv from the
  // key. This will allows us to have a different IVs between keys but keep it
  // deterministic.
  const auto key_iv = DeriveIV(key);
  container->set_iv(key_iv);

  std::vector<uint8_t> final_key = key;
  final_key.insert(final_key.end(), padding.begin(), padding.end());
  container->set_key(EncryptKey(container_key_, key_iv, final_key));
}

void TestLicenseBuilder::AddStubbedContentKey() {
  auto* container = response_.add_key();

  container->set_type(video_widevine::License_KeyContainer_KeyType_CONTENT);
  container->set_id("stubbed-content-key");
  container->set_level(
      video_widevine::License_KeyContainer_SecurityLevel_SW_SECURE_CRYPTO);
  container->set_iv("0000000000000000");

  // We don't bother encrypting the key, it should never be used and there is no
  // way to verify it.
  container->set_key("0000000000000000");
}

void TestLicenseBuilder::AddContentKey(
    video_widevine::License_KeyContainer_SecurityLevel level,
    const std::vector<uint8_t>& key_id,
    const std::vector<uint8_t>& key,
    const std::vector<uint8_t>& padding) {
  DCHECK_GT(key_id.size(), 0u);
  DCHECK_EQ(key.size(), 16u);

  auto* container = response_.add_key();

  container->set_type(video_widevine::License_KeyContainer_KeyType_CONTENT);
  container->set_id(key_id.data(), key_id.size());
  container->set_level(level);

  // To avoid having to define a key iv for each key, derive a key iv from the
  // key. This will allows us to have a different IVs between keys but keep it
  // deterministic.
  const auto key_iv = DeriveIV(key);
  container->set_iv(key_iv);

  std::vector<uint8_t> final_key = key;
  final_key.insert(final_key.end(), padding.begin(), padding.end());
  container->set_key(EncryptKey(container_key_, key_iv, final_key));
}

void TestLicenseBuilder::AddOperatorSessionKey(
    const std::vector<uint8_t>& key_id) {
  DCHECK_GT(key_id.size(), 0u);

  // We only set the type and id because the key should not actually be used.
  auto* container = response_.add_key();
  container->set_type(
      video_widevine::License_KeyContainer_KeyType_OPERATOR_SESSION);
  container->set_id(key_id.data(), key_id.size());
}

void TestLicenseBuilder::SetRemoteAttestation(RemoteAttestation setting) {
  switch (setting) {
    case RemoteAttestation::kUnavailable:
      response_.clear_remote_attestation_verified();
      break;
    case RemoteAttestation::kVerified:
      response_.set_remote_attestation_verified(true);
      break;
    case RemoteAttestation::kUnverified:
      response_.set_remote_attestation_verified(false);
      break;
  }
}

void TestLicenseBuilder::SetVerificationStatus(VerificationStatus setting) {
  switch (setting) {
    case VerificationStatus::kUnavailable:
      response_.clear_platform_verification_status();
      break;
    case VerificationStatus::kHardwareVerified:
      response_.set_platform_verification_status(
          video_widevine::PLATFORM_HARDWARE_VERIFIED);
      break;
    case VerificationStatus::kOther:
      response_.set_platform_verification_status(
          video_widevine::PLATFORM_UNVERIFIED);
      break;
  }
}

void TestLicenseBuilder::Build(const RsaPublicKey& public_key,
                               License* license) const {
  DCHECK(license);

  const std::string message_str = response_.SerializeAsString();

  std::string signing_key = crypto_util::DeriveKey(
      session_key_, crypto_util::kSigningKeyLabel, serialized_request_,
      crypto_util::kSigningKeySizeBits * 2);
  signing_key.resize(crypto_util::kSigningKeySizeBytes);

  const std::string signature_str =
      crypto_util::CreateSignatureHmacSha256(signing_key, message_str);

  std::string session_key_str;
  CHECK(public_key.Encrypt(session_key_, &session_key_str));

  license->request = std::vector<uint8_t>(serialized_request_.begin(),
                                          serialized_request_.end());
  license->message =
      std::vector<uint8_t>(message_str.begin(), message_str.end());
  license->signature =
      std::vector<uint8_t>(signature_str.begin(), signature_str.end());
  license->session_key =
      std::vector<uint8_t>(session_key_str.begin(), session_key_str.end());
}
}  // namespace widevine