DLLHelper can fix imports (does not work at the moment)

1 files changed, 124 insertions, 2 deletions

diff --git a/MemDriverLib/DLLHelper.cpp b/MemDriverLib/DLLHelper.cpp
index fba1f09..bd76552 100644
--- a/MemDriverLib/DLLHelper.cpp
+++ b/MemDriverLib/DLLHelper.cpp
@@ -10,6 +10,68 @@
 #define MakeDelta(cast, x, y) (cast) ((DWORD_PTR)(x) - (DWORD_PTR)(y))
 
 
+static FARPROC GetRemoteProcAddress(HMODULE localMod, HMODULE remoteMod, char *func_name)
+{
+	/*
+	 * Account for potential differences in base address
+	 * of modules in different processes.
+	 */
+	unsigned long delta = MakeDelta(unsigned long, remoteMod, localMod);
+	return MakePtr(FARPROC, GetProcAddress(localMod, func_name), delta);
+}
+
+static HMODULE GetRemoteModuleHandle(char *module_name,
+	std::vector<MODULE_DATA> &modules)
+{
+	SIZE_T remote_module_name_length;
+
+	for (auto& mod : modules) {
+		remote_module_name_length = strnlen(mod.BaseDllName, sizeof mod.BaseDllName);
+		if (strlen(module_name) == remote_module_name_length &&
+			!strncmp(module_name, mod.BaseDllName, remote_module_name_length))
+		{
+			return (HMODULE)mod.DllBase;
+		}
+	}
+
+	return NULL;
+}
+
+static PIMAGE_SECTION_HEADER GetEnclosingSectionHeader(DWORD rva, PIMAGE_NT_HEADERS pNTHeader)
+{
+	PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(pNTHeader);
+	unsigned int i;
+
+	for (i = 0; i < pNTHeader->FileHeader.NumberOfSections; i++, section++)
+	{
+		// This 3 line idiocy is because Watcom's linker actually sets the
+		// Misc.VirtualSize field to 0.  (!!! - Retards....!!!)
+		DWORD size = section->Misc.VirtualSize;
+		if (0 == size)
+			size = section->SizeOfRawData;
+
+		// Is the RVA within this section?
+		if ((rva >= section->VirtualAddress) &&
+			(rva < (section->VirtualAddress + size)))
+			return section;
+	}
+
+	return 0;
+}
+
+static LPVOID GetPtrFromRVA(DWORD rva, IMAGE_NT_HEADERS *pNTHeader, PBYTE imageBase)
+{
+	PIMAGE_SECTION_HEADER pSectionHdr;
+	INT delta;
+
+	pSectionHdr = GetEnclosingSectionHeader(rva, pNTHeader);
+	if (!pSectionHdr)
+		return 0;
+
+	delta = (INT)(pSectionHdr->VirtualAddress - pSectionHdr->PointerToRawData);
+	return (PVOID)(imageBase + rva - delta);
+}
+
 DLLHelper::DLLHelper()
 {
 }
@@ -22,7 +84,7 @@ DLLHelper::~DLLHelper()
 	}
 }
 
-bool DLLHelper::Init(HANDLE targetPID, std::string& fullDllPath) {
+bool DLLHelper::Init(HANDLE targetPID, const char * fullDllPath) {
 	if (!targetPID) {
 		return false;
 	}
@@ -106,9 +168,69 @@ bool DLLHelper::InitTargetMemory()
 	if (!ki.VAlloc(m_TargetPID, &wantedBaseAddr, &wantedSize, PAGE_EXECUTE_READWRITE)) {
 		return false;
 	}
-	if (wantedSize != m_DLLSize) {
+	if (wantedSize < m_DLLSize) {
+		return false;
+	}
+
+	m_TargetBaseAddress = wantedBaseAddr;
+	return true;
+}
+
+bool DLLHelper::FixImports()
+{
+	IMAGE_IMPORT_DESCRIPTOR *impDesc = (IMAGE_IMPORT_DESCRIPTOR *)GetPtrFromRVA(
+		(DWORD)(m_NTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress),
+		m_NTHeader,
+		(PBYTE)m_DLLPtr);
+	char *module_name;
+	std::vector<MODULE_DATA> modules;
+	KInterface& ki = KInterface::getInstance();
+
+	if (!m_TargetPID || !m_TargetBaseAddress || !m_NTHeader ||
+		!m_NTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size)
+	{
+		std::stringstream err_str;
+		err_str << "Pre-requirement failed (PID: " << m_TargetPID << ", BaseAddress: "
+			<< m_TargetBaseAddress << ", NTHeader: " << m_NTHeader;
+		throw std::runtime_error(err_str.str());
+		return false;
+	}
+	if (!ki.Modules(m_TargetPID, modules)) {
 		return false;
 	}
 
+	//   Loop through all the required modules
+	while ((module_name = (char *)GetPtrFromRVA((DWORD)(impDesc->Name), m_NTHeader,
+		(PBYTE)m_TargetBaseAddress)))
+	{
+		HMODULE localMod = LoadLibraryA(module_name);
+		HMODULE remoteMod = GetRemoteModuleHandle(module_name, modules);
+
+		if (!remoteMod) {
+			std::stringstream err_str;
+			err_str << "Module '" << module_name << "' does not exist in the target process "
+				<< m_TargetPID << " and we are not allowed to load it.";
+			throw std::runtime_error(err_str.str());
+			return false;
+		}
+
+		IMAGE_THUNK_DATA *itd =
+			(IMAGE_THUNK_DATA *)GetPtrFromRVA((DWORD)(impDesc->FirstThunk), m_NTHeader,
+			(PBYTE)m_TargetBaseAddress);
+
+		while (itd->u1.AddressOfData)
+		{
+			IMAGE_IMPORT_BY_NAME *iibn;
+			iibn = (IMAGE_IMPORT_BY_NAME *)GetPtrFromRVA((DWORD)(itd->u1.AddressOfData),
+				m_NTHeader, (PBYTE)m_TargetBaseAddress);
+
+			itd->u1.Function = MakePtr(DWORD, GetRemoteProcAddress(localMod,
+				remoteMod, (char *)iibn->Name), 0);
+
+			itd++;
+		}
+		impDesc++;
+	}
+
 	return true;
 }
\ No newline at end of file