// Copyright 2022 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.

#include "license_holder.h"

#include "license_request.h"
#include "message_dumper.h"
#include "oec_device_features.h"
#include "properties.h"
#include "test_base.h"

using video_widevine::SignedMessage;

namespace wvcdm {
namespace {
const std::string kCencMimeType = "cenc";
}  // namespace

LicenseHolder::~LicenseHolder() {}

void LicenseHolder::OpenSession() {
  CdmResponseType status = cdm_engine_->OpenSession(
      config_.key_system(), nullptr, &event_listener_, &session_id_);
  ASSERT_EQ(NO_ERROR, status);
  ASSERT_TRUE(cdm_engine_->IsOpenSession(session_id_));
}

void LicenseHolder::CloseSession() {
  CdmResponseType status = cdm_engine_->CloseSession(session_id_);
  ASSERT_EQ(NO_ERROR, status);
  ASSERT_FALSE(cdm_engine_->IsOpenSession(session_id_));
}

void LicenseHolder::FetchLicense() {
  video_widevine::WidevinePsshData pssh;
  pssh.set_content_id(content_id_);
  const std::string init_data_string = MakePSSH(pssh);
  const InitializationData init_data(kCencMimeType, init_data_string);
  init_data.DumpToLogs();
  CdmKeyRequest key_request;
  ASSERT_NO_FATAL_FAILURE(GenerateKeyRequest(init_data, &key_request))
      << "Failed for " << content_id();
  if (config_.dump_golden_data()) {
    MessageDumper::DumpLicenseRequest(key_request);
  }
  if (log_core_message_) MessageDumper::PrintLicenseRequest(key_request);
  ASSERT_NO_FATAL_FAILURE(GetKeyResponse(key_request))
      << "Failed for " << content_id();
  if (config_.dump_golden_data()) {
    MessageDumper::DumpLicense(key_response_);
  }
  if (pst_length_ > 0) {
    SignedMessage signed_message;
    ASSERT_TRUE(signed_message.ParseFromString(key_response_));

    video_widevine::License license;
    ASSERT_TRUE(license.ParseFromString(signed_message.msg()));

    // Verify that pst exists and matches |pst_length_|.
    ASSERT_TRUE(!license.id().provider_session_token().empty());
    ASSERT_EQ(pst_length_, license.id().provider_session_token().length());
  }
  if (log_core_message_) MessageDumper::PrintLicenseResponse(key_response_);
}

void LicenseHolder::LoadLicense() {
  CdmLicenseType license_type;
  ASSERT_EQ(KEY_ADDED, cdm_engine_->AddKey(session_id_, key_response_,
                                           &license_type, &key_set_id_))
      << "Failed to load license for " << content_id();
  if (can_persist_) {
    ASSERT_EQ(license_type, kLicenseTypeOffline)
        << "Failed for " << content_id();
  } else {
    ASSERT_EQ(license_type, kLicenseTypeStreaming)
        << "Failed for " << content_id();
  }
}

void LicenseHolder::FailLoadLicense() {
  CdmLicenseType license_type;
  ASSERT_NE(KEY_ADDED, cdm_engine_->AddKey(session_id_, key_response_,
                                           &license_type, &key_set_id_))
      << "Unexpected success loading license for " << content_id();
}

void LicenseHolder::ReloadLicense() {
  CdmResponseType status = cdm_engine_->RestoreKey(session_id_, key_set_id_);
  ASSERT_EQ(KEY_ADDED, status)
      << "Failed to reload license for " << content_id();
}

void LicenseHolder::FailReloadLicense() {
  const CdmResponseType status =
      cdm_engine_->RestoreKey(session_id_, key_set_id_);
  ASSERT_NE(KEY_ADDED, status)
      << "Unexpected success loading license for " << content_id();
}

void LicenseHolder::GenerateAndPostRenewalRequest(
    const std::string& renewal_policy_id) {
  event_listener_.set_renewal_needed(false);
  CdmKeyRequest request;
  const CdmResponseType result =
      cdm_engine_->GenerateRenewalRequest(session_id_, &request);
  ASSERT_EQ(KEY_MESSAGE, result) << "Failed for " << content_id();
  if (config_.dump_golden_data()) {
    MessageDumper::DumpRenewalRequest(request);
  }

  if (renewal_info_check_) {
    SignedMessage signed_message;
    ASSERT_TRUE(signed_message.ParseFromString(request.message));

    video_widevine::LicenseRequest renewal_request;
    ASSERT_TRUE(renewal_request.ParseFromString(signed_message.msg()));

    if (client_id_check_) {
      // Check that a renewal request where client ID is always included
      // has a client ID.
      ASSERT_TRUE(renewal_request.has_client_id());
    } else {
      // Check that a renewal request where can_persist = true does not include
      // the client ID.
      ASSERT_FALSE(renewal_request.has_client_id());
    }

    // Check that a renewal request where can_persist = true does not include
    // usage entry info.
    ASSERT_FALSE(renewal_request.content_id()
                     .existing_license()
                     .has_session_usage_table_entry());

    // Check that a renewal request where can_persist = true includes the PST.
    ASSERT_TRUE(renewal_request.content_id()
                    .existing_license()
                    .license_id()
                    .has_provider_session_token());
  }

  const std::string url = MakeUrl(config_.renewal_server(), renewal_policy_id);
  request_in_flight_.reset(new UrlRequest(url));
  ASSERT_TRUE(request_in_flight_->is_connected())
      << "Failed for " << content_id();
  request_in_flight_->PostRequest(request.message);
}

void LicenseHolder::FetchRenewal() {
  ASSERT_NE(request_in_flight_, nullptr) << "Failed for " << content_id();
  ASSERT_NO_FATAL_FAILURE(
      request_in_flight_->AssertOkResponseWithRetry(&request_response_))
      << "Renewal failed for " << content_id();
}

void LicenseHolder::LoadRenewal() {
  LicenseRequest license_request;
  license_request.GetDrmMessage(request_response_, request_message_);
  if (config_.dump_golden_data()) {
    MessageDumper::DumpRenewal(request_message_);
  }
  if (renewal_info_check_) {
    SignedMessage signed_message;
    ASSERT_TRUE(signed_message.ParseFromString(request_message_));

    video_widevine::License renewal;
    ASSERT_TRUE(renewal.ParseFromString(signed_message.msg()));

    // Verify that pst exist.
    ASSERT_TRUE(renewal.id().has_provider_session_token());
  }
  ASSERT_EQ(KEY_ADDED, cdm_engine_->RenewKey(session_id_, request_message_))
      << "Failed for " << content_id();
}

void LicenseHolder::GenerateAndPostReleaseRequest(
    const std::string& release_policy_id) {
  event_listener_.set_renewal_needed(false);
  CdmKeyRequest request;
  CdmAppParameterMap empty_app_parameters;
  video_widevine::WidevinePsshData pssh;
  pssh.set_content_id(content_id_);
  const std::string init_data_string = MakePSSH(pssh);
  const InitializationData init_data(kCencMimeType, init_data_string);
  init_data.DumpToLogs();

  CdmSessionId session_id;
  CdmKeySetId key_set_id;
  CdmResponseType result;
  // For Android when key set IDs are used, the key set ID passed in should have
  // a value and the session ID should be empty.
  if (!Properties::AlwaysUseKeySetIds()) {
    key_set_id = key_set_id_;
    result = cdm_engine_->OpenKeySetSession(key_set_id_, nullptr, nullptr);
    ASSERT_EQ(NO_ERROR, result) << "Failed for " << content_id();
    // For CE CDM, we only need the session ID to be valid.
  } else {
    session_id = session_id_;
  }
  result = cdm_engine_->GenerateKeyRequest(session_id, key_set_id, init_data,
                                           kLicenseTypeRelease,
                                           empty_app_parameters, &request);
  ASSERT_EQ(KEY_MESSAGE, result) << "Failed for " << content_id();
  if (!Properties::AlwaysUseKeySetIds()) {
    cdm_engine_->CloseKeySetSession(key_set_id_);
  }
  const std::string url = MakeUrl(config_.renewal_server(), release_policy_id);
  request_in_flight_.reset(new UrlRequest(url));
  ASSERT_TRUE(request_in_flight_->is_connected())
      << "Failed for " << content_id();
  request_in_flight_->PostRequest(request.message);
}

void LicenseHolder::FetchRelease() {
  ASSERT_NE(request_in_flight_, nullptr) << "Failed for " << content_id();
  ASSERT_NO_FATAL_FAILURE(
      request_in_flight_->AssertOkResponseWithRetry(&request_response_))
      << "Renewal failed for " << content_id();
}

void LicenseHolder::LoadRelease() {
  LicenseRequest license_request;
  license_request.GetDrmMessage(request_response_, request_message_);
}

void LicenseHolder::RemoveLicense() {
  ASSERT_EQ(NO_ERROR, cdm_engine_->RemoveLicense(session_id_))
      << "Failed for " << content_id();
}

CdmResponseType LicenseHolder::Decrypt(const std::string& key_id) {
  constexpr size_t buffer_size = 500;
  const std::vector<uint8_t> input(buffer_size, 0);
  std::vector<uint8_t> output(buffer_size, 0);
  const std::vector<uint8_t> iv(KEY_IV_SIZE, 0);
  CdmDecryptionParametersV16 params(key_id);
  params.is_secure = false;
  CdmDecryptionSample sample(input.data(), output.data(), 0, input.size(), iv);
  CdmDecryptionSubsample subsample(0, input.size());
  sample.subsamples.push_back(subsample);
  params.samples.push_back(sample);
  return cdm_engine_->DecryptV16(session_id_, params);
}

CdmResponseType LicenseHolder::DecryptClearLead(const std::string& key_id,
                                                size_t num_samples) {
  constexpr size_t buffer_size = 500;
  const std::vector<uint8_t> input(buffer_size, 0);
  std::vector<uint8_t> output(buffer_size, 0);
  const std::vector<uint8_t> iv(KEY_IV_SIZE, 0);
  CdmDecryptionParametersV16 params(key_id);
  params.is_secure = false;
  CdmDecryptionSubsample subsample(input.size(), 0);
  for (size_t i = 0; i < num_samples; i++) {
    CdmDecryptionSample sample(input.data(), output.data(), 0, input.size(),
                               iv);
    sample.subsamples.push_back(subsample);
    params.samples.push_back(sample);
  }
  return cdm_engine_->DecryptV16(session_id_, params);
}

void LicenseHolder::DecryptSecure(const KeyId& key_id) {
  ASSERT_TRUE(wvoec::global_features.test_secure_buffers);
  constexpr size_t buffer_size = 500;
  const std::vector<uint8_t> input(buffer_size, 0);
  const std::vector<uint8_t> iv(KEY_IV_SIZE, 0);

  // To create a secure buffer, we need to know the OEMCrypto session id.
  CdmQueryMap query_map;
  cdm_engine_->QueryOemCryptoSessionId(session_id_, &query_map);
  const std::string oec_session_id_string =
      query_map[QUERY_KEY_OEMCRYPTO_SESSION_ID];
  uint32_t oec_session_id = std::stoi(oec_session_id_string);

  int secure_buffer_fid;
  OEMCrypto_DestBufferDesc output_descriptor;
  output_descriptor.type = OEMCrypto_BufferType_Secure;
  output_descriptor.buffer.secure.secure_buffer_length = buffer_size;
  ASSERT_EQ(
      OEMCrypto_AllocateSecureBuffer(oec_session_id, buffer_size,
                                     &output_descriptor, &secure_buffer_fid),
      OEMCrypto_SUCCESS);

  ASSERT_NE(output_descriptor.buffer.secure.secure_buffer, nullptr);
  // It is OK if OEMCrypto changes the maximum size, but there must
  // still be enough room for our data.
  ASSERT_GE(output_descriptor.buffer.secure.secure_buffer_length, buffer_size);
  output_descriptor.buffer.secure.offset = 0;

  // Now create a sample array for the CDM layer.
  CdmDecryptionParametersV16 params(key_id);
  params.is_secure = true;
  CdmDecryptionSample sample(input.data(),
                             output_descriptor.buffer.secure.secure_buffer, 0,
                             input.size(), iv);
  CdmDecryptionSubsample subsample(0, input.size());
  sample.subsamples.push_back(subsample);
  params.samples.push_back(sample);
  CdmResponseType status = cdm_engine_->DecryptV16(session_id_, params);

  // Free the secure buffer before we check the return status.
  EXPECT_EQ(OEMCrypto_FreeSecureBuffer(oec_session_id, &output_descriptor,
                                       secure_buffer_fid),
            OEMCrypto_SUCCESS);

  ASSERT_EQ(status, NO_ERROR);
}

void LicenseHolder::FailDecrypt(const KeyId& key_id,
                                CdmResponseEnum expected_status) {
  CdmResponseType status = Decrypt(key_id);
  // If the server knows we cannot handle the key, it would not have given us
  // the key. In that case, the status should indicate no key.
  if (status != NEED_KEY) {
    // Otherwise, we should have gotten the expected error.
    ASSERT_EQ(expected_status, status);
  }
}

void LicenseHolder::GenerateKeyRequest(const InitializationData& init_data,
                                       CdmKeyRequest* key_request) {
  CdmAppParameterMap empty_app_parameters;
  CdmKeySetId empty_key_set_id;
  CdmLicenseType license_type =
      can_persist_ ? kLicenseTypeOffline : kLicenseTypeStreaming;
  CdmResponseType result = cdm_engine_->GenerateKeyRequest(
      session_id_, empty_key_set_id, init_data, license_type,
      empty_app_parameters, key_request);
  ASSERT_EQ(KEY_MESSAGE, result);
  ASSERT_EQ(kKeyRequestTypeInitial, key_request->type);
}

std::string LicenseHolder::MakeUrl(const std::string& server_url,
                                   const std::string& renewal_policy_id) {
  // For tests, we want to specify the policy, but the UAT server only allows us
  // to set the content id as the video_id. So each policy is matched to a
  // single license with the same name.  The local license server, on the other
  // hand, wants to see the policy id in the url. So we have to guess which
  // format to use based on the name of the server.
  const std::string path = server_url + config_.client_auth();
  if (!renewal_policy_id.empty()) {
    const std::string video_query =
        "video_id=" + content_id() + "&renewal_policy=" + renewal_policy_id;
    // If there is already a parameter, then we don't need to add another
    // question mark.
    return path + ((path.find('?') == std::string::npos) ? '?' : '&') +
           video_query;
  } else {
    return path;
  }
}

void LicenseHolder::GetKeyResponse(const CdmKeyRequest& key_request) {
  // The content id matches the policy id used on UAT.
  const std::string url = MakeUrl(config_.license_server(), "");
  UrlRequest url_request(url);
  ASSERT_TRUE(url_request.is_connected());

  std::string http_response;
  url_request.PostRequest(key_request.message);
  ASSERT_NO_FATAL_FAILURE(url_request.AssertOkResponseWithRetry(&http_response))
      << "Failed for " << content_id();
  LicenseRequest license_request;
  license_request.GetDrmMessage(http_response, key_response_);
}

}  // namespace wvcdm