/*
 * 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_client_property_set.h"
#include "cdm_identifier.h"
#include "wv_cdm_constants.h"
#include "wv_cdm_event_listener.h"
#include "wv_content_decryption_module.h"
#include <fuzzer/FuzzedDataProvider.h>

using namespace wvcdm;
using namespace android;

const std::string kAppId = "com.example.test";
static constexpr int32_t kMaxByte = 256;
static constexpr int32_t kMinSetId = 1;
static constexpr int32_t kMaxSetId = 100;
static constexpr int32_t kMinSize = 1;
static constexpr int32_t kMaxSize = 1024;
static constexpr int32_t kMaxAppParamSize = 10;
static constexpr int32_t kMaxRuns = 100;

const std::string kSecurityLevel[] = {"QUERY_VALUE_SECURITY_LEVEL_L1",
                                      "QUERY_VALUE_SECURITY_LEVEL_L2",
                                      "QUERY_VALUE_SECURITY_LEVEL_L3"};

const std::string kQueryLevel[] = {
    "QUERY_KEY_SECURITY_LEVEL",
    "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",
};
const std::string kInitDataTypes[] = {
    HLS_INIT_DATA_FORMAT,  ISO_BMFF_VIDEO_MIME_TYPE, ISO_BMFF_AUDIO_MIME_TYPE,
    CENC_INIT_DATA_FORMAT, WEBM_VIDEO_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 FuzzCdmClientPropertySet : public CdmClientPropertySet {
public:
  FuzzCdmClientPropertySet(FuzzedDataProvider *fdp) {
    mFdp = fdp;
    mSecurityLevel = fdp->ConsumeBool()
                         ? fdp->ConsumeRandomLengthString(kMaxByte)
                         : fdp->PickValueInArray(kSecurityLevel);
    mUsePrivacyMode = fdp->ConsumeBool();
    mIsSessionSharingEnabled = fdp->ConsumeBool();
    mSessionSharingId = fdp->ConsumeIntegral<int32_t>();
    mAppId =
        fdp->ConsumeBool() ? kAppId : fdp->ConsumeRandomLengthString(kMaxByte);
  }

  const std::string &security_level() const override { return mSecurityLevel; }

  bool use_privacy_mode() const override { return mUsePrivacyMode; }

  const std::string &service_certificate() const override {
    return mRawServiceCertificate;
  }

  void set_service_certificate(const std::string &cert) override {
    if (mFdp->ConsumeBool()) {
      mRawServiceCertificate = cert;
    }
  }

  bool is_session_sharing_enabled() const override {
    return mIsSessionSharingEnabled;
  }

  uint32_t session_sharing_id() const override {
    uint32_t sessionId = 0;
    if (mFdp->ConsumeBool()) {
      sessionId = mSessionSharingId;
    }
    return sessionId;
  }

  bool use_atsc_mode() const override { return false; }

  void set_session_sharing_id(uint32_t id) override {
    if (mFdp->ConsumeBool()) {
      mSessionSharingId = id;
    }
  }

  const std::string &app_id() const override { return mAppId; }

  void enable_privacy_mode() {
    if (mFdp->ConsumeBool()) {
      mUsePrivacyMode = true;
    }
  }

private:
  FuzzedDataProvider *mFdp;
  std::string mSecurityLevel;
  std::string mRawServiceCertificate;
  std::string mAppId;
  uint32_t mSessionSharingId;
  bool mUsePrivacyMode;
  bool mIsSessionSharingEnabled;
};

class ContentDecryptionFuzzer {
public:
  ContentDecryptionFuzzer(const uint8_t *data, size_t size) : mFdp(data, size) {
    mDecryptor = sp<WvContentDecryptionModule>::make();
    FuzzWvCdmEventListener fuzzWvCdmEventListener;
    mFuzzCdmClientPropertySet.reset(new FuzzCdmClientPropertySet(&mFdp));

    if (mFdp.ConsumeBool()) {
      mCdmIdentifier = kDefaultCdmIdentifier;
    } else {
      mCdmIdentifier.spoid = mFdp.ConsumeRandomLengthString(kMaxByte);
      mCdmIdentifier.origin = mFdp.ConsumeRandomLengthString(kMaxByte);
      mCdmIdentifier.app_package_name =
          mFdp.ConsumeRandomLengthString(kMaxByte);
      mCdmIdentifier.unique_id = mFdp.ConsumeIntegral<uint32_t>();
      mCdmIdentifier.user_id = mFdp.ConsumeIntegral<uint32_t>();
    }
    mDecryptor->OpenSession(KEY_SYSTEM, mFuzzCdmClientPropertySet.get(),
                            mCdmIdentifier, &fuzzWvCdmEventListener,
                            &mSessionId);
  };
  ~ContentDecryptionFuzzer() { mDecryptor->CloseSession(mSessionId); }
  void process();

private:
  FuzzedDataProvider mFdp;
  std::unique_ptr<FuzzCdmClientPropertySet> mFuzzCdmClientPropertySet;
  CdmSessionId mSessionId;
  CdmIdentifier mCdmIdentifier;
  sp<WvContentDecryptionModule> mDecryptor;
  void invokeDecryptorSessionAPIs();
  void invokeProvisionAPIs();
};

void ContentDecryptionFuzzer::invokeDecryptorSessionAPIs() {
  CdmKeySetId keySetId = mFdp.ConsumeRandomLengthString(kMaxByte);
  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 invokeDecryptionSessionAPI =
        mFdp.PickValueInArray<const std::function<void()>>({
            [&]() { mDecryptor->IsOpenSession(mSessionId); },
            [&]() {
              CdmAppParameterMap appParameters;
              int32_t maxSize =
                  mFdp.ConsumeIntegralInRange<int32_t>(0, kMaxAppParamSize);
              for (size_t i = 0; i < maxSize; ++i) {
                std::string key = mFdp.ConsumeRandomLengthString(kMaxByte);
                std::string value = mFdp.ConsumeRandomLengthString(kMaxByte);
                std::string cdmKey(key.c_str(), key.size());
                std::string cdmValue(value.c_str(), value.size());
                appParameters[cdmKey] = cdmValue;
              }

              std::string initDataType = mFdp.PickValueInArray(kInitDataTypes);

              CdmInitData initData = mFdp.ConsumeRandomLengthString(kMaxByte);

              CdmLicenseType licenseType =
                  (CdmLicenseType)mFdp.ConsumeIntegralInRange<int32_t>(
                      kLicenseTypeOffline, kLicenseTypeEmbeddedKeyData);

              CdmKeyRequest keyRequest;
              keyRequest.message = mFdp.ConsumeRandomLengthString(kMaxByte);

              keyRequest.type =
                  (CdmKeyRequestType)mFdp.ConsumeIntegralInRange<int32_t>(
                      kKeyRequestTypeUnknown, kKeyRequestTypeUpdate);

              keyRequest.url = mFdp.ConsumeRandomLengthString(kMaxByte);

              FuzzCdmClientPropertySet *propertySet =
                  mFdp.ConsumeBool() ? mFuzzCdmClientPropertySet.get()
                                     : nullptr;

              mDecryptor->GenerateKeyRequest(
                  mSessionId, keySetId, initDataType, initData, licenseType,
                  appParameters, propertySet, mCdmIdentifier,
                  (mFdp.ConsumeBool() ? nullptr : &keyRequest));
            },
            [&]() {
              CdmKeyResponse response =
                  mFdp.ConsumeRandomLengthString(kMaxByte);
              mDecryptor->AddKey(mSessionId, response, &keySetId);
            },
            [&]() { mDecryptor->RestoreKey(mSessionId, keySetId); },
            [&]() { mDecryptor->RemoveKeys(mSessionId); },
            [&]() {
              RequestedSecurityLevel securityLevel =
                  (RequestedSecurityLevel)mFdp.ConsumeIntegral<uint32_t>();
              std::string outputStr;
              std::string queryToken =
                  mFdp.ConsumeBool() ? mFdp.ConsumeRandomLengthString(kMaxByte)
                                     : mFdp.PickValueInArray(kQueryLevel);
              ;
              mDecryptor->QueryStatus(
                  securityLevel, queryToken.c_str(),
                  (mFdp.ConsumeBool() ? nullptr : &outputStr));
            },
            [&]() {
              CdmQueryMap keyInfo;
              mDecryptor->QuerySessionStatus(
                  mSessionId, (mFdp.ConsumeBool() ? nullptr : &keyInfo));
            },
            [&]() {
              CdmQueryMap keyInfo;
              mDecryptor->QueryKeyStatus(
                  mSessionId, (mFdp.ConsumeBool() ? nullptr : &keyInfo));
            },
            [&]() {
              CdmQueryMap keyInfo;
              mDecryptor->QueryOemCryptoSessionId(
                  mSessionId, (mFdp.ConsumeBool() ? nullptr : &keyInfo));
            },
            [&]() {
              CdmSecurityLevel level =
                  (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(
                      kSecurityLevelUninitialized, kSecurityLevelUnknown);
              mDecryptor->IsSecurityLevelSupported(level);
            },

        });
    invokeDecryptionSessionAPI();
  }
}

void ContentDecryptionFuzzer::invokeProvisionAPIs() {

  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 invokeProvisionAPI =
        mFdp.PickValueInArray<const std::function<void()>>({
            [&]() {
              CdmSecurityLevel level =
                  (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(
                      kSecurityLevelUninitialized, kSecurityLevelUnknown);
              mDecryptor->Unprovision(level, mCdmIdentifier);
            },
            [&]() {
              CdmKeyMessage keyMsg;
              CdmCertificateType certificateType =
                  (CdmCertificateType)mFdp.ConsumeIntegral<int32_t>();

              std::string certAuthority =
                  mFdp.ConsumeRandomLengthString(kMaxByte);

              RequestedSecurityLevel securityLevel =
                  (RequestedSecurityLevel)mFdp.ConsumeIntegral<uint32_t>();

              std::string serverUrl = mFdp.ConsumeRandomLengthString(kMaxByte);

              std::string serviceCertificate =
                  mFdp.ConsumeRandomLengthString(kMaxByte);

              mDecryptor->GetProvisioningRequest(
                  certificateType, certAuthority, mCdmIdentifier,
                  serviceCertificate, securityLevel,
                  (mFdp.ConsumeBool() ? nullptr : &keyMsg),
                  (mFdp.ConsumeBool() ? nullptr : &serverUrl));
            },
            [&]() {
              std::string response = mFdp.ConsumeRandomLengthString(kMaxByte);
              RequestedSecurityLevel securityLevel =
                  (RequestedSecurityLevel)mFdp.ConsumeIntegral<uint32_t>();
              std::string cert, wrappedKey;
              mDecryptor->HandleProvisioningResponse(
                  mCdmIdentifier, response, securityLevel,
                  (mFdp.ConsumeBool() ? nullptr : &cert),
                  (mFdp.ConsumeBool() ? nullptr : &wrappedKey));
            },
            [&]() {
              CdmSecurityLevel level =
                  (CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(
                      kSecurityLevelUninitialized, kSecurityLevelUnknown);
              std::string origin = mFdp.ConsumeRandomLengthString(kMaxByte);
              std::string spoid = mFdp.ConsumeRandomLengthString(kMaxByte);
              mDecryptor->IsProvisioned(
                  level, origin, spoid,
                  mFdp.ConsumeBool() /*atsc_mode_enabled*/);
            },
        });
    invokeProvisionAPI();
  }
}

void ContentDecryptionFuzzer::process() {
  auto invokeContentDecryptionAPI =
      mFdp.PickValueInArray<const std::function<void()>>({
          [&]() { invokeDecryptorSessionAPIs(); },
          [&]() { invokeProvisionAPIs(); },
      });
  invokeContentDecryptionAPI();
}

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