// 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 {

//
// Certificate Root of Trust Tests
//
// 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 OEMCryptoSessionTests {
 protected:
  void SetUp() override {
    OEMCryptoSessionTests::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());
    OEMCryptoSessionTests::TearDown();
  }

  Session session_;
};

/** These tests cover all systems that can load a DRM Certificate. That includes
 * Provisioning 2, 3 and 4. */
class OEMCryptoLoadsCertificate : public OEMCryptoUsesCertificate {
 protected:
  void SetUp() override {
    OEMCryptoUsesCertificate::SetUp();
    if (!global_features.loads_certificate) {
      GTEST_SKIP() << "Test for devices that load a DRM certificate only.";
    }
  }

  /** Verify that the specified padding scheme does not work with the DRM
   * key and the function OEMCrypto_GenerateRSASignature. */
  void DisallowForbiddenPaddingDRMKey(RSA_Padding_Scheme scheme, size_t size) {
    Session s;
    ASSERT_NO_FATAL_FAILURE(s.open());
    ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_drm_key_));
    DisallowForbiddenPadding(s.session_id(), scheme, size);
  }

  /** Verify that the specified padding scheme does not work with whichever key
   * is currently loaded into the specified session and the function
   * OEMCrypto_GenerateRSASignature. */
  void DisallowForbiddenPadding(OEMCrypto_SESSION session,
                                RSA_Padding_Scheme scheme, size_t size) {
    OEMCryptoResult sts;
    // Sign a Message
    vector<uint8_t> message(size);
    GetRandBytes(message.data(), message.size());
    size_t signature_length = 256;
    vector<uint8_t> signature(signature_length);
    sts = OEMCrypto_GenerateRSASignature(session, message.data(),
                                         message.size(), signature.data(),
                                         &signature_length, scheme);
    // Allow OEMCrypto to request a full buffer.
    if (sts == OEMCrypto_ERROR_SHORT_BUFFER) {
      ASSERT_NE(static_cast<size_t>(0), signature_length);
      signature.assign(signature_length, 0);
      sts = OEMCrypto_GenerateRSASignature(session, message.data(),
                                           message.size(), signature.data(),
                                           &signature_length, scheme);
    }

    EXPECT_NE(OEMCrypto_SUCCESS, sts)
        << "Signed with forbidden padding scheme=" << (int)scheme
        << ", size=" << (int)size;
    const vector<uint8_t> zero(signature.size(), 0);
    ASSERT_EQ(zero, signature);  // signature should not be computed.
  }

  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);
  }
};

// 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 OEMCryptoLoadsCertificate {
 protected:
  void SetUp() override {
    OEMCryptoLoadsCertificate::SetUp();
    if (global_features.provisioning_method != OEMCrypto_Keybox) {
      GTEST_SKIP() << "Test for Prov 2.0 devices only.";
    }
    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)));
    ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_IsKeyboxValid())
        << "After loading Test keybox, the keybox was still not valid.";
  }
};

// This class is for tests that have an OEM Certificate instead of a keybox.
class OEMCryptoProv30Test : public OEMCryptoLoadsCertificate {
 protected:
  void SetUp() override {
    OEMCryptoLoadsCertificate::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 OEMCryptoLoadsCertificate {
 protected:
  void SetUp() override {
    OEMCryptoLoadsCertificate::SetUp();
    if (global_features.provisioning_method != OEMCrypto_BootCertificateChain) {
      GTEST_SKIP() << "Test for Prov 4.0 devices only.";
    }
  }
};

class OEMCryptoProv40CastTest : public OEMCryptoProv40Test,
                                public testing::WithParamInterface<bool> {
 protected:
  void SetUp() override {
    OEMCryptoProv40Test::SetUp();
    if (!global_features.cast_receiver) {
      GTEST_SKIP() << "Test for cast devices only.";
    }
  }
};

}  // namespace wvoec

#endif  // CDM_OEMCRYPTO_PROVISIONING_TEST_