// Copyright 2023 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//
// Test data for OEMCrypto unit tests.
//
#ifndef CDM_OEMCRYPTO_USAGE_TABLE_TEST_
#define CDM_OEMCRYPTO_USAGE_TABLE_TEST_

#include <gtest/gtest.h>
#include <openssl/hmac.h>
#include <openssl/sha.h>

#include "OEMCryptoCENC.h"
#include "log.h"
#include "oemcrypto_basic_test.h"
#include "oemcrypto_license_test.h"
#include "test_sleep.h"

namespace wvoec {

// This class is for testing the generic crypto functionality.
class OEMCryptoGenericCryptoTest : public OEMCryptoRefreshTest {
 protected:
  // buffer_size_ must be a multiple of encryption block size, 16.  We'll use a
  // reasonable number of blocks for most of the tests.
  OEMCryptoGenericCryptoTest() : buffer_size_(160) {}

  void SetUp() override {
    OEMCryptoRefreshTest::SetUp();
    if (!wvoec::global_features.generic_crypto) {
      GTEST_SKIP() << "Test for devices with generic crypto API only";
    }
    ASSERT_NO_FATAL_FAILURE(license_messages_.SignAndVerifyRequest());
    ASSERT_NO_FATAL_FAILURE(
        license_messages_.CreateResponseWithGenericCryptoKeys());
    InitializeClearBuffer();
  }

  void InitializeClearBuffer() {
    clear_buffer_.assign(buffer_size_, 0);
    for (size_t i = 0; i < clear_buffer_.size(); i++) {
      clear_buffer_[i] = 1 + i % 250;
    }
    for (size_t i = 0; i < wvoec::KEY_IV_SIZE; i++) {
      iv_[i] = i;
    }
  }

  void ResizeBuffer(size_t new_size) {
    buffer_size_ = new_size;
    InitializeClearBuffer();  // Re-initialize the clear buffer.
  }

  void EncryptAndLoadKeys() {
    ASSERT_NO_FATAL_FAILURE(license_messages_.EncryptAndSignResponse());
    ASSERT_EQ(OEMCrypto_SUCCESS, license_messages_.LoadResponse());
  }

  // Encrypt the buffer with the specified key made in
  // CreateResponseWithGenericCryptoKeys.
  void EncryptBuffer(unsigned int key_index, const vector<uint8_t>& in_buffer,
                     vector<uint8_t>* out_buffer) {
    EncryptBufferWithKey(session_.license().keys[key_index].key_data, in_buffer,
                         out_buffer);
  }
  // Encrypt the buffer with the specified key.
  void EncryptBufferWithKey(const uint8_t* key_data,
                            const vector<uint8_t>& in_buffer,
                            vector<uint8_t>* out_buffer) {
    AES_KEY aes_key;
    ASSERT_EQ(0, AES_set_encrypt_key(key_data, AES_BLOCK_SIZE * 8, &aes_key));
    uint8_t iv_buffer[wvoec::KEY_IV_SIZE];
    memcpy(iv_buffer, iv_, wvoec::KEY_IV_SIZE);
    out_buffer->resize(in_buffer.size());
    ASSERT_GT(in_buffer.size(), 0u);
    ASSERT_EQ(0u, in_buffer.size() % AES_BLOCK_SIZE);
    AES_cbc_encrypt(in_buffer.data(), out_buffer->data(), in_buffer.size(),
                    &aes_key, iv_buffer, AES_ENCRYPT);
  }

  // Sign the buffer with the specified key made in
  // CreateResponseWithGenericCryptoKeys.
  void SignBuffer(unsigned int key_index, const vector<uint8_t>& in_buffer,
                  vector<uint8_t>* signature) {
    SignBufferWithKey(session_.license().keys[key_index].key_data, in_buffer,
                      signature);
  }

  // Sign the buffer with the specified key.
  void SignBufferWithKey(const uint8_t* key_data,
                         const vector<uint8_t>& in_buffer,
                         vector<uint8_t>* signature) {
    unsigned int md_len = SHA256_DIGEST_LENGTH;
    signature->resize(SHA256_DIGEST_LENGTH);
    HMAC(EVP_sha256(), key_data, wvoec::MAC_KEY_SIZE, in_buffer.data(),
         in_buffer.size(), signature->data(), &md_len);
  }

  OEMCryptoResult GenericEncrypt(const uint8_t* key_handle,
                                 size_t key_handle_length,
                                 const uint8_t* clear_buffer,
                                 size_t clear_buffer_length, const uint8_t* iv,
                                 OEMCrypto_Algorithm algorithm,
                                 uint8_t* out_buffer) {
    if (ShouldGenerateCorpus()) {
      const std::string file_name =
          GetFileName("oemcrypto_generic_encrypt_fuzz_seed_corpus");
      OEMCrypto_Generic_Api_Fuzz fuzzed_structure;
      fuzzed_structure.cipher_mode = OEMCrypto_CipherMode_CENC;
      fuzzed_structure.algorithm = algorithm;
      // Cipher mode and algorithm.
      AppendToFile(file_name, reinterpret_cast<const char*>(&fuzzed_structure),
                   sizeof(fuzzed_structure));
      AppendToFile(file_name, reinterpret_cast<const char*>(iv),
                   wvoec::KEY_IV_SIZE);
      AppendSeparator(file_name);
      AppendToFile(file_name, reinterpret_cast<const char*>(clear_buffer),
                   clear_buffer_length);
    }
    return OEMCrypto_Generic_Encrypt(key_handle, key_handle_length,
                                     clear_buffer, clear_buffer_length, iv,
                                     algorithm, out_buffer);
  }

  OEMCryptoResult GenericDecrypt(
      const uint8_t* key_handle, size_t key_handle_length,
      const uint8_t* encrypted_buffer, size_t encrypted_buffer_length,
      const uint8_t* iv, OEMCrypto_Algorithm algorithm, uint8_t* out_buffer) {
    if (ShouldGenerateCorpus()) {
      const std::string file_name =
          GetFileName("oemcrypto_generic_decrypt_fuzz_seed_corpus");
      OEMCrypto_Generic_Api_Fuzz fuzzed_structure;
      fuzzed_structure.cipher_mode = OEMCrypto_CipherMode_CENC;
      fuzzed_structure.algorithm = algorithm;
      // Cipher mode and algorithm.
      AppendToFile(file_name, reinterpret_cast<const char*>(&fuzzed_structure),
                   sizeof(fuzzed_structure));
      AppendToFile(file_name, reinterpret_cast<const char*>(iv),
                   wvoec::KEY_IV_SIZE);
      AppendSeparator(file_name);
      AppendToFile(file_name, reinterpret_cast<const char*>(encrypted_buffer),
                   encrypted_buffer_length);
    }
    return OEMCrypto_Generic_Decrypt(key_handle, key_handle_length,
                                     encrypted_buffer, encrypted_buffer_length,
                                     iv, algorithm, out_buffer);
  }

  OEMCryptoResult GenericVerify(const uint8_t* key_handle,
                                size_t key_handle_length,
                                const uint8_t* clear_buffer,
                                size_t clear_buffer_length,
                                OEMCrypto_Algorithm algorithm,
                                const uint8_t* signature,
                                size_t signature_length) {
    if (ShouldGenerateCorpus()) {
      const std::string file_name =
          GetFileName("oemcrypto_generic_verify_fuzz_seed_corpus");
      OEMCrypto_Generic_Api_Fuzz fuzzed_structure;
      fuzzed_structure.cipher_mode = OEMCrypto_CipherMode_CENC;
      fuzzed_structure.algorithm = algorithm;
      // Cipher mode and algorithm.
      AppendToFile(file_name, reinterpret_cast<const char*>(&fuzzed_structure),
                   sizeof(fuzzed_structure));
      AppendToFile(file_name, reinterpret_cast<const char*>(clear_buffer),
                   clear_buffer_length);
      AppendSeparator(file_name);
      AppendToFile(file_name, reinterpret_cast<const char*>(signature),
                   signature_length);
    }
    return OEMCrypto_Generic_Verify(key_handle, key_handle_length, clear_buffer,
                                    clear_buffer_length, algorithm, signature,
                                    signature_length);
  }

  OEMCryptoResult GenericSign(const uint8_t* key_handle,
                              size_t key_handle_length,
                              const uint8_t* clear_buffer,
                              size_t clear_buffer_length,
                              OEMCrypto_Algorithm algorithm, uint8_t* signature,
                              size_t* signature_length) {
    if (ShouldGenerateCorpus()) {
      const std::string file_name =
          GetFileName("oemcrypto_generic_sign_fuzz_seed_corpus");
      OEMCrypto_Generic_Api_Fuzz fuzzed_structure;
      fuzzed_structure.cipher_mode = OEMCrypto_CipherMode_CENC;
      fuzzed_structure.algorithm = algorithm;
      // Cipher mode and algorithm.
      AppendToFile(file_name, reinterpret_cast<const char*>(&fuzzed_structure),
                   sizeof(fuzzed_structure));
      AppendToFile(file_name, reinterpret_cast<const char*>(clear_buffer),
                   clear_buffer_length);
    }
    return OEMCrypto_Generic_Sign(key_handle, key_handle_length, clear_buffer,
                                  clear_buffer_length, algorithm, signature,
                                  signature_length);
  }

  // This asks OEMCrypto to encrypt with the specified key, and expects a
  // failure.
  void BadEncrypt(unsigned int key_index, OEMCrypto_Algorithm algorithm,
                  size_t buffer_length) {
    OEMCryptoResult sts;
    vector<uint8_t> expected_encrypted;
    EncryptBuffer(key_index, clear_buffer_, &expected_encrypted);
    vector<uint8_t> key_handle;
    sts = GetKeyHandleIntoVector(
        session_.session_id(), session_.license().keys[key_index].key_id,
        session_.license().keys[key_index].key_id_length,
        OEMCrypto_CipherMode_CENC, key_handle);
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
    vector<uint8_t> encrypted(buffer_length);
    sts = GenericEncrypt(key_handle.data(), key_handle.size(),
                         clear_buffer_.data(), buffer_length, iv_, algorithm,
                         encrypted.data());
    EXPECT_NE(OEMCrypto_SUCCESS, sts);
    expected_encrypted.resize(buffer_length);
    EXPECT_NE(encrypted, expected_encrypted);
  }

  // This asks OEMCrypto to decrypt with the specified key, and expects a
  // failure.
  void BadDecrypt(unsigned int key_index, OEMCrypto_Algorithm algorithm,
                  size_t buffer_length) {
    OEMCryptoResult sts;
    vector<uint8_t> encrypted;
    EncryptBuffer(key_index, clear_buffer_, &encrypted);
    vector<uint8_t> key_handle;
    sts = GetKeyHandleIntoVector(
        session_.session_id(), session_.license().keys[key_index].key_id,
        session_.license().keys[key_index].key_id_length,
        OEMCrypto_CipherMode_CENC, key_handle);
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
    vector<uint8_t> resultant(encrypted.size());
    sts = GenericDecrypt(key_handle.data(), key_handle.size(), encrypted.data(),
                         buffer_length, iv_, algorithm, resultant.data());
    EXPECT_NE(OEMCrypto_SUCCESS, sts);
    EXPECT_NE(clear_buffer_, resultant);
  }

  // This asks OEMCrypto to sign with the specified key, and expects a
  // failure.
  void BadSign(unsigned int key_index, OEMCrypto_Algorithm algorithm) {
    OEMCryptoResult sts;
    vector<uint8_t> expected_signature;
    SignBuffer(key_index, clear_buffer_, &expected_signature);

    vector<uint8_t> key_handle;
    sts = GetKeyHandleIntoVector(
        session_.session_id(), session_.license().keys[key_index].key_id,
        session_.license().keys[key_index].key_id_length,
        OEMCrypto_CipherMode_CENC, key_handle);
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
    size_t signature_length = (size_t)SHA256_DIGEST_LENGTH;
    vector<uint8_t> signature(SHA256_DIGEST_LENGTH);
    sts = GenericSign(key_handle.data(), key_handle.size(),
                      clear_buffer_.data(), clear_buffer_.size(), algorithm,
                      signature.data(), &signature_length);
    EXPECT_NE(OEMCrypto_SUCCESS, sts);
    EXPECT_NE(signature, expected_signature);
  }

  // This asks OEMCrypto to verify a signature with the specified key, and
  // expects a failure.
  void BadVerify(unsigned int key_index, OEMCrypto_Algorithm algorithm,
                 size_t signature_size, bool alter_data) {
    OEMCryptoResult sts;
    vector<uint8_t> signature;
    SignBuffer(key_index, clear_buffer_, &signature);
    if (alter_data) {
      signature[0] ^= 42;
    }
    if (signature.size() < signature_size) {
      signature.resize(signature_size);
    }

    vector<uint8_t> key_handle;
    sts = GetKeyHandleIntoVector(
        session_.session_id(), session_.license().keys[key_index].key_id,
        session_.license().keys[key_index].key_id_length,
        OEMCrypto_CipherMode_CENC, key_handle);
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
    sts = GenericVerify(key_handle.data(), key_handle.size(),
                        clear_buffer_.data(), clear_buffer_.size(), algorithm,
                        signature.data(), signature_size);
    EXPECT_NE(OEMCrypto_SUCCESS, sts);
  }

  // This must be a multiple of encryption block size.
  size_t buffer_size_;
  vector<uint8_t> clear_buffer_;
  vector<uint8_t> encrypted_buffer_;
  uint8_t iv_[wvoec::KEY_IV_SIZE];
};

class OEMCryptoUsageTableTest : public OEMCryptoGenericCryptoTest {
 public:
  void SetUp() override { OEMCryptoGenericCryptoTest::SetUp(); }

  virtual void ShutDown() {
    ASSERT_NO_FATAL_FAILURE(session_.close());
    ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_Terminate());
  }

  virtual void Restart() {
    OEMCrypto_SetSandbox(kTestSandbox, sizeof(kTestSandbox));
    ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_Initialize());
    (void)OEMCrypto_SetMaxAPIVersion(kCurrentAPI);
    (void)OEMCrypto_EnterTestMode();
    EnsureTestROT();
    ASSERT_NO_FATAL_FAILURE(session_.open());
  }

  void PrintDotsWhileSleep(int64_t total_seconds, int64_t interval_seconds) {
    int64_t dot_time = interval_seconds;
    int64_t elapsed_time = 0;
    const int64_t start_time = wvutil::Clock().GetCurrentTime();
    do {
      wvutil::TestSleep::Sleep(1);
      elapsed_time = wvutil::Clock().GetCurrentTime() - start_time;
      if (elapsed_time >= dot_time) {
        cout << ".";
        cout.flush();
        dot_time += interval_seconds;
      }
    } while (elapsed_time < total_seconds);
    cout << endl;
  }

  OEMCryptoResult LoadUsageTableHeader(
      const vector<uint8_t>& encrypted_usage_header) {
    return OEMCrypto_LoadUsageTableHeader(encrypted_usage_header.data(),
                                          encrypted_usage_header.size());
  }
};

}  // namespace wvoec

#endif  // CDM_OEMCRYPTO_USAGE_TABLE_TEST_