/*
 * 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 "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"

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) {
    fdp_ = fdp;
    security_level_ = fdp->ConsumeBool()
                         ? fdp->ConsumeRandomLengthString(kMaxByte)
                         : fdp->PickValueInArray(kSecurityLevel);
    user_privacy_mode_ = fdp->ConsumeBool();
    is_session_sharing_enabled_ = fdp->ConsumeBool();
    session_sharing_id_ = fdp->ConsumeIntegral<int32_t>();
    app_id_ =
        fdp->ConsumeBool() ? kAppId : fdp->ConsumeRandomLengthString(kMaxByte);
  }

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

  bool use_privacy_mode() const override { return user_privacy_mode_; }

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

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

  bool is_session_sharing_enabled() const override {
    return is_session_sharing_enabled_;
  }

  uint32_t session_sharing_id() const override {
    uint32_t session_id = 0;
    if (fdp_->ConsumeBool()) {
      session_id = session_sharing_id_;
    }
    return session_id;
  }

  bool use_atsc_mode() const override { return false; }

  void set_session_sharing_id(uint32_t id) override {
    if (fdp_->ConsumeBool()) {
      session_sharing_id_ = id;
    }
  }

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

  void enable_privacy_mode() {
    if (fdp_->ConsumeBool()) {
      user_privacy_mode_ = true;
    }
  }

 private:
  FuzzedDataProvider *fdp_;
  std::string security_level_;
  std::string raw_service_certificate_;
  std::string app_id_;
  uint32_t session_sharing_id_;
  bool user_privacy_mode_;
  bool is_session_sharing_enabled_;
};

class ContentDecryptionFuzzer {
 public:
  ContentDecryptionFuzzer(const uint8_t *data, size_t size) : fdp_(data, size) {
    decryptor_ = sp<WvContentDecryptionModule>::make();
    FuzzWvCdmEventListener fuzz_wv_cdm_event_listener_;
    fuzz_cdm_client_property_set_.reset(new FuzzCdmClientPropertySet(&fdp_));

    if (fdp_.ConsumeBool()) {
      cdm_identifier_ = kDefaultCdmIdentifier;
    } else {
      cdm_identifier_.spoid = fdp_.ConsumeRandomLengthString(kMaxByte);
      cdm_identifier_.origin = fdp_.ConsumeRandomLengthString(kMaxByte);
      cdm_identifier_.app_package_name =
          fdp_.ConsumeRandomLengthString(kMaxByte);
      cdm_identifier_.unique_id = fdp_.ConsumeIntegral<uint32_t>();
      cdm_identifier_.user_id = fdp_.ConsumeIntegral<uint32_t>();
    }
    decryptor_->OpenSession(KEY_SYSTEM, fuzz_cdm_client_property_set_.get(),
                            cdm_identifier_, &fuzz_wv_cdm_event_listener_,
                            &session_id_);
  };
  ~ContentDecryptionFuzzer() { decryptor_->CloseSession(session_id_); }
  void Process();

 private:
  FuzzedDataProvider fdp_;
  std::unique_ptr<FuzzCdmClientPropertySet> fuzz_cdm_client_property_set_;
  CdmSessionId session_id_;
  CdmIdentifier cdm_identifier_;
  sp<WvContentDecryptionModule> decryptor_;
  void InvokeDecryptorSessionAPI();
  void InvokeProvisionAPI();
};

void ContentDecryptionFuzzer::InvokeDecryptorSessionAPI() {
  CdmKeySetId key_set_id = fdp_.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 (fdp_.remaining_bytes() > 0 && --runs) {
    auto invoke_decryption_session_API =
        fdp_.PickValueInArray<const std::function<void()>>({
            [&]() { decryptor_->IsOpenSession(session_id_); },
            [&]() {
              CdmAppParameterMap app_parameters;
              int32_t max_size =
                  fdp_.ConsumeIntegralInRange<int32_t>(0, kMaxAppParamSize);
              for (size_t i = 0; i < max_size; ++i) {
                std::string key = fdp_.ConsumeRandomLengthString(kMaxByte);
                std::string value = fdp_.ConsumeRandomLengthString(kMaxByte);
                std::string cdm_key(key.c_str(), key.size());
                std::string cdm_value(value.c_str(), value.size());
                app_parameters[cdm_key] = cdm_value;
              }

              std::string init_data_type = fdp_.PickValueInArray(kInitDataTypes);

              CdmInitData init_data = fdp_.ConsumeRandomLengthString(kMaxByte);

              CdmLicenseType license_type =
                  (CdmLicenseType)fdp_.ConsumeIntegralInRange<int32_t>(
                      kLicenseTypeOffline, kLicenseTypeEmbeddedKeyData);

              CdmKeyRequest key_request;
              key_request.message = fdp_.ConsumeRandomLengthString(kMaxByte);

              key_request.type =
                  (CdmKeyRequestType)fdp_.ConsumeIntegralInRange<int32_t>(
                      kKeyRequestTypeUnknown, kKeyRequestTypeUpdate);

              key_request.url = fdp_.ConsumeRandomLengthString(kMaxByte);

              FuzzCdmClientPropertySet *property_set =
                  fdp_.ConsumeBool() ? fuzz_cdm_client_property_set_.get()
                                     : nullptr;

              decryptor_->GenerateKeyRequest(
                  session_id_, key_set_id, init_data_type, init_data, license_type,
                  app_parameters, property_set, cdm_identifier_,
                  (fdp_.ConsumeBool() ? nullptr : &key_request));
            },
            [&]() {
              CdmKeyResponse response =
                  fdp_.ConsumeRandomLengthString(kMaxByte);
              decryptor_->AddKey(session_id_, response, &key_set_id);
            },
            [&]() { decryptor_->RestoreKey(session_id_, key_set_id); },
            [&]() { decryptor_->RemoveKeys(session_id_); },
            [&]() {
              RequestedSecurityLevel security_level =
                  (RequestedSecurityLevel)fdp_.ConsumeIntegral<uint32_t>();
              std::string output_str;
              std::string query_token =
                  fdp_.ConsumeBool() ? fdp_.ConsumeRandomLengthString(kMaxByte)
                                     : fdp_.PickValueInArray(kQueryLevel);
              decryptor_->QueryStatus(
                  security_level, query_token.c_str(),
                  (fdp_.ConsumeBool() ? nullptr : &output_str));
            },
            [&]() {
              CdmQueryMap key_info;
              decryptor_->QuerySessionStatus(
                  session_id_, (fdp_.ConsumeBool() ? nullptr : &key_info));
            },
            [&]() {
              CdmQueryMap key_info;
              decryptor_->QueryKeyStatus(
                  session_id_, (fdp_.ConsumeBool() ? nullptr : &key_info));
            },
            [&]() {
              CdmQueryMap key_info;
              decryptor_->QueryOemCryptoSessionId(
                  session_id_, (fdp_.ConsumeBool() ? nullptr : &key_info));
            },
            [&]() {
              CdmSecurityLevel level =
                  (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(
                      kSecurityLevelUninitialized, kSecurityLevelUnknown);
              decryptor_->IsSecurityLevelSupported(level);
            },

        });
    invoke_decryption_session_API();
  }
}

void ContentDecryptionFuzzer::InvokeProvisionAPI() {
  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_provision_API =
        fdp_.PickValueInArray<const std::function<void()>>({
            [&]() {
              CdmSecurityLevel level =
                  (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(
                      kSecurityLevelUninitialized, kSecurityLevelUnknown);
              decryptor_->Unprovision(level, cdm_identifier_);
            },
            [&]() {
              CdmKeyMessage key_msg;
              CdmCertificateType certificate_type =
                  (CdmCertificateType)fdp_.ConsumeIntegral<int32_t>();

              std::string cert_authority =
                  fdp_.ConsumeRandomLengthString(kMaxByte);

              RequestedSecurityLevel security_level =
                  (RequestedSecurityLevel)fdp_.ConsumeIntegral<uint32_t>();

              std::string server_url = fdp_.ConsumeRandomLengthString(kMaxByte);

              std::string service_certificate =
                  fdp_.ConsumeRandomLengthString(kMaxByte);

              decryptor_->GetProvisioningRequest(
                  certificate_type, cert_authority, cdm_identifier_,
                  service_certificate, security_level,
                  (fdp_.ConsumeBool() ? nullptr : &key_msg),
                  (fdp_.ConsumeBool() ? nullptr : &server_url));
            },
            [&]() {
              std::string response = fdp_.ConsumeRandomLengthString(kMaxByte);
              RequestedSecurityLevel security_level =
                  (RequestedSecurityLevel)fdp_.ConsumeIntegral<uint32_t>();
              std::string cert, wrapped_key;
              decryptor_->HandleProvisioningResponse(
                  cdm_identifier_, response, security_level,
                  (fdp_.ConsumeBool() ? nullptr : &cert),
                  (fdp_.ConsumeBool() ? nullptr : &wrapped_key));
            },
            [&]() {
              CdmSecurityLevel level =
                  (CdmSecurityLevel)fdp_.ConsumeIntegralInRange<int32_t>(
                      kSecurityLevelUninitialized, kSecurityLevelUnknown);
              std::string origin = fdp_.ConsumeRandomLengthString(kMaxByte);
              std::string spoid = fdp_.ConsumeRandomLengthString(kMaxByte);
              decryptor_->IsProvisioned(
                  level, origin, spoid,
                  fdp_.ConsumeBool() /*atsc_mode_enabled*/);
            },
        });
    invoke_provision_API();
  }
}

void ContentDecryptionFuzzer::Process() {
  auto invoke_content_decryption_API =
      fdp_.PickValueInArray<const std::function<void()>>({
          [&]() { InvokeDecryptorSessionAPI(); },
          [&]() { InvokeProvisionAPI(); },
      });
  invoke_content_decryption_API();
}

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