Added KMemDriver GetProcesses to enumerate all processes from kernel space.

 * fixed missing NUL termination for converted ASCII strings

2 files changed, 93 insertions, 25 deletions

diff --git a/KMemDriver/KMemDriver.c b/KMemDriver/KMemDriver.c
index 0350c81..4c7592f 100644
--- a/KMemDriver/KMemDriver.c
+++ b/KMemDriver/KMemDriver.c
@@ -167,6 +167,60 @@ NTSTATUS DriverEntry(
 	return status;
 }
 
+NTSTATUS GetProcesses(OUT PROCESS_DATA* procs, IN OUT SIZE_T* psiz)
+{
+	SIZE_T const max_siz = *psiz;
+	ULONG mem_needed = 0;
+
+	NTSTATUS status = ZwQuerySystemInformation(0x05, NULL, 0, &mem_needed);
+	if (!NT_SUCCESS(status) && mem_needed == 0) {
+		KDBG("NtQuerySystemInformation(ReturnLength: %lu) failed with 0x%X\n", mem_needed, status);
+		return status;
+	}
+
+	if (mem_needed / sizeof(SYSTEM_PROCESS_INFORMATION) > max_siz / sizeof(PROCESS_DATA)) {
+		KDBG("NtQuerySystemInformation buffer too small\n", status);
+		return STATUS_BUFFER_TOO_SMALL;
+	}
+
+	SYSTEM_PROCESS_INFORMATION* sysinfo = MmAllocateNonCachedMemory(mem_needed);
+	if (sysinfo == NULL) {
+		return STATUS_NO_MEMORY;
+	}
+
+	status = ZwQuerySystemInformation(0x05, sysinfo, mem_needed, NULL);
+	if (!NT_SUCCESS(status)) {
+		KDBG("NtQuerySystemInformation(SystemInformationLength: %lu) failed with 0x%X\n", mem_needed, status);
+		goto free_memory;
+	}
+
+	*psiz = 0;
+	SYSTEM_PROCESS_INFORMATION* cur_proc = sysinfo;
+	while (cur_proc->NextEntryOffset > 0 && *psiz < max_siz) {
+		cur_proc = (SYSTEM_PROCESS_INFORMATION*)((PUCHAR)cur_proc + cur_proc->NextEntryOffset);
+
+		ANSI_STRING name;
+		if (NT_SUCCESS(RtlUnicodeStringToAnsiString(&name, &cur_proc->ImageName, TRUE))) {
+			SIZE_T len = (name.Length >= sizeof(procs->ImageName) ? sizeof(procs->ImageName) - 1 : name.Length);
+			RtlCopyMemory(procs->ImageName, name.Buffer, len);
+			procs->ImageName[len] = '\0';
+			RtlFreeAnsiString(&name);
+		}
+
+		procs->NumberOfThreads = cur_proc->NumberOfThreads;
+		procs->UniqueProcessId = cur_proc->UniqueProcessId;
+		procs->HandleCount = cur_proc->HandleCount;
+
+		*psiz += sizeof(procs[0]);
+		procs++;
+	}
+	KDBG("Found %zu processes\n", *psiz / sizeof(*procs));
+
+free_memory:
+	MmFreeNonCachedMemory(sysinfo, mem_needed);
+	return status;
+}
+
 NTSTATUS WaitForControlProcess(OUT PEPROCESS* ppEProcess)
 {
 	NTSTATUS status;
@@ -418,6 +472,24 @@ NTSTATUS KRThread(IN PVOID pArg)
 						KeWriteVirtualMemory(ctrlPEP, ping, (PVOID)SHMEM_ADDR, &siz);
 						break;
 					}
+					case MEM_PROCESSES: {
+						PKERNEL_PROCESSES_REQUEST procs = (PKERNEL_PROCESSES_REQUEST)shm_buf;
+						PROCESS_DATA* data = (PPROCESS_DATA)(procs + 1);
+						siz = SHMEM_SIZE - sizeof(*procs);
+						//siz = sizeof(*data) * 128;
+						KDBG("Got a PROCESSES request\n");
+
+						procs->ProcessCount = 0;
+						procs->StatusRes = GetProcesses(data, &siz);
+						if (!NT_SUCCESS(procs->StatusRes))
+						{
+							break;
+						}
+
+						procs->ProcessCount = siz / sizeof(*data);
+						KeWriteVirtualMemory(ctrlPEP, procs, (PVOID)SHMEM_ADDR, &siz);
+						break;
+					}
 					case MEM_PAGES: {
 						PKERNEL_PAGE pages = (PKERNEL_PAGE)shm_buf;
 						KDBG("Got a PAGES request for process 0x%X start at address 0x%p\n",
diff --git a/KMemDriver/Memory.c b/KMemDriver/Memory.c
index 2b0621b..86d3f8c 100644
--- a/KMemDriver/Memory.c
+++ b/KMemDriver/Memory.c
@@ -7,8 +7,8 @@
 
 NTSTATUS GetPages(
 	IN PEPROCESS Process,
-	OUT MEMORY_BASIC_INFORMATION *mbiArr,
-	IN SIZE_T mbiArrLen, OUT SIZE_T *mbiUsed,
+	OUT MEMORY_BASIC_INFORMATION* mbiArr,
+	IN SIZE_T mbiArrLen, OUT SIZE_T* mbiUsed,
 	IN PVOID start_addr
 )
 {
@@ -31,8 +31,8 @@ NTSTATUS GetPages(
 	do {
 		mbiReturn = 0;
 		status = ZwQueryVirtualMemory(procHandle, (PVOID)baseAddr,
-			MemoryBasicInformation, mbiArr, sizeof *mbiArr * mbiArrLen, &mbiReturn);
-		mbiLength = mbiReturn / sizeof *mbiArr;
+			MemoryBasicInformation, mbiArr, sizeof * mbiArr * mbiArrLen, &mbiReturn);
+		mbiLength = mbiReturn / sizeof * mbiArr;
 		if (!NT_SUCCESS(status)) {
 			if (status == STATUS_INVALID_PARAMETER)
 				status = STATUS_SUCCESS;
@@ -58,7 +58,7 @@ NTSTATUS GetPages(
 
 NTSTATUS GetModules(
 	IN PEPROCESS Process,
-	OUT PMODULE_DATA pmod, IN OUT SIZE_T *psiz,
+	OUT PMODULE_DATA pmod, IN OUT SIZE_T* psiz,
 	IN SIZE_T start_index,
 	IN BOOLEAN isWow64
 )
@@ -106,20 +106,18 @@ NTSTATUS GetModules(
 			tmpUnicodeStr.Length = ldrEntry32->BaseDllName.Length;
 			tmpUnicodeStr.MaximumLength = ldrEntry32->BaseDllName.MaximumLength;
 			if (NT_SUCCESS(RtlUnicodeStringToAnsiString(&name, &tmpUnicodeStr, TRUE))) {
-				RtlCopyMemory(pmod->BaseDllName, name.Buffer,
-					(name.Length > sizeof pmod->BaseDllName ?
-						sizeof pmod->BaseDllName : name.Length)
-				);
+				SIZE_T len = (name.Length >= sizeof(pmod->BaseDllName) ? sizeof(pmod->BaseDllName) - 1 : name.Length);
+				RtlCopyMemory(pmod->BaseDllName, name.Buffer, len);
+				pmod->BaseDllName[len] = '\0';
 				RtlFreeAnsiString(&name);
 			}
 			tmpUnicodeStr.Buffer = (PWCH)ldrEntry32->FullDllName.Buffer;
 			tmpUnicodeStr.Length = ldrEntry32->FullDllName.Length;
 			tmpUnicodeStr.MaximumLength = ldrEntry32->FullDllName.MaximumLength;
 			if (NT_SUCCESS(RtlUnicodeStringToAnsiString(&name, &tmpUnicodeStr, TRUE))) {
-				RtlCopyMemory(pmod->FullDllPath, name.Buffer,
-					(name.Length > sizeof pmod->FullDllPath ?
-						sizeof pmod->FullDllPath : name.Length)
-				);
+				SIZE_T len = (name.Length >= sizeof(pmod->FullDllPath) ? sizeof(pmod->FullDllPath) - 1 : name.Length);
+				RtlCopyMemory(pmod->FullDllPath, name.Buffer, len);
+				pmod->FullDllPath[len] = '\0';
 				RtlFreeAnsiString(&name);
 			}
 			pmod->DllBase = (PVOID)ldrEntry32->DllBase;
@@ -164,17 +162,15 @@ NTSTATUS GetModules(
 			PLDR_DATA_TABLE_ENTRY ldrEntry = CONTAINING_RECORD(listEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
 			ANSI_STRING name;
 			if (NT_SUCCESS(RtlUnicodeStringToAnsiString(&name, &ldrEntry->BaseDllName, TRUE))) {
-				RtlCopyMemory(pmod->BaseDllName, name.Buffer,
-					(name.Length > sizeof pmod->BaseDllName ?
-						sizeof pmod->BaseDllName : name.Length)
-				);
+				SIZE_T len = (name.Length >= sizeof(pmod->BaseDllName) ? sizeof(pmod->BaseDllName) - 1 : name.Length);
+				RtlCopyMemory(pmod->BaseDllName, name.Buffer, len);
+				pmod->BaseDllName[len] = '\0';
 				RtlFreeAnsiString(&name);
 			}
 			if (NT_SUCCESS(RtlUnicodeStringToAnsiString(&name, &ldrEntry->FullDllName, TRUE))) {
-				RtlCopyMemory(pmod->FullDllPath, name.Buffer,
-					(name.Length > sizeof pmod->FullDllPath ?
-						sizeof pmod->FullDllPath : name.Length)
-				);
+				SIZE_T len = (name.Length >= sizeof(pmod->FullDllPath) ? sizeof(pmod->FullDllPath) - 1 : name.Length);
+				RtlCopyMemory(pmod->FullDllPath, name.Buffer, len);
+				pmod->FullDllPath[len] = '\0';
 				RtlFreeAnsiString(&name);
 			}
 			pmod->DllBase = ldrEntry->DllBase;
@@ -239,7 +235,7 @@ NTSTATUS KeWriteVirtualMemory(
 NTSTATUS KeProtectVirtualMemory(
 	IN HANDLE hProcess, IN PVOID addr,
 	IN SIZE_T siz, IN ULONG new_prot,
-	OUT ULONG *old_prot
+	OUT ULONG* old_prot
 )
 {
 	NTSTATUS status;
@@ -277,7 +273,7 @@ NTSTATUS KeRestoreProtectVirtualMemory(IN HANDLE hProcess,
 	return status;
 }
 
-NTSTATUS AllocMemoryToProcess(IN PEPROCESS pep, IN OUT PVOID *baseAddr, IN OUT SIZE_T *outSize, IN ULONG protect)
+NTSTATUS AllocMemoryToProcess(IN PEPROCESS pep, IN OUT PVOID* baseAddr, IN OUT SIZE_T* outSize, IN ULONG protect)
 {
 	NTSTATUS status;
 	PKAPC_STATE apc;
@@ -339,7 +335,7 @@ NTSTATUS WritePhysicalPage(IN PVOID addr, IN PUCHAR content, IN OUT PSIZE_T cont
 	return MmCopyMemory(vaddr, mm, 4096, MM_COPY_MEMORY_VIRTUAL, content_size_and_transferred);
 }
 
-NTSTATUS ReadPhysicalPage(IN PHYSICAL_ADDRESS * addr, OUT PUCHAR content, IN OUT PSIZE_T content_size_and_transferred)
+NTSTATUS ReadPhysicalPage(IN PHYSICAL_ADDRESS* addr, OUT PUCHAR content, IN OUT PSIZE_T content_size_and_transferred)
 {
 	MM_COPY_ADDRESS mm = { 0 };
 
@@ -348,7 +344,7 @@ NTSTATUS ReadPhysicalPage(IN PHYSICAL_ADDRESS * addr, OUT PUCHAR content, IN OUT
 		return STATUS_UNSUCCESSFUL;
 	}
 
-	mm.PhysicalAddress = *(PHYSICAL_ADDRESS *)addr;
+	mm.PhysicalAddress = *(PHYSICAL_ADDRESS*)addr;
 	return MmCopyMemory(content, mm, 4096, MM_COPY_MEMORY_PHYSICAL, content_size_and_transferred);
 }