MemoryScanner.cpp

#include "MemoryScanner.h"
#include <TlHelp32.h>
#include <Psapi.h>
#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>

#ifndef PAGE_READABLE
#define PAGE_READABLE (PAGE_READONLY | PAGE_READWRITE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE)
#endif

MemoryScanner::MemoryScanner() : m_hProcess(nullptr), m_moduleBase(0) {}

MemoryScanner::~MemoryScanner() {
    if (m_hProcess) {
        CloseHandle(m_hProcess);
    }
}

bool MemoryScanner::Initialize(const std::string& processName) {
    m_processName = processName;
    DWORD processId = GetProcessIdByName(processName);
    if (processId == 0) {
        return false;
    }
    m_hProcess = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, FALSE, processId);
    if (!m_hProcess) {
        return false;
    }
    m_moduleBase = GetModuleBaseAddress(processId, processName);
    if (m_moduleBase == 0) {
        return false;
    }
    return true;
}

std::vector<ScanResult> MemoryScanner::ScanAllSignatures() {
    std::vector<ScanResult> results;
    auto signatures = Signatures::GetAllSignatures();
    for (const auto& sig : signatures) {
        ScanResult result;
        result.signatureName = sig.name;
        result.pattern = sig.pattern;
        if (sig.pattern.find("0x") == 0) {
            result.address = std::stoull(sig.pattern.substr(2), nullptr, 16);
            result.found = true;
        } else {
            result.address = FindPattern(sig.pattern);
            result.found = (result.address != 0);
        }
        results.push_back(result);
    }
    return results;
}

ScanResult MemoryScanner::ScanSignature(const std::string& signatureName) {
    auto allSignatures = Signatures::GetAllSignatures();
    for (const auto& sig : allSignatures) {
        if (sig.name.compare(signatureName) == 0) {
            ScanResult result;
            result.signatureName = sig.name;
            result.pattern = sig.pattern;
            if (sig.pattern.find("0x") == 0) {
                result.address = std::stoull(sig.pattern.substr(2), nullptr, 16);
                result.found = true;
            } else {
                result.address = FindPattern(sig.pattern);
                result.found = (result.address != 0);
            }
            return result;
        }
    }
    return ScanResult{};
}

bool MemoryScanner::DumpResultsToFile(const std::string& filename, const std::vector<ScanResult>& results) {
    std::ofstream file(filename);
    if (!file.is_open()) {
        return false;
    }
    for (const auto& result : results) {
        if (result.found) {
            file << result.signatureName << " = 0x" << std::hex << result.address << "\n";
        } else {
            file << result.signatureName << " = 0x0\n";
        }
    }
    file.close();
    return true;
}

uintptr_t MemoryScanner::FindPattern(const std::string& pattern) {
    auto pattern_to_ints = [](const std::string& pat) {
        std::vector<int> bytes;
        const char* p = pat.c_str();
        const char* end = p + pat.size();
        for (const char* cur = p; cur < end; ++cur) {
            if (*cur == '?') {
                // wildcard "?" or "??"
                if (cur + 1 < end && *(cur + 1) == '?') ++cur;
                bytes.push_back(-1);
            } else if (*cur == ' ') {
                continue;
            } else {
                bytes.push_back(strtoul(cur, const_cast<char**>(&cur), 16));
            }
        }
        return bytes;
    };

    std::vector<int> sig = pattern_to_ints(pattern);
    if (sig.empty()) return 0;

    MODULEINFO moduleInfo;
    if (!GetModuleInformation(m_hProcess, (HMODULE)m_moduleBase, &moduleInfo, sizeof(moduleInfo))) {
        return 0;
    }

    const uintptr_t base = m_moduleBase;
    const uintptr_t size = moduleInfo.SizeOfImage;
    if (size == 0) return 0;

    const SIZE_T pageSize = 4096;
    PSAPI_WORKING_SET_EX_INFORMATION wsInfo{};

    for (uintptr_t pageStart = base; pageStart < base + size; pageStart += pageSize) {
        wsInfo.VirtualAddress = reinterpret_cast<PVOID>(pageStart);
        if (!QueryWorkingSetEx(m_hProcess, &wsInfo, sizeof(wsInfo))) {
            continue;
        }
        if (!wsInfo.VirtualAttributes.Valid) {
            continue;
        }
        unsigned long protect = wsInfo.VirtualAttributes.Win32Protection;
        if (!(protect == PAGE_EXECUTE_READ || protect == PAGE_EXECUTE_READWRITE || protect == PAGE_READONLY || protect == PAGE_READWRITE)) {
            continue;
        }

        SIZE_T toRead = static_cast<SIZE_T>(std::min<uintptr_t>(pageSize, (base + size) - pageStart));
        std::vector<uint8_t> buffer(toRead);
        SIZE_T bytesRead = 0;
        if (!::ReadProcessMemory(m_hProcess, reinterpret_cast<LPCVOID>(pageStart), buffer.data(), toRead, &bytesRead) || bytesRead == 0) {
            continue;
        }
        buffer.resize(bytesRead);

        const size_t sigLen = sig.size();
        if (bytesRead < sigLen) continue;

        for (size_t i = 0; i + sigLen <= buffer.size(); ++i) {
            bool match = true;
            for (size_t j = 0; j < sigLen; ++j) {
                int b = sig[j];
                if (b != -1 && static_cast<uint8_t>(b) != buffer[i + j]) {
                    match = false;
                    break;
                }
            }
            if (match) {
                return pageStart + i;
            }
        }
    }
    return 0;
}

std::vector<uint8_t> MemoryScanner::PatternToBytes(const std::string& pattern) {
    std::vector<uint8_t> bytes;
    std::istringstream iss(pattern);
    std::string token;
    while (iss >> token) {
        if (token == "?") {
            bytes.push_back(0);
        } else {
            bytes.push_back((uint8_t)std::stoul(token, nullptr, 16));
        }
    }
    return bytes;
}

bool MemoryScanner::PatternMatch(const uint8_t* data, const std::vector<uint8_t>& pattern) {
    for (size_t i = 0; i < pattern.size(); i++) {
        if (pattern[i] != 0 && data[i] != pattern[i]) {
            return false;
        }
    }
    return true;
}

bool MemoryScanner::ReadProcessMemory(uintptr_t address, void* buffer, size_t size) {
    SIZE_T bytesRead;
    return ::ReadProcessMemory(m_hProcess, (LPCVOID)address, buffer, size, &bytesRead) && bytesRead == size;
}

std::vector<uint8_t> MemoryScanner::ReadMemoryRegion(uintptr_t address, size_t size) {
    std::vector<uint8_t> buffer(size);
    SIZE_T bytesRead;
    if (::ReadProcessMemory(m_hProcess, (LPCVOID)address, buffer.data(), size, &bytesRead)) {
        buffer.resize(bytesRead);
        return buffer;
    }
    return std::vector<uint8_t>();
}

DWORD MemoryScanner::GetProcessIdByName(const std::string& processName) {
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnapshot == INVALID_HANDLE_VALUE) {
        return 0;
    }
    PROCESSENTRY32W pe32;
    pe32.dwSize = sizeof(PROCESSENTRY32W);
    if (Process32FirstW(hSnapshot, &pe32)) {
        do {
            int size_needed = WideCharToMultiByte(CP_UTF8, 0, pe32.szExeFile, -1, NULL, 0, NULL, NULL);
            std::string exeName(size_needed, 0);
            WideCharToMultiByte(CP_UTF8, 0, pe32.szExeFile, -1, &exeName[0], size_needed, NULL, NULL);
            if (!exeName.empty() && exeName.back() == '\0') exeName.pop_back();
            if (processName.compare(exeName) == 0) {
                CloseHandle(hSnapshot);
                return pe32.th32ProcessID;
            }
        } while (Process32NextW(hSnapshot, &pe32));
    }
    CloseHandle(hSnapshot);
    return 0;
}

uintptr_t MemoryScanner::GetModuleBaseAddress(DWORD processId, const std::string& moduleName) {
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32, processId);
    if (hSnapshot == INVALID_HANDLE_VALUE) {
        return 0;
    }
    MODULEENTRY32W me32;
    me32.dwSize = sizeof(MODULEENTRY32W);
    if (Module32FirstW(hSnapshot, &me32)) {
        do {
            int size_needed = WideCharToMultiByte(CP_UTF8, 0, me32.szModule, -1, NULL, 0, NULL, NULL);
            std::string modName(size_needed, 0);
            WideCharToMultiByte(CP_UTF8, 0, me32.szModule, -1, &modName[0], size_needed, NULL, NULL);
            if (!modName.empty() && modName.back() == '\0') modName.pop_back();
            if (moduleName.compare(modName) == 0) {
                CloseHandle(hSnapshot);
                return (uintptr_t)me32.modBaseAddr;
            }
        } while (Module32NextW(hSnapshot, &me32));
    }
    CloseHandle(hSnapshot);
    return 0;
}