Merge pull request #135 from namazso/master

Almost only-RPM module enumeration (closes #131)

Author: KN4CK3R

NativeCore/Windows/EnumerateRemoteSectionsAndModules.cpp

@@ -1,10 +1,120 @@
 #include <windows.h>
+#include <winternl.h>
 #include <tlhelp32.h>
 #include <vector>
 #include <algorithm>
 
 #include "NativeCore.hpp"
 
+static DWORD GetRemotePeb(HANDLE process, PPEB* ppeb)
+{
+	const auto ntdll = GetModuleHandle(TEXT("ntdll"));
+	if (!ntdll)
+		return ERROR_MOD_NOT_FOUND;
+
+	using tRtlNtStatusToDosError = ULONG (NTAPI *)(
+			_In_ NTSTATUS Status
+		);
+	const auto pRtlNtStatusToDosError = tRtlNtStatusToDosError(GetProcAddress(ntdll, "RtlNtStatusToDosError"));
+	if (!pRtlNtStatusToDosError)
+		return ERROR_NOT_FOUND;
+
+	using tNtQueryInformationProcess = NTSTATUS (NTAPI *)(
+			_In_ HANDLE ProcessHandle,
+			_In_ PROCESSINFOCLASS ProcessInformationClass,
+			_Out_writes_bytes_(ProcessInformationLength) PVOID ProcessInformation,
+			_In_ ULONG ProcessInformationLength,
+			_Out_opt_ PULONG ReturnLength
+		);
+
+	const auto pNtQueryInformationProcess = tNtQueryInformationProcess(GetProcAddress(ntdll, "NtQueryInformationProcess"));
+	if (!pNtQueryInformationProcess)
+		return ERROR_NOT_FOUND;
+
+	PROCESS_BASIC_INFORMATION pbi;
+	const auto status = pNtQueryInformationProcess(process, ProcessBasicInformation, &pbi, sizeof(pbi), nullptr);
+	if (!NT_SUCCESS(status))
+		return pRtlNtStatusToDosError(status);
+
+	*ppeb = pbi.PebBaseAddress;
+	
+	return ERROR_SUCCESS;
+}
+
+template <typename Proc>
+static DWORD EnumerateRemoteModulesNative(HANDLE process, Proc proc)
+{
+	PPEB ppeb;
+	const auto error = GetRemotePeb(process, &ppeb);
+	if (error != ERROR_SUCCESS)
+		return error;
+	
+	PPEB_LDR_DATA ldr;
+	auto success = ReadRemoteMemory(process, &ppeb->Ldr, &ldr, 0, sizeof(ldr));
+	if (!success)
+		return ERROR_READ_FAULT; // we seem to swallow the error anyways, might aswell give a distinctive one back
+
+	const auto list_head = &ldr->InMemoryOrderModuleList; // remote address
+	PLIST_ENTRY list_current; // remote address
+	success = ReadRemoteMemory(process, &list_head->Flink, &list_current, 0, sizeof(list_current));
+	if (!success)
+		return ERROR_READ_FAULT;
+	
+	while (list_current != list_head)
+	{
+		// TODO: error handling - what do we do if module list changed? We can't un-call the callback
+
+		LDR_DATA_TABLE_ENTRY mod;
+		success = ReadRemoteMemory(process, CONTAINING_RECORD(list_current, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks), &mod, 0, sizeof(mod));
+		if (!success)
+			return ERROR_SUCCESS; // return success here to prevent running the other one
+
+		EnumerateRemoteModuleData data = {};
+		data.BaseAddress = mod.DllBase;
+		data.Size = *(ULONG*)&mod.Reserved2[1]; // instead of undocced member could read ImageSize from headers
+		const auto path_len = std::min(sizeof(RC_UnicodeChar) * (PATH_MAXIMUM_LENGTH - 1), size_t(mod.FullDllName.Length));
+		success = ReadRemoteMemory(process, mod.FullDllName.Buffer, data.Path, 0, int(path_len));
+		if (!success)
+			return ERROR_SUCCESS; // return success here to prevent running the other one
+		
+		// UNICODE_STRING is not guaranteed to be null terminated
+		data.Path[path_len / 2] = 0;
+		
+		proc(&data);
+		
+		list_current = mod.InMemoryOrderLinks.Flink;
+	}
+	
+	return ERROR_SUCCESS;
+}
+
+template <typename Proc>
+static DWORD EnumerateRemoteModulesWinapi(HANDLE process, Proc proc)
+{
+	const auto handle = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetProcessId(process));
+	if (handle == INVALID_HANDLE_VALUE)
+		return GetLastError();
+	
+	MODULEENTRY32W me32 = {};
+	me32.dwSize = sizeof(MODULEENTRY32W);
+	if (Module32FirstW(handle, &me32))
+	{
+		do
+		{
+			EnumerateRemoteModuleData data = {};
+			data.BaseAddress = me32.modBaseAddr;
+			data.Size = me32.modBaseSize;
+			std::memcpy(data.Path, me32.szExePath, std::min(MAX_PATH, PATH_MAXIMUM_LENGTH));
+
+			proc(&data);
+		} while (Module32NextW(handle, &me32));
+	}
+
+	CloseHandle(handle);
+
+	return ERROR_SUCCESS;
+}
+
 void RC_CallConv EnumerateRemoteSectionsAndModules(RC_Pointer process, EnumerateRemoteSectionsCallback callbackSection, EnumerateRemoteModulesCallback callbackModule)
 {
 	if (callbackSection == nullptr && callbackModule == nullptr)
@@ -55,48 +165,34 @@ void RC_CallConv EnumerateRemoteSectionsAndModules(RC_Pointer process, Enumerate
 		address = reinterpret_cast<size_t>(memInfo.BaseAddress) + memInfo.RegionSize;
 	}
 
-	const auto handle = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetProcessId(process));
-	if (handle != INVALID_HANDLE_VALUE)
+	const auto moduleEnumerator = [&](EnumerateRemoteModuleData* data)
 	{
-		MODULEENTRY32W me32 = {};
-		me32.dwSize = sizeof(MODULEENTRY32W);
-		if (Module32FirstW(handle, &me32))
+		if (callbackModule != nullptr)
+			callbackModule(data);
+
+		if (callbackSection != nullptr)
 		{
-			do
-			{
-				if (callbackModule != nullptr)
+			auto it = std::lower_bound(std::begin(sections), std::end(sections), static_cast<LPVOID>(data->BaseAddress), [&sections](const auto& lhs, const LPVOID& rhs)
 				{
-					EnumerateRemoteModuleData data = {};
-					data.BaseAddress = me32.modBaseAddr;
-					data.Size = me32.modBaseSize;
-					std::memcpy(data.Path, me32.szExePath, std::min(MAX_PATH, PATH_MAXIMUM_LENGTH));
+					return lhs.BaseAddress < rhs;
+				});
 
-					callbackModule(&data);
-				}
+			IMAGE_DOS_HEADER DosHdr = {};
+			IMAGE_NT_HEADERS NtHdr = {};
+
+			ReadRemoteMemory(process, data->BaseAddress, &DosHdr, 0, sizeof(IMAGE_DOS_HEADER));
+			ReadRemoteMemory(process, PUCHAR(data->BaseAddress) + DosHdr.e_lfanew, &NtHdr, 0, sizeof(IMAGE_NT_HEADERS));
+
+			std::vector<IMAGE_SECTION_HEADER> sectionHeaders(NtHdr.FileHeader.NumberOfSections);
+			ReadRemoteMemory(process, PUCHAR(data->BaseAddress) + DosHdr.e_lfanew + sizeof(IMAGE_NT_HEADERS), sectionHeaders.data(), 0, NtHdr.FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER));
+			for (auto i = 0; i < NtHdr.FileHeader.NumberOfSections; ++i)
+			{
+				auto&& sectionHeader = sectionHeaders[i];
 
-				if (callbackSection != nullptr)
+				const auto sectionAddress = reinterpret_cast<size_t>(data->BaseAddress) + sectionHeader.VirtualAddress;
+				for (auto j = it; j != std::end(sections); ++j)
 				{
-					auto it = std::lower_bound(std::begin(sections), std::end(sections), static_cast<LPVOID>(me32.modBaseAddr), [&sections](const auto& lhs, const LPVOID& rhs)
-					{
-						return lhs.BaseAddress < rhs;
-					});
-
-					IMAGE_DOS_HEADER DosHdr = {};
-					IMAGE_NT_HEADERS NtHdr = {};
-
-					ReadRemoteMemory(process, me32.modBaseAddr, &DosHdr, 0, sizeof(IMAGE_DOS_HEADER));
-					ReadRemoteMemory(process, me32.modBaseAddr + DosHdr.e_lfanew, &NtHdr, 0, sizeof(IMAGE_NT_HEADERS));
-
-					std::vector<IMAGE_SECTION_HEADER> sectionHeaders(NtHdr.FileHeader.NumberOfSections);
-					ReadRemoteMemory(process, me32.modBaseAddr + DosHdr.e_lfanew + sizeof(IMAGE_NT_HEADERS), sectionHeaders.data(), 0, NtHdr.FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER));
-					for (auto i = 0; i < NtHdr.FileHeader.NumberOfSections; ++i)
-					{
-						auto&& sectionHeader = sectionHeaders[i];
-
-						const auto sectionAddress = reinterpret_cast<size_t>(me32.modBaseAddr) + sectionHeader.VirtualAddress;
-						for (auto j = it; j != std::end(sections); ++j)
-						{
-							if (sectionAddress >= reinterpret_cast<size_t>(j->BaseAddress) 
+					if (sectionAddress >= reinterpret_cast<size_t>(j->BaseAddress) 
 								&& sectionAddress < reinterpret_cast<size_t>(j->BaseAddress) + static_cast<size_t>(j->Size)
 								&& sectionHeader.VirtualAddress + sectionHeader.Misc.VirtualSize <= me32.modBaseSize )
 							{
@@ -121,21 +217,20 @@ void RC_CallConv EnumerateRemoteSectionsAndModules(RC_Pointer process, Enumerate
 
 								break;
 							}
-						}
-
-					}
 				}
-			} while (Module32NextW(handle, &me32));
-		}
 
-		CloseHandle(handle);
+			}
+		}
+	};
+	
+	if(EnumerateRemoteModulesNative(process, moduleEnumerator) != ERROR_SUCCESS)
+		EnumerateRemoteModulesWinapi(process, moduleEnumerator);
 
-		if (callbackSection != nullptr)
+	if (callbackSection != nullptr)
+	{
+		for (auto&& section : sections)
 		{
-			for (auto&& section : sections)
-			{
-				callbackSection(&section);
-			}
+			callbackSection(&section);
 		}
 	}
 }