// 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_PROVISIONING_TEST_
#define CDM_OEMCRYPTO_PROVISIONING_TEST_

#include <gtest/gtest.h>

#include "OEMCryptoCENC.h"
#include "oec_extra_test_keys.h"
#include "oemcrypto_basic_test.h"
#include "oemcrypto_license_test.h"
#include "oemcrypto_resource_test.h"

namespace wvoec {

// Tests using this class are only used for devices with a keybox.  They are not
// run for devices with an OEM Certificate.
class OEMCryptoKeyboxTest : public OEMCryptoClientTest {
  void SetUp() override {
    OEMCryptoClientTest::SetUp();
    OEMCryptoResult sts = OEMCrypto_IsKeyboxValid();
    // If the production keybox is valid, use it for these tests.  Most of the
    // other tests will use a test keybox anyway, but it's nice to check the
    // device ID for the real keybox if we can.
    if (sts == OEMCrypto_SUCCESS) return;
    printf("Production keybox is NOT valid.  All tests use test keybox.\n");
    ASSERT_EQ(
        OEMCrypto_SUCCESS,
        OEMCrypto_LoadTestKeybox(reinterpret_cast<const uint8_t*>(&kTestKeybox),
                                 sizeof(kTestKeybox)));
  }
};

// This class is for tests that have an OEM Certificate instead of a keybox.
class OEMCryptoProv30Test : public OEMCryptoClientTest {
  void SetUp() override {
    OEMCryptoClientTest::SetUp();
    if (global_features.provisioning_method != OEMCrypto_OEMCertificate) {
      GTEST_SKIP() << "Test for Prov 3.0 devices only.";
    }
  }
};

// This class is for tests that have boot certificate chain instead of a keybox.
class OEMCryptoProv40Test : public OEMCryptoClientTest {
  void SetUp() override {
    OEMCryptoClientTest::SetUp();
    if (global_features.provisioning_method != OEMCrypto_BootCertificateChain) {
      GTEST_SKIP() << "Test for Prov 4.0 devices only.";
    }
  }
};

class OEMCryptoProv40CastTest : public OEMCryptoClientTest,
                                public testing::WithParamInterface<bool> {
  void SetUp() override {
    OEMCryptoClientTest::SetUp();
    if (!global_features.cast_receiver) {
      GTEST_SKIP() << "Test for cast devices only.";
    }
    if (global_features.provisioning_method != OEMCrypto_BootCertificateChain) {
      GTEST_SKIP() << "Test for Prov 4.0 devices only.";
    }
  }
};

//
// Certificate Root of Trust Tests
//
class OEMCryptoLoadsCertificate : public OEMCryptoSessionTestKeyboxTest {
 protected:
  void TestPrepareProvisioningRequestForHugeBufferLengths(
      const std::function<void(size_t, ProvisioningRoundTrip*)> f,
      bool check_status) {
    auto oemcrypto_function = [&](size_t message_length) {
      Session s;
      s.open();
      if (global_features.provisioning_method == OEMCrypto_OEMCertificate) {
        s.LoadOEMCert(true);
      } else {
        s.GenerateDerivedKeysFromKeybox(keybox_);
      }
      ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
      f(message_length, &provisioning_messages);
      return provisioning_messages
          .SignAndCreateRequestWithCustomBufferLengths();
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, check_status);
  }

  void TestLoadProvisioningForHugeBufferLengths(
      const std::function<void(size_t, ProvisioningRoundTrip*)> f,
      bool check_status, bool update_core_message_substring_values) {
    auto oemcrypto_function = [&](size_t message_length) {
      Session s;
      ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
      provisioning_messages.PrepareSession(keybox_);
      provisioning_messages.SignAndVerifyRequest();
      provisioning_messages.CreateDefaultResponse();
      if (update_core_message_substring_values) {
        // Make provisioning message big enough so that updated core message
        // substring offset and length values from tests are still able to read
        // valid data from provisioning_message buffer rather than some garbage
        // data.
        provisioning_messages.set_message_size(
            sizeof(provisioning_messages.response_data()) + message_length);
      }
      f(message_length, &provisioning_messages);
      provisioning_messages
          .EncryptAndSignResponseWithoutUpdatingEncPrivateKeyLength();
      OEMCryptoResult result = provisioning_messages.LoadResponse();
      s.close();
      return result;
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, check_status);
  }

  void TestLoadProvisioningForOutOfRangeSubstringOffsetAndLengths(
      const std::function<void(size_t, ProvisioningRoundTrip*)> f) {
    Session s;
    ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
    provisioning_messages.PrepareSession(keybox_);
    provisioning_messages.SignAndVerifyRequest();
    provisioning_messages.CreateDefaultResponse();
    size_t message_length = sizeof(provisioning_messages.response_data());
    f(message_length, &provisioning_messages);
    provisioning_messages
        .EncryptAndSignResponseWithoutUpdatingEncPrivateKeyLength();
    OEMCryptoResult result = provisioning_messages.LoadResponse();
    s.close();
    // Verifying error is not due to signature failure which can be caused due
    // to test code.
    ASSERT_NE(OEMCrypto_ERROR_SIGNATURE_FAILURE, result);
    ASSERT_NE(OEMCrypto_SUCCESS, result);
  }
};

// These tests are run by all L1 devices that load and use certificates.  It is
// also run by a few L3 devices that use a baked in certificate, but cannot load
// a certificate.
class OEMCryptoUsesCertificate : public OEMCryptoLoadsCertificate {
 protected:
  void SetUp() override {
    OEMCryptoLoadsCertificate::SetUp();
    ASSERT_NO_FATAL_FAILURE(session_.open());
    if (global_features.derive_key_method ==
        DeviceFeatures::LOAD_TEST_RSA_KEY) {
      ASSERT_NO_FATAL_FAILURE(session_.SetRsaPublicKeyFromPrivateKeyInfo(
          encoded_rsa_key_.data(), encoded_rsa_key_.size()));
    } else {
      InstallTestDrmKey(&session_);
    }
  }

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

  Session session_;
};

}  // namespace wvoec

#endif  // CDM_OEMCRYPTO_PROVISIONING_TEST_