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

// These tests perform various end-to-end actions similar to what an application
// would do. They verify that policies specified on UAT are honored on the
// device.

#include <openssl/err.h>
#include <openssl/rsa.h>
#include <openssl/x509.h>
#include <string>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "cdm_engine.h"
#include "license_holder.h"
#include "log.h"
#include "oec_device_features.h"
#include "properties.h"
#include "provisioning_holder.h"
#include "test_base.h"
#include "test_printers.h"
#include "test_sleep.h"
#include "wv_cdm_types.h"

const uint64_t kLargePstLength = 127;

namespace wvcdm {
// Core Policy Integration Test
class CorePIGTest : public WvCdmTestBaseWithEngine {
 protected:
  void SetUp() override {
    WvCdmTestBase::SetUp();
    EnsureProvisioned();
  }
};

/*
 * Check the version of the response returned by the server.
 */
TEST_F(CorePIGTest, PrintClientAndServerVersionNumber) {
  LicenseHolder holder("CDM_OfflineNoNonce", &cdm_engine_, config_);
  holder.set_can_persist(true);
  const KeyId key_id = "0000000000000000";

  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.set_log_core_message(true));
  ASSERT_NO_FATAL_FAILURE(holder.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder.LoadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());

  if (wvoec::global_features.provisioning_method == OEMCrypto_DrmCertificate) {
    GTEST_SKIP() << "Test not for DRM certificate provisioning method.";
  }
  std::string level;
  ASSERT_EQ(
      NO_ERROR,
      cdm_engine_.QueryStatus(kLevelDefault, QUERY_KEY_SECURITY_LEVEL, &level)
          .code());
  CdmSecurityLevel security_level = level == QUERY_VALUE_SECURITY_LEVEL_L1
                                        ? kSecurityLevelL1
                                        : kSecurityLevelL3;
  std::string prov_model;
  ASSERT_EQ(NO_ERROR,
            cdm_engine_.QueryStatus(kLevelDefault, QUERY_KEY_PROVISIONING_MODEL,
                                    &prov_model));

  ASSERT_NO_FATAL_FAILURE(cdm_engine_.Unprovision(security_level));
  ASSERT_FALSE(cdm_engine_.IsProvisioned(security_level));
  ProvisioningHolder provisioner(&cdm_engine_, config_);
  ASSERT_NO_FATAL_FAILURE(provisioner.set_log_core_message(true));
  if (prov_model == QUERY_VALUE_BOOT_CERTIFICATE_CHAIN) {
    ASSERT_NO_FATAL_FAILURE(provisioner.Provision(binary_provisioning_));
  }
  ASSERT_NO_FATAL_FAILURE(provisioner.Provision(binary_provisioning_));
}

/**
 * An offline license with nonce not required.
 */
TEST_F(CorePIGTest, OfflineNoNonce) {
  LicenseHolder holder("CDM_OfflineNoNonce", &cdm_engine_, config_);
  holder.set_can_persist(true);
  const KeyId key_id = "0000000000000000";

  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder.LoadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
  // Should be able to close the previous session, open a new session,
  // and reload the license.
  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.ReloadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
}

/**
 * An offline license with nonce and provider session token.
 */
TEST_F(CorePIGTest, OfflineWithPST) {
  LicenseHolder holder("CDM_OfflineWithPST", &cdm_engine_, config_);
  holder.set_can_persist(true);
  const KeyId key_id = "0000000000000000";

  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder.LoadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
  // Should be able to close the previous session, open a new session,
  // and reload the license.
  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.ReloadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
}

/**
 * An offline license with nonce and large provider session token.
 */
TEST_F(CorePIGTest, OfflineWithLargePST) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Test for usage table devices only.";
  }
  LicenseHolder holder("CDM_OfflineWithPST_Length_127", &cdm_engine_, config_);
  holder.set_can_persist(true);
  const KeyId key_id = "0000000000000000";

  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  holder.set_pst_length(kLargePstLength);
  ASSERT_NO_FATAL_FAILURE(holder.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder.LoadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
  // Should be able to close the previous session, open a new session,
  // and reload the license.
  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.ReloadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
}

TEST_F(CorePIGTest, OfflineMultipleLicensesWithDefrag) {
  const KeyId key_id = "0000000000000000";

  // 1. Open a session, load license, close session
  LicenseHolder holder1("CDM_OfflineWithPST", &cdm_engine_, config_);
  holder1.set_can_persist(true);
  ASSERT_NO_FATAL_FAILURE(holder1.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder1.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder1.LoadLicense());
  ASSERT_NO_FATAL_FAILURE(holder1.CloseSession());

  // 2. Open a session, load license, keep session open
  LicenseHolder holder2("CDM_OfflineWithPST", &cdm_engine_, config_);
  holder2.set_can_persist(true);
  ASSERT_NO_FATAL_FAILURE(holder2.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder2.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder2.LoadLicense());

  // 3. Remove first license
  ASSERT_NO_FATAL_FAILURE(holder1.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder1.ReloadLicense());
  ASSERT_NO_FATAL_FAILURE(holder1.RemoveLicense());
  ASSERT_NO_FATAL_FAILURE(holder1.CloseSession());

  // 4. Open a session, load license
  LicenseHolder holder3("CDM_OfflineWithPST", &cdm_engine_, config_);
  holder3.set_can_persist(true);
  ASSERT_NO_FATAL_FAILURE(holder3.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder3.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder3.LoadLicense());
  EXPECT_EQ(NO_ERROR, holder3.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder3.CloseSession());

  ASSERT_NO_FATAL_FAILURE(holder2.CloseSession());

  // Ensure first offline license can no longer be used
  ASSERT_NO_FATAL_FAILURE(holder1.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder1.FailReloadLicense());
  EXPECT_NE(NO_ERROR, holder1.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder1.CloseSession());

  // Ensure second and third offline licenses can be used
  ASSERT_NO_FATAL_FAILURE(holder2.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder2.ReloadLicense());
  EXPECT_EQ(NO_ERROR, holder2.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder2.RemoveLicense());
  ASSERT_NO_FATAL_FAILURE(holder2.CloseSession());

  ASSERT_NO_FATAL_FAILURE(holder3.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder3.ReloadLicense());
  EXPECT_EQ(NO_ERROR, holder3.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder3.RemoveLicense());
  ASSERT_NO_FATAL_FAILURE(holder3.CloseSession());
}

/**
 * This test verifies that the system can download and install license with a
 * key that requires secure buffers. It also verifies that we cannot decrypt to
 * a non-secure buffer using this key, but that we can decrypt to a secure
 * buffer, if the test harness supports secure buffers.
 */
TEST_F(CorePIGTest, OfflineHWSecureRequired) {
  LicenseHolder holder("CDM_OfflineHWSecureRequired", &cdm_engine_, config_);
  holder.set_can_persist(true);
  const KeyId sw_key_id = "0000000000000000";
  const KeyId hw_key_id = "0000000000000001";

  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder.LoadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(sw_key_id));
  ASSERT_NO_FATAL_FAILURE(
      holder.FailDecrypt(hw_key_id, DECRYPT_ERROR));
  // Next, if possible, we try to decrypt to a secure buffer, and verify
  // success.
  if (wvoec::global_features.test_secure_buffers) {
    ASSERT_NO_FATAL_FAILURE(holder.DecryptSecure(hw_key_id));
  } else {
    LOGI("Test harness cannot create secure buffers. test skipped.");
  }
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());

  // Should be able to close the previous session, open a new session,
  // and reload the license.
  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.ReloadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(sw_key_id));
  ASSERT_NO_FATAL_FAILURE(
      holder.FailDecrypt(hw_key_id, DECRYPT_ERROR));
  // Next, if possible, we try to decrypt to a secure buffer, and verify
  // success.
  if (wvoec::global_features.test_secure_buffers) {
    ASSERT_NO_FATAL_FAILURE(holder.DecryptSecure(hw_key_id));
  } else {
    LOGI("Test harness cannot create secure buffers. test skipped.");
  }
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
}

/**
 * Should be able to request license, perform playback, generate a license
 * release, and receive the release response.
 */
TEST_F(CorePIGTest, LicenseRelease1) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Test for usage table devices only.";
  }
  LicenseHolder holder("CDM_UnlimitedStreaming_can_persist", &cdm_engine_,
                       config_);
  holder.set_can_persist(true);
  const KeyId key_id = "0000000000000000";

  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder.LoadLicense());
  EXPECT_EQ(NO_ERROR, holder.Decrypt(key_id));
  // For Android where AlwaysUseKeySetIds() is false, the CDM engine generates
  // a session separately. Thus, we close the session and only for CE CDM reopen
  // it for the license release.
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
  if (Properties::AlwaysUseKeySetIds()) {
    ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
    ASSERT_NO_FATAL_FAILURE(holder.ReloadLicense());
  }
  ASSERT_NO_FATAL_FAILURE(holder.GenerateAndPostReleaseRequest(
      "CDM_UnlimitedStreaming_can_persist"));
  EXPECT_NE(NO_ERROR, holder.Decrypt(key_id));
  ASSERT_NO_FATAL_FAILURE(holder.FetchRelease());
  ASSERT_NO_FATAL_FAILURE(holder.LoadRelease());
  EXPECT_NE(NO_ERROR, holder.Decrypt(key_id));
  // For CE CDM, we can close the session after we have gotten the release.
  if (Properties::AlwaysUseKeySetIds()) {
    ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
  }
}

/**
 * Should be able to request license, wait some time, generate a license
 * release, and receive the release response.
 */
TEST_F(CorePIGTest, LicenseRelease2) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Test for usage table devices only.";
  }
  LicenseHolder holder("CDM_UnlimitedStreaming_can_persist", &cdm_engine_,
                       config_);
  holder.set_can_persist(true);
  const KeyId key_id = "0000000000000000";

  ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
  ASSERT_NO_FATAL_FAILURE(holder.FetchLicense());
  ASSERT_NO_FATAL_FAILURE(holder.LoadLicense());
  wvutil::TestSleep::Sleep(10);
  // For Android where AlwaysUseKeySetIds() is false, the CDM engine generates
  // a session separately. Thus, we close the session and only for CE CDM reopen
  // it for the license release.
  ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
  if (Properties::AlwaysUseKeySetIds()) {
    ASSERT_NO_FATAL_FAILURE(holder.OpenSession());
    ASSERT_NO_FATAL_FAILURE(holder.ReloadLicense());
  }
  ASSERT_NO_FATAL_FAILURE(holder.GenerateAndPostReleaseRequest(
      "CDM_UnlimitedStreaming_can_persist"));
  ASSERT_NO_FATAL_FAILURE(holder.FetchRelease());
  ASSERT_NO_FATAL_FAILURE(holder.LoadRelease());
  // For CE CDM, we can close the session after we have gotten the release.
  if (Properties::AlwaysUseKeySetIds()) {
    ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
  }
}

TEST_F(CorePIGTest, CastReceiverProvisioningUsingCdm) {
  const std::string digest_hex_str =
      // digest info header
      "3021300906052b0e03021a05000414"
      // sha1 of kMessage
      "d2662f893aaec72f3ca6decc2aa942f3949e8b21";
  const auto digest = wvutil::a2b_hex(digest_hex_str);

  if (!wvoec::global_features.cast_receiver) {
    GTEST_SKIP() << "OEMCrypto does not support CAST Receiver functionality";
  }

  // Provision x509 cert for CAST Receiver.
  ProvisioningHolder provisioner(&cdm_engine_, config_);
  provisioner.Provision(kCertificateX509, binary_provisioning_);

  // cdm_engine_.SignRsa
  std::string signature_str;
  const std::string digest_str(digest.begin(), digest.end());
  ASSERT_EQ(NO_ERROR, cdm_engine_.SignRsa(provisioner.wrapped_key(), digest_str,
                                          &signature_str, kSign_PKCS1_Block1));

  // Verify the generated signature
  const std::vector<uint8_t> signature(signature_str.begin(), signature_str.end());
  LOGI("digest.size(): %zu, signature.size(): %zu", digest.size(),
       signature.size());

  const std::string cert = provisioner.certificate();
  const char* const cert_str_ptr = cert.c_str();
  LOGI("cert: %s", cert_str_ptr);

  // Extract the public key from the x509 cert chain
  std::unique_ptr<BIO, void (*)(BIO*)> bio(BIO_new(BIO_s_mem()), BIO_free_all);
  ASSERT_NE(bio, nullptr);
  ASSERT_GT(BIO_puts(bio.get(), cert_str_ptr), 0);
  std::unique_ptr<X509, void (*)(X509*)> x509(
      PEM_read_bio_X509(bio.get(), nullptr, nullptr, nullptr), X509_free);
  ASSERT_NE(x509, nullptr);
  std::unique_ptr<EVP_PKEY, void (*)(EVP_PKEY*)> pubkey(
        X509_get_pubkey(x509.get()), EVP_PKEY_free);
  ASSERT_NE(pubkey, nullptr);

  // remove digest info header for verification
  // SHA1 is 20 bytes long
  const std::vector<uint8_t> sha1_digest(digest.begin() + digest.size() - 20, digest.end());

  // Modified from openssl example
  // https://www.openssl.org/docs/man3.0/man3/EVP_PKEY_verify_init.html
  // Set RSA padding as RSA_PKCS1_PADDING and digest algo to SHA1.
  const unsigned char* const md = sha1_digest.data();
  const unsigned char* const sig = signature.data();
  const size_t mdlen = sha1_digest.size();
  const size_t siglen = signature.size();

  std::unique_ptr<EVP_PKEY_CTX, void (*)(EVP_PKEY_CTX*)> ctx(
      EVP_PKEY_CTX_new(pubkey.get(), nullptr /* no engine */), EVP_PKEY_CTX_free);

  ASSERT_NE(ctx, nullptr);
  ASSERT_GT(EVP_PKEY_verify_init(ctx.get()), 0);
  ASSERT_GT(EVP_PKEY_CTX_set_rsa_padding(ctx.get(), RSA_PKCS1_PADDING), 0);
  ASSERT_GT(EVP_PKEY_CTX_set_signature_md(ctx.get(), EVP_sha1()), 0);

  /* Perform operation */
  EXPECT_EQ(1, EVP_PKEY_verify(ctx.get(), sig, siglen, md, mdlen));
}
}  // namespace wvcdm