From 81f1e4319fb88e4ced8d41253fa2a281759a1eb4 Mon Sep 17 00:00:00 2001 From: noita-player <56001276+noita-player@users.noreply.github.com> Date: Fri, 6 Mar 2026 23:21:10 -0800 Subject: [PATCH] Add process.info MCP tool for PEB/TEB enumeration Expose PEB address via provider interface and query it in the ProcessMemory plugin using NtQueryInformationProcess. The new process.info MCP tool returns the PEB VA and enumerates TEBs by querying thread information via NtQuerySystemInformation and NtQueryInformationThread for each thread in the target process. --- plugins/ProcessMemory/ProcessMemoryPlugin.cpp | 117 ++++++++++++++++++ plugins/ProcessMemory/ProcessMemoryPlugin.h | 3 + src/mcp/mcp_bridge.cpp | 50 ++++++++ src/mcp/mcp_bridge.h | 1 + src/providers/provider.h | 7 ++ 5 files changed, 178 insertions(+) diff --git a/plugins/ProcessMemory/ProcessMemoryPlugin.cpp b/plugins/ProcessMemory/ProcessMemoryPlugin.cpp index b185f41..d97d6fe 100644 --- a/plugins/ProcessMemory/ProcessMemoryPlugin.cpp +++ b/plugins/ProcessMemory/ProcessMemoryPlugin.cpp @@ -19,6 +19,60 @@ #include #include #include + +typedef struct _UNICODE_STRING { USHORT Length, MaximumLength; PWSTR Buffer; } UNICODE_STRING; +typedef struct _CLIENT_ID { HANDLE UniqueProcess; HANDLE UniqueThread; } CLIENT_ID; +typedef struct _SYSTEM_THREAD_INFORMATION { + LARGE_INTEGER KernelTime, UserTime, CreateTime; + ULONG WaitTime; PVOID StartAddress; CLIENT_ID ClientId; + LONG Priority, BasePriority; ULONG ContextSwitches, ThreadState, WaitReason; +} SYSTEM_THREAD_INFORMATION; +typedef struct _SYSTEM_PROCESS_INFORMATION { + ULONG NextEntryOffset; // 0x000 + ULONG NumberOfThreads; // 0x004 + LARGE_INTEGER WorkingSetPrivateSize; // 0x008 + ULONG HardFaultCount; // 0x010 + ULONG NumberOfThreadsHighWatermark; // 0x014 + ULONGLONG CycleTime; // 0x018 + LARGE_INTEGER CreateTime; // 0x020 + LARGE_INTEGER UserTime; // 0x028 + LARGE_INTEGER KernelTime; // 0x030 + UNICODE_STRING ImageName; // 0x038 + LONG BasePriority; // 0x048 + HANDLE UniqueProcessId; // 0x050 + PVOID InheritedFromUniqueProcessId; // 0x058 + ULONG HandleCount; // 0x060 + ULONG SessionId; // 0x064 + ULONG_PTR UniqueProcessKey; // 0x068 + SIZE_T PeakVirtualSize; // 0x070 + SIZE_T VirtualSize; // 0x078 + ULONG PageFaultCount; // 0x080 + ULONG _pad0; // 0x084 + SIZE_T PeakWorkingSetSize; // 0x088 + SIZE_T WorkingSetSize; // 0x090 + SIZE_T QuotaPeakPagedPoolUsage; // 0x098 + SIZE_T QuotaPagedPoolUsage; // 0x0A0 + SIZE_T QuotaPeakNonPagedPoolUsage; // 0x0A8 + SIZE_T QuotaNonPagedPoolUsage; // 0x0B0 + SIZE_T PagefileUsage; // 0x0B8 + SIZE_T PeakPagefileUsage; // 0x0C0 + SIZE_T PrivatePageCount; // 0x0C8 + LARGE_INTEGER ReadOperationCount; // 0x0D0 + LARGE_INTEGER WriteOperationCount; // 0x0D8 + LARGE_INTEGER OtherOperationCount; // 0x0E0 + LARGE_INTEGER ReadTransferCount; // 0x0E8 + LARGE_INTEGER WriteTransferCount; // 0x0F0 + LARGE_INTEGER OtherTransferCount; // 0x0F8 +} SYSTEM_PROCESS_INFORMATION; // sizeof = 0x100 +typedef struct alignas(8) _THREAD_BASIC_INFORMATION { + NTSTATUS ExitStatus; // 0x00 + ULONG _pad; // 0x04 + PVOID TebBaseAddress; // 0x08 + CLIENT_ID ClientId; // 0x10 + ULONG_PTR AffinityMask; // 0x20 + LONG Priority; // 0x28 + LONG BasePriority; // 0x2C +} THREAD_BASIC_INFORMATION; #elif defined(__linux__) #include #include @@ -61,6 +115,17 @@ ProcessMemoryProvider::ProcessMemoryProvider(uint32_t pid, const QString& proces BOOL isWow64 = FALSE; if (IsWow64Process(m_handle, &isWow64) && isWow64) m_pointerSize = 4; + // Query PEB address via NtQueryInformationProcess + { + typedef NTSTATUS(NTAPI* NtQIP_t)(HANDLE, ULONG, PVOID, ULONG, PULONG); + static NtQIP_t pNtQIP = (NtQIP_t)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQueryInformationProcess"); + if (pNtQIP) { + struct { PVOID r1; PVOID PebBaseAddress; PVOID r2[2]; ULONG_PTR pid; PVOID r3; } pbi = {}; + ULONG retLen = 0; + if (pNtQIP(m_handle, /*ProcessBasicInformation*/0, &pbi, sizeof(pbi), &retLen) >= 0 && pbi.PebBaseAddress) + m_peb = (uint64_t)(uintptr_t)pbi.PebBaseAddress; + } + } cacheModules(); } } @@ -426,6 +491,58 @@ int ProcessMemoryProvider::size() const #endif } + +QVector ProcessMemoryProvider::tebs() const +{ +#ifdef _WIN32 + QVector result; + if (!m_handle || !m_peb) return result; + + typedef NTSTATUS(NTAPI* NtQSI_t)(ULONG, PVOID, ULONG, PULONG); + typedef NTSTATUS(NTAPI* NtQIT_t)(HANDLE, ULONG, PVOID, ULONG, PULONG); + static auto pNtQSI = (NtQSI_t)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQuerySystemInformation"); + static auto pNtQIT = (NtQIT_t)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQueryInformationThread"); + if (!pNtQSI || !pNtQIT) return result; + + // Enumerate threads via SystemProcessInformation (class 5) + ULONG retLen = 0; + ULONG bufSize = 1 << 20; + QByteArray buf(bufSize, 0); + NTSTATUS qsiSt; + for (int attempt = 0; attempt < 8; ++attempt) { + qsiSt = pNtQSI(5, buf.data(), bufSize, &retLen); + if ((uint32_t)qsiSt != 0xC0000004u) break; + bufSize *= 2; + buf.resize(bufSize); + } + if (qsiSt < 0) return result; + + // Walk process entries to find ours + auto* proc = (SYSTEM_PROCESS_INFORMATION*)buf.data(); + for (;;) { + if ((uintptr_t)proc->UniqueProcessId == m_pid) { + auto* threads = (SYSTEM_THREAD_INFORMATION*)((char*)proc + sizeof(*proc)); + for (ULONG i = 0; i < proc->NumberOfThreads; ++i) { + DWORD tid = (DWORD)(uintptr_t)threads[i].ClientId.UniqueThread; + HANDLE hThread = OpenThread(THREAD_QUERY_LIMITED_INFORMATION, FALSE, tid); + if (!hThread) continue; + THREAD_BASIC_INFORMATION tbi = {}; + ULONG tbiLen = 0; + NTSTATUS qitSt = pNtQIT(hThread, 0, &tbi, sizeof(tbi), &tbiLen); + if (qitSt >= 0 && tbi.TebBaseAddress) + result.append({(uint64_t)(uintptr_t)tbi.TebBaseAddress, tid}); + CloseHandle(hThread); + } + break; + } + if (!proc->NextEntryOffset) break; + proc = (SYSTEM_PROCESS_INFORMATION*)((char*)proc + proc->NextEntryOffset); + } + return result; +#else + return {}; +#endif +} // ────────────────────────────────────────────────────────────────────────── // ProcessMemoryPlugin implementation // ────────────────────────────────────────────────────────────────────────── diff --git a/plugins/ProcessMemory/ProcessMemoryPlugin.h b/plugins/ProcessMemory/ProcessMemoryPlugin.h index 1bf56ef..03a6301 100644 --- a/plugins/ProcessMemory/ProcessMemoryPlugin.h +++ b/plugins/ProcessMemory/ProcessMemoryPlugin.h @@ -41,6 +41,8 @@ public: // Process-specific helpers uint32_t pid() const { return m_pid; } void refreshModules() { m_modules.clear(); cacheModules(); } + uint64_t peb() const override { return m_peb; } + QVector tebs() const override; private: void cacheModules(); @@ -56,6 +58,7 @@ private: bool m_writable; uint64_t m_base; int m_pointerSize = 8; + uint64_t m_peb = 0; struct ModuleInfo { QString name; diff --git a/src/mcp/mcp_bridge.cpp b/src/mcp/mcp_bridge.cpp index 40ec41a..f8f0fc6 100644 --- a/src/mcp/mcp_bridge.cpp +++ b/src/mcp/mcp_bridge.cpp @@ -447,6 +447,22 @@ QJsonObject McpBridge::handleToolsList(const QJsonValue& id) { }} }); + + // process.info + tools.append(QJsonObject{ + {"name", "process.info"}, + {"description", "Returns PEB address and enumerates all Thread Environment Blocks (TEBs) for the attached process. " + "TEBs are discovered via NtQuerySystemInformation and NtQueryInformationThread. " + "Each TEB entry includes: address, threadId. " + "Requires a live process provider with PEB support."}, + {"inputSchema", QJsonObject{ + {"type", "object"}, + {"properties", QJsonObject{ + {"tabIndex", QJsonObject{{"type", "integer"}, + {"description", "MDI tab index (0-based). Omit for active tab."}}} + }} + }} + }); return okReply(id, QJsonObject{{"tools", tools}}); } @@ -472,6 +488,7 @@ QJsonObject McpBridge::handleToolsCall(const QJsonValue& id, const QJsonObject& else if (toolName == "ui.action") result = toolUiAction(args); else if (toolName == "tree.search") result = toolTreeSearch(args); else if (toolName == "node.history") result = toolNodeHistory(args); + else if (toolName == "process.info") result = toolProcessInfo(args); else return errReply(id, -32601, "Unknown tool: " + toolName); m_mainWindow->clearMcpStatus(); @@ -1327,6 +1344,39 @@ QJsonObject McpBridge::toolNodeHistory(const QJsonObject& args) { QJsonDocument(result).toJson(QJsonDocument::Compact))); } + +// ════════════════════════════════════════════════════════════════════ +// TOOL: process.info — PEB address + TEB enumeration +// ════════════════════════════════════════════════════════════════════ + +QJsonObject McpBridge::toolProcessInfo(const QJsonObject& args) { + auto* tab = resolveTab(args); + if (!tab) return makeTextResult("No active tab", true); + + auto* prov = tab->doc->provider.get(); + if (!prov) return makeTextResult("No data source attached", true); + if (!prov->isLive()) return makeTextResult("Not a live provider", true); + + uint64_t pebAddr = prov->peb(); + if (!pebAddr) return makeTextResult("PEB not available for this provider", true); + + QJsonObject out; + out["peb"] = "0x" + QString::number(pebAddr, 16).toUpper(); + + auto tebList = prov->tebs(); + QJsonArray tebArr; + for (const auto& t : tebList) { + tebArr.append(QJsonObject{ + {"address", "0x" + QString::number(t.tebAddress, 16).toUpper()}, + {"threadId", (qint64)t.threadId} + }); + } + + out["tebs"] = tebArr; + out["tebCount"] = tebArr.size(); + return makeTextResult(QString::fromUtf8(QJsonDocument(out).toJson(QJsonDocument::Indented))); +} + // ════════════════════════════════════════════════════════════════════ // Notifications (call from MainWindow/Controller hooks) // ════════════════════════════════════════════════════════════════════ diff --git a/src/mcp/mcp_bridge.h b/src/mcp/mcp_bridge.h index b2151b6..8c5f1c5 100644 --- a/src/mcp/mcp_bridge.h +++ b/src/mcp/mcp_bridge.h @@ -60,6 +60,7 @@ private: QJsonObject toolUiAction(const QJsonObject& args); QJsonObject toolTreeSearch(const QJsonObject& args); QJsonObject toolNodeHistory(const QJsonObject& args); + QJsonObject toolProcessInfo(const QJsonObject& args); // Helpers QJsonObject makeTextResult(const QString& text, bool isError = false); diff --git a/src/providers/provider.h b/src/providers/provider.h index 787647b..67df724 100644 --- a/src/providers/provider.h +++ b/src/providers/provider.h @@ -73,6 +73,13 @@ public: // Default: returns empty (scan engine falls back to [0, size())). virtual QVector enumerateRegions() const { return {}; } + // Process Environment Block address (x64 PEB VA in target process). + // Only meaningful for live process providers. Returns 0 if unavailable. + virtual uint64_t peb() const { return 0; } + + struct ThreadInfo { uint64_t tebAddress; uint32_t threadId; }; + virtual QVector tebs() const { return {}; } + // --- Derived convenience (non-virtual, never override) --- bool isValid() const { return size() > 0; }