/*
 * 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 <cdm_engine.h>
#include <cdm_engine_factory.h>
#include <fuzzer/FuzzedDataProvider.h>
#include <stdlib.h>
#include <wv_cdm_constants.h>
#include <wv_cdm_event_listener.h>
#include <wv_cdm_types.h>

#include <fstream>
#include <iostream>
#include <string>

#include "vendor_widevine_fuzz_helper.h"

const int32_t kMinByte = 0;
const int32_t kMaxByte = 256;
const int32_t kMinEnum = 0;
const int32_t kMaxEnum = 4;
const int32_t kMinValue = 0;
const int32_t kMaxValue = 1;
const int32_t kMinKey = 10;
const int32_t kMaxKey = 1000;
const int32_t kMinSize = 1;
const int32_t kMaxSize = 100;
const char *kKeySystem = "widevine";
const char *kKeySetIdPrefix = "ksid";
const char *kEmptyString = "";
const std::string kUsageFilePrefixL1 = "/data/vendor/mediadrm/IDM0/L1/";
const std::string kUsageFilePrefixL3 = "/data/vendor/mediadrm/IDM0/L3/";
const std::string kCertFile = kUsageFilePrefixL3 + "cert.bin";
const std::string kCertFile1 = kUsageFilePrefixL3 + "cert1.bin";
static constexpr int32_t kMaxRuns = 500;

using namespace wvutil;
using namespace wvcdm::metrics;
using namespace drm_metrics;
using namespace wvcdm;

std::string kQueryToken[] = {QUERY_KEY_SECURITY_LEVEL,
                             QUERY_KEY_PLAY_ALLOWED,
                             QUERY_KEY_CURRENT_HDCP_LEVEL,
                             QUERY_KEY_MAX_HDCP_LEVEL,
                             QUERY_KEY_USAGE_SUPPORT,
                             QUERY_KEY_NUMBER_OF_OPEN_SESSIONS,
                             QUERY_KEY_MAX_NUMBER_OF_SESSIONS,
                             QUERY_KEY_OEMCRYPTO_API_VERSION,
                             QUERY_KEY_CURRENT_SRM_VERSION,
                             QUERY_KEY_SRM_UPDATE_SUPPORT,
                             QUERY_KEY_WVCDM_VERSION,
                             QUERY_KEY_RESOURCE_RATING_TIER,
                             QUERY_KEY_OEMCRYPTO_BUILD_INFORMATION,
                             QUERY_KEY_DECRYPT_HASH_SUPPORT,
                             QUERY_KEY_PROVISIONING_MODEL,
                             QUERY_KEY_MAX_USAGE_TABLE_ENTRIES,
                             QUERY_KEY_OEMCRYPTO_API_MINOR_VERSION,
                             QUERY_KEY_ANALOG_OUTPUT_CAPABILITIES,
                             QUERY_KEY_CAN_DISABLE_ANALOG_OUTPUT,
                             QUERY_KEY_WATERMARKING_SUPPORT,
                             QUERY_KEY_PRODUCTION_READY,
                             QUERY_KEY_SYSTEM_ID,
                             QUERY_KEY_DEBUG_BOOT_CERTIFICATE_CHAIN,
                             QUERY_KEY_DEVICE_ID,
                             QUERY_KEY_PROVISIONING_ID};

std::string kInitType[] = {ISO_BMFF_VIDEO_MIME_TYPE, WEBM_VIDEO_MIME_TYPE,
                           HLS_INIT_DATA_FORMAT,     CENC_INIT_DATA_FORMAT,
                           ISO_BMFF_AUDIO_MIME_TYPE, WEBM_AUDIO_MIME_TYPE,
                           WEBM_INIT_DATA_FORMAT};

class FuzzWvCdmEventListener : public WvCdmEventListener {
 public:
  void OnSessionRenewalNeeded(const CdmSessionId & /*session_id*/) {}
  void OnSessionKeysChange(const CdmSessionId & /*session_id*/,
                           const CdmKeyStatusMap & /*keys_status*/,
                           bool /*has_new_usable_key*/) {}
  void OnExpirationUpdate(const CdmSessionId & /*session_id*/,
                          int64_t /*new_expiry_time_seconds*/) {}
};

class CdmEngineFuzzer {
 public:
  CdmEngineFuzzer(const uint8_t *data, size_t size) : fdp_(data, size){};
  void Process();
  void Init(CdmEngine *cdm_engine);
  void CreateUsageFiles(std::string app_id_);
  std::vector<std::string> CreateVector(bool enable_key_set_id);

 private:
  FuzzedDataProvider fdp_;
  std::ofstream out_cert_file_1_;
  std::ofstream out_cert_file_2_;
  std::ofstream out_usage_file_1_;
  std::ofstream out_usage_file_2_;
  std::string app_id_;
  std::string key_id_;
  std::string sign_;
  std::string default_url_;
  std::unique_ptr<wvutil::FileSystem> file_system_ =
      std::make_unique<wvutil::FileSystem>();
  std::unique_ptr<FuzzCdmClientPropertySet> fuzz_cdm_client_property_ =
      std::make_unique<FuzzCdmClientPropertySet>(&fdp_);
  CdmSessionId session_id_;
  CdmKeySystem key_system_;
  CdmKeySetId key_set_id_;
};

void CdmEngineFuzzer::CreateUsageFiles(std::string app_id_) {
  /**
   * These files are being fetched by the GetUsageInfo() &
   * RemoveAllUsageInfo() APIs, since they were not created
   * by any other API calls in the fuzzer, we are creating them.
   */
  DeviceFiles device_files(file_system_.get());
  std::string file_name = device_files.GetUsageInfoFileName(app_id_);
  std::string create_file_1 = kUsageFilePrefixL1 + file_name;
  std::string create_file_2 = kUsageFilePrefixL3 + file_name;
  out_usage_file_1_.open(create_file_1.c_str());
  out_usage_file_2_.open(create_file_2.c_str());
}

std::vector<std::string> CdmEngineFuzzer::CreateVector(bool enable_key_set_id) {
  int32_t size = fdp_.ConsumeIntegralInRange(kMinSize, kMaxSize);
  std::vector<std::string> vector;
  for (int i = 0; i < size; ++i) {
    vector.push_back(enable_key_set_id ? key_set_id_
                                    : fdp_.ConsumeRandomLengthString(kMaxByte));
  }
  return vector;
}

void CdmEngineFuzzer::Init(CdmEngine *cdm_engine) {
  app_id_ =
      std::to_string(fdp_.ConsumeIntegralInRange<uint32_t>(kMinSize, kMaxSize));
  key_id_ =
      std::to_string(fdp_.ConsumeIntegralInRange<uint32_t>(kMinSize, kMaxSize));
  sign_ = fdp_.ConsumeRandomLengthString(kMaxByte);

  key_system_ = fdp_.ConsumeBool() ? kKeySystem
                                  : fdp_.ConsumeRandomLengthString(kMaxByte);

  FuzzWvCdmEventListener fuzz_wv_cdm_event_listener;
  cdm_engine->OpenSession(
      key_system_,
      fdp_.ConsumeBool() ? fuzz_cdm_client_property_.get() : nullptr,
      fdp_.ConsumeBool() ? &fuzz_wv_cdm_event_listener : nullptr, &session_id_);

  std::string request;
  default_url_ = fdp_.ConsumeRandomLengthString(kMaxByte);
  cdm_engine->GetProvisioningRequest(
      (CdmCertificateType)fdp_.ConsumeIntegralInRange<int32_t>(kMinValue,
                                                               kMaxValue),
      fdp_.ConsumeRandomLengthString(kMaxByte) /* cert_authority */,
      fdp_.ConsumeRandomLengthString(kMaxByte) /* service_certificate */,
      (RequestedSecurityLevel)fdp_.ConsumeIntegralInRange<uint32_t>(kMinValue,
                                                                    kMaxValue),
      fdp_.ConsumeBool() ? &request : nullptr, &default_url_);

  std::string cert;
  std::string wrapped_key;
  CdmProvisioningResponse response_message =
      fdp_.ConsumeRandomLengthString(kMaxByte);
  cdm_engine->HandleProvisioningResponse(
      response_message,
      (RequestedSecurityLevel)fdp_.ConsumeIntegralInRange<uint32_t>(kMinValue,
                                                                    kMaxValue),
      fdp_.ConsumeBool() ? &cert : nullptr,
      fdp_.ConsumeBool() ? &wrapped_key : nullptr);

  CdmAppParameterMap app_parameters;
  app_parameters.insert({fdp_.ConsumeRandomLengthString(kMaxByte),
                        fdp_.ConsumeRandomLengthString(kMaxByte)});
  CdmInitData data = fdp_.ConsumeRandomLengthString(kMaxByte);
  const std::string oec_version =
      std::to_string(fdp_.ConsumeIntegralInRange<uint32_t>(kMinSize, kMaxSize));
  InitializationData init_data(fdp_.PickValueInArray<std::string>(kInitType),
                               data, oec_version);

  key_set_id_ =
      kKeySetIdPrefix +
      std::to_string(fdp_.ConsumeIntegralInRange<uint32_t>(kMinKey, kMaxKey));

  CdmLicenseType generate_lic_type =
      static_cast<CdmLicenseType>(fdp_.ConsumeIntegralInRange(
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeOffline),
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeEmbeddedKeyData)));
  CdmKeyRequest key_request;
  key_request.message = fdp_.ConsumeRandomLengthString(kMaxByte);
  key_request.type = static_cast<CdmKeyRequestType>(fdp_.ConsumeIntegralInRange(
      static_cast<uint32_t>(CdmKeyRequestType::kKeyRequestTypeUnknown),
      static_cast<uint32_t>(CdmKeyRequestType::kKeyRequestTypeUpdate)));
  key_request.url = fdp_.ConsumeBool()
                       ? default_url_
                       : fdp_.ConsumeRandomLengthString(kMaxByte);
  cdm_engine->GenerateKeyRequest(session_id_,
                                fdp_.ConsumeBool() ? key_set_id_ : kEmptyString,
                                init_data, generate_lic_type, app_parameters,
                                fdp_.ConsumeBool() ? &key_request : nullptr);

  CdmLicenseType add_lic_type =
      static_cast<CdmLicenseType>(fdp_.ConsumeIntegralInRange(
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeOffline),
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeEmbeddedKeyData)));
  CdmKeyResponse key_data =
      std::to_string(fdp_.ConsumeIntegralInRange<uint32_t>(kMinKey, kMaxKey));
  cdm_engine->AddKey(session_id_, fdp_.ConsumeBool() ? key_data : kEmptyString,
                    &add_lic_type, &key_set_id_);
}

void CdmEngineFuzzer::Process() {
  CdmEngine *cdm_engine = CdmEngineFactory::CreateCdmEngine(file_system_.get());

  CdmEngineFuzzer::Init(cdm_engine);

   int32_t runs = kMaxRuns;
  /* Limited the while loop to prevent a timeout caused by the
  /* CryptoSession constructor, which took time to initialize
  /* OEMCrypto in each iteration.*/
  while (fdp_.remaining_bytes() > 0 && --runs) {
    auto invoke_cdm_engine_API = fdp_.PickValueInArray<
        const std::function<void()>>({
        [&]() { cdm_engine->OnTimerEvent(); },
        [&]() { cdm_engine->IsOpenSession(session_id_); },
        [&]() { cdm_engine->SetDefaultOtaKeyboxFallbackDurationRules(); },
        [&]() { cdm_engine->SetFastOtaKeyboxFallbackDurationRules(); },
        [&]() {
          cdm_engine->OpenSession(
              key_system_, nullptr /* property_set */,
              fdp_.ConsumeRandomLengthString(kMaxByte) /* forced_session_id */,
              nullptr /* event_listener */);
        },
        [&]() {
          cdm_engine->StoreAtscLicense(
              static_cast<RequestedSecurityLevel>(fdp_.ConsumeIntegralInRange(
                  static_cast<int32_t>(RequestedSecurityLevel::kLevelDefault),
                  static_cast<int32_t>(RequestedSecurityLevel::kLevel3))),
              key_set_id_,
              fdp_.ConsumeRandomLengthString(
                  kMaxByte) /* serialized_license_data */);
        },
        [&]() { cdm_engine->RestoreKey(session_id_, key_set_id_); },
        [&]() {
          cdm_engine->SetSessionServiceCertificate(
              session_id_, fdp_.ConsumeRandomLengthString(
                              kMaxByte) /* service_certificate */);
        },
        [&]() {
          std::string query_response;
          cdm_engine->QueryStatus(
              static_cast<RequestedSecurityLevel>(fdp_.ConsumeIntegralInRange(
                  static_cast<int32_t>(RequestedSecurityLevel::kLevelDefault),
                  static_cast<int32_t>(RequestedSecurityLevel::kLevel3))),
              fdp_.PickValueInArray<std::string>(kQueryToken),
              fdp_.ConsumeBool() ? &query_response : nullptr);
        },
        [&]() {
          CdmQueryMap query_response;
          cdm_engine->QuerySessionStatus(session_id_, &query_response);
        },
        [&]() {
          CdmKeyRequest key_request;
          key_request.message = fdp_.ConsumeRandomLengthString(kMaxByte);
          key_request.type =
              static_cast<CdmKeyRequestType>(fdp_.ConsumeIntegralInRange(
                  static_cast<uint32_t>(
                      CdmKeyRequestType::kKeyRequestTypeUnknown),
                  static_cast<uint32_t>(
                      CdmKeyRequestType::kKeyRequestTypeUpdate)));
          key_request.url = default_url_;

          cdm_engine->GenerateRenewalRequest(
              session_id_, fdp_.ConsumeBool() ? &key_request : nullptr);

          CdmKeyResponse key_data =
              fdp_.ConsumeBool()
                  ? std::to_string(
                        fdp_.ConsumeIntegralInRange<uint32_t>(kMinKey, kMaxKey))
                  : kEmptyString;
          cdm_engine->RenewKey(session_id_, key_data);
        },
        [&]() { cdm_engine->IsReleaseSession(session_id_); },
        [&]() { cdm_engine->IsOfflineSession(session_id_); },
        [&]() {
          out_cert_file_1_.open(kCertFile.c_str());
          out_cert_file_2_.open(kCertFile1.c_str());
          cdm_engine->RemoveOfflineLicense(
              key_set_id_,
              static_cast<CdmSecurityLevel>(fdp_.ConsumeIntegralInRange(
                  static_cast<int32_t>(
                      CdmSecurityLevel::kSecurityLevelUninitialized),
                  static_cast<int32_t>(
                      CdmSecurityLevel::kSecurityLevelUnknown))));
        },
        [&]() {
          CdmQueryMap query_response;
          cdm_engine->QueryKeyStatus(session_id_, &query_response);
        },
        [&]() {
          cdm_engine->IsSecurityLevelSupported(
              (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum));
        },
        [&]() {
          CdmKeyAllowedUsage key_usage;
          cdm_engine->QueryKeyAllowedUsage(
              session_id_, key_id_, fdp_.ConsumeBool() ? &key_usage : nullptr);
        },
        [&]() {
          CdmKeyAllowedUsage key_usage;
          cdm_engine->QueryKeyAllowedUsage(
              key_id_, fdp_.ConsumeBool() ? &key_usage : nullptr);
        },
        [&]() {
          CdmQueryMap query_response;
          cdm_engine->QueryOemCryptoSessionId(session_id_, &query_response);
        },
        [&]() {
          std::vector<std::string> key_set_ids = CreateVector(true);
          cdm_engine->ListStoredLicenses(
              (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              fdp_.ConsumeBool() ? &key_set_ids : nullptr);
        },
        [&]() {
          std::vector<std::string> key_set_ids = CreateVector(true);
          std::vector<std::string> provider_session_tokens = CreateVector(false);

          cdm_engine->ListUsageIds(
              app_id_,
              (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              fdp_.ConsumeBool() ? &key_set_ids : nullptr,
              fdp_.ConsumeBool() ? &provider_session_tokens : nullptr);
        },
        [&]() {
          cdm_engine->DeleteUsageRecord(
              app_id_,
              (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              key_set_id_);
        },
        [&]() {
          CdmOfflineLicenseState license_state =
              static_cast<CdmOfflineLicenseState>(fdp_.ConsumeIntegralInRange(
                  static_cast<int32_t>(
                      CdmOfflineLicenseState::kLicenseStateActive),
                  static_cast<int32_t>(
                      CdmOfflineLicenseState::kLicenseStateUnknown)));
          cdm_engine->GetOfflineLicenseState(
              key_set_id_,
              (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              fdp_.ConsumeBool() ? &license_state : nullptr);
        },
        [&]() {
          cdm_engine->SetPlaybackId(session_id_,
                                   fdp_.ConsumeRandomLengthString(kMaxByte));
        },
        [&]() {
          CdmUsageReport usage_report;
          CreateUsageFiles(app_id_);
          int error_detail;
          cdm_engine->GetUsageInfo(app_id_,
                                  fdp_.ConsumeBool() ? &error_detail : nullptr,
                                  fdp_.ConsumeBool() ? &usage_report : nullptr);
        },
        [&]() {
          CreateUsageFiles(app_id_);
          cdm_engine->RemoveAllUsageInfo(app_id_);
        },
        [&]() {
          cdm_engine->ReleaseUsageInfo(
              fdp_.ConsumeRandomLengthString(kMaxByte) /* message */);
        },
        [&]() {
          std::string release_message = fdp_.ConsumeRandomLengthString(kMaxByte);
          cdm_engine->LoadUsageSession(
              key_set_id_, fdp_.ConsumeBool() ? &release_message : nullptr);
        },
        [&]() {
          CdmDecryptionParametersV16 parameters(key_id_);
          cdm_engine->DecryptV16(session_id_, parameters);
        },
        [&]() {
          std::string out = fdp_.ConsumeRandomLengthString(kMaxByte);
          std::string iv_GE = fdp_.ConsumeRandomLengthString(kMaxByte);
          cdm_engine->GenericEncrypt(
              session_id_, fdp_.ConsumeRandomLengthString(kMaxByte), key_id_,
              iv_GE,
              (CdmEncryptionAlgorithm)fdp_.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              fdp_.ConsumeBool() ? &out : nullptr);
        },
        [&]() {
          std::string out = fdp_.ConsumeRandomLengthString(kMaxByte);
          std::string iv_GE = fdp_.ConsumeRandomLengthString(kMaxByte);
          cdm_engine->GenericDecrypt(
              session_id_, fdp_.ConsumeRandomLengthString(kMaxByte), key_id_,
              iv_GE,
              (CdmEncryptionAlgorithm)fdp_.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              fdp_.ConsumeBool() ? &out : nullptr);
        },
        [&]() {
          cdm_engine->GenericSign(
              session_id_, fdp_.ConsumeRandomLengthString(kMaxByte), key_id_,
              (CdmSigningAlgorithm)fdp_.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              fdp_.ConsumeBool() ? &sign_ : nullptr);

          cdm_engine->GenericVerify(
              session_id_, fdp_.ConsumeRandomLengthString(kMaxByte), key_id_,
              (CdmSigningAlgorithm)fdp_.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              sign_);
        },
        [&]() {
          cdm_engine->SetDebugIgnoreKeyboxCount(
              fdp_.ConsumeIntegral<uint32_t>());
        },
        [&]() {
          std::string sess_id = fdp_.ConsumeRandomLengthString(kMaxByte);
          std::string hash;
          std::string hash_string =
              fdp_.ConsumeRandomLengthString(kMaxByte) + "," +
              fdp_.ConsumeRandomLengthString(kMaxByte) + "," +
              fdp_.ConsumeRandomLengthString(kMaxByte);
          uint32_t frame_num = fdp_.ConsumeIntegral<uint32_t>();
          cdm_engine->ParseDecryptHashString(
              hash_string, fdp_.ConsumeBool() ? &sess_id : nullptr,
              fdp_.ConsumeBool() ? &frame_num : nullptr,
              fdp_.ConsumeBool() ? &hash : nullptr);

          cdm_engine->SetDecryptHash(session_id_,
                                    fdp_.ConsumeIntegral<uint32_t>(), hash);

          std::string hash_error;
          cdm_engine->GetDecryptHashError(
              session_id_, fdp_.ConsumeBool() ? &hash_error : nullptr);
        },
        [&]() {
          std::string session_id = fdp_.ConsumeRandomLengthString(kMaxByte);
          cdm_engine->FindSessionForKey(
              key_id_, fdp_.ConsumeBool() ? &session_id : nullptr);
        },
        [&]() {
          cdm_engine->NotifyResolution(
              session_id_, fdp_.ConsumeIntegral<uint32_t>() /* width */,
              fdp_.ConsumeIntegral<uint32_t>() /* height */);
        },
        [&]() {
          cdm_engine->ValidateServiceCertificate(
              fdp_.ConsumeRandomLengthString(kMaxByte) /* cert */);
        },
        [&]() {
          CdmUsageReport usage_report;
          CreateUsageFiles(app_id_);
          CdmSecureStopId ssid = fdp_.ConsumeRandomLengthString(kMaxByte);
          int error_detail = fdp_.ConsumeIntegral<int32_t>();
          cdm_engine->GetUsageInfo(app_id_, ssid,
                                  fdp_.ConsumeBool() ? &error_detail : nullptr,
                                  fdp_.ConsumeBool() ? &usage_report : nullptr);
        },
        [&]() {
          CdmSecureStopId ssid = fdp_.ConsumeRandomLengthString(kMaxByte);
          cdm_engine->RemoveUsageInfo(app_id_, ssid);
        },
        [&]() { cdm_engine->RemoveLicense(session_id_); },
        [&]() { cdm_engine->IsKeyLoaded(key_id_); },
    });
    invoke_cdm_engine_API();
  }

  cdm_engine->Unprovision((CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(
      kMinEnum, kMaxEnum));

  cdm_engine->RemoveKeys(session_id_);

  out_cert_file_1_.close();
  out_cert_file_2_.close();
  out_usage_file_1_.close();
  out_usage_file_2_.close();

  delete cdm_engine;
}

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