mirror of
https://github.com/NohamR/Reclass.git
synced 2026-05-10 19:59:21 +00:00
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.
This commit is contained in:
@@ -19,6 +19,60 @@
|
||||
#include <tlhelp32.h>
|
||||
#include <psapi.h>
|
||||
#include <shellapi.h>
|
||||
|
||||
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 <climits>
|
||||
#include <sys/types.h>
|
||||
@@ -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<rcx::Provider::ThreadInfo> ProcessMemoryProvider::tebs() const
|
||||
{
|
||||
#ifdef _WIN32
|
||||
QVector<ThreadInfo> 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
|
||||
// ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -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<ThreadInfo> 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;
|
||||
|
||||
@@ -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)
|
||||
// ════════════════════════════════════════════════════════════════════
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -73,6 +73,13 @@ public:
|
||||
// Default: returns empty (scan engine falls back to [0, size())).
|
||||
virtual QVector<MemoryRegion> 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<ThreadInfo> tebs() const { return {}; }
|
||||
|
||||
// --- Derived convenience (non-virtual, never override) ---
|
||||
|
||||
bool isValid() const { return size() > 0; }
|
||||
|
||||
Reference in New Issue
Block a user