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

#include <vector>

#include "FuzzedDataProvider.h"
#include "OEMCryptoCENC.h"
#include "oec_device_features.h"
#include "oemcrypto_corpus_generator_helper.h"
#include "oemcrypto_session_tests_helper.h"

namespace wvoec {
// Forward-declare the libFuzzer's mutator callback. Mark it weak so that
// the program links successfully even outside of --config=asan-fuzzer
// (apparently the only config in which LLVM uses our custom mutator).
extern "C" size_t LLVMFuzzerMutate(uint8_t* Data, size_t Size, size_t MaxSize)
    __attribute__((weak));

const size_t KB = 1024;
// Maximum signature length. If fuzzed signature length is greater that this,
// this value will be used for signature length.
const size_t MAX_FUZZ_SIGNATURE_LENGTH = 5 * KB;
// Initial setup to create a valid OEMCrypto state such as initializing crypto
// firmware/hardware, installing golden key box etc. in order to fuzz
// OEMCrypto APIs.
class InitializeFuzz : public SessionUtil {
 public:
  InitializeFuzz() {
    wvoec::global_features.Initialize();
    OEMCrypto_SetSandbox(kTestSandbox, sizeof(kTestSandbox));
    OEMCrypto_Initialize();
    EnsureTestKeys();
  }

  ~InitializeFuzz() { OEMCrypto_Terminate(); }
};

class OEMCryptoLicenseAPIFuzz : public InitializeFuzz {
 public:
  OEMCryptoLicenseAPIFuzz() : license_messages_(&session_) {
    session_.open();
    InstallTestRSAKey(&session_);
    session_.GenerateNonce();
  }

  ~OEMCryptoLicenseAPIFuzz() { session_.close(); }

  LicenseRoundTrip& license_messages() { return license_messages_; }

  Session* session() { return &session_; }

  void LoadLicense();

 private:
  Session session_;
  LicenseRoundTrip license_messages_;
};

class OEMCryptoProvisioningAPIFuzz : public InitializeFuzz {
 public:
  OEMCryptoProvisioningAPIFuzz()
      : provisioning_messages_(&session_, encoded_rsa_key_) {
    // Opens a session and Generates Nonce.
    provisioning_messages_.PrepareSession(keybox_);
  }

  ~OEMCryptoProvisioningAPIFuzz() { session_.close(); }

  void LoadProvisioning();
  ProvisioningRoundTrip& provisioning_messages() {
    return provisioning_messages_;
  }
  Session* session() { return &session_; }

 private:
  Session session_;
  ProvisioningRoundTrip provisioning_messages_;
};

// Initial setup to create a valid state such as creating session, installing
// golden key box etc. in order to fuzz Load Renewal API.
class OEMCryptoRenewalAPIFuzz : public OEMCryptoLicenseAPIFuzz {
 public:
  OEMCryptoRenewalAPIFuzz() : renewal_messages_(&license_messages()) {}

  RenewalRoundTrip& renewal_messages() { return renewal_messages_; }

 private:
  RenewalRoundTrip renewal_messages_;
};

class LicenseWithUsageEntryFuzz : public InitializeFuzz {
 public:
  LicenseWithUsageEntryFuzz() : license_messages_(&session_) {
    license_messages_.set_pst("my_pst");
  }

  void CreateUsageTableHeader();
  LicenseRoundTrip& license_messages() { return license_messages_; }
  const vector<uint8_t>& encrypted_usage_header() {
    return encrypted_usage_header_;
  }
  void LoadLicense();

 private:
  vector<uint8_t> encrypted_usage_header_;
  LicenseRoundTrip license_messages_;
  Session session_;
};

// Convert data to valid enum value.
template <typename T>
void ConvertDataToValidEnum(T max_enum_value, T* t) {
  FuzzedDataProvider fuzzed_enum_data(reinterpret_cast<uint8_t*>(t), sizeof(T));
  *t = static_cast<T>(fuzzed_enum_data.ConsumeIntegralInRange<uint32_t>(
      0, static_cast<uint32_t>(max_enum_value)));
}

// Redirect printf and log statements from oemcrypto functions to a file to
// reduce noise
void RedirectStdoutToFile();

// Function to split fuzzer input using delimiter "-_^_".
std::vector<std::vector<uint8_t>> SplitInput(const uint8_t* data, size_t size);
// Check the status and exit fuzzer if arguments do not match. This is usually
// called to check status of APIs which are called to setup state for fuzzers.
void CheckStatusAndExitFuzzerOnFailure(OEMCryptoResult result,
                                       OEMCryptoResult expected_status);
}  // namespace wvoec

#endif  // OEMCRYPTO_FUZZ_HELPER_H_