Commit 5ceef5ea authored by Kilo Code's avatar Kilo Code

Add process tree monitoring capability

parent fd78e572
......@@ -26,6 +26,17 @@
#include <tlhelp32.h>
#include <psapi.h>
// Structure to track monitored processes
typedef struct {
DWORD pid;
HANDLE hProcess;
char processName[MAX_PATH];
} MonitoredProcess;
#define MAX_PROCESSES 256
MonitoredProcess monitoredProcesses[MAX_PROCESSES];
int processCount = 0;
// #pragma comment(lib, "ws2_32.lib")
// #pragma comment(lib, "iphlpapi.lib")
......@@ -41,6 +52,79 @@ int is_internal_ip(DWORD ip) {
return 0;
}
// Function to add a process to the monitoring list
void AddMonitoredProcess(DWORD pid, HANDLE hProcess, const char* processName) {
if (processCount >= MAX_PROCESSES) return;
monitoredProcesses[processCount].pid = pid;
monitoredProcesses[processCount].hProcess = hProcess;
strcpy(monitoredProcesses[processCount].processName, processName);
processCount++;
}
// Function to remove a terminated process from monitoring
void RemoveMonitoredProcess(DWORD pid) {
for (int i = 0; i < processCount; i++) {
if (monitoredProcesses[i].pid == pid) {
CloseHandle(monitoredProcesses[i].hProcess);
// Shift remaining processes
for (int j = i; j < processCount - 1; j++) {
monitoredProcesses[j] = monitoredProcesses[j + 1];
}
processCount--;
break;
}
}
}
// Function to check if a process is already being monitored
int IsProcessMonitored(DWORD pid) {
for (int i = 0; i < processCount; i++) {
if (monitoredProcesses[i].pid == pid) return 1;
}
return 0;
}
// Function to inject DLL into a process
BOOL InjectDLLIntoProcess(DWORD pid) {
HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
if (!hProcess) return FALSE;
if (!InjectDLL(hProcess, "ssl_hook.dll")) {
CloseHandle(hProcess);
return FALSE;
}
CloseHandle(hProcess);
return TRUE;
}
// Function to find and monitor child processes
void MonitorChildProcesses(DWORD parentPid) {
HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
if (hSnapshot == INVALID_HANDLE_VALUE) return;
PROCESSENTRY32 pe32;
pe32.dwSize = sizeof(PROCESSENTRY32);
if (Process32First(hSnapshot, &pe32)) {
do {
if (pe32.th32ParentProcessID == parentPid && !IsProcessMonitored(pe32.th32ProcessID)) {
// Found a child process, inject DLL and add to monitoring
if (InjectDLLIntoProcess(pe32.th32ProcessID)) {
HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pe32.th32ProcessID);
if (hProcess) {
AddMonitoredProcess(pe32.th32ProcessID, hProcess, pe32.szExeFile);
printf("Now monitoring child process: %s (PID: %lu)\n", pe32.szExeFile, pe32.th32ProcessID);
}
}
}
} while (Process32Next(hSnapshot, &pe32));
}
CloseHandle(hSnapshot);
}
BOOL InjectDLL(HANDLE hProcess, const char* dllPath) {
LPVOID pRemoteBuf = VirtualAllocEx(hProcess, NULL, strlen(dllPath) + 1, MEM_COMMIT, PAGE_READWRITE);
if (!pRemoteBuf) return FALSE;
......@@ -127,7 +211,12 @@ int main() {
ResumeThread(pi.hThread);
DWORD pid = pi.dwProcessId;
DWORD rootPid = pi.dwProcessId;
// Add the root process to monitoring
AddMonitoredProcess(rootPid, pi.hProcess, program);
printf("Now monitoring root process: %s (PID: %lu)\n", program, rootPid);
Sleep(2000); // Wait for process to potentially establish connections
char internal_log_path[512];
......@@ -138,7 +227,7 @@ int main() {
FILE* internal_log = fopen(internal_log_path, "a");
FILE* external_log = fopen(external_log_path, "a");
// Initial snapshot of connections
// Initial snapshot of connections for all monitored processes
PMIB_TCPTABLE_OWNER_PID tcpTable;
DWORD size = 0;
GetExtendedTcpTable(NULL, &size, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0);
......@@ -146,21 +235,26 @@ int main() {
if (GetExtendedTcpTable(tcpTable, &size, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0) == NO_ERROR) {
for (DWORD i = 0; i < tcpTable->dwNumEntries; i++) {
MIB_TCPROW_OWNER_PID row = tcpTable->table[i];
if (row.dwOwningPid == pid) {
DWORD localIP = row.dwLocalAddr;
DWORD remoteIP = row.dwRemoteAddr;
char log_entry[256];
sprintf(log_entry, "Connection: Local %lu.%lu.%lu.%lu:%u -> Remote %lu.%lu.%lu.%lu:%u State:%lu\n",
(localIP >> 24) & 0xFF, (localIP >> 16) & 0xFF, (localIP >> 8) & 0xFF, localIP & 0xFF, ntohs(row.dwLocalPort),
(remoteIP >> 24) & 0xFF, (remoteIP >> 16) & 0xFF, (remoteIP >> 8) & 0xFF, remoteIP & 0xFF, ntohs(row.dwRemotePort),
row.dwState);
if (is_internal_ip(remoteIP)) {
fprintf(internal_log, "%s", log_entry);
} else {
fprintf(external_log, "%s", log_entry);
// Check if this connection belongs to any monitored process
for (int p = 0; p < processCount; p++) {
if (row.dwOwningPid == monitoredProcesses[p].pid) {
DWORD localIP = row.dwLocalAddr;
DWORD remoteIP = row.dwRemoteAddr;
char log_entry[256];
sprintf(log_entry, "[%s] Connection: Local %lu.%lu.%lu.%lu:%u -> Remote %lu.%lu.%lu.%lu:%u State:%lu\n",
monitoredProcesses[p].processName,
(localIP >> 24) & 0xFF, (localIP >> 16) & 0xFF, (localIP >> 8) & 0xFF, localIP & 0xFF, ntohs(row.dwLocalPort),
(remoteIP >> 24) & 0xFF, (remoteIP >> 16) & 0xFF, (remoteIP >> 8) & 0xFF, remoteIP & 0xFF, ntohs(row.dwRemotePort),
row.dwState);
if (is_internal_ip(remoteIP)) {
fprintf(internal_log, "%s", log_entry);
} else {
fprintf(external_log, "%s", log_entry);
}
printf("%s", log_entry);
break;
}
printf("%s", log_entry);
}
}
}
......@@ -171,33 +265,52 @@ int main() {
// Continuous monitoring loop
DWORD lastStatusTime = GetTickCount();
const DWORD STATUS_INTERVAL = 5000; // 5 seconds
DWORD lastChildCheckTime = GetTickCount();
const DWORD CHILD_CHECK_INTERVAL = 2000; // 2 seconds
while (1) {
DWORD currentTime = GetTickCount();
// Check for new child processes periodically
if (currentTime - lastChildCheckTime >= CHILD_CHECK_INTERVAL) {
// Check all monitored processes for new children
for (int p = 0; p < processCount; p++) {
MonitorChildProcesses(monitoredProcesses[p].pid);
}
lastChildCheckTime = currentTime;
}
// Periodic status logging
if (currentTime - lastStatusTime >= STATUS_INTERVAL) {
// Get current TCP connections for the monitored process
// Aggregate data from all monitored processes
DWORD totalSocketCount = 0;
DWORD totalThreadCount = 0;
DWORD totalMemoryUsage = 0;
// Get current TCP connections for all monitored processes
size = 0;
GetExtendedTcpTable(NULL, &size, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0);
tcpTable = (PMIB_TCPTABLE_OWNER_PID)malloc(size);
DWORD socketCount = 0;
if (GetExtendedTcpTable(tcpTable, &size, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0) == NO_ERROR) {
for (DWORD i = 0; i < tcpTable->dwNumEntries; i++) {
if (tcpTable->table[i].dwOwningPid == pid) {
socketCount++;
for (int p = 0; p < processCount; p++) {
if (tcpTable->table[i].dwOwningPid == monitoredProcesses[p].pid) {
totalSocketCount++;
break;
}
}
}
}
// Get process information
HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
DWORD threadCount = 0;
DWORD memoryUsage = 0;
if (hProcess) {
// Get process information for all monitored processes
for (int p = 0; p < processCount; p++) {
DWORD threadCount = 0;
DWORD memoryUsage = 0;
PROCESS_MEMORY_COUNTERS pmc;
if (GetProcessMemoryInfo(hProcess, &pmc, sizeof(pmc))) {
if (GetProcessMemoryInfo(monitoredProcesses[p].hProcess, &pmc, sizeof(pmc))) {
memoryUsage = pmc.WorkingSetSize / 1024; // KB
totalMemoryUsage += memoryUsage;
}
// Get thread count
......@@ -207,14 +320,14 @@ int main() {
te32.dwSize = sizeof(THREADENTRY32);
if (Thread32First(hSnapshot, &te32)) {
do {
if (te32.th32OwnerProcessID == pid) {
if (te32.th32OwnerProcessID == monitoredProcesses[p].pid) {
threadCount++;
}
} while (Thread32Next(hSnapshot, &te32));
}
CloseHandle(hSnapshot);
}
CloseHandle(hProcess);
totalThreadCount += threadCount;
}
// Get system time for logging
......@@ -222,25 +335,35 @@ int main() {
GetSystemTime(&st);
// Print status to console
printf("[%04d-%02d-%02d %02d:%02d:%02d] STATUS - Sockets: %lu, Threads: %lu, Memory: %lu KB\n",
printf("[%04d-%02d-%02d %02d:%02d:%02d] STATUS - Processes: %d, Sockets: %lu, Threads: %lu, Memory: %lu KB\n",
st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond,
socketCount, threadCount, memoryUsage);
processCount, totalSocketCount, totalThreadCount, totalMemoryUsage);
free(tcpTable);
lastStatusTime = currentTime;
}
// Check if process is still running
DWORD exitCode;
if (GetExitCodeProcess(pi.hProcess, &exitCode) && exitCode != STILL_ACTIVE) {
printf("Monitored process has terminated.\n");
// Check for terminated processes and remove them
for (int p = processCount - 1; p >= 0; p--) {
DWORD exitCode;
if (GetExitCodeProcess(monitoredProcesses[p].hProcess, &exitCode) && exitCode != STILL_ACTIVE) {
printf("Process %s (PID: %lu) has terminated.\n", monitoredProcesses[p].processName, monitoredProcesses[p].pid);
RemoveMonitoredProcess(monitoredProcesses[p].pid);
}
}
// Exit if no processes are left to monitor
if (processCount == 0) {
printf("All monitored processes have terminated.\n");
break;
}
Sleep(1000); // Check every second
}
CloseHandle(pi.hProcess);
CloseHandle(pi.hThread);
// Clean up any remaining process handles
for (int p = 0; p < processCount; p++) {
CloseHandle(monitoredProcesses[p].hProcess);
}
return 0;
}
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment