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

#ifndef WVCDM_CORE_TEST_BASE_H_
#define WVCDM_CORE_TEST_BASE_H_

#include <gtest/gtest.h>

#include <memory>
#include <string>
#include <vector>

#include "cdm_engine.h"
#include "config_test_env.h"
#include "create_test_file_system.h"
#include "crypto_session.h"
#include "metrics_collections.h"
#include "oemcrypto_types.h"
#include "string_conversions.h"

namespace wvcdm {
// This is the base class for Widevine CDM integration tests.  It's main use is
// to configure OEMCrypto to use a test keybox.
class WvCdmTestBase : public ::testing::Test {
 public:
  WvCdmTestBase();
  ~WvCdmTestBase() override {}
  void SetUp() override;
  virtual std::string binary_key_id() const {
    return wvutil::a2bs_hex(config_.key_id());
  }

  // Returns true if the test program should continue, if false, the caller
  // should exit.  This should be called by main() to allow the user to pass in
  // command line switches.  The |extra_help_text| parameter can be used to
  // append platform-specific information to the usage information printed when
  // invalid flags are detected. For instance, a platform might add information
  // about platform-specific flags that were already parsed before calling
  // Initialize().
  static bool Initialize(int argc, const char* const argv[],
                         const std::string& extra_help_text = std::string());
  // Install a test keybox, if appropriate.
  static void InstallTestRootOfTrust();

  // Send provisioning request to the server and handle response.
  virtual void Provision();
  // Calls Provision() if not already provisioned.
  virtual void EnsureProvisioned();

  // Locate the portion of the server's provisioning response message that is
  // between the strings jason_start_substr and json_end_substr. Returns the
  // string through *result. If the start substring match fails, assume the
  // entire string represents a serialized protobuf mesaage and return true with
  // the entire string. If the end_substring match fails, return false with an
  // empty *result.
  bool ExtractSignedMessage(const std::string& response, std::string* result);

  // Fill a buffer with some nonconstant data of the given size. The first
  // byte will be set to <init> to help you find the buffer when debugging.
  static void StripeBuffer(std::vector<uint8_t>* buffer, size_t size,
                           uint8_t init);

  // Helper method for doing cryptography.
  static std::string Aes128CbcEncrypt(std::vector<uint8_t> key,
                                      const std::vector<uint8_t>& clear,
                                      std::vector<uint8_t> iv);
  // Helper method for doing cryptography.
  static std::string Aes128CbcDecrypt(std::vector<uint8_t> key,
                                      const std::vector<uint8_t>& clear,
                                      std::vector<uint8_t> iv);
  // Helper method for doing cryptography.
  static std::string SignHMAC(const std::string& message,
                              const std::vector<uint8_t>& key);

  // The default test configuration.  This is created in Initialize() and
  // influenced by command line arguments before any tests are created.
  static std::unique_ptr<ConfigTestEnv> default_config_;

  // If the tests should use the QA test keybox.
  static bool use_qa_test_keybox_;

  // Configuration for an individual test.  This is initialized to be the
  // default configuration, but can be modified by the test itself.
  ConfigTestEnv config_;

  // This should be set by test subclasses BEFORE calling SetUp -- i.e. in the
  // tests's constructor.
  bool binary_provisioning_;
};

// This just makes the constructor public so that we can create one with dummy
// metrics and file system.
class TestCdmEngine : public CdmEngine {
 public:
  TestCdmEngine(wvutil::FileSystem* file_system,
                std::shared_ptr<metrics::EngineMetrics> metrics)
      : CdmEngine(file_system, metrics) {}
  const CdmSession* GetCdmSession(std::string sessionId) const;
};

class WvCdmTestBaseWithEngine : public WvCdmTestBase {
 public:
  WvCdmTestBaseWithEngine()
      : file_system_(CreateTestFileSystem()),
        dummy_engine_metrics_(new metrics::EngineMetrics()),
        cdm_engine_(file_system_.get(), dummy_engine_metrics_) {}

 protected:
  std::unique_ptr<wvutil::FileSystem> file_system_;
  std::shared_ptr<metrics::EngineMetrics> dummy_engine_metrics_;
  TestCdmEngine cdm_engine_;
};

struct TestCryptoSessionConfig {
  // Disables newly created TestCryptoSession instances from installing
  // a test keybox.
  bool disable_test_keybox = false;
};

class TestCryptoSession : public CryptoSession {
 public:
  explicit TestCryptoSession(metrics::CryptoMetrics* crypto_metrics);
  TestCryptoSession(metrics::CryptoMetrics* crypto_metrics,
                    const TestCryptoSessionConfig* config);
  // This intercepts nonce flood errors, which is useful for tests that request
  // many nonces and are not time critical.
  CdmResponseType GenerateNonce(uint32_t* nonce) override;

 private:
  // Called once when TestCryptoSession is constructed.
  void MaybeInstallTestKeybox();
  bool IsTestKeyboxNeeded();

  // An un-owned pointer to the config.
  const TestCryptoSessionConfig* const config_ = nullptr;
};

// Given a PSSH data structure, this makes a PSSH string for use in
// generating a license request.
std::string MakePSSH(const video_widevine::WidevinePsshData& header);
// Given a serialized PSSH data, this generates a full PSSH string.
std::string MakePSSH(const std::string& serialized_header);

}  // namespace wvcdm

#endif  // WVCDM_CORE_TEST_BASE_H_