// Copyright 2023 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//
// Test data for OEMCrypto unit tests.
//
#ifndef CDM_OEMCRYPTO_LICENSE_TEST_
#define CDM_OEMCRYPTO_LICENSE_TEST_

#include <gtest/gtest.h>

#include "OEMCryptoCENC.h"
#include "clock.h"
#include "log.h"
#include "oemcrypto_basic_test.h"
#include "oemcrypto_corpus_generator_helper.h"
#include "oemcrypto_resource_test.h"
#include "wvcrc32.h"

namespace wvoec {

using ::testing::WithParamInterface;

// Used for testing oemcrypto APIs with huge buffers.
typedef const std::function<OEMCryptoResult(size_t)> oemcrypto_function;
void TestHugeLengthDoesNotCrashAPI(oemcrypto_function f,
                                   size_t start_buffer_length,
                                   size_t end_buffer_length, bool check_status);
void TestHugeLengthDoesNotCrashAPI(oemcrypto_function f, bool check_status);

void TestMaxKeys(SessionUtil* util, size_t num_keys_per_session);

class OEMCryptoSessionTests : public OEMCryptoClientTest {
 public:
  vector<uint8_t> encrypted_usage_header_;

 protected:
  OEMCryptoSessionTests() {}

  void SetUp() override {
    OEMCryptoClientTest::SetUp();
    if (wvoec::global_features.derive_key_method ==
        wvoec::DeviceFeatures::NO_METHOD) {
      GTEST_SKIP() << "Test for devices that can derive session keys only.";
    }
    EnsureTestROT();
    if (global_features.usage_table) {
      CreateUsageTableHeader();
    }
  }

  void CreateUsageTableHeader(bool expect_success = true) {
    size_t header_buffer_length = 0;
    OEMCryptoResult sts =
        OEMCrypto_CreateUsageTableHeader(nullptr, &header_buffer_length);
    if (expect_success) {
      ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
    } else {
      ASSERT_NE(OEMCrypto_SUCCESS, sts);
      if (sts != OEMCrypto_ERROR_SHORT_BUFFER) return;
    }
    encrypted_usage_header_.resize(header_buffer_length);
    sts = OEMCrypto_CreateUsageTableHeader(encrypted_usage_header_.data(),
                                           &header_buffer_length);
    if (expect_success) {
      ASSERT_EQ(OEMCrypto_SUCCESS, sts);
      encrypted_usage_header_.resize(header_buffer_length);
    } else {
      ASSERT_NE(OEMCrypto_SUCCESS, sts);
    }
  }

  void TestPrepareLicenseRequestForHugeBufferLengths(
      const std::function<void(size_t, LicenseRoundTrip*)> f,
      bool check_status) {
    auto oemcrypto_function = [&](size_t message_length) {
      Session s;
      s.open();
      InstallTestDrmKey(&s);
      LicenseRoundTrip license_messages(&s);
      f(message_length, &license_messages);
      OEMCryptoResult result =
          license_messages.SignAndCreateRequestWithCustomBufferLengths();
      s.close();
      return result;
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, check_status);
  }

  OEMCryptoResult LoadLicense(Session& s, LicenseRoundTrip& license_messages) {
    InstallTestDrmKey(&s);
    license_messages.SignAndVerifyRequest();
    license_messages.CreateDefaultResponse();
    license_messages.EncryptAndSignResponse();
    return license_messages.LoadResponse();
  }
};

// This class is for testing a single license with the default API version
// of 16.
class OEMCryptoLicenseTestAPI16 : public OEMCryptoSessionTests {
 public:
  OEMCryptoLicenseTestAPI16()
      : license_api_version_(kCurrentAPI), license_messages_(&session_) {}

  void SetUp() override {
    OEMCryptoSessionTests::SetUp();
    ASSERT_NO_FATAL_FAILURE(session_.open());
    ASSERT_NO_FATAL_FAILURE(InstallTestDrmKey(&session_));
  }

  void TearDown() override {
    ASSERT_NO_FATAL_FAILURE(session_.close());
    OEMCryptoSessionTests::TearDown();
  }

 protected:
  Session session_;
  uint32_t license_api_version_;
  LicenseRoundTrip license_messages_;
};

// This class is used to test a license that is from a server with the specified
// version parameter. Up to two versions old.
class OEMCryptoLicenseTest : public OEMCryptoLicenseTestAPI16,
                             public WithParamInterface<uint32_t> {
 protected:
  void SetUp() override {
    // The only difference between this class and its parent is that we use a
    // different license api:
    license_api_version_ = GetParam();
    license_messages_.set_api_version(license_api_version_);
    OEMCryptoLicenseTestAPI16::SetUp();
  }

  void LoadLicense() {
    license_messages_.SignAndVerifyRequest();
    license_messages_.CreateDefaultResponse();
    license_messages_.EncryptAndSignResponse();
    license_messages_.LoadResponse();
  }

  void TestDecryptCENCForHugeBufferLengths(
      const std::function<void(size_t, OEMCrypto_SampleDescription*)> f,
      bool check_status) {
    LoadLicense();
    auto oemcrypto_function = [&](size_t message_length) {
      vector<uint8_t> key_handle;
      GetKeyHandleIntoVector(session_.session_id(),
                             session_.license().keys[0].key_id,
                             session_.license().keys[0].key_id_length,
                             OEMCrypto_CipherMode_CENC, key_handle);

      size_t input_buffer_size = 1;
      vector<uint8_t> in_buffer(input_buffer_size + message_length);
      vector<uint8_t> out_buffer(in_buffer.size());

      OEMCrypto_SampleDescription sample_description;
      OEMCrypto_SubSampleDescription subsample_description;
      GenerateSimpleSampleDescription(
          in_buffer, out_buffer, &sample_description, &subsample_description);

      OEMCrypto_SubSampleDescription* sub_samples =
          const_cast<OEMCrypto_SubSampleDescription*>(
              sample_description.subsamples);
      // Actual tests modifies either of these fields to match clear + encrypted
      // = in_buffer.size().
      sub_samples[0].num_bytes_clear = 0;
      sub_samples[0].num_bytes_encrypted = input_buffer_size;

      // Create the pattern description (always 0,0 for CTR)
      OEMCrypto_CENCEncryptPatternDesc pattern = {0, 0};
      int secure_fd = 0;
      f(message_length, &sample_description);
      if (sample_description.buffers.output_descriptor.type ==
          OEMCrypto_BufferType_Secure) {
        OEMCryptoResult sts = OEMCrypto_AllocateSecureBuffer(
            session_.session_id(), in_buffer.size(),
            &sample_description.buffers.output_descriptor, &secure_fd);
        if (sts != OEMCrypto_SUCCESS) {
          LOGI("Secure buffers are not supported.");
          return sts;
        }
      }
      // Try to decrypt the data
      OEMCryptoResult result =
          OEMCrypto_DecryptCENC(key_handle.data(), key_handle.size(),
                                &sample_description, 1, &pattern);
      if (sample_description.buffers.output_descriptor.type ==
          OEMCrypto_BufferType_Secure) {
        OEMCrypto_FreeSecureBuffer(
            session_.session_id(),
            &sample_description.buffers.output_descriptor, secure_fd);
      }
      return result;
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, check_status);
  }

  void TestDecryptCENCForOutOfRangeOffsetsAndLengths(
      const std::function<void(OEMCrypto_SampleDescription*)> f,
      bool update_secure_buffer) {
    LoadLicense();
    vector<uint8_t> key_handle;
    GetKeyHandleIntoVector(session_.session_id(),
                           session_.license().keys[0].key_id,
                           session_.license().keys[0].key_id_length,
                           OEMCrypto_CipherMode_CENC, key_handle);

    vector<uint8_t> in_buffer(256);
    vector<uint8_t> out_buffer(in_buffer.size());

    OEMCrypto_SampleDescription sample_description;
    OEMCrypto_SubSampleDescription subsample_description;
    GenerateSimpleSampleDescription(in_buffer, out_buffer, &sample_description,
                                    &subsample_description);

    // Create the pattern description (always 0,0 for CTR)
    OEMCrypto_CENCEncryptPatternDesc pattern = {0, 0};
    int secure_fd = 0;
    if (update_secure_buffer) {
      OEMCryptoResult sts = OEMCrypto_AllocateSecureBuffer(
          session_.session_id(), in_buffer.size(),
          &sample_description.buffers.output_descriptor, &secure_fd);
      if (sts != OEMCrypto_SUCCESS) {
        LOGI("Secure buffers are not supported.");
        return;
      }
    }
    f(&sample_description);
    // Try to decrypt the data
    OEMCryptoResult result = OEMCrypto_DecryptCENC(
        key_handle.data(), key_handle.size(), &sample_description, 1, &pattern);
    if (update_secure_buffer) {
      OEMCrypto_FreeSecureBuffer(session_.session_id(),
                                 &sample_description.buffers.output_descriptor,
                                 secure_fd);
    }
    ASSERT_NE(OEMCrypto_SUCCESS, result);
  }
};

// Test usage table functionality.
class LicenseWithUsageEntry {
 public:
  LicenseWithUsageEntry(const std::string& pst = "my_pst")
      : session_(),
        license_messages_(&session_),
        generic_crypto_(false),
        time_license_received_(0),
        time_first_decrypt_(0),
        time_last_decrypt_(0),
        active_(true) {
    license_messages_.set_pst(pst);
  }

  void MakeAndLoadOnline(OEMCryptoSessionTests* test) {
    MakeAndLoad(test,
                wvoec::kControlNonceEnabled | wvoec::kControlNonceRequired);
  }

  // If status in not a nullptr, then creating a new entry is allowed to fail,
  // and its error code is stored in status.
  void MakeOfflineAndClose(OEMCryptoSessionTests* test,
                           OEMCryptoResult* status = nullptr) {
    MakeAndLoad(test, wvoec::kControlNonceOrEntry, status);
    if (status != nullptr && *status != OEMCrypto_SUCCESS) {
      ASSERT_NO_FATAL_FAILURE(session_.close());
      return;
    }
    ASSERT_NO_FATAL_FAILURE(
        session_.UpdateUsageEntry(&(test->encrypted_usage_header_)));
    ASSERT_NO_FATAL_FAILURE(GenerateVerifyReport(kUnused));
    ASSERT_NO_FATAL_FAILURE(session_.close());
  }

  // If status in not a nullptr, then creating a new entry is allowed to fail,
  // and its error code is stored in status.
  void MakeAndLoad(SessionUtil* util, uint32_t control,
                   OEMCryptoResult* status = nullptr) {
    license_messages_.set_control(control);
    ASSERT_NO_FATAL_FAILURE(session_.open());
    ASSERT_NO_FATAL_FAILURE(util->InstallTestDrmKey(&session_));
    ASSERT_NO_FATAL_FAILURE(license_messages_.SignAndVerifyRequest());
    if (generic_crypto_) {
      ASSERT_NO_FATAL_FAILURE(
          license_messages_.CreateResponseWithGenericCryptoKeys());
    } else {
      ASSERT_NO_FATAL_FAILURE(license_messages_.CreateDefaultResponse());
    }
    ASSERT_NO_FATAL_FAILURE(license_messages_.EncryptAndSignResponse());
    ASSERT_NO_FATAL_FAILURE(session_.CreateNewUsageEntry(status));
    if (status != nullptr && *status != OEMCrypto_SUCCESS) return;
    ASSERT_EQ(OEMCrypto_SUCCESS, license_messages_.LoadResponse());
    time_license_received_ = wvutil::Clock().GetCurrentTime();
  }

  void OpenAndReload(SessionUtil* util) {
    ASSERT_NO_FATAL_FAILURE(session_.open());
    ASSERT_NO_FATAL_FAILURE(session_.ReloadUsageEntry());
    ASSERT_NO_FATAL_FAILURE(util->InstallTestDrmKey(&session_));
    ASSERT_NO_FATAL_FAILURE(session_.GenerateDerivedKeysFromSessionKey());
    ASSERT_EQ(OEMCrypto_SUCCESS, license_messages_.LoadResponse());
  }

  // Test decrypt, and update the decrypt times for the pst report.
  void TestDecryptCTR(bool select_key_first = true,
                      OEMCryptoResult expected_result = OEMCrypto_SUCCESS) {
    session_.TestDecryptCTR(select_key_first, expected_result);
    time_last_decrypt_ = wvutil::Clock().GetCurrentTime();
    if (time_first_decrypt_ == 0) time_first_decrypt_ = time_last_decrypt_;
  }

  void DeactivateUsageEntry() {
    active_ = false;
    if (ShouldGenerateCorpus()) {
      const std::string file_name =
          GetFileName("oemcrypto_deactivate_usage_entry_fuzz_seed_corpus");
      AppendToFile(file_name, pst().c_str(), pst().length());
    }
    ASSERT_EQ(
        OEMCrypto_SUCCESS,
        OEMCrypto_DeactivateUsageEntry(
            session_.session_id(),
            reinterpret_cast<const uint8_t*>(pst().c_str()), pst().length()));
  }

  void GenerateVerifyReport(OEMCrypto_Usage_Entry_Status status) {
    ASSERT_NO_FATAL_FAILURE(session_.GenerateReport(pst()));
    Test_PST_Report expected(pst(), status);
    ASSERT_NO_FATAL_FAILURE(
        session_.VerifyReport(std::move(expected), time_license_received_,
                              time_first_decrypt_, time_last_decrypt_));
    // The PST report was signed above.  Below we verify that the entire message
    // that is sent to the server will be signed by the right mac keys.
    RenewalRoundTrip renewal_messages(&license_messages_);
    renewal_messages.set_is_release(!active_);
    ASSERT_NO_FATAL_FAILURE(renewal_messages.SignAndVerifyRequest());
  }

  void ReloadUsageEntry() {
    session_.ReloadUsageEntry();
    session_.set_mac_keys(license_messages_.response_data().mac_keys);
  }

  const std::string& pst() const { return license_messages_.pst(); }
  void set_pst(const std::string& pst) { license_messages_.set_pst(pst); }
  LicenseRoundTrip& license_messages() { return license_messages_; }
  Session& session() { return session_; }
  void set_generic_crypto(bool generic_crypto) {
    generic_crypto_ = generic_crypto;
  }

 private:
  Session session_;
  LicenseRoundTrip license_messages_;
  bool generic_crypto_;
  int64_t time_license_received_;
  int64_t time_first_decrypt_;
  int64_t time_last_decrypt_;
  bool active_;
};

class OEMCryptoRefreshTest : public OEMCryptoLicenseTest {
 protected:
  void SetUp() override {
    OEMCryptoLicenseTest::SetUp();
    // These values allow us to run a few simple duration tests or just start
    // playback right away. All times are in seconds since the license was
    // signed.
    // Soft expiry false means timers are strictly enforce.
    timer_limits_.soft_enforce_rental_duration = true;
    timer_limits_.soft_enforce_playback_duration = false;
    // Playback may begin immediately.
    timer_limits_.earliest_playback_start_seconds = 0;
    // First playback may be within the first two seconds.
    timer_limits_.rental_duration_seconds = kDuration;
    // Once started, playback may last two seconds without a renewal.
    timer_limits_.initial_renewal_duration_seconds = kDuration;
    // Total playback is not limited.
    timer_limits_.total_playback_duration_seconds = 0;
  }

  void LoadLicense() {
    license_messages_.core_response().timer_limits = timer_limits_;
    ASSERT_NO_FATAL_FAILURE(license_messages_.SignAndVerifyRequest());
    ASSERT_NO_FATAL_FAILURE(license_messages_.CreateDefaultResponse());
    ASSERT_NO_FATAL_FAILURE(license_messages_.EncryptAndSignResponse());
    ASSERT_EQ(OEMCrypto_SUCCESS, license_messages_.LoadResponse());
  }

  void MakeRenewalRequest(RenewalRoundTrip* renewal_messages) {
    ASSERT_NO_FATAL_FAILURE(renewal_messages->SignAndVerifyRequest());
    ASSERT_NO_FATAL_FAILURE(renewal_messages->CreateDefaultResponse());
  }

  void LoadRenewal(RenewalRoundTrip* renewal_messages,
                   OEMCryptoResult expected_result) {
    ASSERT_NO_FATAL_FAILURE(renewal_messages->EncryptAndSignResponse());
    ASSERT_EQ(expected_result, renewal_messages->LoadResponse());
  }

  void MakeReleaseRequest(ReleaseRoundTrip* release_messages) {
    ASSERT_NO_FATAL_FAILURE(release_messages->SignAndVerifyRequest());
    ASSERT_NO_FATAL_FAILURE(release_messages->CreateDefaultResponse());
  }

  void LoadRelease(ReleaseRoundTrip* release_messages,
                   OEMCryptoResult expected_result) {
    ASSERT_NO_FATAL_FAILURE(release_messages->EncryptAndSignResponse());
    ASSERT_EQ(expected_result, release_messages->LoadResponse());
  }

  ODK_TimerLimits timer_limits_;
};

// This class is for the refresh tests that should only be run on licenses with
// a core message.
class OEMCryptoRefreshTestAPI16 : public OEMCryptoRefreshTest {};

}  // namespace wvoec

#endif  // CDM_OEMCRYPTO_LICENSE_TEST_