/*
 * Copyright (C) 2023 The Android Open Source Project
 *
 * Licensed under the Apache license, Version 2.0 (the "license");
 * you may not use this file except in compliance with the license.
 * You may obtain a copy of the license at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the license is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the license for the specific language governing permissions and
 * limitations under the license.
 *
 */

#include <fuzzer/FuzzedDataProvider.h>

#include "crypto_session.h"
#include "entitlement_key_session.h"
#include "properties.h"
#include "string_conversions.h"
#include "wv_cdm_constants.h"

using namespace wvcdm;

constexpr int32_t kStringLength = 20;
constexpr int32_t kMaxSamplesCount = 1000;
constexpr int32_t kMinSamplesCount = 1;
constexpr int32_t kMaxByte = 1000;
constexpr int32_t kMinByte = 0;
constexpr int32_t kIvSize = 16;
constexpr int32_t kDefaultSampleSize = 0;

class CryptoSessionFuzzer {
 public:
  CryptoSessionFuzzer(const uint8_t* data, size_t size) : fdp_(data, size){};
  void Process();
  void FillDecryptParams(CdmDecryptionParametersV16& params);
  void SetKeyParams(CryptoKey& input_key);
  void SetUp(CryptoSession* crypto_session, const std::string& key_id,
             std::string& message, std::string& signature,
             std::string& core_message);
  void InitializeStripeBuffer(std::vector<uint8_t>* buffer, size_t size);

 private:
  FuzzedDataProvider fdp_;
};

void CryptoSessionFuzzer::InitializeStripeBuffer(std::vector<uint8_t>* buffer,
                                                 size_t size) {
  buffer->assign(size, 0);
  for (size_t i = 0; i < size; ++i) {
    (*buffer)[i] = fdp_.ConsumeIntegral<uint8_t>();
  }
}

void CryptoSessionFuzzer::SetKeyParams(CryptoKey& input_key) {
  input_key.set_key_data(fdp_.ConsumeRandomLengthString(kStringLength));
  input_key.set_key_data_iv(fdp_.ConsumeRandomLengthString(kStringLength));
  input_key.set_key_control(fdp_.ConsumeRandomLengthString(kStringLength));
  input_key.set_key_control_iv(fdp_.ConsumeRandomLengthString(kStringLength));
  input_key.set_cipher_mode(fdp_.ConsumeBool() ? kCipherModeCbc
                                              : kCipherModeCtr);
  input_key.set_track_label(fdp_.ConsumeRandomLengthString(kStringLength));
  input_key.set_entitlement_key_id(
      fdp_.ConsumeRandomLengthString(kStringLength));
}

void CryptoSessionFuzzer::FillDecryptParams(
    CdmDecryptionParametersV16& params) {
  params.cipher_mode = fdp_.ConsumeBool() ? kCipherModeCbc : kCipherModeCtr;
  params.observe_legacy_fields = fdp_.ConsumeBool();
  uint32_t no_of_samples =
      fdp_.ConsumeIntegralInRange<uint32_t>(kMinSamplesCount, kMaxSamplesCount);

  while (--no_of_samples) {
    size_t sample_size = 0;
    uint32_t no_of_sub_samples = fdp_.ConsumeIntegralInRange<uint32_t>(
        kMinSamplesCount, kMaxSamplesCount);
    std::vector<CdmDecryptionSubsample> sub_sample_vector;
    while (--no_of_sub_samples) {
      size_t clear_bytes =
          fdp_.ConsumeIntegralInRange<size_t>(kMinByte, kMaxByte);
      size_t protected_bytes =
          fdp_.ConsumeIntegralInRange<size_t>(kMinByte, kMaxByte);
      CdmDecryptionSubsample sub_sample(clear_bytes, protected_bytes);
      sample_size += clear_bytes + protected_bytes;
      sub_sample_vector.push_back(sub_sample);
    }
    std::vector<uint8_t> iv;
    iv = fdp_.ConsumeBytes<uint8_t>(kIvSize);
    if (iv.size() != kIvSize) {
      iv.resize(kIvSize, 0);
    }
    void* decrypt_buffer;
    std::vector<uint8_t> e_buffer;
    for (int i = 0; i < sample_size; ++i) {
      e_buffer.push_back(
          fdp_.ConsumeIntegralInRange<uint8_t>(kMinByte, kMaxByte));
    }
    size_t decrypt_buffer_offset =
        fdp_.ConsumeIntegralInRange<uint8_t>(kMinByte, kMaxByte);
    CdmDecryptionSample sample(e_buffer.data(), decrypt_buffer,
                               decrypt_buffer_offset, sample_size, iv);
    sample.subsamples = sub_sample_vector;
    params.samples.push_back(sample);
  }
}

void CryptoSessionFuzzer::SetUp(CryptoSession* crypto_session,
                                const std::string& key_id, std::string& message,
                                std::string& signature,
                                std::string& core_message) {
  bool should_specify_algorithm;
  OEMCrypto_SignatureHashAlgorithm algorithm;
  CdmLicenseKeyType key_type_1 = kLicenseKeyTypeEntitlement;
  crypto_session->PrepareAndSignProvisioningRequest(
      message, &core_message, &signature, should_specify_algorithm, algorithm);
  crypto_session->PrepareAndSignLicenseRequest(
      message, &core_message, &signature, should_specify_algorithm, algorithm);
  crypto_session->LoadLicense(message, core_message, signature, key_type_1);

  CryptoKey input_key;
  input_key.set_key_id(key_id);
  CryptoSessionFuzzer::SetKeyParams(input_key);
  std::vector<CryptoKey> keys = {input_key};
  crypto_session->LoadEntitledContentKeys(keys);
}

void CryptoSessionFuzzer::Process() {
  metrics::CryptoMetrics crypto_metrics;
  std::unique_ptr<CryptoSession> crypto_session(
      CryptoSession::MakeCryptoSession(&crypto_metrics));

  CdmLicenseKeyType key_type;
  OEMCrypto_SignatureHashAlgorithm algorithm;
  Properties::Init();

  RequestedSecurityLevel requested_security_level =
      fdp_.ConsumeBool() ? kLevelDefault : kLevel3;
  crypto_session->Open(requested_security_level);

  while (fdp_.remaining_bytes()) {
    auto invoke_crypto_session_API = fdp_.PickValueInArray<
        const std::function<void()>>({
        [&]() {
          std::string token = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string additional_token =
              fdp_.ConsumeRandomLengthString(kStringLength);
          crypto_session->GetProvisioningToken(
              fdp_.ConsumeBool() ? &token : nullptr,
              fdp_.ConsumeBool() ? &additional_token : nullptr);
        },
        [&]() {
          crypto_session->SetSystemId(
              fdp_.ConsumeIntegral<uint32_t>() /*system_id*/);
        },
        [&]() {
          std::string provisioning_id;
          crypto_session->GetProvisioningId(&provisioning_id);
          crypto_session->GetExternalDeviceUniqueId(&provisioning_id);
        },
        [&]() {
          crypto_session->LoadLicense(
              fdp_.ConsumeRandomLengthString(kStringLength) /*signed_message*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*core_message*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*signature*/,
              key_type);
        },
        [&]() {
          std::string core_message =
              fdp_.ConsumeRandomLengthString(kStringLength);
          std::string signature = fdp_.ConsumeRandomLengthString(kStringLength);
          crypto_session->PrepareAndSignRenewalRequest(
              fdp_.ConsumeRandomLengthString(kStringLength) /*message*/,
              fdp_.ConsumeBool() ? &core_message : nullptr,
              fdp_.ConsumeBool() ? &signature : nullptr);
        },
        [&]() {
          crypto_session->LoadRenewal(
              fdp_.ConsumeRandomLengthString(kStringLength) /*signed_message*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*core_message*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*signature*/);
        },
        [&]() {
          const std::string& message =
              fdp_.ConsumeRandomLengthString(kStringLength);
          std::string core_message =
              fdp_.ConsumeRandomLengthString(kStringLength);
          std::string signature = fdp_.ConsumeRandomLengthString(kStringLength);
          OEMCrypto_SignatureHashAlgorithm algorithm;
          bool should_specify_algorithm = fdp_.ConsumeBool();
          crypto_session->PrepareAndSignProvisioningRequest(
              message, fdp_.ConsumeBool() ? &core_message : nullptr,
              fdp_.ConsumeBool() ? &signature : nullptr, should_specify_algorithm,
              algorithm);
        },
        [&]() {
          if (fdp_.ConsumeBool()) {
            const CryptoWrappedKey private_key;
            crypto_session->LoadCertificatePrivateKey(private_key);
          } else {
            CryptoWrappedKey::Type type =
                (CryptoWrappedKey::Type)fdp_.ConsumeIntegral<int32_t>();
            std::string key = fdp_.ConsumeRandomLengthString(kStringLength);
            const CryptoWrappedKey private_key(type, key);
            crypto_session->LoadCertificatePrivateKey(private_key);
          }
        },
        [&]() {
          std::string public_key;
          std::string public_key_signature;
          std::string wrapped_private_key;
          CryptoWrappedKey::Type key_type;
          crypto_session->GenerateCertificateKeyPair(
              &public_key, &public_key_signature, &wrapped_private_key, &key_type);
        },
        [&]() {
          if (fdp_.ConsumeBool()) {
            const CryptoWrappedKey private_key;
            crypto_session->LoadOemCertificatePrivateKey(private_key);
          } else {
            CryptoWrappedKey::Type type =
                (CryptoWrappedKey::Type)fdp_.ConsumeIntegral<int32_t>();
            std::string key = fdp_.ConsumeRandomLengthString(kStringLength);
            const CryptoWrappedKey private_key(type, key);
            crypto_session->LoadOemCertificatePrivateKey(private_key);
          }
        },
        [&]() {
          size_t out;
          RequestedSecurityLevel security_level =
              (RequestedSecurityLevel)fdp_.ConsumeIntegral<int32_t>();
          crypto_session->GetNumberOfOpenSessions(security_level, &out);
        },
        [&]() {
          std::string info = fdp_.ConsumeRandomLengthString(kStringLength);
          crypto_session->GetBuildInformation(&info);
        },
        [&]() {
          CdmWatermarkingSupport support =
              (CdmWatermarkingSupport)fdp_.ConsumeIntegral<int32_t>();
          crypto_session->GetWatermarkingSupport(&support);
        },
        [&]() {
          CdmProductionReadiness readiness =
              (CdmProductionReadiness)fdp_.ConsumeIntegral<int32_t>();
          crypto_session->GetProductionReadiness(&readiness);
        },
        [&]() {
          RequestedSecurityLevel security_level =
              (RequestedSecurityLevel)fdp_.ConsumeIntegral<int32_t>();
          size_t number_of_entries = fdp_.ConsumeIntegral<size_t>();
          crypto_session->GetMaximumUsageTableEntries(security_level,
                                                     &number_of_entries);
        },
        [&]() {
          RequestedSecurityLevel security_level =
              (RequestedSecurityLevel)fdp_.ConsumeIntegral<int32_t>();
          uint32_t decrypt_hash_support = fdp_.ConsumeIntegral<uint32_t>();
          crypto_session->GetDecryptHashSupport(security_level,
                                               &decrypt_hash_support);
        },
        [&]() {
          bool has_support = fdp_.ConsumeBool();
          crypto_session->HasUsageTableSupport(&has_support);
        },
        [&]() {
          crypto_session->DeactivateUsageInformation(
              fdp_.ConsumeRandomLengthString(
                  kStringLength) /*provider_session_token*/);
        },
        [&]() {
          std::string usage_report =
              fdp_.ConsumeRandomLengthString(kStringLength);
          CryptoSession::UsageDurationStatus usage_duration_status;
          int64_t seconds_since_started = fdp_.ConsumeIntegral<int64_t>();
          int64_t seconds_since_last_played = fdp_.ConsumeIntegral<int64_t>();
          crypto_session->GenerateUsageReport(
              fdp_.ConsumeRandomLengthString(
                  kStringLength) /*provider_session_token*/,
              &usage_report, &usage_duration_status, &seconds_since_started,
              &seconds_since_last_played);
        },
        [&]() {
          uint32_t nonce = fdp_.ConsumeIntegral<uint32_t>();
          crypto_session->GenerateNonce(&nonce);
        },
        [&]() {
          std::string signed_message, core_message, signature;
          std::string wrapped_private_key;
          crypto_session->LoadProvisioning(
              signed_message, core_message, signature,
              fdp_.ConsumeBool() ? &wrapped_private_key : nullptr);
        },
        [&]() {
          crypto_session->SetDecryptHash(
              fdp_.ConsumeIntegral<int16_t>() /*frame_number*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*hash*/);
        },
        [&]() {
          CdmEncryptionAlgorithm algorithm;
          std::string out_buffer;
          crypto_session->GenericDecrypt(
              fdp_.ConsumeRandomLengthString(kStringLength) /*in_buffer*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*key_id*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*iv*/, algorithm,
              &out_buffer);
        },
        [&]() {
          CdmSigningAlgorithm algorithm = fdp_.ConsumeBool()
                                              ? kSigningAlgorithmHmacSha256
                                              : kSigningAlgorithmUnknown;
          std::string signature;
          crypto_session->GenericSign(
              fdp_.ConsumeRandomLengthString(kStringLength) /*message*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*key_id*/,
              algorithm, &signature);
        },
        [&]() {
          CdmSigningAlgorithm algorithm = kSigningAlgorithmHmacSha256;
          const std::string signature;
          crypto_session->GenericVerify(
              fdp_.ConsumeRandomLengthString(kStringLength) /*message*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*key_id*/,
              algorithm, signature);
        },
        [&]() {
          UsageTableHeader usage_table_header;
          crypto_session->CreateUsageTableHeader(requested_security_level,
                                                &usage_table_header);

          crypto_session->LoadUsageTableHeader(requested_security_level,
                                              usage_table_header);
          uint32_t new_entry_count;
          crypto_session->ShrinkUsageTableHeader(
              requested_security_level, new_entry_count, &usage_table_header);
        },
        [&]() {
          const UsageTableHeader usage_table_header;
          crypto_session->LoadUsageTableHeader(requested_security_level,
                                              usage_table_header);
        },
        [&]() {
          uint32_t entry_number;
          crypto_session->CreateUsageEntry(&entry_number);
          crypto_session->MoveUsageEntry(entry_number);
        },
        [&]() {
          const UsageEntry usage_entry;
          crypto_session->LoadUsageEntry(
              fdp_.ConsumeIntegral<uint32_t>() /*entry_number*/, usage_entry);
        },
        [&]() {
          UsageTableHeader usage_table_header;
          UsageEntry usage_entry;
          crypto_session->UpdateUsageEntry(&usage_table_header, &usage_entry);
        },
        [&]() {
          bool can_support_output = fdp_.ConsumeBool();
          bool can_disable_output = fdp_.ConsumeBool();
          bool can_support_cgmsa = fdp_.ConsumeBool();
          crypto_session->GetAnalogOutputCapabilities(
              &can_support_output, &can_disable_output, &can_support_cgmsa);
        },
        [&]() {
          crypto_session->SetDebugIgnoreKeyboxCount(
              fdp_.ConsumeIntegral<uint32_t>() /*count*/);
        },
        [&]() { crypto_session->GetOkpFallbackPolicy(); },
        [&]() {
          std::string request = fdp_.ConsumeRandomLengthString(kStringLength);
          crypto_session->PrepareOtaProvisioningRequest(
              fdp_.ConsumeBool() /*use_test_key*/, &request);
        },
        [&]() {
          crypto_session->LoadOtaProvisioning(
              fdp_.ConsumeBool() /*use_test_key*/,
              fdp_.ConsumeRandomLengthString(kStringLength) /*response*/);
        },
        [&]() {
          uint32_t tier = fdp_.ConsumeIntegral<uint32_t>();
          crypto_session->GetResourceRatingTier(&tier);
        },
        [&]() {
          size_t max;
          crypto_session->GetMaxNumberOfSessions(requested_security_level, &max);
        },
        [&]() {
          std::string key_data = fdp_.ConsumeRandomLengthString(kStringLength);
          crypto_session->GetTokenFromKeybox(requested_security_level, &key_data);
        },
        [&]() {
          std::string bcc = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string additional_signature =
              fdp_.ConsumeRandomLengthString(kStringLength);
          crypto_session->GetBootCertificateChain(&bcc, &additional_signature);
        },
        [&]() {
          crypto_session->GenerateDerivedKeys(
              fdp_.ConsumeRandomLengthString(kStringLength) /*message*/);
        },
        [&]() {
          CdmDecryptionParametersV16 params_2;
          crypto_session->Decrypt(params_2);
        },
        [&]() {
          const std::string& message =
              fdp_.ConsumeRandomLengthString(kStringLength);
          std::string core_message, signature;
          bool should_specify_algorithm;
          OEMCrypto_SignatureHashAlgorithm algorithm;

          crypto_session->PrepareAndSignProvisioningRequest(
              message, &core_message, &signature, should_specify_algorithm,
              algorithm);
          crypto_session->PrepareAndSignLicenseRequest(
              message, &core_message, &signature, should_specify_algorithm,
              algorithm);
        },
        [&]() {
          std::string key_id = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string message = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string signature, core_message;
          SetUp(crypto_session.get(), key_id, message, signature, core_message);

          std::vector<uint8_t> in_vector;
          std::vector<uint8_t> iv_vector;
          std::string in_buffer;
          std::string iv;
          CdmEncryptionAlgorithm algorithm_encrypt =
              kEncryptionAlgorithmAesCbc128;
          std::string out_buffer;
          InitializeStripeBuffer(&in_vector, CONTENT_KEY_SIZE * 15);
          in_buffer = std::string(in_vector.begin(), in_vector.end());
          InitializeStripeBuffer(&iv_vector, KEY_IV_SIZE);
          iv = std::string(iv_vector.begin(), iv_vector.end());

          crypto_session->GenericEncrypt(in_buffer, key_id, iv, algorithm_encrypt,
                                        &out_buffer);
        },
        [&]() {
          std::string key_id = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string message = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string signature, core_message;
          SetUp(crypto_session.get(), key_id, message, signature, core_message);

          std::vector<uint8_t> in_vector;
          std::vector<uint8_t> iv_vector;
          std::string in_buffer;
          std::string iv;
          CdmEncryptionAlgorithm algorithm_encrypt =
              kEncryptionAlgorithmAesCbc128;
          std::string out_buffer;
          InitializeStripeBuffer(&in_vector, CONTENT_KEY_SIZE * 15);
          in_buffer = std::string(in_vector.begin(), in_vector.end());
          InitializeStripeBuffer(&iv_vector, KEY_IV_SIZE);
          iv = std::string(iv_vector.begin(), iv_vector.end());

          crypto_session->GenericDecrypt(in_buffer, key_id, iv, algorithm_encrypt,
                                        &out_buffer);
        },
        [&]() {
          std::string key_id = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string message = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string signature, core_message;
          SetUp(crypto_session.get(), key_id, message, signature, core_message);

          CdmSigningAlgorithm algorithm_2 = kSigningAlgorithmHmacSha256;
          crypto_session->GenericSign(message, key_id, algorithm_2, &signature);
        },
        [&]() {
          std::string key_id = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string message = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string signature, core_message;
          SetUp(crypto_session.get(), key_id, message, signature, core_message);

          CdmSigningAlgorithm algorithm_2 = kSigningAlgorithmHmacSha256;
          crypto_session->GenericVerify(message, key_id, algorithm_2, signature);
        },
        [&]() {
          const std::string key_id =
              fdp_.ConsumeRandomLengthString(kStringLength);
          std::string message = fdp_.ConsumeRandomLengthString(kStringLength);
          std::string signature, core_message;
          SetUp(crypto_session.get(), key_id, message, signature, core_message);
          CdmDecryptionParametersV16 params(key_id);
          CryptoSessionFuzzer::FillDecryptParams(params);
          crypto_session->Decrypt(params);
        },
    });
    invoke_crypto_session_API();
  }
  crypto_session->Close();
}

extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
  CryptoSessionFuzzer crypto_session_fuzzer(data, size);
  crypto_session_fuzzer.Process();
  return 0;
}