/*
 * 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 "policy_engine.h"
#include "wv_cdm_event_listener.h"
#include <fuzzer/FuzzedDataProvider.h>

using namespace wvcdm;
using namespace wvcdm::metrics;
using namespace video_widevine;

static constexpr int32_t kLicenseIDVersion = 1;
static constexpr int32_t kMaxByte = 256;
static constexpr int32_t kMinResolution = 1;
static constexpr int32_t kMaxResolution = 1024;
static constexpr int32_t kMaxEntitledKey = 5;
const KeyId kSigningKeyId = "signing_key";
const KeyId kEntitlementKeyId = "entitlementkeyid";
const char *kRenewalServerUrl =
    "https://test.google.com/license/GetCencLicense";

class FuzzWvCdmEventListener : public WvCdmEventListener {
  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 PolicyEngineFuzzer {
public:
  PolicyEngineFuzzer(const uint8_t *data, size_t size) : mFdp(data, size){};
  void process();

private:
  FuzzedDataProvider mFdp;
  KeyId mKeyId;
  License mLicense;
  void setUp();
};

void policySetBool(std::function<void(bool)> function,
                   FuzzedDataProvider *fdp) {
  if (fdp->ConsumeBool()) {
    function(fdp->ConsumeBool());
  }
}

void policySetInt64(std::function<void(int64_t)> function,
                    FuzzedDataProvider *fdp) {
  if (fdp->ConsumeBool()) {
    function(fdp->ConsumeIntegral<int64_t>());
  }
}

void PolicyEngineFuzzer::setUp() {

  License::KeyContainer *key = mLicense.add_key();
  int32_t keyType = mFdp.ConsumeIntegralInRange<uint8_t>(
      License::KeyContainer::SIGNING, License::KeyContainer::ENTITLEMENT);
  key->set_type((License_KeyContainer_KeyType)keyType);
  mKeyId = mFdp.ConsumeRandomLengthString(kMaxByte);

  LicenseIdentification *id = mLicense.mutable_id();
  if (mFdp.ConsumeBool()) {
    id->set_version(kLicenseIDVersion);
  }
  if (mFdp.ConsumeBool()) {
    id->set_type(mFdp.ConsumeBool() ? STREAMING : OFFLINE);
  }
  if (mFdp.ConsumeBool()) {
      id->set_provider_session_token(mFdp.ConsumeRandomLengthString(kMaxByte));
  }

  if (mFdp.ConsumeBool()) {
    mLicense.set_license_start_time(mFdp.ConsumeIntegral<int64_t>());
  }

  if (mFdp.ConsumeBool()) {
    License_KeyContainer_OutputProtection *outputProtection =
        key->mutable_required_protection();
    License_KeyContainer_OutputProtection_HDCP hdcpLevel =
        (License_KeyContainer_OutputProtection_HDCP)mFdp
            .ConsumeIntegralInRange<uint8_t>(HDCP_NONE, HDCP_NO_DIGITAL_OUTPUT);
    outputProtection->set_hdcp(hdcpLevel);
  }

  if (mFdp.ConsumeBool()) {
    const std::string ivString = mFdp.ConsumeRandomLengthString(kMaxByte);
    key->set_iv(ivString);
  }

  switch (keyType) {
     case License::KeyContainer::SIGNING: {
        if (mFdp.ConsumeBool()) {
          mKeyId = mFdp.ConsumeBool() ? kSigningKeyId : mKeyId;
          key->set_id(mKeyId);
        }
        break;
     }
     case License::KeyContainer::CONTENT: {
        if (mFdp.ConsumeBool()) {
          key->set_id(mKeyId);
        }
        if (mFdp.ConsumeBool()) {
          int32_t level = mFdp.ConsumeIntegralInRange<uint8_t>(
              License::KeyContainer::SW_SECURE_CRYPTO,
              License::KeyContainer::HW_SECURE_ALL);
          key->set_level((License_KeyContainer_SecurityLevel)level);
        }

        License_Policy *policy = mLicense.mutable_policy();
        policySetBool(
            std::bind(&License_Policy::set_can_play, policy, std::placeholders::_1),
            &mFdp);
        policySetBool(std::bind(&License_Policy::set_can_persist, policy,
                                std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_can_renew, policy,
                                std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_renew_with_usage, policy,
                                std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_soft_enforce_rental_duration,
                                policy, std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_soft_enforce_playback_duration,
                                policy, std::placeholders::_1), &mFdp);

        policySetInt64(std::bind(&License_Policy::set_rental_duration_seconds,
                                 policy, std::placeholders::_1), &mFdp);
        policySetInt64(std::bind(&License_Policy::set_playback_duration_seconds,
                                 policy, std::placeholders::_1), &mFdp);
        policySetInt64(std::bind(&License_Policy::set_license_duration_seconds,
                                 policy, std::placeholders::_1), &mFdp);
        policySetInt64(
            std::bind(&License_Policy::set_renewal_recovery_duration_seconds,
                      policy, std::placeholders::_1), &mFdp);
        policySetInt64(
            std::bind(&License_Policy::set_renewal_retry_interval_seconds, policy,
                      std::placeholders::_1), &mFdp);
        policySetInt64(
            std::bind(&License_Policy::set_renewal_delay_seconds, policy,
                      std::placeholders::_1), &mFdp);

        if (mFdp.ConsumeBool()) {
          video_widevine::License::Policy::TimerDelayBase timeDelayBase =
              (video_widevine::License::Policy::TimerDelayBase)mFdp.ConsumeIntegralInRange<uint8_t>(
              License_Policy_TimerDelayBase_TIMER_DELAY_BASE_UNSPECIFIED ,
              License_Policy_TimerDelayBase_FIRST_DECRYPT);
          policy->set_initial_renewal_delay_base(timeDelayBase);
        }

        if (mFdp.ConsumeBool()) {
          std::string serverUrl = mFdp.ConsumeBool()
                                      ? kRenewalServerUrl
                                      : mFdp.ConsumeRandomLengthString(kMaxByte);
          policy->set_renewal_server_url(serverUrl);
        }
        break;
     }
     case License::KeyContainer::KEY_CONTROL: {
         key->set_type(License::KeyContainer::KEY_CONTROL);
         // No control bits for this type of key container used for license
         break;
     }
     case License::KeyContainer::OPERATOR_SESSION: {
        key->set_type(License::KeyContainer::OPERATOR_SESSION);

        if (mFdp.ConsumeBool()) {
          key->set_id(mKeyId);
        }

        if (mFdp.ConsumeBool()) {
          int32_t level = mFdp.ConsumeIntegralInRange<uint8_t>(
              License::KeyContainer::SW_SECURE_CRYPTO,
              License::KeyContainer::HW_SECURE_ALL);
          key->set_level((License_KeyContainer_SecurityLevel)level);
        }

        if (mFdp.ConsumeBool()) {
          policySetBool(std::bind(&License_Policy::set_can_persist, mLicense.mutable_policy(),
                                std::placeholders::_1), &mFdp);
        }

        License::KeyContainer::OperatorSessionKeyPermissions *permissions =
            key->mutable_operator_session_key_permissions();
        policySetBool(
            std::bind(&License::KeyContainer::OperatorSessionKeyPermissions::
                          set_allow_encrypt,
                      permissions, std::placeholders::_1), &mFdp);
        policySetBool(
            std::bind(&License::KeyContainer::OperatorSessionKeyPermissions::
                          set_allow_decrypt,
                      permissions, std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License::KeyContainer::
                                    OperatorSessionKeyPermissions::set_allow_sign,
                                permissions, std::placeholders::_1), &mFdp);
        policySetBool(
            std::bind(&License::KeyContainer::OperatorSessionKeyPermissions::
                          set_allow_signature_verify,
                      permissions, std::placeholders::_1), &mFdp);
        break;
     }
     case License::KeyContainer::ENTITLEMENT:
     default: {
        if (mFdp.ConsumeBool()) {
          mKeyId = mFdp.ConsumeBool() ? kEntitlementKeyId : mKeyId;
          key->set_id(mKeyId);
        }
        if (mFdp.ConsumeBool()) {
          int32_t level = mFdp.ConsumeIntegralInRange<uint8_t>(
              License::KeyContainer::SW_SECURE_CRYPTO,
              License::KeyContainer::HW_SECURE_ALL);
          key->set_level((License_KeyContainer_SecurityLevel)level);
        }

        License_Policy *policy = mLicense.mutable_policy();
        policySetBool(
            std::bind(&License_Policy::set_can_play, policy, std::placeholders::_1),
            &mFdp);
        policySetBool(std::bind(&License_Policy::set_can_persist, policy,
                                std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_can_renew, policy,
                                std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_renew_with_usage, policy,
                                std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_soft_enforce_rental_duration,
                                policy, std::placeholders::_1), &mFdp);
        policySetBool(std::bind(&License_Policy::set_soft_enforce_playback_duration,
                                policy, std::placeholders::_1), &mFdp);

        policySetInt64(std::bind(&License_Policy::set_rental_duration_seconds,
                                 policy, std::placeholders::_1), &mFdp);
        policySetInt64(std::bind(&License_Policy::set_playback_duration_seconds,
                                 policy, std::placeholders::_1), &mFdp);
        policySetInt64(std::bind(&License_Policy::set_license_duration_seconds,
                                 policy, std::placeholders::_1), &mFdp);
        policySetInt64(
            std::bind(&License_Policy::set_renewal_recovery_duration_seconds,
                      policy, std::placeholders::_1), &mFdp);
        policySetInt64(
            std::bind(&License_Policy::set_renewal_retry_interval_seconds, policy,
                      std::placeholders::_1), &mFdp);
        policySetInt64(
            std::bind(&License_Policy::set_renewal_delay_seconds, policy,
                      std::placeholders::_1), &mFdp);

        if (mFdp.ConsumeBool()) {
          video_widevine::License::Policy::TimerDelayBase timeDelayBase = (video_widevine::License::Policy::TimerDelayBase)mFdp.ConsumeIntegralInRange<uint8_t>(
              License_Policy_TimerDelayBase_TIMER_DELAY_BASE_UNSPECIFIED ,
              License_Policy_TimerDelayBase_FIRST_DECRYPT);
          policy->set_initial_renewal_delay_base(timeDelayBase);
        }
        break;
     }
  }
}

void PolicyEngineFuzzer::process() {
  setUp();
  FuzzWvCdmEventListener fuzzWvCdmEventListener;
  CdmSessionId sessionId = mFdp.ConsumeRandomLengthString(kMaxByte);
  metrics::CryptoMetrics crypto_metrics;
  std::unique_ptr<CryptoSession> cryptoSession(
      wvcdm::CryptoSession::MakeCryptoSession(&crypto_metrics));
  std::unique_ptr<PolicyEngine> policyEngine(new PolicyEngine(
      sessionId, &fuzzWvCdmEventListener, cryptoSession.get()));

  while (mFdp.remaining_bytes()) {
    auto invokePolicyEngineAPI =
        mFdp.PickValueInArray<const std::function<void()>>({
            [&]() { policyEngine->CanDecryptContent(mKeyId); },
            [&]() { policyEngine->CanUseKeyForSecurityLevel(mKeyId); },
            [&]() { policyEngine->OnTimerEvent(); },
            [&]() {
              policyEngine->SetLicense(
                  mLicense, mFdp.ConsumeBool() /*defer_license_state_update*/);
            },
            [&]() {
              int32_t maxEntitledKey =
                  mFdp.ConsumeIntegralInRange<uint8_t>(1, kMaxEntitledKey);
              std::vector<WidevinePsshData_EntitledKey> entitled_keys(
                  maxEntitledKey);
              for (int i = 0; i < maxEntitledKey; ++i) {
                entitled_keys[i].set_entitlement_key_id(
                    mFdp.ConsumeRandomLengthString(kMaxByte));
                entitled_keys[i].set_key_id(
                    mFdp.ConsumeBool()
                        ? mFdp.ConsumeRandomLengthString(kMaxByte)
                        : mKeyId);
              }
              policyEngine->SetEntitledLicenseKeys(entitled_keys);
            },
            [&]() { policyEngine->SetLicenseForRelease(mLicense); },
            [&]() { policyEngine->BeginDecryption(); },
            [&]() { policyEngine->DecryptionEvent(); },
            [&]() {
              policyEngine->UpdateLicense(
                  mLicense, mFdp.ConsumeBool() /*defer_license_state_update*/);
            },
            [&]() {
              policyEngine->UpdateLicenseState(
                  mFdp.ConsumeIntegral<int64_t>() /*current_time*/);
            },
            [&]() {
              int32_t width = mFdp.ConsumeIntegralInRange<int32_t>(
                  kMinResolution, kMaxResolution);
              int32_t height = mFdp.ConsumeIntegralInRange<int32_t>(
                  kMinResolution, kMaxResolution);
              policyEngine->NotifyResolution(width, height);
            },
            [&]() { policyEngine->NotifySessionExpiration(); },
            [&]() {
              CdmQueryMap query_info;
              policyEngine->Query(&query_info);
            },
            [&]() {
              CdmKeyAllowedUsage key_usage;
              policyEngine->QueryKeyAllowedUsage(
                  mKeyId, (mFdp.ConsumeBool() ? nullptr : &key_usage));
            },
            [&]() {
              policyEngine->RestorePlaybackTimes(
                  mFdp.ConsumeIntegral<int64_t>() /*playback_start_time*/,
                  mFdp.ConsumeIntegral<int64_t>() /*last_playback_time*/,
                  mFdp.ConsumeIntegral<int64_t>() /*grace_period_end_time*/);
            },
            [&]() {
              policyEngine->HasLicenseOrRentalOrPlaybackDurationExpired(
                  mFdp.ConsumeIntegral<int64_t>());
            },
        });
    invokePolicyEngineAPI();
  }
}

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