#include "oemcrypto_session_tests_helper.h"

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

using namespace std;
using namespace wvoec;

namespace wvoec {

void SessionUtil::CreateWrappedDRMKey() {
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    // Have the device create a wrapped key.
    CreateProv4DRMKey();
  } else {
    // Create a wrapped RSA key from encoded_rsa_key_.
    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_drm_key_ = provisioning_messages.wrapped_rsa_key();
    drm_key_type_ = OEMCrypto_RSA_Private_Key;
    drm_public_key_.clear();
  }
}

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::EnsureTestROT() {
  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 DeviceFeatures::PRELOADED_RSA_KEY:
      // There is already a 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::InstallTestDrmKey(Session* s) {
  if (global_features.loads_certificate) {
    if (wrapped_drm_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(CreateWrappedDRMKey());
    }
    // Load the wrapped drm test key.
    ASSERT_NO_FATAL_FAILURE(
        s->LoadWrappedDrmKey(drm_key_type_, wrapped_drm_key_));
    if (drm_public_key_.size() > 0) {
      ASSERT_NO_FATAL_FAILURE(s->SetPublicKeyFromSubjectPublicKey(
          drm_key_type_, drm_public_key_.data(), drm_public_key_.size()));
    } else {
      ASSERT_NO_FATAL_FAILURE(s->SetPublicKeyFromPrivateKeyInfo(
          drm_key_type_, encoded_rsa_key_.data(), encoded_rsa_key_.size()));
    }
  } else {
    // Test RSA key should be loaded.
    ASSERT_NO_FATAL_FAILURE(s->SetTestRsaPublicKey());
  }
}

// Generate OEM key pair, craft a provisioning 4.0 OEM cert request, sign it
// with the OEM private key and verify the signature. Finally, install OEM
// private to session s.
void SessionUtil::CreateProv4OEMKey(Session* s) {
  ASSERT_NE(s, nullptr);
  if (global_features.provisioning_method != OEMCrypto_BootCertificateChain) {
    FAIL() << "Provisioning 4.0 is required.";
  }
  Provisioning40RoundTrip provisioning_messages(s);
  // Generate key pair.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.PrepareSession(true));
  // Need OEM public key to verify the signed request.
  ASSERT_NO_FATAL_FAILURE(s->SetPublicKeyFromSubjectPublicKey(
      provisioning_messages.oem_key_type(),
      provisioning_messages.oem_public_key().data(),
      provisioning_messages.oem_public_key().size()));
  // Save the generated keys, which will be used by DRM cert provisioning later.
  wrapped_oem_key_ = provisioning_messages.wrapped_oem_key();
  oem_public_key_ = provisioning_messages.oem_public_key();
  oem_key_type_ = provisioning_messages.oem_key_type();
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  // Install OEM private key into the session.
  ASSERT_EQ(OEMCrypto_SUCCESS, provisioning_messages.LoadOEMCertResponse());
}

// Generate DRM key pair, craft a provisioning 4.0 DRM cert request, sign it
// with the OEM private key and verify the signature. Finally, install DRM
// private to session s. An OEM cert needs to be installed first. It is also
// done in this function.
void SessionUtil::CreateProv4DRMKey() {
  if (global_features.provisioning_method != OEMCrypto_BootCertificateChain) {
    FAIL() << "Provisioning 4.0 is required.";
  }
  // Provision OEM key first.
  if (wrapped_oem_key_.size() == 0) {
    Session oem_session;
    ASSERT_NO_FATAL_FAILURE(oem_session.open());
    ASSERT_NO_FATAL_FAILURE(CreateProv4OEMKey(&oem_session));
  }
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            OEMCrypto_InstallOemPrivateKey(
                s.session_id(), oem_key_type_,
                reinterpret_cast<const uint8_t*>(wrapped_oem_key_.data()),
                wrapped_oem_key_.size()));
  ASSERT_NO_FATAL_FAILURE(s.SetPublicKeyFromSubjectPublicKey(
      oem_key_type_, oem_public_key_.data(), oem_public_key_.size()));

  // Provision DRM key.
  Provisioning40RoundTrip provisioning_messages(&s);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.PrepareSession(false));
  // Need DRM public key to verify DRM request signature.
  ASSERT_NO_FATAL_FAILURE(s.SetPublicKeyFromSubjectPublicKey(
      provisioning_messages.drm_key_type(),
      provisioning_messages.drm_public_key().data(),
      provisioning_messages.drm_public_key().size()));
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_EQ(OEMCrypto_SUCCESS, provisioning_messages.LoadDRMCertResponse());
  wrapped_drm_key_ = provisioning_messages.wrapped_drm_key();
  drm_key_type_ = provisioning_messages.drm_key_type();
  drm_public_key_ = provisioning_messages.drm_public_key();
}

// Requires stage 1 prov4 to be complete, ie OEM key is available
void SessionUtil::CreateProv4CastKey(Session* s,
                                     bool load_drm_before_prov_req) {
  if (global_features.provisioning_method != OEMCrypto_BootCertificateChain) {
    FAIL() << "Provisioning 4.0 is required.";
  }

  Provisioning40CastRoundTrip prov_cast(s, encoded_rsa_key_);

  // Calls GenerateCertificateKeyPair(). Generated keys stored in
  // prov_cast.drm_public_key_ and prov_cast.wrapped_drm_key_
  ASSERT_NO_FATAL_FAILURE(prov_cast.PrepareSession());

  // Can choose to load DRM key before preparing the provisioning request, or
  // after
  if (load_drm_before_prov_req) {
    ASSERT_NO_FATAL_FAILURE(prov_cast.LoadDRMPrivateKey());
  }
  ASSERT_NO_FATAL_FAILURE(s->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());
  if (!load_drm_before_prov_req) {
    ASSERT_NO_FATAL_FAILURE(prov_cast.LoadDRMPrivateKey());
  }

  // Generate derived keys in order to verify and decrypt response.
  // We are cheating a little bit here since this GenerateDerivedKeys helper
  // simulates work on both client side (calls
  // OEMCrypto_GenerateDerivedKeysFromSessionKey) and server side (sets
  // key_deriver() keys used to create response)
  ASSERT_NO_FATAL_FAILURE(s->GenerateDerivedKeysFromSessionKey());

  // Response is provisioning 2 with CAST key
  ASSERT_NO_FATAL_FAILURE(prov_cast.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(prov_cast.EncryptAndSignResponse());

  // Should parse and load successfully
  ASSERT_EQ(OEMCrypto_SUCCESS, prov_cast.LoadResponse());
}

}  // namespace wvoec