DLLHelper: added Has* functions to check availability of some data dirs, implemented copy PE header and sections function, fixed invalid VAlloc request size

1 files changed, 69 insertions, 5 deletions

diff --git a/MemDriverLib/DLLHelper.cpp b/MemDriverLib/DLLHelper.cpp
index 0688434..279fc58 100644
--- a/MemDriverLib/DLLHelper.cpp
+++ b/MemDriverLib/DLLHelper.cpp
@@ -158,17 +158,17 @@ bool DLLHelper::VerifyHeader()
 
 bool DLLHelper::InitTargetMemory()
 {
-	if (!m_DLLPtr || !m_DLLSize) {
+	if (!m_DLLPtr || !m_NTHeader) {
 		return false;
 	}
 
 	PVOID wantedBaseAddr = m_TargetBaseAddress;
-	SIZE_T wantedSize = m_DLLSize;
+	SIZE_T wantedSize = m_NTHeader->OptionalHeader.SizeOfImage;
 	KInterface& ki = KInterface::getInstance();
 	if (!ki.VAlloc(m_TargetPID, &wantedBaseAddr, &wantedSize, PAGE_EXECUTE_READWRITE)) {
 		return false;
 	}
-	if (wantedSize < m_DLLSize) {
+	if (wantedSize < m_NTHeader->OptionalHeader.SizeOfImage) {
 		return false;
 	}
 
@@ -190,8 +190,12 @@ bool DLLHelper::FixImports()
 		!m_NTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size)
 	{
 		std::stringstream err_str;
-		err_str << "Pre-requirement failed (PID: " << m_TargetPID << ", BaseAddress: "
+		err_str << "FixImports pre-requirement failed [PID: " << m_TargetPID << ", BaseAddress: "
 			<< m_TargetBaseAddress << ", NTHeader: " << m_NTHeader;
+		if (m_NTHeader) {
+			err_str << " ImportTableSize: " << m_NTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size;
+		}
+		err_str << "]";
 		throw std::runtime_error(err_str.str());
 		return false;
 	}
@@ -246,8 +250,12 @@ bool DLLHelper::FixRelocs()
 		!m_NTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size)
 	{
 		std::stringstream err_str;
-		err_str << "Pre-requirement failed (PID: " << m_TargetPID << ", BaseAddress: "
+		err_str << "FixRelocs pre-requirement failed [PID: " << m_TargetPID << ", BaseAddress: "
 			<< m_TargetBaseAddress << ", NTHeader: " << m_NTHeader;
+		if (m_NTHeader) {
+			err_str << " RelocTableSize: " << m_NTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size;
+		}
+		err_str << "]";
 		throw std::runtime_error(err_str.str());
 		return false;
 	}
@@ -283,4 +291,60 @@ bool DLLHelper::FixRelocs()
 	}
 
 	return true;
+}
+
+bool DLLHelper::CopyHeaderAndSections()
+{
+	IMAGE_SECTION_HEADER *header;
+	unsigned int nBytes = 0;
+	unsigned int virtualSize = 0;
+	unsigned int n = 0;
+	KInterface& ki = KInterface::getInstance();
+
+	if (!m_TargetPID || !m_TargetBaseAddress || !m_NTHeader)
+	{
+		std::stringstream err_str;
+		err_str << "CopyHeaderAndSections pre-requirement failed [PID: " << m_TargetPID << ", BaseAddress: "
+			<< m_TargetBaseAddress << ", NTHeader: " << m_NTHeader << "]";
+		throw std::runtime_error(err_str.str());
+		return false;
+	}
+
+	if (!ki.WPM(m_TargetPID, m_TargetBaseAddress, m_DLLPtr,
+		m_NTHeader->FileHeader.SizeOfOptionalHeader +
+		sizeof(m_NTHeader->FileHeader) +
+		sizeof(m_NTHeader->Signature), NULL))
+	{
+		std::stringstream err_str;
+		err_str << "CopyHeaderAndSections failed [PID: " << m_TargetPID << ", BaseAddress: "
+			<< m_TargetBaseAddress << ", NTHeader: " << m_NTHeader << "]";
+		throw std::runtime_error(err_str.str());
+		return false;
+	}
+
+	header = IMAGE_FIRST_SECTION(m_NTHeader);
+	for (unsigned int i = 0; m_NTHeader->FileHeader.NumberOfSections; i++)
+	{
+		if (nBytes >= m_NTHeader->OptionalHeader.SizeOfImage)
+			break;
+
+		if (!ki.WPM(m_TargetPID, MakePtr(LPVOID, m_TargetBaseAddress, header->VirtualAddress),
+			MakePtr(BYTE *, m_DLLPtr, header->PointerToRawData), header->SizeOfRawData, NULL))
+		{
+			std::stringstream err_str;
+			err_str << "CopyHeaderAndSections failed [PID: " << m_TargetPID << ", BaseAddress: "
+				<< m_TargetBaseAddress << ", NTHeader: " << m_NTHeader
+				<< ", Section: " << header->Name << ", VA: " << header->VirtualAddress
+				<< ", Size: " << header->SizeOfRawData << "]";
+			throw std::runtime_error(err_str.str());
+			return false;
+		}
+
+		virtualSize = header->VirtualAddress;
+		header++;
+		virtualSize = header->VirtualAddress - virtualSize;
+		nBytes += virtualSize;
+	}
+
+	return true;
 }
\ No newline at end of file