// 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_CAST_TEST_
#define CDM_OEMCRYPTO_CAST_TEST_

#include <vector>

#include <gtest/gtest.h>

#include "OEMCryptoCENC.h"
#include "oemcrypto_provisioning_test.h"
#include "oemcrypto_session_tests_helper.h"
#include "oemcrypto_usage_table_test.h"

namespace wvoec {

const char* HDCPCapabilityAsString(OEMCrypto_HDCP_Capability value);

// This test attempts to use alternate algorithms for loaded device certs.
class OEMCryptoLoadsCertificateAlternates : public OEMCryptoLoadsCertificate {
 protected:
  // The message to be signed by OEMCrypto_GenerateRSASignature() starts with a
  // constant digest info prefix followed by a SHA-1 hash of the message.
  void PrepareCastDigestedMessage(const std::vector<uint8_t>& message,
                                  std::vector<uint8_t>& digest) {
    // The application will compute the SHA-1 Hash of the message, so this
    // test must do that also.
    uint8_t hash[SHA_DIGEST_LENGTH];
    if (!SHA1(message.data(), message.size(), hash)) {
      dump_boringssl_error();
      FAIL() << "boringssl error creating SHA1 hash.";
    }
    // The application will prepend the digest info to the hash.
    // SHA-1 digest info prefix = 0x30 0x21 0x30 ...
    static const std::vector<uint8_t> prefix =
        wvutil::a2b_hex("3021300906052b0e03021a05000414");
    digest.insert(digest.end(), prefix.begin(), prefix.end());
    digest.insert(digest.end(), hash, hash + SHA_DIGEST_LENGTH);
  }

  void TestSignature(RSA_Padding_Scheme scheme, size_t size) {
    Session s;
    ASSERT_NO_FATAL_FAILURE(s.open());
    ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_drm_key_));

    vector<uint8_t> licenseRequest(size);
    GetRandBytes(licenseRequest.data(), licenseRequest.size());
    vector<uint8_t> digested_message;
    ASSERT_NO_FATAL_FAILURE(
        PrepareCastDigestedMessage(licenseRequest, digested_message));
    size_t signature_length = 0;
    OEMCryptoResult sts = OEMCrypto_GenerateRSASignature(
        s.session_id(), digested_message.data(), digested_message.size(),
        nullptr, &signature_length, scheme);
    ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
    ASSERT_NE(static_cast<size_t>(0), signature_length);

    std::vector<uint8_t> signature(signature_length, 0);
    sts = OEMCrypto_GenerateRSASignature(
        s.session_id(), digested_message.data(), digested_message.size(),
        signature.data(), &signature_length, scheme);

    ASSERT_EQ(OEMCrypto_SUCCESS, sts)
        << "Failed to sign with padding scheme=" << (int)scheme
        << ", size=" << size;
    signature.resize(signature_length);
    ASSERT_NO_FATAL_FAILURE(s.SetRsaPublicKeyFromPrivateKeyInfo(
        encoded_rsa_key_.data(), encoded_rsa_key_.size()));
    ASSERT_NO_FATAL_FAILURE(s.VerifyRsaSignature(
        digested_message, signature.data(), signature_length, scheme));
  }

  // If force is true, we assert that the key loads successfully.
  void LoadCastCertificateKey(bool force) {
    if (!wvoec::global_features.cast_receiver) {
      GTEST_SKIP() << "Cast not supported";
    }
    // Padding scheme used to sign cast data.
    constexpr uint32_t schemes = kSign_PKCS1_Block1;
    // prov 2 or prov 3
    if (global_features.provisioning_method == OEMCrypto_Keybox ||
        global_features.provisioning_method == OEMCrypto_OEMCertificate) {
      Session s;
      ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
      provisioning_messages.set_allowed_schemes(schemes);
      provisioning_messages.PrepareSession(keybox_);
      ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
      ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
      ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
      OEMCryptoResult sts = provisioning_messages.LoadResponse();
      key_loaded_ = (OEMCrypto_SUCCESS == sts);
      if (key_loaded_) {
        uint8_t* ptr = provisioning_messages.response_data().rsa_key;
        size_t len = provisioning_messages.response_data().rsa_key_length;
        encoded_rsa_key_ = std::vector<uint8_t>(ptr, ptr + len);
        wrapped_drm_key_ = provisioning_messages.wrapped_rsa_key();
        drm_key_type_ = OEMCrypto_RSA_Private_Key;
        EXPECT_GT(wrapped_drm_key_.size(), 0u);
        EXPECT_EQ(nullptr, find(wrapped_drm_key_, encoded_rsa_key_));
      }
      if (force) {
        EXPECT_EQ(OEMCrypto_SUCCESS, sts);
      }
    } else if (global_features.provisioning_method ==
               OEMCrypto_BootCertificateChain) {
      Session s1;
      ASSERT_NO_FATAL_FAILURE(s1.open());
      ASSERT_NO_FATAL_FAILURE(CreateProv4OEMKey(&s1));

      Session s2;
      ASSERT_NO_FATAL_FAILURE(s2.open());
      ASSERT_EQ(OEMCrypto_InstallOemPrivateKey(s2.session_id(), oem_key_type_,
                                               wrapped_oem_key_.data(),
                                               wrapped_oem_key_.size()),
                OEMCrypto_SUCCESS);
      Provisioning40CastRoundTrip prov_cast(&s2, encoded_rsa_key_);
      prov_cast.set_allowed_schemes(schemes);
      ASSERT_NO_FATAL_FAILURE(prov_cast.PrepareSession());
      ASSERT_NO_FATAL_FAILURE(prov_cast.LoadDRMPrivateKey());

      ASSERT_NO_FATAL_FAILURE(s2.SetPublicKeyFromSubjectPublicKey(
          prov_cast.drm_key_type(), prov_cast.drm_public_key().data(),
          prov_cast.drm_public_key().size()));
      ASSERT_NO_FATAL_FAILURE(prov_cast.SignAndVerifyRequest());
      ASSERT_NO_FATAL_FAILURE(s2.GenerateDerivedKeysFromSessionKey());
      ASSERT_NO_FATAL_FAILURE(prov_cast.CreateDefaultResponse());
      ASSERT_NO_FATAL_FAILURE(prov_cast.EncryptAndSignResponse());
      OEMCryptoResult sts = prov_cast.LoadResponse();
      key_loaded_ = (OEMCrypto_SUCCESS == sts);
      if (key_loaded_) {
        uint8_t* ptr = prov_cast.response_data().rsa_key;
        size_t len = prov_cast.response_data().rsa_key_length;
        encoded_rsa_key_ = std::vector<uint8_t>(ptr, ptr + len);
        wrapped_drm_key_ = prov_cast.wrapped_rsa_key();
        drm_key_type_ = OEMCrypto_RSA_Private_Key;
        EXPECT_GT(wrapped_drm_key_.size(), 0u);
        EXPECT_EQ(nullptr, find(wrapped_drm_key_, encoded_rsa_key_));
      }
      if (force) {
        EXPECT_EQ(OEMCrypto_SUCCESS, sts);
      }
    } else {
      FAIL() << "Unsupported provisioning method";
    }
  }

  bool key_loaded_ = false;
};

// Used to test the different HDCP versions.  This test is parameterized by the
// required HDCP version in the key control block.
class OEMCryptoSessionTestLoadCasKeysWithHDCP : public OEMCryptoSessionTests,
                                                public WithParamInterface<int> {
 protected:
  void LoadCasKeysWithHDCP(OEMCrypto_HDCP_Capability version) {
    OEMCryptoResult sts;
    OEMCrypto_HDCP_Capability current, maximum;
    sts = OEMCrypto_GetHDCPCapability(&current, &maximum);
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
    Session s;
    ASSERT_NO_FATAL_FAILURE(s.open());
    ASSERT_NO_FATAL_FAILURE(InstallTestDrmKey(&s));
    LicenseRoundTrip license_messages(&s);
    license_messages.set_control((version << wvoec::kControlHDCPVersionShift) |
                                 wvoec::kControlObserveHDCP |
                                 wvoec::kControlHDCPRequired);
    license_messages.set_license_type(OEMCrypto_EntitlementLicense);
    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());

    uint32_t key_session_id;
    sts = OEMCrypto_CreateEntitledKeySession(s.session_id(), &key_session_id);
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
    EntitledMessage entitled_message_1(&license_messages);
    entitled_message_1.FillKeyArray();
    entitled_message_1.SetEntitledKeySession(key_session_id);

    if (((version <= HDCP_V2_3 || current >= HDCP_V1_0) && version > current) ||
        (current == HDCP_V1 && version >= HDCP_V1_0)) {
      ASSERT_NO_FATAL_FAILURE(entitled_message_1.LoadCasKeys(
          /*load_even=*/true, /*load_odd=*/true,
          OEMCrypto_ERROR_INSUFFICIENT_HDCP))
          << "Failed when current HDCP = " << HDCPCapabilityAsString(current)
          << ", maximum HDCP = " << HDCPCapabilityAsString(maximum)
          << ", license HDCP = " << HDCPCapabilityAsString(version);
    } else {
      ASSERT_NO_FATAL_FAILURE(entitled_message_1.LoadCasKeys(
          /*load_even=*/true, /*load_odd=*/true, OEMCrypto_SUCCESS))
          << "Failed when current HDCP = " << HDCPCapabilityAsString(current)
          << ", maximum HDCP = " << HDCPCapabilityAsString(maximum)
          << ", license HDCP = " << HDCPCapabilityAsString(version);
    }
  }
};
}  // namespace wvoec

#endif  // CDM_OEMCRYPTO_CAST_TEST_