/*
 * 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 <policy_engine.h>
#include <system_fuzz_helper.h>
#include <system_id_extractor.h>
#include <wv_cdm_event_listener.h>
#include <wv_cdm_types.h>

using namespace wvcdm;
using namespace wvutil;
using namespace wvcdm::metrics;

static constexpr int32_t kMinByte = 0;
static constexpr int32_t kMaxByte = 256;
static constexpr int32_t kMaxEnum = 4;
static constexpr int32_t kMinEnum = 0;
static constexpr int32_t kMaxRequestedSecurityLevel = 1;
static constexpr int32_t kMinRequestedSecurityLevel = 0;

class FuzzDeviceFiles : public DeviceFiles {
public:
  FuzzDeviceFiles(wvutil::FileSystem *fs, FuzzedDataProvider *fdp)
      : DeviceFiles(fs), mFdp(fdp) {}

  bool Init(CdmSecurityLevel) override { return mFdp->ConsumeBool(); }

  DeviceFiles::CertificateState
  RetrieveOemCertificate(std::string *certificate,
                         CryptoWrappedKey * /*wrapped_private_key*/) override {
    *certificate = mFdp->ConsumeRandomLengthString(kMaxByte);
    return (DeviceFiles::CertificateState)mFdp->ConsumeIntegralInRange<int32_t>(
        kMinEnum, kMaxEnum);
  }

private:
  FuzzedDataProvider *mFdp;
};

class FuzzFile : public File {
public:
  FuzzFile(FuzzedDataProvider *fdp) : mFdp(fdp) {}

  ssize_t Read(char *buffer, size_t bytes) {
    std::vector<char> initData = mFdp->ConsumeBytes<char>(
        mFdp->ConsumeIntegralInRange<size_t>(kMinByte, bytes));
    buffer = initData.data();
    return initData.size();
  };

  ssize_t Write(const char * /* buffer*/, size_t bytes) {
    return mFdp->ConsumeIntegralInRange<ssize_t>(kMinByte, bytes);
  };

private:
  FuzzedDataProvider *mFdp;
};

class FuzzFileSystem : public FileSystem {
public:
  FuzzFileSystem(FuzzedDataProvider *fdp) : mFdp(fdp) {}

  std::unique_ptr<File> Open(const std::string &, int) {
    return std::unique_ptr<File>(new FuzzFile(mFdp));
  }

  bool Exists(const std::string &) { return mFdp->ConsumeBool(); }

  bool Remove(const std::string &) { return mFdp->ConsumeBool(); }

  ssize_t FileSize(const std::string &) {
    return mFdp->ConsumeIntegralInRange<ssize_t>(kMinByte, kMaxByte);
  }

  bool List(const std::string &, std::vector<std::string> *) {
    return mFdp->ConsumeBool();
  }

private:
  FuzzedDataProvider *mFdp;
};

class FuzzCryptoSession : public CryptoSession {
public:
  FuzzCryptoSession(metrics::CryptoMetrics *metrics, FuzzedDataProvider *fdp)
      : CryptoSession(metrics) {
    mFdp = fdp;
  }

  ~FuzzCryptoSession() override {}

  bool IsOpen() override { return mFdp->ConsumeBool(); }

  bool GetCachedSystemId(uint32_t * /*system_id*/) override {
    return mFdp->ConsumeBool();
  }

  void SetSystemId(uint32_t /*system_id*/) override {}
  CdmSecurityLevel GetSecurityLevel(RequestedSecurityLevel /*level*/) override {
    return (CdmSecurityLevel)mFdp->ConsumeIntegralInRange(kMinEnum, kMaxEnum);
  }

  CdmResponseType
  GetProvisioningMethod(RequestedSecurityLevel /*level*/,
                        CdmClientTokenType *tokenType) override {
    *tokenType =
        (CdmClientTokenType)mFdp->ConsumeIntegralInRange(kMinEnum, kMaxEnum);
    return CdmResponseType(
        mFdp->PickValueInArray<CdmResponseEnum>(kCdmResponseEnums));
  }

  CdmResponseType GetTokenFromKeybox(RequestedSecurityLevel /*level*/,
                                     std::string *keyData) override {
    *keyData = mFdp->ConsumeRandomLengthString(kMaxByte);
    return CdmResponseType(
        mFdp->PickValueInArray<CdmResponseEnum>(kCdmResponseEnums));
  }

  CdmResponseType GetTokenFromOemCert(RequestedSecurityLevel /*level*/,
                                      std::string *oemCert) override {
    *oemCert = mFdp->ConsumeRandomLengthString(kMaxByte);
    return CdmResponseType(
        mFdp->PickValueInArray<CdmResponseEnum>(kCdmResponseEnums));
  }
  void SetSystemIdBase(uint32_t /* system_id*/) {}

private:
  FuzzedDataProvider *mFdp;
};

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

private:
  FuzzedDataProvider mFdp;
};

void SystemIdExtractorFuzzer::process() {

  FuzzFileSystem fuzzFileSystem(&mFdp);
  metrics::CryptoMetrics cryptoMetrics;
  FuzzDeviceFiles deviceFiles(&fuzzFileSystem, &mFdp);

  RequestedSecurityLevel securityLevel =
      (RequestedSecurityLevel)mFdp.ConsumeIntegralInRange(
          kMinRequestedSecurityLevel, kMaxRequestedSecurityLevel);

  std::unique_ptr<FuzzCryptoSession> cryptoSessions =
      std::make_unique<FuzzCryptoSession>(&cryptoMetrics, &mFdp);

  std::unique_ptr<SystemIdExtractor> systemIdExtractor =
      std::make_unique<SystemIdExtractor>(securityLevel, cryptoSessions.get(),
                                          &fuzzFileSystem);

  while (mFdp.remaining_bytes()) {
    auto invokeSystemIdExtractorAPI = mFdp.PickValueInArray<
        const std::function<void()>>({
        [&]() {
          uint32_t systemId;
          systemId = mFdp.ConsumeIntegralInRange<uint32_t>(kMinByte, kMaxByte);
          systemIdExtractor->ExtractSystemId(&systemId);
        },
        [&]() {
          uint32_t systemId;
          systemId = mFdp.ConsumeIntegralInRange<uint32_t>(kMinByte, kMaxByte);
          const std::string keyData = mFdp.ConsumeRandomLengthString(kMaxByte);
          systemIdExtractor->ExtractSystemIdFromKeyboxData(keyData, &systemId);
        },
        [&]() {
          uint32_t systemId;
          systemId = mFdp.ConsumeIntegralInRange<uint32_t>(kMinByte, kMaxByte);
          const std::string oemCert = mFdp.ConsumeRandomLengthString(kMaxByte);
          systemIdExtractor->ExtractSystemIdFromOemCert(oemCert, &systemId);
        },
        [&]() { systemIdExtractor->SetDeviceFilesForTesting(&deviceFiles); },
    });
    invokeSystemIdExtractorAPI();
  }
}

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