/*
 * 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 "vendor_widevine_fuzz_helper.h"
#include <cdm_engine.h>
#include <cdm_engine_factory.h>
#include <fstream>
#include <fuzzer/FuzzedDataProvider.h>
#include <iostream>
#include <stdlib.h>
#include <string>
#include <wv_cdm_constants.h>
#include <wv_cdm_event_listener.h>
#include <wv_cdm_types.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) : mFdp(data, size){};
  void process();
  void init(CdmEngine *cdmEngine);
  void createUsageFile(std::string mAppId);
  std::vector<std::string> createVector(bool enableKeySetId);

private:
  FuzzedDataProvider mFdp;
  std::ofstream mOutCertFile1;
  std::ofstream mOutCertFile2;
  std::ofstream mOutUsageFile1;
  std::ofstream mOutUsageFile2;
  std::string mAppId;
  std::string mKeyId;
  std::string mSign;
  std::string mDefaultUrl;
  std::unique_ptr<wvutil::FileSystem> mFileSystem =
      std::make_unique<wvutil::FileSystem>();
  std::unique_ptr<FuzzCdmClientPropertySet> mFuzzCdmClientPropertySet =
      std::make_unique<FuzzCdmClientPropertySet>(&mFdp);
  CdmSessionId mSessionId;
  CdmKeySystem mKeySystem;
  CdmKeySetId mKeySetId;
};

void CdmEngineFuzzer::createUsageFile(std::string mAppId) {
  /**
   * 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 deviceFiles(mFileSystem.get());
  std::string fileName = deviceFiles.GetUsageInfoFileName(mAppId);
  std::string createFile1 = kUsageFilePrefixL1 + fileName;
  std::string createFile2 = kUsageFilePrefixL3 + fileName;
  mOutUsageFile1.open(createFile1.c_str());
  mOutUsageFile2.open(createFile2.c_str());
}

std::vector<std::string> CdmEngineFuzzer::createVector(bool enableKeySetId) {
  int32_t size = mFdp.ConsumeIntegralInRange(kMinSize, kMaxSize);
  std::vector<std::string> vector;
  for (int i = 0; i < size; ++i) {
    vector.push_back(enableKeySetId ? mKeySetId
                                    : mFdp.ConsumeRandomLengthString(kMaxByte));
  }
  return vector;
}

void CdmEngineFuzzer::init(CdmEngine *cdmEngine) {
  mAppId =
      std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(kMinSize, kMaxSize));
  mKeyId =
      std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(kMinSize, kMaxSize));
  mSign = mFdp.ConsumeRandomLengthString(kMaxByte);

  mKeySystem = mFdp.ConsumeBool() ? kKeySystem
                                  : mFdp.ConsumeRandomLengthString(kMaxByte);

  FuzzWvCdmEventListener fuzzWvCdmEventListener;
  cdmEngine->OpenSession(
      mKeySystem,
      mFdp.ConsumeBool() ? mFuzzCdmClientPropertySet.get() : nullptr,
      mFdp.ConsumeBool() ? &fuzzWvCdmEventListener : nullptr, &mSessionId);

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

  std::string cert;
  std::string wrappedKey;
  CdmProvisioningResponse responseMessage =
      mFdp.ConsumeRandomLengthString(kMaxByte);
  cdmEngine->HandleProvisioningResponse(
      responseMessage,
      (RequestedSecurityLevel)mFdp.ConsumeIntegralInRange<uint32_t>(kMinValue,
                                                                    kMaxValue),
      mFdp.ConsumeBool() ? &cert : nullptr,
      mFdp.ConsumeBool() ? &wrappedKey : nullptr);

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

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

  CdmLicenseType generateLicType =
      static_cast<CdmLicenseType>(mFdp.ConsumeIntegralInRange(
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeOffline),
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeEmbeddedKeyData)));
  CdmKeyRequest keyRequest;
  keyRequest.message = mFdp.ConsumeRandomLengthString(kMaxByte);
  keyRequest.type = static_cast<CdmKeyRequestType>(mFdp.ConsumeIntegralInRange(
      static_cast<uint32_t>(CdmKeyRequestType::kKeyRequestTypeUnknown),
      static_cast<uint32_t>(CdmKeyRequestType::kKeyRequestTypeUpdate)));
  keyRequest.url = mFdp.ConsumeBool()
                       ? mDefaultUrl
                       : mFdp.ConsumeRandomLengthString(kMaxByte);
  cdmEngine->GenerateKeyRequest(mSessionId,
                                mFdp.ConsumeBool() ? mKeySetId : kEmptyString,
                                init_data, generateLicType, appParameters,
                                mFdp.ConsumeBool() ? &keyRequest : nullptr);

  CdmLicenseType addLicType =
      static_cast<CdmLicenseType>(mFdp.ConsumeIntegralInRange(
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeOffline),
          static_cast<int32_t>(CdmLicenseType::kLicenseTypeEmbeddedKeyData)));
  CdmKeyResponse keyData =
      std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(kMinKey, kMaxKey));
  cdmEngine->AddKey(mSessionId, mFdp.ConsumeBool() ? keyData : kEmptyString,
                    &addLicType, &mKeySetId);
}

void CdmEngineFuzzer::process() {
  CdmEngine *cdmEngine = CdmEngineFactory::CreateCdmEngine(mFileSystem.get());

  CdmEngineFuzzer::init(cdmEngine);

   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 (mFdp.remaining_bytes() > 0 && --runs) {
    auto invokeCdmEngineAPI = mFdp.PickValueInArray<
        const std::function<void()>>({
        [&]() { cdmEngine->OnTimerEvent(); },
        [&]() { cdmEngine->IsOpenSession(mSessionId); },
        [&]() { cdmEngine->SetDefaultOtaKeyboxFallbackDurationRules(); },
        [&]() { cdmEngine->SetFastOtaKeyboxFallbackDurationRules(); },
        [&]() {
          cdmEngine->OpenSession(
              mKeySystem, nullptr /* property_set */,
              mFdp.ConsumeRandomLengthString(kMaxByte) /* forced_session_id */,
              nullptr /* event_listener */);
        },
        [&]() {
          cdmEngine->StoreAtscLicense(
              static_cast<RequestedSecurityLevel>(mFdp.ConsumeIntegralInRange(
                  static_cast<int32_t>(RequestedSecurityLevel::kLevelDefault),
                  static_cast<int32_t>(RequestedSecurityLevel::kLevel3))),
              mKeySetId,
              mFdp.ConsumeRandomLengthString(
                  kMaxByte) /* serialized_license_data */);
        },
        [&]() { cdmEngine->RestoreKey(mSessionId, mKeySetId); },
        [&]() {
          cdmEngine->SetSessionServiceCertificate(
              mSessionId, mFdp.ConsumeRandomLengthString(
                              kMaxByte) /* service_certificate */);
        },
        [&]() {
          std::string queryResponse;
          cdmEngine->QueryStatus(
              static_cast<RequestedSecurityLevel>(mFdp.ConsumeIntegralInRange(
                  static_cast<int32_t>(RequestedSecurityLevel::kLevelDefault),
                  static_cast<int32_t>(RequestedSecurityLevel::kLevel3))),
              mFdp.PickValueInArray<std::string>(kQueryToken),
              mFdp.ConsumeBool() ? &queryResponse : nullptr);
        },
        [&]() {
          CdmQueryMap queryResponse;
          cdmEngine->QuerySessionStatus(mSessionId, &queryResponse);
        },
        [&]() {
          CdmKeyRequest keyRequest;
          keyRequest.message = mFdp.ConsumeRandomLengthString(kMaxByte);
          keyRequest.type =
              static_cast<CdmKeyRequestType>(mFdp.ConsumeIntegralInRange(
                  static_cast<uint32_t>(
                      CdmKeyRequestType::kKeyRequestTypeUnknown),
                  static_cast<uint32_t>(
                      CdmKeyRequestType::kKeyRequestTypeUpdate)));
          keyRequest.url = mDefaultUrl;

          cdmEngine->GenerateRenewalRequest(
              mSessionId, mFdp.ConsumeBool() ? &keyRequest : nullptr);

          CdmKeyResponse keyData =
              mFdp.ConsumeBool()
                  ? std::to_string(
                        mFdp.ConsumeIntegralInRange<uint32_t>(kMinKey, kMaxKey))
                  : kEmptyString;
          cdmEngine->RenewKey(mSessionId, keyData);
        },
        [&]() { cdmEngine->IsReleaseSession(mSessionId); },
        [&]() { cdmEngine->IsOfflineSession(mSessionId); },
        [&]() {
          mOutCertFile1.open(kCertFile.c_str());
          mOutCertFile2.open(kCertFile1.c_str());
          cdmEngine->RemoveOfflineLicense(
              mKeySetId,
              static_cast<CdmSecurityLevel>(mFdp.ConsumeIntegralInRange(
                  static_cast<int32_t>(
                      CdmSecurityLevel::kSecurityLevelUninitialized),
                  static_cast<int32_t>(
                      CdmSecurityLevel::kSecurityLevelUnknown))));
        },
        [&]() {
          CdmQueryMap queryResponse;
          cdmEngine->QueryKeyStatus(mSessionId, &queryResponse);
        },
        [&]() {
          cdmEngine->IsSecurityLevelSupported(
              (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum));
        },
        [&]() {
          CdmKeyAllowedUsage keyUsage;
          cdmEngine->QueryKeyAllowedUsage(
              mSessionId, mKeyId, mFdp.ConsumeBool() ? &keyUsage : nullptr);
        },
        [&]() {
          CdmKeyAllowedUsage keyUsage;
          cdmEngine->QueryKeyAllowedUsage(mKeyId, mFdp.ConsumeBool() ? &keyUsage
                                                                     : nullptr);
        },
        [&]() {
          CdmQueryMap queryResponse;
          cdmEngine->QueryOemCryptoSessionId(mSessionId, &queryResponse);
        },
        [&]() {
          std::vector<std::string> keySetIds = createVector(true);
          cdmEngine->ListStoredLicenses(
              (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              mFdp.ConsumeBool() ? &keySetIds : nullptr);
        },
        [&]() {
          std::vector<std::string> keySetIds = createVector(true);
          std::vector<std::string> providerSessionTokens = createVector(false);

          cdmEngine->ListUsageIds(
              mAppId,
              (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              mFdp.ConsumeBool() ? &keySetIds : nullptr,
              mFdp.ConsumeBool() ? &providerSessionTokens : nullptr);
        },
        [&]() {
          cdmEngine->DeleteUsageRecord(
              mAppId,
              (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              mKeySetId);
        },
        [&]() {
          CdmOfflineLicenseState licenseState =
              static_cast<CdmOfflineLicenseState>(mFdp.ConsumeIntegralInRange(
                  static_cast<int32_t>(
                      CdmOfflineLicenseState::kLicenseStateActive),
                  static_cast<int32_t>(
                      CdmOfflineLicenseState::kLicenseStateUnknown)));
          cdmEngine->GetOfflineLicenseState(
              mKeySetId,
              (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(kMinEnum,
                                                                     kMaxEnum),
              mFdp.ConsumeBool() ? &licenseState : nullptr);
        },
        [&]() {
          cdmEngine->SetPlaybackId(mSessionId,
                                   mFdp.ConsumeRandomLengthString(kMaxByte));
        },
        [&]() {
          CdmUsageReport usageReport;
          createUsageFile(mAppId);
          int errorDetail;
          cdmEngine->GetUsageInfo(mAppId,
                                  mFdp.ConsumeBool() ? &errorDetail : nullptr,
                                  mFdp.ConsumeBool() ? &usageReport : nullptr);
        },
        [&]() {
          createUsageFile(mAppId);
          cdmEngine->RemoveAllUsageInfo(mAppId);
        },
        [&]() {
          cdmEngine->ReleaseUsageInfo(
              mFdp.ConsumeRandomLengthString(kMaxByte) /* message */);
        },
        [&]() {
          std::string releaseMessage = mFdp.ConsumeRandomLengthString(kMaxByte);
          cdmEngine->LoadUsageSession(
              mKeySetId, mFdp.ConsumeBool() ? &releaseMessage : nullptr);
        },
        [&]() {
          CdmDecryptionParametersV16 parameters(mKeyId);
          cdmEngine->DecryptV16(mSessionId, parameters);
        },
        [&]() {
          std::string out = mFdp.ConsumeRandomLengthString(kMaxByte);
          std::string ivGE = mFdp.ConsumeRandomLengthString(kMaxByte);
          cdmEngine->GenericEncrypt(
              mSessionId, mFdp.ConsumeRandomLengthString(kMaxByte), mKeyId,
              ivGE,
              (CdmEncryptionAlgorithm)mFdp.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              mFdp.ConsumeBool() ? &out : nullptr);
        },
        [&]() {
          std::string out = mFdp.ConsumeRandomLengthString(kMaxByte);
          std::string ivGE = mFdp.ConsumeRandomLengthString(kMaxByte);
          cdmEngine->GenericDecrypt(
              mSessionId, mFdp.ConsumeRandomLengthString(kMaxByte), mKeyId,
              ivGE,
              (CdmEncryptionAlgorithm)mFdp.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              mFdp.ConsumeBool() ? &out : nullptr);
        },
        [&]() {
          cdmEngine->GenericSign(
              mSessionId, mFdp.ConsumeRandomLengthString(kMaxByte), mKeyId,
              (CdmSigningAlgorithm)mFdp.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              mFdp.ConsumeBool() ? &mSign : nullptr);

          cdmEngine->GenericVerify(
              mSessionId, mFdp.ConsumeRandomLengthString(kMaxByte), mKeyId,
              (CdmSigningAlgorithm)mFdp.ConsumeIntegralInRange<int32_t>(
                  kMinValue, kMaxValue),
              mSign);
        },
        [&]() {
          cdmEngine->SetDebugIgnoreKeyboxCount(
              mFdp.ConsumeIntegral<uint32_t>());
        },
        [&]() {
          std::string sessId = mFdp.ConsumeRandomLengthString(kMaxByte);
          std::string hash;
          std::string hashString =
              mFdp.ConsumeRandomLengthString(kMaxByte) + "," +
              mFdp.ConsumeRandomLengthString(kMaxByte) + "," +
              mFdp.ConsumeRandomLengthString(kMaxByte);
          uint32_t frameNum = mFdp.ConsumeIntegral<uint32_t>();
          cdmEngine->ParseDecryptHashString(
              hashString, mFdp.ConsumeBool() ? &sessId : nullptr,
              mFdp.ConsumeBool() ? &frameNum : nullptr,
              mFdp.ConsumeBool() ? &hash : nullptr);

          cdmEngine->SetDecryptHash(mSessionId,
                                    mFdp.ConsumeIntegral<uint32_t>(), hash);

          std::string hashError;
          cdmEngine->GetDecryptHashError(
              mSessionId, mFdp.ConsumeBool() ? &hashError : nullptr);
        },
        [&]() {
          std::string mSessionId = mFdp.ConsumeRandomLengthString(kMaxByte);
          cdmEngine->FindSessionForKey(mKeyId, mFdp.ConsumeBool() ? &mSessionId
                                                                  : nullptr);
        },
        [&]() {
          cdmEngine->NotifyResolution(
              mSessionId, mFdp.ConsumeIntegral<uint32_t>() /* width */,
              mFdp.ConsumeIntegral<uint32_t>() /* height */);
        },
        [&]() {
          cdmEngine->ValidateServiceCertificate(
              mFdp.ConsumeRandomLengthString(kMaxByte) /* cert */);
        },
        [&]() {
          CdmUsageReport usageReport;
          createUsageFile(mAppId);
          CdmSecureStopId ssid = mFdp.ConsumeRandomLengthString(kMaxByte);
          int errorDetail = mFdp.ConsumeIntegral<int32_t>();
          cdmEngine->GetUsageInfo(mAppId, ssid,
                                  mFdp.ConsumeBool() ? &errorDetail : nullptr,
                                  mFdp.ConsumeBool() ? &usageReport : nullptr);
        },
        [&]() {
          CdmSecureStopId ssid = mFdp.ConsumeRandomLengthString(kMaxByte);
          cdmEngine->RemoveUsageInfo(mAppId, ssid);
        },
        [&]() { cdmEngine->RemoveLicense(mSessionId); },
        [&]() { cdmEngine->IsKeyLoaded(mKeyId); },
    });
    invokeCdmEngineAPI();
  }

  cdmEngine->Unprovision((CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(
      kMinEnum, kMaxEnum));

  cdmEngine->RemoveKeys(mSessionId);

  mOutCertFile1.close();
  mOutCertFile2.close();
  mOutUsageFile1.close();
  mOutUsageFile2.close();

  delete cdmEngine;
}

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