/*
 * 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 "crypto_session.h"
#include "entitlement_key_session.h"
#include "properties.h"
#include "string_conversions.h"
#include "wv_cdm_constants.h"
#include <fuzzer/FuzzedDataProvider.h>

using namespace wvcdm;

constexpr int32_t kStringLength = 20;
constexpr int32_t kMaxSamplesCount = 1000;
constexpr int32_t kMinSamplesCount = 1;
constexpr int32_t kMaxByte = 1000;
constexpr int32_t kMinByte = 0;
constexpr int32_t kIvSize = 16;
constexpr int32_t kDefaultSampleSize = 0;

class CryptoSessionFuzzer {
public:
    CryptoSessionFuzzer(const uint8_t* data, size_t size) : mFdp(data, size) {};
    void process();
    void FillDecryptParams(CdmDecryptionParametersV16& params);
    void SetKeyParams(CryptoKey& inputKey);
    void setUp(CryptoSession* cryptoSession, const std::string& keyId,
        std::string& message, std::string& signature,
        std::string& coreMessage);
    void initializeStripeBuffer(std::vector<uint8_t>* buffer, size_t size);

private:
    FuzzedDataProvider mFdp;
};

void CryptoSessionFuzzer::initializeStripeBuffer(std::vector<uint8_t>* buffer,
    size_t size) {
    buffer->assign(size, 0);
    for (size_t i = 0; i < size; ++i) {
        (*buffer)[i] = mFdp.ConsumeIntegral<uint8_t>();
    }
}

void CryptoSessionFuzzer::SetKeyParams(CryptoKey& inputKey) {
    inputKey.set_key_data(mFdp.ConsumeRandomLengthString(kStringLength));
    inputKey.set_key_data_iv(mFdp.ConsumeRandomLengthString(kStringLength));
    inputKey.set_key_control(mFdp.ConsumeRandomLengthString(kStringLength));
    inputKey.set_key_control_iv(mFdp.ConsumeRandomLengthString(kStringLength));
    inputKey.set_cipher_mode(mFdp.ConsumeBool() ? kCipherModeCbc
        : kCipherModeCtr);
    inputKey.set_track_label(mFdp.ConsumeRandomLengthString(kStringLength));
    inputKey.set_entitlement_key_id(
        mFdp.ConsumeRandomLengthString(kStringLength));
}

void CryptoSessionFuzzer::FillDecryptParams(
    CdmDecryptionParametersV16& params) {
    params.cipher_mode = mFdp.ConsumeBool() ? kCipherModeCbc : kCipherModeCtr;
    params.observe_legacy_fields = mFdp.ConsumeBool();
    uint32_t noOfSamples =
        mFdp.ConsumeIntegralInRange<uint32_t>(kMinSamplesCount, kMaxSamplesCount);

    while (--noOfSamples) {
        size_t sampleSize = 0;
        uint32_t noOfSubSamples = mFdp.ConsumeIntegralInRange<uint32_t>(
            kMinSamplesCount, kMaxSamplesCount);
        std::vector<CdmDecryptionSubsample> subSampleVector;
        while (--noOfSubSamples) {
            size_t clearBytes =
                mFdp.ConsumeIntegralInRange<size_t>(kMinByte, kMaxByte);
            size_t protectedBytes =
                mFdp.ConsumeIntegralInRange<size_t>(kMinByte, kMaxByte);
            CdmDecryptionSubsample subsample(clearBytes, protectedBytes);
            sampleSize += clearBytes + protectedBytes;
            subSampleVector.push_back(subsample);
        }
        std::vector<uint8_t> iv;
        iv = mFdp.ConsumeBytes<uint8_t>(kIvSize);
        if (iv.size() != kIvSize) {
            iv.resize(kIvSize, 0);
        }
        void* decryptBuffer;
        std::vector<uint8_t> eBuffer;
        for (int i = 0; i < sampleSize; ++i) {
            eBuffer.push_back(
                mFdp.ConsumeIntegralInRange<uint8_t>(kMinByte, kMaxByte));
        }
        size_t decryptBufferOffset =
            mFdp.ConsumeIntegralInRange<uint8_t>(kMinByte, kMaxByte);
        CdmDecryptionSample sample(eBuffer.data(), decryptBuffer,
            decryptBufferOffset, sampleSize, iv);
        sample.subsamples = subSampleVector;
        params.samples.push_back(sample);
    }
}

void CryptoSessionFuzzer::setUp(CryptoSession* cryptoSession,
    const std::string& keyId, std::string& message,
    std::string& signature,
    std::string& coreMessage) {
    bool shouldSpecifyAlgorithm;
    OEMCrypto_SignatureHashAlgorithm algorithm;
    CdmLicenseKeyType keyType1 = kLicenseKeyTypeEntitlement;
    cryptoSession->PrepareAndSignProvisioningRequest(
        message, &coreMessage, &signature, shouldSpecifyAlgorithm, algorithm);
    cryptoSession->PrepareAndSignLicenseRequest(
        message, &coreMessage, &signature, shouldSpecifyAlgorithm, algorithm);
    cryptoSession->LoadLicense(message, coreMessage, signature, keyType1);

    CryptoKey inputKey;
    inputKey.set_key_id(keyId);
    CryptoSessionFuzzer::SetKeyParams(inputKey);
    std::vector<CryptoKey> keys = { inputKey };
    cryptoSession->LoadEntitledContentKeys(keys);
}

void CryptoSessionFuzzer::process() {
    metrics::CryptoMetrics cryptoMetrics;
    std::unique_ptr<CryptoSession> cryptoSession(
        CryptoSession::MakeCryptoSession(&cryptoMetrics));

    CdmLicenseKeyType keyType;
    OEMCrypto_SignatureHashAlgorithm algorithm;
    Properties::Init();

    RequestedSecurityLevel requestedSecurityLevel =
        mFdp.ConsumeBool() ? kLevelDefault : kLevel3;
    cryptoSession->Open(requestedSecurityLevel);

    while (mFdp.remaining_bytes()) {
        auto invokeCryptoSessionAPI = mFdp.PickValueInArray<
            const std::function<void()>>({
            [&]() {
              std::string token = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string additionalToken =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              cryptoSession->GetProvisioningToken(
                  mFdp.ConsumeBool() ? &token : nullptr,
                  mFdp.ConsumeBool() ? &additionalToken : nullptr);
            },
            [&]() {
              cryptoSession->SetSystemId(
                  mFdp.ConsumeIntegral<uint32_t>() /*system_id*/);
            },
            [&]() {
              std::string provisioningId;
              cryptoSession->GetProvisioningId(&provisioningId);
              cryptoSession->GetExternalDeviceUniqueId(&provisioningId);
            },
            [&]() {
              cryptoSession->LoadLicense(
                  mFdp.ConsumeRandomLengthString(kStringLength) /*signed_message*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*coreMessage*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*signature*/,
                  keyType);
            },
            [&]() {
              std::string coreMessage =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              std::string signature = mFdp.ConsumeRandomLengthString(kStringLength);
              cryptoSession->PrepareAndSignRenewalRequest(
                  mFdp.ConsumeRandomLengthString(kStringLength) /*message*/,
                  mFdp.ConsumeBool() ? &coreMessage : nullptr,
                  mFdp.ConsumeBool() ? &signature : nullptr);
            },
            [&]() {
              cryptoSession->LoadRenewal(
                  mFdp.ConsumeRandomLengthString(kStringLength) /*signed_message*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*coreMessage*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*signature*/);
            },
            [&]() {
              const std::string& message =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              std::string coreMessage =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              std::string signature = mFdp.ConsumeRandomLengthString(kStringLength);
              OEMCrypto_SignatureHashAlgorithm algorithm;
              bool shouldSpecifyAlgorithm = mFdp.ConsumeBool();
              cryptoSession->PrepareAndSignProvisioningRequest(
                  message, mFdp.ConsumeBool() ? &coreMessage : nullptr,
                  mFdp.ConsumeBool() ? &signature : nullptr, shouldSpecifyAlgorithm,
                  algorithm);
            },
            [&]() {
              if (mFdp.ConsumeBool()) {
               const CryptoWrappedKey privateKey;
               cryptoSession->LoadCertificatePrivateKey(privateKey);
              } else {
               CryptoWrappedKey::Type type =
                 (CryptoWrappedKey::Type)mFdp.ConsumeIntegral<int32_t>();
               std::string key = mFdp.ConsumeRandomLengthString(kStringLength);
               const CryptoWrappedKey privateKey(type, key);
               cryptoSession->LoadCertificatePrivateKey(privateKey);
              }
            },
            [&]() {
              std::string publicKey;
              std::string publicKeySignature;
              std::string wrappedPrivateKey;
              CryptoWrappedKey::Type keyType;
              cryptoSession->GenerateCertificateKeyPair(
                  &publicKey, &publicKeySignature, &wrappedPrivateKey, &keyType);
            },
            [&]() {
              if (mFdp.ConsumeBool()) {
               const CryptoWrappedKey privateKey;
               cryptoSession->LoadOemCertificatePrivateKey(privateKey);
              } else {
               CryptoWrappedKey::Type type =
                 (CryptoWrappedKey::Type)mFdp.ConsumeIntegral<int32_t>();
               std::string key = mFdp.ConsumeRandomLengthString(kStringLength);
               const CryptoWrappedKey privateKey(type, key);
               cryptoSession->LoadOemCertificatePrivateKey(privateKey);
              }
            },
            [&]() {
              size_t out;
              RequestedSecurityLevel security_level =
                 (RequestedSecurityLevel)mFdp.ConsumeIntegral<int32_t>();
              cryptoSession->GetNumberOfOpenSessions(security_level, &out);
            },
            [&]() {
              std::string info = mFdp.ConsumeRandomLengthString(kStringLength);
              cryptoSession->GetBuildInformation(&info);
            },
            [&]() {
              CdmWatermarkingSupport support =
                 (CdmWatermarkingSupport)mFdp.ConsumeIntegral<int32_t>();
              cryptoSession->GetWatermarkingSupport(&support);
            },
           [&]() {
              CdmProductionReadiness readiness =
                 (CdmProductionReadiness)mFdp.ConsumeIntegral<int32_t>();
               cryptoSession->GetProductionReadiness(&readiness);
            },
            [&]() {
               RequestedSecurityLevel security_level =
                 (RequestedSecurityLevel)mFdp.ConsumeIntegral<int32_t>();
               size_t number_of_entries = mFdp.ConsumeIntegral<size_t>();
               cryptoSession->GetMaximumUsageTableEntries(security_level,
                                                     &number_of_entries);
            },
            [&]() {
              RequestedSecurityLevel security_level =
                 (RequestedSecurityLevel)mFdp.ConsumeIntegral<int32_t>();
              uint32_t decrypt_hash_support = mFdp.ConsumeIntegral<uint32_t>();
              cryptoSession->GetDecryptHashSupport(security_level,
                                               &decrypt_hash_support);
            },
            [&]() {
              bool hasSupport = mFdp.ConsumeBool();
              cryptoSession->HasUsageTableSupport(&hasSupport);
            },
            [&]() {
              cryptoSession->DeactivateUsageInformation(
                  mFdp.ConsumeRandomLengthString(
                      kStringLength) /*provider_session_token*/);
            },
            [&]() {
              std::string usageReport =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              CryptoSession::UsageDurationStatus usageDurationStatus;
              int64_t secondsSinceStarted = mFdp.ConsumeIntegral<int64_t>();
              int64_t secondsSinceLastPlayed = mFdp.ConsumeIntegral<int64_t>();
              cryptoSession->GenerateUsageReport(
                  mFdp.ConsumeRandomLengthString(
                      kStringLength) /*provider_session_token*/,
                  &usageReport, &usageDurationStatus, &secondsSinceStarted,
                  &secondsSinceLastPlayed);
            },
            [&]() {
              uint32_t nonce = mFdp.ConsumeIntegral<uint32_t>();
              cryptoSession->GenerateNonce(&nonce);
            },
            [&]() {
              std::string signedMessage, coreMessage, signature;
              std::string wrappedPrivateKey;
              cryptoSession->LoadProvisioning(
                  signedMessage, coreMessage, signature,
                  mFdp.ConsumeBool() ? &wrappedPrivateKey : nullptr);
            },
            [&]() {
              cryptoSession->SetDecryptHash(
                  mFdp.ConsumeIntegral<int16_t>() /*frame_number*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*hash*/);
            },
            [&]() {
              CdmEncryptionAlgorithm algorithm;
              std::string outBuffer;
              cryptoSession->GenericDecrypt(
                  mFdp.ConsumeRandomLengthString(kStringLength) /*in_buffer*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*key_id*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*iv*/, algorithm,
                  &outBuffer);
            },
            [&]() {
              CdmSigningAlgorithm algorithm = mFdp.ConsumeBool()
                                                  ? kSigningAlgorithmHmacSha256
                                                  : kSigningAlgorithmUnknown;
              std::string signature;
              cryptoSession->GenericSign(
                  mFdp.ConsumeRandomLengthString(kStringLength) /*message*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*key_id*/,
                  algorithm, &signature);
            },
            [&]() {
              CdmSigningAlgorithm algorithm = kSigningAlgorithmHmacSha256;
              const std::string signature;
              cryptoSession->GenericVerify(
                  mFdp.ConsumeRandomLengthString(kStringLength) /*message*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*key_id*/,
                  algorithm, signature);
            },
            [&]() {
              UsageTableHeader usageTableHeader;
              cryptoSession->CreateUsageTableHeader(requestedSecurityLevel,
                                                    &usageTableHeader);

              cryptoSession->LoadUsageTableHeader(requestedSecurityLevel,
                                                  usageTableHeader);
              uint32_t newEntryCount;
              cryptoSession->ShrinkUsageTableHeader(
                  requestedSecurityLevel, newEntryCount, &usageTableHeader);
            },
            [&]() {
              const UsageTableHeader usageTableHeader;
              cryptoSession->LoadUsageTableHeader(requestedSecurityLevel,
                                                  usageTableHeader);
            },
            [&]() {
              uint32_t entryNumber;
              cryptoSession->CreateUsageEntry(&entryNumber);
              cryptoSession->MoveUsageEntry(entryNumber);
            },
            [&]() {
              const UsageEntry usageEntry;
              cryptoSession->LoadUsageEntry(
                  mFdp.ConsumeIntegral<uint32_t>() /*entry_number*/, usageEntry);
            },
            [&]() {
              UsageTableHeader usageTableHeader;
              UsageEntry usageEntry;
              cryptoSession->UpdateUsageEntry(&usageTableHeader, &usageEntry);
            },
            [&]() {
              bool canSupportOutput = mFdp.ConsumeBool();
              bool canDisableOutput = mFdp.ConsumeBool();
              bool canSupportCgmsA = mFdp.ConsumeBool();
              cryptoSession->GetAnalogOutputCapabilities(
                  &canSupportOutput, &canDisableOutput, &canSupportCgmsA);
            },
            [&]() {
              cryptoSession->SetDebugIgnoreKeyboxCount(
                  mFdp.ConsumeIntegral<uint32_t>() /*count*/);
            },
            [&]() {
              cryptoSession->SetAllowTestKeybox(mFdp.ConsumeBool() /*allow*/);
            },
            [&]() { cryptoSession->GetOkpFallbackPolicy(); },
            [&]() {
              std::string request = mFdp.ConsumeRandomLengthString(kStringLength);
              cryptoSession->PrepareOtaProvisioningRequest(
                  mFdp.ConsumeBool() /*use_test_key*/, &request);
            },
            [&]() {
              cryptoSession->LoadOtaProvisioning(
                  mFdp.ConsumeBool() /*use_test_key*/,
                  mFdp.ConsumeRandomLengthString(kStringLength) /*response*/);
            },
            [&]() {
              uint32_t tier = mFdp.ConsumeIntegral<uint32_t>();
              cryptoSession->GetResourceRatingTier(&tier);
            },
            [&]() {
              size_t max;
              cryptoSession->GetMaxNumberOfSessions(requestedSecurityLevel, &max);
            },
            [&]() {
              std::string keyData = mFdp.ConsumeRandomLengthString(kStringLength);
              cryptoSession->GetTokenFromKeybox(requestedSecurityLevel, &keyData);
            },
            [&]() {
              std::string bcc = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string additionalSignature =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              cryptoSession->GetBootCertificateChain(&bcc, &additionalSignature);
            },
            [&]() {
              cryptoSession->GenerateDerivedKeys(
                  mFdp.ConsumeRandomLengthString(kStringLength) /*message*/);
            },
            [&]() {
              CdmDecryptionParametersV16 params2;
              cryptoSession->Decrypt(params2);
            },
            [&]() {
              const std::string& message =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              std::string coreMessage, signature;
              bool shouldSpecifyAlgorithm;
              OEMCrypto_SignatureHashAlgorithm algorithm;

              cryptoSession->PrepareAndSignProvisioningRequest(
                  message, &coreMessage, &signature, shouldSpecifyAlgorithm,
                  algorithm);
              cryptoSession->PrepareAndSignLicenseRequest(
                  message, &coreMessage, &signature, shouldSpecifyAlgorithm,
                  algorithm);
            },
            [&]() {
              std::string keyId = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string message = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string signature, coreMessage;
              setUp(cryptoSession.get(), keyId, message, signature, coreMessage);

              std::vector<uint8_t> inVector;
              std::vector<uint8_t> ivVector;
              std::string inBuffer;
              std::string iv;
              CdmEncryptionAlgorithm algorithmEncrypt =
                  kEncryptionAlgorithmAesCbc128;
              std::string outBuffer;
              initializeStripeBuffer(&inVector, CONTENT_KEY_SIZE * 15);
              inBuffer = std::string(inVector.begin(), inVector.end());
              initializeStripeBuffer(&ivVector, KEY_IV_SIZE);
              iv = std::string(ivVector.begin(), ivVector.end());

              cryptoSession->GenericEncrypt(inBuffer, keyId, iv, algorithmEncrypt,
                                            &outBuffer);
            },
            [&]() {
              std::string keyId = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string message = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string signature, coreMessage;
              setUp(cryptoSession.get(), keyId, message, signature, coreMessage);

              std::vector<uint8_t> inVector;
              std::vector<uint8_t> ivVector;
              std::string inBuffer;
              std::string iv;
              CdmEncryptionAlgorithm algorithmEncrypt =
                  kEncryptionAlgorithmAesCbc128;
              std::string outBuffer;
              initializeStripeBuffer(&inVector, CONTENT_KEY_SIZE * 15);
              inBuffer = std::string(inVector.begin(), inVector.end());
              initializeStripeBuffer(&ivVector, KEY_IV_SIZE);
              iv = std::string(ivVector.begin(), ivVector.end());

              cryptoSession->GenericDecrypt(inBuffer, keyId, iv, algorithmEncrypt,
                                            &outBuffer);
            },
            [&]() {
              std::string keyId = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string message = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string signature, coreMessage;
              setUp(cryptoSession.get(), keyId, message, signature, coreMessage);

              CdmSigningAlgorithm algorithm2 = kSigningAlgorithmHmacSha256;
              cryptoSession->GenericSign(message, keyId, algorithm2, &signature);
            },
            [&]() {
              std::string keyId = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string message = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string signature, coreMessage;
              setUp(cryptoSession.get(), keyId, message, signature, coreMessage);

              CdmSigningAlgorithm algorithm2 = kSigningAlgorithmHmacSha256;
              cryptoSession->GenericVerify(message, keyId, algorithm2, signature);
            },
            [&]() {
              const std::string keyId =
                  mFdp.ConsumeRandomLengthString(kStringLength);
              std::string message = mFdp.ConsumeRandomLengthString(kStringLength);
              std::string signature, coreMessage;
              setUp(cryptoSession.get(), keyId, message, signature, coreMessage);
              CdmDecryptionParametersV16 params(keyId);
              CryptoSessionFuzzer::FillDecryptParams(params);
              cryptoSession->Decrypt(params);
            },
                });
        invokeCryptoSessionAPI();
    }
    cryptoSession->Close();
}

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