// Copyright 2020 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.

#include "OEMCryptoCENC.h"
#include "oemcrypto_fuzz_helper.h"

namespace wvoec {

void RedirectStdoutToFile() { freopen("log.txt", "a", stdout); }

std::vector<FuzzedData> SplitFuzzedData(const uint8_t* data, size_t size) {
  std::vector<FuzzedData> result;
  const uint8_t* const end = data + size;
  // Using memmem to find separator.
  while (
      const uint8_t* const separator = reinterpret_cast<const uint8_t*>(memmem(
          data, end - data, kFuzzDataSeparator, sizeof(kFuzzDataSeparator)))) {
    result.push_back({data, static_cast<size_t>(separator - data)});
    data = separator + sizeof(kFuzzDataSeparator);
  }
  if (data < end) {
    result.push_back({data, static_cast<size_t>(end - data)});
  }
  return result;
}

void InitializeFuzz(SessionUtil& session_util) {
  wvoec::global_features.Initialize();
  OEMCrypto_SetSandbox(kTestSandbox, sizeof(kTestSandbox));
  OEMCrypto_Initialize();
  OEMCrypto_EnterTestMode();
  session_util.EnsureTestROT();
}

void SessionFuzz::Initialize() {
  InitializeFuzz(session_util_);
  session_.open();
}

void SessionFuzz::Terminate() {
  session_.close();
  OEMCrypto_Terminate();
}

void OEMCryptoLicenseAPIFuzz::Initialize() {
  session_fuzz_.Initialize();
  session_fuzz_.InstallTestDrmKey();
}

void OEMCryptoLicenseAPIFuzz::Terminate() {
  session_fuzz_.Terminate();
}

void OEMCryptoLicenseAPIFuzz::LoadLicense(bool generic_crypto_keys) {
  license_messages_.SignAndVerifyRequest();
  if (generic_crypto_keys) {
    license_messages_.CreateResponseWithGenericCryptoKeys();
  } else {
    license_messages_.CreateDefaultResponse();
  }
  license_messages_.EncryptAndSignResponse();
  OEMCryptoResult sts = license_messages_.LoadResponse();
  CheckStatusAndExitFuzzerOnFailure(sts, OEMCrypto_SUCCESS);
}

void OEMCryptoProvisioningAPIFuzz::Initialize() {
  InitializeFuzz(session_util_);

  // Opens a session and Generates Nonce.
  provisioning_messages_.PrepareSession(session_util_.keybox_);
}

void OEMCryptoProvisioningAPIFuzz::Terminate() {
  session_.close();
  OEMCrypto_Terminate();
}

void OEMCryptoProvisioningAPIFuzz::LoadProvisioning() {
  provisioning_messages_.SignAndVerifyRequest();
  provisioning_messages_.CreateDefaultResponse();
  provisioning_messages_.EncryptAndSignResponse();
  OEMCryptoResult sts = provisioning_messages_.LoadResponse();
  CheckStatusAndExitFuzzerOnFailure(sts, OEMCrypto_SUCCESS);
}

void LicenseWithUsageEntryFuzz::CreateUsageTableHeader() {
  size_t header_buffer_length = 0;
  OEMCryptoResult sts =
      OEMCrypto_CreateUsageTableHeader(nullptr, &header_buffer_length);
  encrypted_usage_header_.resize(header_buffer_length);
  sts = OEMCrypto_CreateUsageTableHeader(encrypted_usage_header_.data(),
                                         &header_buffer_length);
  CheckStatusAndExitFuzzerOnFailure(sts, OEMCrypto_SUCCESS);
}

void LicenseWithUsageEntryFuzz::LoadLicense() {
  license_messages_.SignAndVerifyRequest();
  license_messages_.CreateDefaultResponse();
  license_messages_.EncryptAndSignResponse();
  OEMCryptoResult sts = license_messages_.LoadResponse();
  CheckStatusAndExitFuzzerOnFailure(sts, OEMCrypto_SUCCESS);
}

void CheckStatusAndExitFuzzerOnFailure(OEMCryptoResult result,
                                       OEMCryptoResult expected_status) {
  if (result != expected_status) {
    abort();
  }
}

}  // namespace wvoec