#define NOMINMAX

#include <windows.h>
#include <tlhelp32.h>
#include <psapi.h>
#include <string>
#include <vector>
#include <fstream>
#include <iostream>
#include <cstdarg>
#include <memory>
#include <cstdint>
#include <algorithm>
#include <sstream>
#include <iomanip>
#include <map>

// Struct to ensure HANDLEs are properly closed
struct HandleCloser {
    void operator()(HANDLE handle) const {
        if (handle && handle != INVALID_HANDLE_VALUE) {
            CloseHandle(handle);
        }
    }
};
using UniqueHandle = std::unique_ptr<std::remove_pointer<HANDLE>::type, HandleCloser>;

// Function to log messages to a log file only
void log_message(std::ofstream& log_file, const char* format, ...) {
    va_list args;
    va_start(args, format);
    char buffer[2048]; // Increased buffer size for longer messages
    vsnprintf(buffer, sizeof(buffer), format, args);
    va_end(args);

    log_file << buffer;
    log_file.flush();
}

// Function to retrieve the Process ID (PID) by process name
DWORD get_process_id_by_name(const std::string& process_name, std::ofstream& log_file) {
    UniqueHandle hSnapshot(CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0));
    if (!hSnapshot || hSnapshot.get() == INVALID_HANDLE_VALUE) {
        log_message(log_file, "CreateToolhelp32Snapshot failed with error: %lu\n", GetLastError());
        return 0;
    }

    PROCESSENTRY32 pe32;
    pe32.dwSize = sizeof(PROCESSENTRY32);
    if (!Process32First(hSnapshot.get(), &pe32)) {
        log_message(log_file, "Process32First failed with error: %lu\n", GetLastError());
        return 0;
    }

    do {
        std::wstring ws(pe32.szExeFile);
        std::string exeFile(ws.begin(), ws.end());
        if (_stricmp(exeFile.c_str(), process_name.c_str()) == 0) {
            return pe32.th32ProcessID;
        }
    } while (Process32Next(hSnapshot.get(), &pe32));

    return 0;
}

// Function to retrieve the base address of the main module of the process
uintptr_t get_base_address(DWORD pid, const std::string& process_name, std::ofstream& log_file) {
    UniqueHandle hSnapshot(CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32, pid));
    if (!hSnapshot || hSnapshot.get() == INVALID_HANDLE_VALUE) {
        log_message(log_file, "CreateToolhelp32Snapshot for modules failed with error: %lu\n", GetLastError());
        return 0;
    }

    MODULEENTRY32 me32;
    me32.dwSize = sizeof(MODULEENTRY32);
    if (!Module32First(hSnapshot.get(), &me32)) {
        log_message(log_file, "Module32First failed with error: %lu\n", GetLastError());
        return 0;
    }

    do {
        // Convert module name to std::string for comparison
        std::wstring ws(me32.szModule);
        std::string moduleName(ws.begin(), ws.end());

        // Compare the module name with the process name
        if (_stricmp(moduleName.c_str(), process_name.c_str()) == 0) {
            log_message(log_file, "Main module %s found at base address 0x%llX\n", moduleName.c_str(),
                reinterpret_cast<unsigned long long>(me32.modBaseAddr));
            return reinterpret_cast<uintptr_t>(me32.modBaseAddr);
        }
    } while (Module32Next(hSnapshot.get(), &me32));

    log_message(log_file, "Main module not found in the module list.\n");
    return 0;
}

// Struct to store section information
struct Section {
    std::string name;
    uintptr_t start;
    uintptr_t end;
};

// Function to parse PE headers and extract sections
bool get_sections(HANDLE hProcess, uintptr_t base_addr, std::vector<Section>& sections, std::ofstream& log_file) {
    // Read the DOS header
    IMAGE_DOS_HEADER dosHeader;
    SIZE_T bytesRead;
    if (!ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(base_addr), &dosHeader, sizeof(dosHeader), &bytesRead) || bytesRead != sizeof(dosHeader)) {
        log_message(log_file, "Failed to read DOS header. Error: %lu\n", GetLastError());
        return false;
    }

    if (dosHeader.e_magic != IMAGE_DOS_SIGNATURE) {
        log_message(log_file, "Invalid DOS signature.\n");
        return false;
    }

    uintptr_t ntHeaderAddr = base_addr + dosHeader.e_lfanew;

    DWORD signature;
    if (!ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(ntHeaderAddr), &signature, sizeof(signature), &bytesRead) || bytesRead != sizeof(signature)) {
        log_message(log_file, "Failed to read NT headers signature. Error: %lu\n", GetLastError());
        return false;
    }

    if (signature != IMAGE_NT_SIGNATURE) {
        log_message(log_file, "Invalid NT signature.\n");
        return false;
    }

    IMAGE_FILE_HEADER fileHeader;
    if (!ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(ntHeaderAddr + sizeof(DWORD)), &fileHeader, sizeof(fileHeader), &bytesRead) || bytesRead != sizeof(fileHeader)) {
        log_message(log_file, "Failed to read IMAGE_FILE_HEADER. Error: %lu\n", GetLastError());
        return false;
    }

    WORD sizeOfOptionalHeader = fileHeader.SizeOfOptionalHeader;

    uintptr_t optionalHeaderAddr = ntHeaderAddr + sizeof(DWORD) + sizeof(IMAGE_FILE_HEADER);

    WORD magic;
    if (!ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(optionalHeaderAddr), &magic, sizeof(magic), &bytesRead) || bytesRead != sizeof(magic)) {
        log_message(log_file, "Failed to read Optional Header Magic. Error: %lu\n", GetLastError());
        return false;
    }

    bool is64 = false;
    if (magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
        is64 = true;
    }
    else if (magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
        is64 = false;
    }
    else {
        log_message(log_file, "Unknown Optional Header Magic: 0x%X.\n", magic);
        return false;
    }

    // Read the entire Optional Header
    std::vector<BYTE> optionalHeaderData(sizeOfOptionalHeader);
    if (!ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(optionalHeaderAddr), optionalHeaderData.data(), sizeOfOptionalHeader, &bytesRead) || bytesRead != sizeOfOptionalHeader) {
        log_message(log_file, "Failed to read Optional Header. Error: %lu\n", GetLastError());
        return false;
    }

    uintptr_t sectionHeaderAddr = optionalHeaderAddr + sizeOfOptionalHeader;

    for (int i = 0; i < fileHeader.NumberOfSections; ++i) {
        IMAGE_SECTION_HEADER sectionHeader;
        if (!ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(sectionHeaderAddr + i * sizeof(IMAGE_SECTION_HEADER)), &sectionHeader, sizeof(sectionHeader), &bytesRead) || bytesRead != sizeof(sectionHeader)) {
            log_message(log_file, "Failed to read IMAGE_SECTION_HEADER for section %d. Error: %lu\n", i + 1, GetLastError());
            return false;
        }

        Section sec;
        char secName[9] = { 0 };
        memcpy(secName, sectionHeader.Name, 8);
        sec.name = std::string(secName);

        sec.start = base_addr + sectionHeader.VirtualAddress;
        sec.end = sec.start + sectionHeader.Misc.VirtualSize;

        sections.push_back(sec);
        log_message(log_file, "Section %s: Start=0x%llX, End=0x%llX (Relative: Start=0x%llX, End=0x%llX)\n",
            sec.name.c_str(),
            static_cast<unsigned long long>(sec.start),
            static_cast<unsigned long long>(sec.end),
            static_cast<unsigned long long>(sec.start - base_addr),
            static_cast<unsigned long long>(sec.end - base_addr));
    }

    return true;
}

// Function to copy the memory of the target process to shadow memory
bool copy_memory_to_shadow(HANDLE hProcess, uintptr_t base_addr, std::vector<Section>& sections, std::vector<char>& shadow_memory, std::ofstream& log_file) {
    for (const auto& section : sections) {
        uintptr_t addr = section.start;
        while (addr < section.end) {
            MEMORY_BASIC_INFORMATION mbi;
            if (!VirtualQueryEx(hProcess, reinterpret_cast<LPCVOID>(addr), &mbi, sizeof(mbi))) {
                log_message(log_file, "VirtualQueryEx failed at address 0x%llX with error: %lu\n",
                    static_cast<unsigned long long>(addr), GetLastError());
                break;
            }

            if (mbi.State == MEM_COMMIT &&
                (mbi.Protect & (PAGE_READWRITE | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY | PAGE_READONLY | PAGE_EXECUTE_READ))) {

                SIZE_T regionSize = mbi.RegionSize;
                if (addr + regionSize > section.end) {
                    regionSize = section.end - addr;
                }

                std::vector<char> buffer(regionSize);
                SIZE_T bytesReadMem;
                if (ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(addr), buffer.data(), regionSize, &bytesReadMem) && bytesReadMem > 0) {
                    shadow_memory.insert(shadow_memory.end(), buffer.begin(), buffer.end());
                }

                addr += regionSize;
            }
            else {
                addr += mbi.RegionSize;
            }
        }
    }

    return true;
}

// Function to search for a string within the shadow memory
uintptr_t find_string_in_shadow_memory(const std::vector<char>& shadow_memory, const std::string& string_to_search, std::ofstream& log_file) {
    size_t string_length = string_to_search.length();
    for (size_t i = 0; i < shadow_memory.size() - string_length + 1; ++i) {
        if (std::memcmp(shadow_memory.data() + i, string_to_search.c_str(), string_length) == 0) {
            return i;
        }
    }

    return 0;
}

// Function to enable debug privileges
bool EnableDebugPrivileges() {
    HANDLE hToken;
    LUID luid;
    TOKEN_PRIVILEGES tp;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        std::cerr << "OpenProcessToken failed with error: " << GetLastError() << std::endl;
        return false;
    }

    if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) {
        std::cerr << "LookupPrivilegeValue failed with error: " << GetLastError() << std::endl;
        CloseHandle(hToken);
        return false;
    }

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    if (!AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(tp), NULL, NULL)) {
        std::cerr << "AdjustTokenPrivileges failed with error: " << GetLastError() << std::endl;
        CloseHandle(hToken);
        return false;
    }

    CloseHandle(hToken);
    return true;
}

int main() {
    if (!EnableDebugPrivileges()) {
        std::cerr << "Failed to enable debug privileges." << std::endl;
        return 1;
    }

    std::string process_name = "r5apex.exe";

    std::ofstream log_file("memscan_detailed.log", std::ios::out);
    if (!log_file.is_open()) {
        std::cerr << "Could not open log file for writing\n";
        return 1;
    }

    DWORD pid = get_process_id_by_name(process_name, log_file);
    if (pid == 0) {
        log_message(log_file, "Process %s not found.\n", process_name.c_str());
        std::cerr << "Process not found. Check the process name and ensure it's running.\n";
        return 1;
    }
    log_message(log_file, "Found process %s with PID: %u\n", process_name.c_str(), pid);
    std::cout << "Found process " << process_name << " with PID: " << pid << "\n";

    uintptr_t base_addr = get_base_address(pid, process_name, log_file);
    if (base_addr == 0) {
        log_message(log_file, "Failed to get base address for process ID %u.\n", pid);
        std::cerr << "Failed to get base address.\n";
        return 1;
    }
    log_message(log_file, "Base address of %s: 0x%llX\n", process_name.c_str(), static_cast<unsigned long long>(base_addr));
    std::cout << "Base address of " << process_name << ": 0x" << std::hex << base_addr << std::dec << "\n";

    UniqueHandle hProcess(OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, FALSE, pid));
    if (!hProcess) {
        log_message(log_file, "Failed to open process with PID %u. Error: %lu\n", pid, GetLastError());
        std::cerr << "Failed to open process. Ensure you have the necessary permissions.\n";
        return 1;
    }
    log_message(log_file, "Successfully opened process handle.\n");

    std::vector<Section> all_sections;
    if (!get_sections(hProcess.get(), base_addr, all_sections, log_file)) {
        log_message(log_file, "Failed to retrieve sections.\n");
        std::cerr << "Failed to retrieve sections.\n";
        return 1;
    }

    for (const auto& section : all_sections) {
        log_message(log_file, "Section %s: Start=0x%llX, End=0x%llX (Relative: Start=0x%llX, End=0x%llX)\n",
            section.name.c_str(),
            static_cast<unsigned long long>(section.start),
            static_cast<unsigned long long>(section.end),
            static_cast<unsigned long long>(section.start - base_addr),
            static_cast<unsigned long long>(section.end - base_addr));
    }

    // Copy the memory of the target process to shadow memory
    std::vector<char> shadow_memory;
    if (!copy_memory_to_shadow(hProcess.get(), base_addr, all_sections, shadow_memory, log_file)) {
        log_message(log_file, "Failed to copy memory to shadow memory.\n");
        std::cerr << "Failed to copy memory to shadow memory.\n";
        return 1;
    }

    // Map to store found strings and their relative offsets
    std::map<std::string, uintptr_t> foundStrings;

    // Search for useful strings in the shadow memory
    for (const auto& section : all_sections) {
        if (section.name == ".rdata" || section.name == ".data") {
            uintptr_t addr = section.start;
            while (addr < section.end) {
                MEMORY_BASIC_INFORMATION mbi;
                if (!VirtualQueryEx(hProcess.get(), reinterpret_cast<LPCVOID>(addr), &mbi, sizeof(mbi))) {
                    log_message(log_file, "VirtualQueryEx failed at address 0x%llX with error: %lu\n",
                        static_cast<unsigned long long>(addr), GetLastError());
                    break;
                }

                if (mbi.State == MEM_COMMIT &&
                    (mbi.Protect & (PAGE_READWRITE | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY | PAGE_READONLY | PAGE_EXECUTE_READ))) {

                    SIZE_T regionSize = mbi.RegionSize;
                    if (addr + regionSize > section.end) {
                        regionSize = section.end - addr;
                    }

                    std::vector<char> buffer(regionSize);
                    SIZE_T bytesReadMem;
                    if (ReadProcessMemory(hProcess.get(), reinterpret_cast<LPCVOID>(addr), buffer.data(), regionSize, &bytesReadMem) && bytesReadMem > 0) {
                        for (SIZE_T i = 0; i < bytesReadMem; ++i) {
                            if (std::isprint(buffer[i]) && buffer[i] != ' ') {
                                std::stringstream ss;
                                ss << buffer[i];
                                for (SIZE_T j = i + 1; j < bytesReadMem && std::isprint(buffer[j]) && buffer[j] != ' '; ++j) {
                                    ss << buffer[j];
                                }
                                std::string str = ss.str();
                                if (str.length() > 3) { // Consider strings longer than 3 characters as useful
                                    uintptr_t relativeAddress = addr + i - base_addr;
                                    foundStrings[str] = relativeAddress;
                                    log_message(log_file, "Found string \"%s\" at relative address 0x%llX (absolute address: 0x%llX) in section %s\n",
                                        str.c_str(),
                                        static_cast<unsigned long long>(relativeAddress),
                                        static_cast<unsigned long long>(addr + i),
                                        section.name.c_str());
                                }
                                i += str.length() - 1;
                            }
                        }
                    }

                    addr += regionSize;
                }
                else {
                    addr += mbi.RegionSize;
                }
            }
        }
    }

    // Output the results in the desired format
    log_message(log_file, "\nFound strings:\n");
    std::cout << "\nFound strings:\n";
    for (const auto& str : foundStrings) {
        log_message(log_file, "    \"%s\" => 0x%llX\n", str.first.c_str(), static_cast<unsigned long long>(str.second));
        std::cout << "    \"" << str.first << "\" => 0x" << std::hex << str.second << "\n";
    }

    std::cout << "Press Enter to exit...";
    std::cin.get();

    return 0;
}