scanner.hpp

#pragma once

#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <Windows.h>

#include "Pattern.hpp"
#include "SIMDCache.hpp"
#include "ThreadPool.hpp"
#include "StringScanner.hpp"
#include "MemoryUtils.hpp"

#include <vector>
#include <memory>
#include <unordered_map>
#include <mutex>
#include <queue>
#include <functional>

class MemoryScanner {
public:
    struct ScanRegion {
        uintptr_t baseAddress;
        size_t size;
        uint32_t protection;
        bool isGuarded;
        bool isNoAccess;

        bool IsReadable() const {
            return !isNoAccess &&
                (protection & (PAGE_READONLY | PAGE_READWRITE | PAGE_EXECUTE_READ |
                    PAGE_EXECUTE_READWRITE)) != 0;
        }
    };

    struct ScanJob {
        Pattern pattern;
        std::function<void(uintptr_t)> callback;
        bool scanStrings;
        std::string stringPattern;
    };

    MemoryScanner(HANDLE processHandle = nullptr, size_t threadCount = 0)
        : processHandle(processHandle ? processHandle : GetCurrentProcess()),
        scanInProgress(false) {

        if (!processHandle) {
            processHandle = GetCurrentProcess();
        }

        size_t poolSize = threadCount ? threadCount : std::thread::hardware_concurrency();

        threadPool = std::make_unique<ThreadPool>(poolSize);

        cachePool.resize(poolSize);
        for (auto& cache : cachePool) {
            cache = std::make_unique<SIMDCache>();
            cache->Initialize(1024 * 1024); // 1MB cache per thread
        }
    }

    void EnableDllAutoScan(bool enable) {
        autoScanDlls = enable;
        if (enable) {
            MonitorDllLoads();
        }
    }

    void AddScanJob(const ScanJob& job) {
        std::lock_guard<std::mutex> lock(jobMutex);
        scanJobs.push(job);
    }

    void ProcessJobs() {
        std::lock_guard<std::mutex> lock(jobMutex);
        while (!scanJobs.empty()) {
            auto job = scanJobs.front();
            scanJobs.pop();

            if (job.scanStrings && !job.stringPattern.empty()) {
                ScanStrings({ job.stringPattern }, false);
            }
            else {
                PerformScan(job.pattern);
            }
        }
    }


    std::vector<uintptr_t> PerformScan(const Pattern& pattern) {

        std::vector<uintptr_t> allResults;
        std::mutex resultsMutex;

        auto regions = GetScanRegions();

        regions.erase(std::remove_if(regions.begin(), regions.end(),
            [](const ScanRegion& region) {
                return region.isGuarded || region.isNoAccess || !region.IsReadable();
            }), regions.end());

        if (regions.empty() || !pattern.IsValid()) {
            return allResults;
        }

        auto patternBytes = pattern.GetPatternBytes();
        auto mask = pattern.GetMask();

        size_t regionsPerThread = std::max<size_t>(1, regions.size() / cachePool.size());

        for (size_t i = 0; i < cachePool.size(); ++i) {
            size_t startIdx = i * regionsPerThread;
            size_t endIdx = (i == cachePool.size() - 1) ? regions.size() : startIdx + regionsPerThread;

            if (startIdx >= regions.size()) break;

            threadPool->Enqueue([this, &patternBytes, &mask, &regions,
                startIdx, endIdx, &allResults,
                &resultsMutex, i]() {

                    auto& cache = cachePool[i];
                    std::vector<uintptr_t> threadResults;

                    for (size_t j = startIdx; j < endIdx; ++j) {
                        const auto& region = regions[j];

                        if (cache->ReadMemory(reinterpret_cast<void*>(region.baseAddress),
                            region.size, processHandle)) {

                            auto patternResults = cache->ScanPattern(patternBytes, mask);

                            for (auto offset : patternResults) {
                                threadResults.push_back(region.baseAddress + offset);
                            }
                        }
                    }

                    if (!threadResults.empty()) {
                        std::lock_guard<std::mutex> lock(resultsMutex);
                        allResults.insert(allResults.end(),
                            threadResults.begin(), threadResults.end());
                    }
                });
        }

        WaitForCompletion();

        return allResults;
    }

    std::vector<std::pair<std::string, uintptr_t>>
        ScanStrings(const std::vector<std::string>& strings,
            bool caseSensitive = true) {

        std::vector<std::pair<std::string, uintptr_t>> allResults;
        std::mutex resultsMutex;

        auto regions = GetScanRegions();

        for (size_t i = 0; i < cachePool.size(); ++i) {
            size_t regionsPerThread = std::max<size_t>(1, regions.size() / cachePool.size());
            size_t startIdx = i * regionsPerThread;
            size_t endIdx = (i == cachePool.size() - 1) ? regions.size() : startIdx + regionsPerThread;

            threadPool->Enqueue([this, &strings, caseSensitive,
                &regions, startIdx, endIdx,
                &allResults, &resultsMutex, i]() {

                    auto& cache = cachePool[i];
                    std::vector<std::pair<std::string, uintptr_t>> threadResults;

                    for (size_t j = startIdx; j < endIdx; ++j) {
                        const auto& region = regions[j];

                        if (cache->ReadMemory(reinterpret_cast<void*>(region.baseAddress),
                            region.size, processHandle)) {

                            auto stringResults = StringScanner::SearchMultiple(
                                cache->GetData(), cache->GetSize(),
                                strings, caseSensitive);

                            for (const auto& result : stringResults) {
                                threadResults.emplace_back(
                                    result.first,
                                    region.baseAddress + result.second
                                );
                            }
                        }
                    }

                    if (!threadResults.empty()) {
                        std::lock_guard<std::mutex> lock(resultsMutex);
                        allResults.insert(allResults.end(),
                            threadResults.begin(), threadResults.end());
                    }
                });
        }

        WaitForCompletion();
        return allResults;
    }

    size_t GetPendingJobCount() const {
        std::lock_guard<std::mutex> lock(jobMutex);
        return scanJobs.size();
    }

    bool IsScanning() const {
        return threadPool->GetJobCount() > 0;
    }

    void WaitForCompletion() {
        while (threadPool->GetJobCount() > 0) {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }

    // Get thread count
    size_t GetThreadCount() const {
        return threadPool->GetThreadCount();
    }

private:
    HANDLE processHandle;
    std::unique_ptr<ThreadPool> threadPool;
    std::vector<std::unique_ptr<SIMDCache>> cachePool;
    std::queue<ScanJob> scanJobs;
    mutable std::mutex jobMutex;
    std::atomic<bool> scanInProgress;
    bool autoScanDlls = false;

    std::vector<ScanRegion> GetScanRegions() {
        std::vector<ScanRegion> regions;

        uintptr_t address = 0;
        MEMORY_BASIC_INFORMATION mbi;

        while (VirtualQueryEx(processHandle, reinterpret_cast<void*>(address),
            &mbi, sizeof(mbi))) {

            if (mbi.State == MEM_COMMIT && mbi.RegionSize > 0) {
                regions.push_back({
                    reinterpret_cast<uintptr_t>(mbi.BaseAddress),
                    mbi.RegionSize,
                    mbi.Protect,
                    (mbi.Protect & PAGE_GUARD) != 0,
                    mbi.Protect == PAGE_NOACCESS
                    });
            }

            address += mbi.RegionSize;
            if (address >= 0x7FFFFFFF0000) break; // Stop at high memory
        }

        return regions;
    }

    void MonitorDllLoads() {

        std::thread([this]() {
            while (autoScanDlls) {
                // check or any othrer dll
                // This would use debugging apis
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
            }
            }).detach();
    }
};