#include "oemcrypto_session_tests_helper.h"

#include <gtest/gtest.h>
#include "oec_test_data.h"

using namespace std;
using namespace wvoec;

namespace wvoec {

// Make this function available when in Fuzz mode because we are not inheriting
// from OEMCryptoClientTest.
const uint8_t* find(const vector<uint8_t>& message,
                    const vector<uint8_t>& substring) {
  vector<uint8_t>::const_iterator pos = search(
      message.begin(), message.end(), substring.begin(), substring.end());
  if (pos == message.end()) {
    return nullptr;
  }
  return &(*pos);
}

// This creates a wrapped RSA key.
void SessionUtil::CreateWrappedRSAKey() {
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  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());
  ASSERT_EQ(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  wrapped_rsa_key_ = provisioning_messages.wrapped_rsa_key();
}

void SessionUtil::InstallKeybox(const wvoec::WidevineKeybox& keybox,
                                bool good) {
  uint8_t wrapped[sizeof(wvoec::WidevineKeybox)];
  size_t length = sizeof(wvoec::WidevineKeybox);
  keybox_ = keybox;
  ASSERT_EQ(OEMCrypto_SUCCESS,
            OEMCrypto_WrapKeybox(reinterpret_cast<const uint8_t*>(&keybox),
                                 sizeof(keybox), wrapped, &length, nullptr, 0));
  OEMCryptoResult sts = OEMCrypto_InstallKeybox(wrapped, sizeof(keybox));
  if (good) {
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
  } else {
    // Can return error now, or return error on IsKeyboxValid.
  }
}

void SessionUtil::EnsureTestKeys() {
  switch (global_features.derive_key_method) {
    case DeviceFeatures::LOAD_TEST_KEYBOX:
      keybox_ = kTestKeybox;
      ASSERT_EQ(
          OEMCrypto_SUCCESS,
          OEMCrypto_LoadTestKeybox(reinterpret_cast<const uint8_t*>(&keybox_),
                                   sizeof(keybox_)));
      break;
    case DeviceFeatures::LOAD_TEST_RSA_KEY:
      ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_LoadTestRSAKey());
      break;
    case DeviceFeatures::TEST_PROVISION_30:
      // Can use oem certificate to install test rsa key.
      break;
    case wvoec::DeviceFeatures::TEST_PROVISION_40:
      // OEM certificate is retrieved from the server.
      break;
    default:
      FAIL() << "Cannot run test without test keybox or RSA key installed.";
  }
}

// This makes sure that the derived keys (encryption key and two mac keys)
// are installed in OEMCrypto and in the test session.
void SessionUtil::InstallTestRSAKey(Session* s) {
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    const size_t buffer_size = 5000;  // Make sure it is large enough.
    std::vector<uint8_t> public_key(buffer_size);
    size_t public_key_size = buffer_size;
    std::vector<uint8_t> public_key_signature(buffer_size);
    size_t public_key_signature_size = buffer_size;
    std::vector<uint8_t> wrapped_private_key(buffer_size);
    size_t wrapped_private_key_size = buffer_size;
    OEMCrypto_PrivateKeyType key_type;
    // Assume OEM cert has been loaded.
    ASSERT_EQ(
        OEMCrypto_SUCCESS,
        OEMCrypto_GenerateCertificateKeyPair(
            s->session_id(), public_key.data(), &public_key_size,
            public_key_signature.data(), &public_key_signature_size,
            wrapped_private_key.data(), &wrapped_private_key_size, &key_type));
    // Assume the public key has been verified by the server and the DRM cert is
    // returned.
    ASSERT_EQ(OEMCrypto_SUCCESS,
              OEMCrypto_LoadDRMPrivateKey(s->session_id(), key_type,
                                          wrapped_private_key.data(),
                                          wrapped_private_key_size));
    ASSERT_NO_FATAL_FAILURE(
        s->SetRsaPublicKey(public_key.data(), public_key_size));
    return;
  }

  if (global_features.loads_certificate) {
    if (wrapped_rsa_key_.size() == 0) {
      // If we don't have a wrapped key yet, create one.
      // This wrapped key will be shared by all sessions in the test.
      ASSERT_NO_FATAL_FAILURE(CreateWrappedRSAKey());
    }
    // Load the wrapped rsa test key.
    ASSERT_NO_FATAL_FAILURE(s->InstallRSASessionTestKey(wrapped_rsa_key_));
  }
  // Test RSA key should be loaded.
  ASSERT_NO_FATAL_FAILURE(s->PreparePublicKey());
}

}  // namespace wvoec