/*
 * 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 <device_files.h>
#include <file_store.h>
#include <fuzzer/FuzzedDataProvider.h>
#include <license_protocol.pb.h>
#include <string_conversions.h>

using namespace wvcdm;
using namespace wvutil;
using namespace video_widevine;

static constexpr uint16_t kMaxByte = 256;
static constexpr uint16_t kMaxVectorSize = 1000;
static constexpr uint8_t kMinVectorSize = 0;
static constexpr uint8_t kMinKeySetId = 0;
static constexpr uint8_t kMaxKeySetId = 100;

class DeviceFilesFuzzer {
public:
  DeviceFilesFuzzer(const uint8_t *data, size_t size) : mFdp(data, size){};
  void process();

private:
  FuzzedDataProvider mFdp;
};

void createLicenseData(DeviceFiles::CdmLicenseData &licenseData,
                       FuzzedDataProvider &mFdp) {
  std::map<std::string, std::string> appParameters;
  appParameters[mFdp.ConsumeBytesAsString(kMaxByte)] =
      mFdp.ConsumeBytesAsString(kMaxByte);

  licenseData.key_set_id = std::to_string(
      mFdp.ConsumeIntegralInRange<uint32_t>(kMinKeySetId, kMaxKeySetId));
  licenseData.state =
      (CdmOfflineLicenseState)mFdp.ConsumeIntegralInRange<int32_t>(
          kLicenseStateActive, kLicenseStateUnknown);
  licenseData.pssh_data = mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.license_request = mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.license = mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.license_renewal_request =
      mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.license_renewal = mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.release_server_url = mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.playback_start_time = mFdp.ConsumeIntegral<int64_t>();
  licenseData.last_playback_time = mFdp.ConsumeIntegral<int64_t>();
  licenseData.grace_period_end_time = mFdp.ConsumeIntegral<int64_t>();
  licenseData.app_parameters = appParameters;
  licenseData.usage_entry = mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.usage_entry_index = mFdp.ConsumeIntegral<uint32_t>();
  licenseData.drm_certificate = mFdp.ConsumeRandomLengthString(kMaxByte);
  licenseData.wrapped_private_key = CryptoWrappedKey(
      (CryptoWrappedKey::Type)mFdp.ConsumeIntegralInRange<int32_t>(
          CryptoWrappedKey::Type::kUninitialized, CryptoWrappedKey::Type::kEcc),
      mFdp.ConsumeRandomLengthString(kMaxByte));
}

void createUsageData(DeviceFiles::CdmUsageData &usageData,
                     FuzzedDataProvider &mFdp) {
  usageData.provider_session_token = mFdp.ConsumeRandomLengthString(kMaxByte);
  usageData.license_request = mFdp.ConsumeRandomLengthString(kMaxByte);
  usageData.license = mFdp.ConsumeRandomLengthString(kMaxByte);
  usageData.key_set_id = std::to_string(
      mFdp.ConsumeIntegralInRange<uint32_t>(kMinKeySetId, kMaxKeySetId));
  usageData.usage_entry = mFdp.ConsumeRandomLengthString(kMaxByte);
  usageData.usage_entry_index = mFdp.ConsumeIntegral<uint32_t>();
  usageData.drm_certificate = mFdp.ConsumeRandomLengthString(kMaxByte);
  usageData.wrapped_private_key = CryptoWrappedKey(
      (CryptoWrappedKey::Type)mFdp.ConsumeIntegralInRange<int32_t>(
          CryptoWrappedKey::Type::kUninitialized, CryptoWrappedKey::Type::kEcc),
      mFdp.ConsumeRandomLengthString(kMaxByte));
}

void createCertificateAndKey(std::string &certificate, CryptoWrappedKey &key,
                             FuzzedDataProvider &mFdp) {
  DrmCertificate drmCertificate;
  drmCertificate.set_type(
      (DrmCertificate_Type)mFdp.ConsumeIntegralInRange<int32_t>(
          DrmCertificate_Type::DrmCertificate_Type_ROOT,
          DrmCertificate_Type::DrmCertificate_Type_PROVISIONER));
  drmCertificate.set_serial_number(
      std::to_string(mFdp.ConsumeIntegral<int64_t>()));
  drmCertificate.set_system_id(mFdp.ConsumeIntegral<uint32_t>());
  drmCertificate.set_creation_time_seconds(mFdp.ConsumeIntegral<int32_t>());
  drmCertificate.set_expiration_time_seconds(mFdp.ConsumeIntegral<int32_t>());
  std::string setDrmCertificate;
  drmCertificate.SerializeToString(&setDrmCertificate);

  SignedDrmCertificate signedDrmCertificate;
  if (mFdp.ConsumeBool()) {
    signedDrmCertificate.set_drm_certificate(setDrmCertificate);
  }

  signedDrmCertificate.SerializeToString(&certificate);

  key = CryptoWrappedKey(
      (CryptoWrappedKey::Type)mFdp.ConsumeIntegralInRange<int32_t>(
          CryptoWrappedKey::Type::kUninitialized, CryptoWrappedKey::Type::kEcc),
      mFdp.ConsumeBytesAsString(kMaxByte));
}

void DeviceFilesFuzzer::process() {
  FileSystem fileSystem;
  DeviceFiles deviceFiles(mFdp.ConsumeBool() ? &fileSystem : nullptr);
  deviceFiles.Reset((CdmSecurityLevel)mFdp.ConsumeIntegralInRange<int32_t>(
      kSecurityLevelUninitialized, kSecurityLevelUnknown));

  while (mFdp.remaining_bytes()) {
    auto invokeDeviceFilesAPI = mFdp.PickValueInArray<
        const std::function<void()>>(
        {[&]() {
           std::string certificate;
           CryptoWrappedKey wrappedPrivateKey;
           createCertificateAndKey(certificate, wrappedPrivateKey, mFdp);
           deviceFiles.StoreCertificate(mFdp.ConsumeBool() ? certificate : "",
                                        wrappedPrivateKey);
         },

         [&]() {
           bool AtscModeEnabled = mFdp.ConsumeBool();
           std::string certificate;
           CryptoWrappedKey privateKey;
           std::string serialNumber;
           uint32_t systemId;

           deviceFiles.RetrieveCertificate(
               AtscModeEnabled, mFdp.ConsumeBool() ? &certificate : nullptr,
               &privateKey, mFdp.ConsumeBool() ? &serialNumber : nullptr,
               mFdp.ConsumeBool() ? &systemId : nullptr);
         },

         [&]() {
           std::string certificate;
           CryptoWrappedKey privateKey;
           std::string serialNumber;
           uint32_t systemId;

           deviceFiles.RetrieveLegacyCertificate(&certificate, &privateKey,
                                                 &serialNumber, &systemId);
         },

         [&]() { deviceFiles.RemoveCertificate(); },

         [&]() {
           std::string certificate;
           CryptoWrappedKey wrappedPrivateKey;
           createCertificateAndKey(certificate, wrappedPrivateKey, mFdp);

           deviceFiles.StoreOemCertificate(
               mFdp.ConsumeBool() ? certificate : "", wrappedPrivateKey);
         },

         [&]() {
           std::string certificate;
           CryptoWrappedKey privateKey;

           deviceFiles.RetrieveOemCertificate(&certificate, &privateKey);
         },

         [&]() { deviceFiles.HasOemCertificate(); },

         [&]() { deviceFiles.RemoveOemCertificate(); },

         [&]() {
           DeviceFiles::CdmLicenseData licenseData;
           createLicenseData(licenseData, mFdp);
           DeviceFiles::ResponseType result;
           deviceFiles.StoreLicense(licenseData,
                                    mFdp.ConsumeBool() ? &result : nullptr);
         },

         [&]() {
           DeviceFiles::CdmLicenseData licenseData;
           DeviceFiles::ResponseType result;

           deviceFiles.RetrieveLicense(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* key_set_id */
               mFdp.ConsumeBool() ? &licenseData : nullptr,
               mFdp.ConsumeBool() ? &result : nullptr);
         },

         [&]() {
           deviceFiles.StoreAtscLicense(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* key_set_id */
               mFdp.ConsumeRandomLengthString(
                   kMaxByte) /* serialized_license_data */
           );
         },

         [&]() {
           deviceFiles.DeleteLicense(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)) /* key_set_id */
           );
         },

         [&]() {
           std::vector<std::string> keySetIds;

           deviceFiles.ListLicenses(&keySetIds);
         },

         [&]() { deviceFiles.DeleteAllFiles(); },

         [&]() { deviceFiles.DeleteAllLicenses(); },

         [&]() {
           deviceFiles.LicenseExists(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)) /* key_set_id */
           );
         },

         [&]() {
           deviceFiles.ReserveLicenseId(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)) /* key_set_id */
           );
         },

         [&]() {
           deviceFiles.UnreserveLicenseId(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)) /* key_set_id */
           );
         },

         [&]() {
           deviceFiles.StoreUsageInfo(
               mFdp.ConsumeRandomLengthString(
                   kMaxByte), /* provider_session_token */
               mFdp.ConsumeRandomLengthString(kMaxByte), /* key_request */
               mFdp.ConsumeRandomLengthString(kMaxByte), /* key_response */
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)),         /* key_set_id */
               mFdp.ConsumeRandomLengthString(kMaxByte), /* usage_entry */
               mFdp.ConsumeIntegral<uint32_t>(),         /* usage_entry_index */
               mFdp.ConsumeRandomLengthString(kMaxByte), /* drm_certificate */
               CryptoWrappedKey(
                   (CryptoWrappedKey::Type)mFdp.ConsumeIntegralInRange<int32_t>(
                       CryptoWrappedKey::Type::kUninitialized,
                       CryptoWrappedKey::Type::kEcc),       /* type */
                   mFdp.ConsumeRandomLengthString(kMaxByte) /* key */
                   ) /* wrapped_private_key */
           );
         },

         [&]() {
           std::vector<std::string> keySetIds;
           std::vector<std::string> providerSessionTokens;

           deviceFiles.ListUsageIds(
               mFdp.ConsumeRandomLengthString(kMaxByte), /* app_id */
               mFdp.ConsumeBool() ? &keySetIds : nullptr,
               mFdp.ConsumeBool() ? &providerSessionTokens : nullptr);
         },

         [&]() {
           std::string providerSessionToken;

           deviceFiles.GetProviderSessionToken(
               mFdp.ConsumeRandomLengthString(kMaxByte), /* app_id */
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* key_set_id */
               &providerSessionToken);
         },

         [&]() {
           deviceFiles.DeleteUsageInfo(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)) /* key_set_id */
           );
         },

         [&]() {
           uint32_t keySetIdsSize = mFdp.ConsumeIntegralInRange<uint32_t>(
               kMinVectorSize, kMaxVectorSize);
           std::vector<std::string> keySetIds;
           for (int i = 0; i < keySetIdsSize; ++i)
             keySetIds.push_back(
                 std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                     kMinKeySetId, kMaxKeySetId)) /* key_set_id */
             );

           deviceFiles.DeleteMultipleUsageInfoByKeySetIds(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               keySetIds);
         },

         [&]() {
           std::vector<std::string> providerSessionTokens;

           deviceFiles.DeleteAllUsageInfoForApp(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               &providerSessionTokens);
         },

         [&]() { deviceFiles.DeleteAllUsageInfo(); },

         [&]() {
           CdmKeyMessage licenseRequest;
           CdmKeyResponse license;
           UsageEntry usageEntry;
           UsageEntryIndex usageEntryIndex;
           std::string drmCertificate;
           CryptoWrappedKey wrappedPrivateKey;

           deviceFiles.RetrieveUsageInfo(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* provider_session_token */
               &licenseRequest, &license, &usageEntry, &usageEntryIndex,
               &drmCertificate, &wrappedPrivateKey);
         },

         [&]() {
           std::string providerSessionToken;
           CdmKeyMessage licenseRequest;
           CdmKeyResponse license;
           UsageEntry usageEntry;
           UsageEntryIndex usageEntryIndex;
           std::string drmCertificate;
           CryptoWrappedKey wrappedPrivateKey;

           deviceFiles.RetrieveUsageInfoByKeySetId(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* key_set_id */
               &providerSessionToken, &licenseRequest, &license, &usageEntry,
               &usageEntryIndex, &drmCertificate, &wrappedPrivateKey);
         },

         [&]() {
           std::vector<std::string> usageInfoFileNames;

           deviceFiles.ListUsageInfoFiles(&usageInfoFileNames);
         },

         [&]() {
           std::vector<DeviceFiles::CdmUsageData> usageData;

           deviceFiles.RetrieveUsageInfo(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               &usageData);
         },

         [&]() {
           std::string providerSessionToken =
               mFdp.ConsumeRandomLengthString(kMaxByte);
           DeviceFiles::CdmUsageData usageData;

           deviceFiles.RetrieveUsageInfo(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               providerSessionToken, &usageData);
         },

         [&]() {
           uint32_t usageDataSize = mFdp.ConsumeIntegralInRange<uint32_t>(
               kMinVectorSize, kMaxVectorSize);
           std::vector<DeviceFiles::CdmUsageData> usageData;
           for (int i = 0; i < usageDataSize; ++i) {
             DeviceFiles::CdmUsageData unit;
             createUsageData(unit, mFdp);
             usageData.push_back(unit);
           }

           deviceFiles.StoreUsageInfo(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               usageData);
         },

         [&]() {
           DeviceFiles::CdmUsageData usageData;
           createUsageData(usageData, mFdp);
           deviceFiles.UpdateUsageInfo(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* usage_info_file_name */
               usageData);
         },

         [&]() {
           uint32_t mediaSegmentIvSize = mFdp.ConsumeIntegralInRange<uint32_t>(
               kMinVectorSize, kMaxVectorSize);
           std::vector<uint8_t> mediaSegmentIv;
           for (int i = 0; i < mediaSegmentIvSize; ++i)
             mediaSegmentIv.push_back(mFdp.ConsumeIntegral<uint8_t>());

           deviceFiles.StoreHlsAttributes(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* key_set_id */
               (CdmHlsMethod)mFdp.ConsumeIntegralInRange<int32_t>(
                   kHlsMethodNone, kHlsMethodSampleAes), /* method */
               mediaSegmentIv);
         },

         [&]() {
           CdmHlsMethod method;
           std::vector<uint8_t> mediaSegmentIv;

           deviceFiles.RetrieveHlsAttributes(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)), /* key_set_id */
               &method, &mediaSegmentIv);
         },

         [&]() {
           deviceFiles.DeleteHlsAttributes(
               std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                   kMinKeySetId, kMaxKeySetId)) /* key_set_id */
           );
         },

         [&]() {
           uint32_t usageEntryInfoListSize =
               mFdp.ConsumeIntegralInRange<uint32_t>(kMinVectorSize,
                                                     kMaxVectorSize);
           std::vector<CdmUsageEntryInfo> usageEntryInfoList;
           for (int i = 0; i < usageEntryInfoListSize; ++i) {
             CdmUsageEntryInfo unit;

             unit.storage_type =
                 (CdmUsageEntryStorageType)mFdp.ConsumeIntegralInRange<int32_t>(
                     kStorageLicense, kStorageTypeUnknown);
             unit.key_set_id = mFdp.ConsumeRandomLengthString(kMaxByte);
             unit.usage_info_file_name =
                 std::to_string(mFdp.ConsumeIntegralInRange<uint32_t>(
                     kMinKeySetId, kMaxKeySetId));
             unit.last_use_time = mFdp.ConsumeIntegral<int64_t>();
             unit.offline_license_expiry_time = mFdp.ConsumeIntegral<int64_t>();
             usageEntryInfoList.push_back(unit);
           }

           deviceFiles.StoreUsageTableInfo(
               mFdp.ConsumeRandomLengthString(
                   kMaxByte), /* usage_table_header */
               usageEntryInfoList);
         },

         [&]() {
           UsageTableHeader usageTableHeader;
           std::vector<CdmUsageEntryInfo> usageEntryInfoList;
           bool lruUpgrade;

           deviceFiles.RetrieveUsageTableInfo(&usageTableHeader,
                                              &usageEntryInfoList, &lruUpgrade);
         },

         [&]() { deviceFiles.DeleteUsageTableInfo(); },

         [&]() {
           okp::SystemFallbackInfo info;
           info.SetState((okp::SystemState)mFdp.ConsumeIntegralInRange<int32_t>(
               (int)okp::SystemState::kUnknown,
               (int)okp::SystemState::kProvisioned));
           info.SetFirstCheckedTime(mFdp.ConsumeIntegral<int64_t>());
           info.SetBackoffStartTime(mFdp.ConsumeIntegral<int64_t>());
           info.SetBackoffDuration(mFdp.ConsumeIntegral<int64_t>());
           info.SetProvisioningTime(mFdp.ConsumeIntegral<int64_t>());

           deviceFiles.StoreOkpInfo(info);
         },

         [&]() {
           okp::SystemFallbackInfo info;

           deviceFiles.RetrieveOkpInfo(&info);
         },

         [&]() { deviceFiles.DeleteOkpInfo(); }});

    invokeDeviceFilesAPI();
  }
}

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