Add memory scan based on detected module

Memory/Memory.cpp

@@ -77,9 +77,10 @@ void Memory::RestoreBytes(void *address)
     }
 }
 
-bool Memory::WaitForModule(const std::string& module_name, int timeoutMs = 15000, int intervalMs = 500)
+MODULEINFO Memory::WaitForModule(const std::string& module_name, int timeoutMs, int intervalMs)
 {
     const HANDLE hProc = GetCurrentProcess();
+    MODULEINFO modInfo{};
 
     for (int waited = 0; waited < timeoutMs; waited += intervalMs)
     {
@@ -91,11 +92,10 @@ bool Memory::WaitForModule(const std::string& module_name, int timeoutMs = 15000
             for (unsigned int i = 0; i < (cbNeeded / sizeof(HMODULE)); ++i)
             {
                 char modName[MAX_PATH];
-                if (GetModuleBaseNameA(hProc, hMods[i], modName, sizeof(modName)))
+                if (GetModuleBaseNameA(hProc, hMods[i], modName, sizeof(modName))) {
-                {
+                    if (_stricmp(modName, module_name.c_str()) == 0) {
-                    if (_stricmp(modName, module_name.c_str()) == 0)
+                        if (GetModuleInformation(hProc, hMods[i], &modInfo, sizeof(modInfo)))
-                    {
+                            return modInfo;
-                        return true;
                     }
                 }
             }
@@ -105,7 +105,7 @@ bool Memory::WaitForModule(const std::string& module_name, int timeoutMs = 15000
     }
 
     if (_log) _log->warn("Timeout: module '{}' not found in process after {} ms.", module_name, timeoutMs);
-    return false;
+    return MODULEINFO{};
 }
 
 std::string Memory::ByteToHexEscaped(const BYTE byte) {
@@ -121,14 +121,38 @@ uint8_t* Memory::AOBScan(
     DWORD protect_flags = PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_READWRITE | PAGE_EXECUTE_WRITECOPY,
     std::shared_ptr<spdlog::logger> log) {
 
+    bool found = false;
+
     _log = log;
-    if (!WaitForModule(module_name))
+    HANDLE hProc = GetCurrentProcess();
-    {
+    MODULEINFO modInfo{};
-        if (log) log->warn("Skipping AOB scan because module '{}' is unavailable.", module_name);
+    HMODULE targetModule = nullptr;
+    // Get module when name is specidifed
+    if (!(module_name.empty() || module_name == "*")) {
+        if (_log) _log->info("Module name: {}", module_name);
+        MODULEINFO modinfo = WaitForModule(module_name);
+        if (modinfo.lpBaseOfDll == nullptr) {
+            if (_log) _log->warn("Skipping AOB scan because module '{}' is unavailable.", module_name);
             return nullptr;
         }
+    }
+    // Fallback to determine module loaded
+    if (!found || module_name.empty() || module_name == "*") {
+        char exeBuf[MAX_PATH] = { 0 };
+        DWORD exeLen = GetModuleFileNameA(nullptr, exeBuf, MAX_PATH);
+        std::string exeName = (exeLen > 0) ? std::string(exeBuf, exeBuf + exeLen) : std::string();
+        size_t pos = exeName.find_last_of("\\/");
 
-    // Convert signature to bytes
+        if (pos != std::string::npos) exeName = exeName.substr(pos + 1);
+        if (_log && !exeName.empty()) _log->info("Module name: {}", exeName);
+
+        targetModule = GetModuleHandleA(nullptr);
+        if (!targetModule || !GetModuleInformation(hProc, targetModule, &modInfo, sizeof(modInfo))) {
+            if (_log) _log->error("Failed to find main module.");
+            return nullptr;
+        }
+    }
+    // Convert AOB string into vector bytes
     std::vector<int> pattern_bytes;
     std::istringstream stream(signature);
     std::string byte_str;
@@ -140,57 +164,45 @@ uint8_t* Memory::AOBScan(
             pattern_bytes.push_back(static_cast<int>(std::strtol(byte_str.c_str(), nullptr, 16)));
     }
 
-    HMODULE hMods[1024];
+    if (pattern_bytes.empty()) {
-    DWORD cbNeeded;
+        if (_log) _log->warn("Empty AOB pattern passed.");
-    HANDLE hProc = GetCurrentProcess();
-
-    if (!EnumProcessModules(hProc, hMods, sizeof(hMods), &cbNeeded))
-    {
-        spdlog::error("EnumProcessModules failed.");
         return nullptr;
     }
-
+    // Logging scanning area
-    for (unsigned int i = 0; i < (cbNeeded / sizeof(HMODULE)); ++i)
-    {
-        char modName[MAX_PATH];
-        if (GetModuleBaseNameA(hProc, hMods[i], modName, sizeof(modName)))
-        {
-            if (_stricmp(modName, module_name.c_str()) == 0)
-            {
-                MODULEINFO modInfo;
-                if (!GetModuleInformation(hProc, hMods[i], &modInfo, sizeof(modInfo)))
-                {
-                    if (log) log->error("GetModuleInformation failed for '{}'", module_name);
-                    return nullptr;
-                }
-
     uint8_t* base = reinterpret_cast<uint8_t*>(modInfo.lpBaseOfDll);
     size_t size = modInfo.SizeOfImage;
-                if (log) log->info("Scanning memory region: 0x{:X} - 0x{:X}", reinterpret_cast<uintptr_t>(base), reinterpret_cast<uintptr_t>(base + size));
+    if (_log) _log->info("Scanning memory region: 0x{:X} - 0x{:X}",
+        reinterpret_cast<uintptr_t>(base), reinterpret_cast<uintptr_t>(base + size));
 
+    // Memory scan
     MEMORY_BASIC_INFORMATION mbi{};
     for (uint8_t* current = base; current < base + size;)
     {
         if (!VirtualQuery(current, &mbi, sizeof(mbi)))
             break;
 
-                    if ((mbi.State & MEM_COMMIT) && (mbi.Protect & protect_flags))
+        bool isCommitted = (mbi.State & MEM_COMMIT) != 0;
-                    {
+        bool hasAccess = (mbi.Protect & protect_flags) != 0;
-                        for (size_t i = 0; i <= mbi.RegionSize - pattern_bytes.size(); ++i)
+        bool isNoAccess = (mbi.Protect & PAGE_NOACCESS) != 0;
+        bool isGuard = (mbi.Protect & PAGE_GUARD) != 0;
+
+        if (isCommitted && hasAccess && !isNoAccess && !isGuard) {
+            uint8_t* regionBase = reinterpret_cast<uint8_t*>(mbi.BaseAddress);
+            size_t regionSize = mbi.RegionSize;
+
+            for (size_t i = 0; i <= regionSize - pattern_bytes.size(); ++i)
             {
                 bool match = true;
                 for (size_t j = 0; j < pattern_bytes.size(); ++j)
                 {
-                                if (pattern_bytes[j] != -1 && current[i + j] != static_cast<uint8_t>(pattern_bytes[j]))
+                    if (pattern_bytes[j] != -1 && regionBase[i + j] != static_cast<uint8_t>(pattern_bytes[j])) {
-                                {
                         match = false;
                         break;
                     }
                 }
 
-                            if (match)
+                if (match) {
-                            {
+                    uint8_t* result = regionBase + i;
-                                uint8_t* result = current + i;
                     return result;
                 }
             }
@@ -198,14 +210,6 @@ uint8_t* Memory::AOBScan(
 
         current = reinterpret_cast<uint8_t*>(mbi.BaseAddress) + mbi.RegionSize;
     }
-
-                if (log) log->warn("No AOB match found in module '{}'.", module_name);
-                return nullptr;
-            }
-        }
-    }
-
-    if (log) log->warn("Module '{}' unexpectedly disappeared during scan.", module_name);
     return nullptr;
 }