added InitTargetMemory for allocating a memory block in the target process (cause we need the base address ASAP)

2 files changed, 32 insertions, 5 deletions

diff --git a/MemDriverLib/DLLHelper.cpp b/MemDriverLib/DLLHelper.cpp
index 0be5f08..fba1f09 100644
--- a/MemDriverLib/DLLHelper.cpp
+++ b/MemDriverLib/DLLHelper.cpp
@@ -1,5 +1,6 @@
 #include "stdafx.h"
 #include "DLLHelper.h"
+#include "KInterface.h"
 
 #include <sstream>
 #include <Windows.h>
@@ -21,7 +22,11 @@ DLLHelper::~DLLHelper()
 	}
 }
 
-bool DLLHelper::Init(std::string& fullDllPath) {
+bool DLLHelper::Init(HANDLE targetPID, std::string& fullDllPath) {
+	if (!targetPID) {
+		return false;
+	}
+	m_TargetPID = targetPID;
 	m_DLLPath = fullDllPath;
 
 	HANDLE hFile = CreateFileA(m_DLLPath.c_str(),
@@ -87,4 +92,23 @@ bool DLLHelper::VerifyHeader()
 	}
 
 	return true;
+}
+
+bool DLLHelper::InitTargetMemory()
+{
+	if (!m_DLLPtr || !m_DLLSize) {
+		return false;
+	}
+
+	PVOID wantedBaseAddr = m_TargetBaseAddress;
+	SIZE_T wantedSize = m_DLLSize;
+	KInterface& ki = KInterface::getInstance();
+	if (!ki.VAlloc(m_TargetPID, &wantedBaseAddr, &wantedSize, PAGE_EXECUTE_READWRITE)) {
+		return false;
+	}
+	if (wantedSize != m_DLLSize) {
+		return false;
+	}
+
+	return true;
 }
\ No newline at end of file
diff --git a/MemDriverLib/DLLHelper.h b/MemDriverLib/DLLHelper.h
index 0a6176f..3384ec6 100644
--- a/MemDriverLib/DLLHelper.h
+++ b/MemDriverLib/DLLHelper.h
@@ -8,14 +8,17 @@ public:
 	DLLHelper();
 	~DLLHelper();
 
-	bool Init(std::string& fullDllPath);
+	bool Init(HANDLE targetPID, std::string& fullDllPath);
 	bool VerifyHeader();
+	bool InitTargetMemory();
 
 private:
+	HANDLE m_TargetPID = 0;
 	std::string m_DLLPath;
-	DWORD m_DLLSize;
+	DWORD m_DLLSize = 0;
 	UINT8 *m_DLLPtr = nullptr;
-	IMAGE_DOS_HEADER *m_DOSHeader;
-	IMAGE_NT_HEADERS *m_NTHeader;
+	IMAGE_DOS_HEADER *m_DOSHeader = nullptr;
+	IMAGE_NT_HEADERS *m_NTHeader = nullptr;
+	PVOID m_TargetBaseAddress = nullptr;
 };