Update PE.cpp

m0n0ph1 · Sep 12, 2023 · f0e0885 · f0e0885
1 parent 585d413
commit f0e0885
Showing 1 changed file with 50 additions and 72 deletions.
diff --git a/sourcecode/ProcessHollowing/PE.cpp b/sourcecode/ProcessHollowing/PE.cpp
@@ -3,93 +3,71 @@
 #include "internals.h"
 #include "pe.h"
 
-DWORD FindRemotePEB(HANDLE hProcess)
-{
-	HMODULE hNTDLL = LoadLibraryA("ntdll");
-
-	if (!hNTDLL)
-		return 0;
+HMODULE hNTDLL = nullptr;
+_NtQueryInformationProcess ntQueryInformationProcess = nullptr;
 
-	FARPROC fpNtQueryInformationProcess = GetProcAddress
-		(
-		hNTDLL,
-		"NtQueryInformationProcess"
-		);
-
-	if (!fpNtQueryInformationProcess)
-		return 0;
+bool InitializeNtQueryInformationProcess()
+{
+    hNTDLL = LoadLibraryA("ntdll");
+    if (!hNTDLL)
+        return false;
 
-	_NtQueryInformationProcess ntQueryInformationProcess = 
-		(_NtQueryInformationProcess)fpNtQueryInformationProcess;
+    FARPROC fpNtQueryInformationProcess = GetProcAddress(hNTDLL, "NtQueryInformationProcess");
+    if (!fpNtQueryInformationProcess)
+        return false;
 
-	PROCESS_BASIC_INFORMATION* pBasicInfo = 
-		new PROCESS_BASIC_INFORMATION();
+    ntQueryInformationProcess = (_NtQueryInformationProcess)fpNtQueryInformationProcess;
+    return true;
+}
 
-	DWORD dwReturnLength = 0;
+DWORD FindRemotePEB(HANDLE hProcess)
+{
+    if(!ntQueryInformationProcess)
+    {
+        if(!InitializeNtQueryInformationProcess())
+            return 0;
+    }
 
-	ntQueryInformationProcess
-		(
-		hProcess, 
-		0, 
-		pBasicInfo, 
-		sizeof(PROCESS_BASIC_INFORMATION), 
-		&dwReturnLength
-		);
+    PROCESS_BASIC_INFORMATION basicInfo = {0};
+    DWORD dwReturnLength = 0;
 
-	return pBasicInfo->PebBaseAddress;
+    ntQueryInformationProcess(hProcess, 0, &basicInfo, sizeof(basicInfo), &dwReturnLength);
+    return basicInfo.PebBaseAddress;
 }
 
 PEB* ReadRemotePEB(HANDLE hProcess)
 {
-	DWORD dwPEBAddress = FindRemotePEB(hProcess);
-
-	PEB* pPEB = new PEB();
+    DWORD dwPEBAddress = FindRemotePEB(hProcess);
+    if(!dwPEBAddress)
+        return nullptr;
 
-	BOOL bSuccess = ReadProcessMemory
-		(
-		hProcess,
-		(LPCVOID)dwPEBAddress,
-		pPEB,
-		sizeof(PEB),
-		0
-		);
+    PEB* pPEB = new PEB();
 
-	if (!bSuccess)
-		return 0;
+    if(!ReadProcessMemory(hProcess, (LPCVOID)dwPEBAddress, pPEB, sizeof(PEB), nullptr))
+    {
+        delete pPEB;
+        return nullptr;
+    }
 
-	return pPEB;
+    return pPEB;
 }
 
 PLOADED_IMAGE ReadRemoteImage(HANDLE hProcess, LPCVOID lpImageBaseAddress)
 {
-	BYTE* lpBuffer = new BYTE[BUFFER_SIZE];
-
-	BOOL bSuccess = ReadProcessMemory
-		(
-		hProcess,
-		lpImageBaseAddress,
-		lpBuffer,
-		BUFFER_SIZE,
-		0
-		);
-
-	if (!bSuccess)
-		return 0;	
-
-	PIMAGE_DOS_HEADER pDOSHeader = (PIMAGE_DOS_HEADER)lpBuffer;
-
-	PLOADED_IMAGE pImage = new LOADED_IMAGE();
-
-	pImage->FileHeader = 
-		(PIMAGE_NT_HEADERS32)(lpBuffer + pDOSHeader->e_lfanew);
-
-	pImage->NumberOfSections = 
-		pImage->FileHeader->FileHeader.NumberOfSections;
-
-	pImage->Sections = 
-		(PIMAGE_SECTION_HEADER)(lpBuffer + pDOSHeader->e_lfanew + 
-		sizeof(IMAGE_NT_HEADERS32));
-
-	return pImage;
+    BYTE* lpBuffer = new BYTE[BUFFER_SIZE];
+    if(!ReadProcessMemory(hProcess, lpImageBaseAddress, lpBuffer, BUFFER_SIZE, nullptr))
+    {
+        delete[] lpBuffer;
+        return nullptr;	
+    }
+
+    PIMAGE_DOS_HEADER pDOSHeader = (PIMAGE_DOS_HEADER)lpBuffer;
+    PLOADED_IMAGE pImage = new LOADED_IMAGE();
+
+    pImage->FileHeader = (PIMAGE_NT_HEADERS32)(lpBuffer + pDOSHeader->e_lfanew);
+    pImage->NumberOfSections = pImage->FileHeader->FileHeader.NumberOfSections;
+    pImage->Sections = (PIMAGE_SECTION_HEADER)(lpBuffer + pDOSHeader->e_lfanew + sizeof(IMAGE_NT_HEADERS32));
+
+    delete[] lpBuffer; // Avoid memory leak
+    return pImage;
 }
-