Driver: PsTree: Copy info from user-mode process

This commit is contained in:
Nodir Temirkhodjaev 2022-02-09 13:05:00 +03:00
parent 8d2ad778f0
commit cac77fae22

View File

@ -389,15 +389,12 @@ FORT_API PFORT_PSNAME fort_pstree_get_proc_name(
static HANDLE OpenProcessById(DWORD processId) static HANDLE OpenProcessById(DWORD processId)
{ {
if (processId == 0 || processId == 4)
return NULL; // skip "System Idle Process" (0) and "System" (4) processes
NTSTATUS status; NTSTATUS status;
PEPROCESS peProcess = NULL; PEPROCESS peProcess = NULL;
status = PsLookupProcessByProcessId((HANDLE) (ptrdiff_t) processId, &peProcess); status = PsLookupProcessByProcessId((HANDLE) (ptrdiff_t) processId, &peProcess);
if (!NT_SUCCESS(status) || peProcess == NULL) { if (!NT_SUCCESS(status) || peProcess == NULL) {
LOG("PsTree: Lookup Process Error: %x\n", status); LOG("PsTree: Lookup Process Error: pid=%d %x\n", processId, status);
return NULL; return NULL;
} }
@ -407,7 +404,7 @@ static HANDLE OpenProcessById(DWORD processId)
ObDereferenceObject(peProcess); ObDereferenceObject(peProcess);
if (!NT_SUCCESS(status) || processHandle == NULL) { if (!NT_SUCCESS(status) || processHandle == NULL) {
LOG("PsTree: Open Process Object Error: %x\n", status); LOG("PsTree: Open Process Object Error: pid=%d %x\n", processId, status);
return NULL; return NULL;
} }
@ -444,13 +441,21 @@ static NTSTATUS GetCurrentProcessPathArgs(PUNICODE_STRING path, PUNICODE_STRING
if (params == NULL) if (params == NULL)
return STATUS_INVALID_ADDRESS; return STATUS_INVALID_ADDRESS;
path->Length = params->ImagePathName.Length; const USHORT pathLength = params->ImagePathName.Length;
path->MaximumLength = params->ImagePathName.Length; if (pathLength != 0 && pathLength <= path->MaximumLength - sizeof(WCHAR)) {
path->Buffer = GetUnicodeStringBuffer(&params->ImagePathName, params); PCWCHAR userBuffer = GetUnicodeStringBuffer(&params->ImagePathName, params);
RtlCopyMemory(path->Buffer, userBuffer, pathLength);
path->Buffer[pathLength / sizeof(WCHAR)] = L'\0';
path->Length = pathLength;
}
commandLine->Length = params->CommandLine.Length; const USHORT commandLineLength = params->CommandLine.Length;
commandLine->MaximumLength = params->CommandLine.Length; if (commandLineLength != 0 && commandLineLength <= commandLine->MaximumLength - sizeof(WCHAR)) {
commandLine->Buffer = GetUnicodeStringBuffer(&params->CommandLine, params); PCWCHAR userBuffer = GetUnicodeStringBuffer(&params->CommandLine, params);
RtlCopyMemory(commandLine->Buffer, userBuffer, commandLineLength);
commandLine->Buffer[commandLineLength / sizeof(WCHAR)] = L'\0';
commandLine->Length = commandLineLength;
}
return STATUS_SUCCESS; return STATUS_SUCCESS;
} }
@ -467,42 +472,57 @@ static void fort_pstree_attach_process(PSYSTEM_PROCESSES processEntry, HANDLE pr
return; return;
} }
WCHAR pathBuffer[256];
UNICODE_STRING path = {
.Length = 0, .MaximumLength = sizeof(pathBuffer), .Buffer = pathBuffer
};
WCHAR commandLineBuffer[512];
UNICODE_STRING commandLine = {
.Length = 0, .MaximumLength = sizeof(commandLineBuffer), .Buffer = commandLineBuffer
};
// Copy info from user-mode process to stack
KAPC_STATE apcState; KAPC_STATE apcState;
KeStackAttachProcess(process, &apcState); KeStackAttachProcess(process, &apcState);
{ {
UNICODE_STRING path;
UNICODE_STRING commandLine;
status = GetCurrentProcessPathArgs(&path, &commandLine); status = GetCurrentProcessPathArgs(&path, &commandLine);
if (NT_SUCCESS(status)) {
const HANDLE processId = (HANDLE) (ptrdiff_t) processEntry->ProcessId;
const HANDLE parentProcessId = (HANDLE) (ptrdiff_t) processEntry->ParentProcessId;
PS_CREATE_NOTIFY_INFO createInfo = { .ParentProcessId = parentProcessId,
.ImageFileName = &path,
.CommandLine = &commandLine };
fort_pstree_notify(/*process=*/NULL, processId, &createInfo);
} else {
LOG("PsTree: Query Process Error: pid=%d %x\n", processEntry->ProcessId, status);
}
} }
KeUnstackDetachProcess(&apcState); KeUnstackDetachProcess(&apcState);
ObDereferenceObject(process); ObDereferenceObject(process);
// Process the info
if (!NT_SUCCESS(status)) {
LOG("PsTree: Query Process Error: pid=%d %x\n", processEntry->ProcessId, status);
} else if (path.Length != 0 && commandLine.Length != 0) {
const HANDLE processId = (HANDLE) (ptrdiff_t) processEntry->ProcessId;
const HANDLE parentProcessId = (HANDLE) (ptrdiff_t) processEntry->ParentProcessId;
PS_CREATE_NOTIFY_INFO createInfo = {
.ParentProcessId = parentProcessId, .ImageFileName = &path, .CommandLine = &commandLine
};
fort_pstree_notify(/*process=*/NULL, processId, &createInfo);
}
} }
static void fort_pstree_enum_processes_loop(PSYSTEM_PROCESSES processEntry) static void fort_pstree_enum_processes_loop(PSYSTEM_PROCESSES processEntry)
{ {
for (;;) { for (;;) {
const DWORD processId = (DWORD) processEntry->ProcessId; const DWORD processId = (DWORD) processEntry->ProcessId;
const DWORD parentProcessId = (DWORD) processEntry->ParentProcessId;
if (processId == 0 || processId == 4 || parentProcessId == 4) {
// skip System (sub)processes
} else {
const HANDLE processHandle = OpenProcessById(processId); const HANDLE processHandle = OpenProcessById(processId);
if (processHandle != NULL) { if (processHandle != NULL) {
fort_pstree_attach_process(processEntry, processHandle); fort_pstree_attach_process(processEntry, processHandle);
ZwClose(processHandle); ZwClose(processHandle);
} }
}
if (processEntry->NextEntryOffset == 0) if (processEntry->NextEntryOffset == 0)
break; break;