Local-PE-Injection

Local-PE-Injection

LocalPE-Injection.c

/*
	PE Injection Steps:
		1. Allocate memory to hold injected PE file.
		2. Copy over the PE sections to the allocated memory
		3. Address the PE relocations
		3. Fix PE IAT (Import Address Table)
		4. Set correct permission for each PE section
		5. Execute the PE entry point
*/

#include <Windows.h>
#include <stdio.h>
#include <winternl.h>

/*
	Reads a file from disk into a heap-allocated buffer

	cFileName -> Name of the PE file to read
	ppAddress -> Receives a pointer to the heap buffer containing the file data
	pdwFileSize -> Receives the size of the file in bytes
*/
BOOL ReadFileDisk(_In_ LPCSTR cFileName, _Out_ PBYTE* ppAddress, _Out_ PDWORD pdwFileSize) {
	
	HANDLE hFile = INVALID_HANDLE_VALUE;
	PBYTE pBuffer = NULL;
	DWORD dwFileSize = 0x00, dwNumBytesRead = 0x00;

	// Verify parameters are filled
	if (!cFileName || !ppAddress || !pdwFileSize)
		return FALSE;


	*ppAddress = NULL;
	*pdwFileSize = 0;

	// Open a handle to the existing file with read access
	hFile = CreateFileA(cFileName, GENERIC_READ, 0x00, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (hFile == INVALID_HANDLE_VALUE) {
		printf("[!] CreateFileA Failed: %d \n", GetLastError());
		goto _CleanUp;
	}

	// Get filesize of file to be read
	dwFileSize = GetFileSize(hFile, NULL);
	if (dwFileSize == INVALID_FILE_SIZE) {
		printf("[!] GetFileSize Failed: %d \n", GetLastError());
		goto _CleanUp;
	}

	// Allocate memory for the file
	pBuffer = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, dwFileSize);
	if (!pBuffer) {
		printf("[!] HeapAlloc Failed: %d \n", GetLastError());
		goto _CleanUp;
	}

	// Reading file
	if (!ReadFile(hFile, pBuffer, dwFileSize, &dwNumBytesRead, NULL) || dwFileSize != dwNumBytesRead) {
		printf("[!] ReadFile Failed: %d \n", GetLastError());
		goto _CleanUp;
	}

	// Assign output parameters
	*ppAddress = pBuffer;
	*pdwFileSize = dwFileSize;


_CleanUp:
	if (hFile != INVALID_HANDLE_VALUE)
		CloseHandle(hFile);
	if (!*ppAddress && pBuffer)
		HeapFree(GetProcessHeap(), 0, pBuffer);
	return ((*ppAddress != NULL) && (*pdwFileSize != 0x00)) ? TRUE : FALSE;

}

/*
	PE header structure — populated by ParsePEStruct(). Caches pointers to all relevant PE headers and data directory entries, so callers don't have to recompute them from the raw buffer repeatedly.
*/

typedef struct _PE_HDRS
{
	PBYTE                    pFileBuffer;
	DWORD                    dwFileSize;

	PIMAGE_NT_HEADERS        pImgNtHdr;
	PIMAGE_SECTION_HEADER    pImgSecHdr;

	PIMAGE_DATA_DIRECTORY    pEntryImportDataDir;
	PIMAGE_DATA_DIRECTORY    pEntryBaseRelocDataDir;
	PIMAGE_DATA_DIRECTORY    pEntryTLSDataDir;
	PIMAGE_DATA_DIRECTORY    pEntryExceptionDataDir;
	PIMAGE_DATA_DIRECTORY    pEntryExportDataDir;

	BOOL                     bIsDLLFile;

} PE_HDR, *PPE_HDR;

/*
	Parses a raw PE file buffer and populates a PE_HDR structure with pointers to the relevant headers and data directory entries

	pPEHdr      -> Pointer to the PE_HDR structure to populate
	pFileBuffer -> Pointer to the raw PE file buffer in memory
	dwFileSize  -> Size of the raw PE file buffer (in bytes)
*/
BOOL ParsePEStruct(_Out_ PPE_HDR pPEHdr, _In_ PBYTE pFileBuffer, _In_ DWORD dwFileSize) {

	// Verify parameters are filled
	if (!pPEHdr || !pFileBuffer || !dwFileSize)
		return FALSE;

	// Assign function parameters to struct
	pPEHdr->pFileBuffer = pFileBuffer;
	pPEHdr->dwFileSize = dwFileSize;
	pPEHdr->pImgNtHdr = (PIMAGE_NT_HEADERS)(pFileBuffer + ((PIMAGE_DOS_HEADER)pFileBuffer)->e_lfanew);

	// Check if NT Header is valid
	if (pPEHdr->pImgNtHdr->Signature != IMAGE_NT_SIGNATURE)
		return FALSE;

	// Populate Elements to struct
	pPEHdr->bIsDLLFile = (pPEHdr->pImgNtHdr->FileHeader.Characteristics & IMAGE_FILE_DLL) ? TRUE : FALSE;
	pPEHdr->pImgSecHdr = IMAGE_FIRST_SECTION(pPEHdr->pImgNtHdr);
	pPEHdr->pEntryImportDataDir = &pPEHdr->pImgNtHdr->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
	pPEHdr->pEntryBaseRelocDataDir = &pPEHdr->pImgNtHdr->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
	pPEHdr->pEntryTLSDataDir = &pPEHdr->pImgNtHdr->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS];
	pPEHdr->pEntryExceptionDataDir = &pPEHdr->pImgNtHdr->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
	pPEHdr->pEntryExportDataDir = &pPEHdr->pImgNtHdr->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];

	return TRUE;
}

/* 
	Resolves all imported function addresses for a PE image and patches them into the IAT (Import Address Thunk / FirstThunk array).
	Walks the import descriptor table, loads each referenced DLL via LoadLibraryA and resolves each function either by ordinal or by name via GetProcAddress.

	pEntryImportDataDir -> Pointer to the import data directory entry from the PE optional header (DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT])
	pPeBaseAddr -> Base address of the PE image mapped into memory, all RVAs in the import descriptors are relative to this
*/
BOOL InitIAT(_In_ PIMAGE_DATA_DIRECTORY pEntryImportDataDir, _In_ PBYTE pPeBaseAddr) {

	// Check parameters filled
	if (!pEntryImportDataDir || !pPeBaseAddr)
		return FALSE;

	// Current import descriptor one entry per imported DLL
	PIMAGE_IMPORT_DESCRIPTOR pImgDescriptor = NULL;

	// Walk the import descriptor array
	for (SIZE_T i = 0; i < pEntryImportDataDir->Size; i += sizeof(IMAGE_IMPORT_DESCRIPTOR)) {
		
		// Resolve current import descriptor from its RVA
		pImgDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)(pPeBaseAddr + pEntryImportDataDir->VirtualAddress + i);

		// A fully zeroed descriptor marks the end of the import descriptor table
		if (pImgDescriptor->OriginalFirstThunk == NULL && pImgDescriptor->FirstThunk == NULL)
			break;

		// OriginalFirstThunk (INT) is the unmodified import name table
	    // Some linkers omit it fall back to FirstThunk (IAT) in that case
		LPSTR cDllName = (LPSTR)(pPeBaseAddr + pImgDescriptor->Name);
		ULONG_PTR uOriginalFirstThunkRVA = pImgDescriptor->OriginalFirstThunk;
		ULONG_PTR uFirstThunkRVA = pImgDescriptor->FirstThunk;
		SIZE_T sImgThunkSize = 0x00;
		HMODULE hModule = NULL;

		// Load the DLL into the current process so we can resolve its export(s)
		hModule = LoadLibraryA(cDllName);
		if (!hModule) {
			printf("[!] LoadLibraryA Failed: %d \n", GetLastError());
			return FALSE;
		}

		// Walk the thunk arrays for this DLL one entry per imported function
		while (TRUE) {

			// Get pointers to the first thunk and original first thunk data
			PIMAGE_THUNK_DATA			pOriginalFirstThunk = (PIMAGE_THUNK_DATA)(pPeBaseAddr + uOriginalFirstThunkRVA + sImgThunkSize);
			PIMAGE_THUNK_DATA			pFirstThunk = (PIMAGE_THUNK_DATA)(pPeBaseAddr + uFirstThunkRVA + sImgThunkSize);
			PIMAGE_IMPORT_BY_NAME		pImgImportByName = NULL;
			ULONG_PTR					pFuncAddress = 0;

			// A zeroed thunk entry marks the end of this DLL's import list
			if (pOriginalFirstThunk->u1.Function == NULL && pFirstThunk->u1.Function == 0)
				break;
			
			// High bit of Ordinal field set function is imported by ordinal, not by name
			if (IMAGE_SNAP_BY_ORDINAL(pOriginalFirstThunk->u1.Ordinal)) {
				if (!(pFuncAddress = (ULONG_PTR)GetProcAddress(hModule, IMAGE_ORDINAL(pOriginalFirstThunk->u1.Ordinal)))) {
					printf("[!] Could Not Import !%s#%d \n", cDllName, (int)pOriginalFirstThunk->u1.Ordinal);
					FreeLibrary(hModule);
					return FALSE;
				}
			}

			// Ordinal flag not set resolve by function name via IMAGE_IMPORT_BY_NAME
			else {
				pImgImportByName = (PIMAGE_IMPORT_BY_NAME)(pPeBaseAddr + pOriginalFirstThunk->u1.AddressOfData);
				if (!(pFuncAddress = (ULONG_PTR)GetProcAddress(hModule, pImgImportByName->Name))) {
					printf("[!] Could Not Import !%s.%s \n", cDllName, pImgImportByName->Name);
					FreeLibrary(hModule);
					return FALSE;
				}
			}

			// Patch the IAT entry with the resolved runtime address of the function
			pFirstThunk->u1.Function = (ULONGLONG)pFuncAddress;

			// Advance by one thunk entry (sizeof IMAGE_THUNK_DATA = pointer size)
			sImgThunkSize += sizeof(IMAGE_THUNK_DATA);

		}
	}

	return TRUE;
}

typedef struct _BASE_RELOCATION_ENTRY {
	WORD	Offset : 12;  // Specifies where the base relocation is to be applied.
	WORD	Type : 4;   // Indicates the type of base relocation to be applied.
} BASE_RELOCATION_ENTRY, * PBASE_RELOCATION_ENTRY;

/*
	Reallocation is necsarry to all entries within each block, to adjust hardcoded address within an executable image when its loaded at an different from its preferred base address

	pEntryBaseRelocDataDir -> Pointer to data dir for the relocations of the PE image
	pPeBaseAddr -> Address of PE image currently loaded in memory
	pPrefAddr -> The preferred base address of the PE image
*/


BOOL ApplyReloc(_In_ PIMAGE_DATA_DIRECTORY pEntryBaseRelocDataDir, _In_ ULONG_PTR pPeBaseAddr, _In_ ULONG_PTR pPrefAddr) {

	// Verify parameters are filled
	if (!pEntryBaseRelocDataDir || !pPeBaseAddr || !pPrefAddr)
		return FALSE;

	// Pointer to the beginning of the base relocation block
	PIMAGE_BASE_RELOCATION pImgBaseRelocation = (pPeBaseAddr + pEntryBaseRelocDataDir->VirtualAddress);

	// The difference between the current PE image base address and its preferable base address
	ULONG_PTR uDeltaOffset = pPeBaseAddr - pPrefAddr;
	
	// Pointer to individual base relocation entries.
	PBASE_RELOCATION_ENTRY pBaseRelocEntry = NULL;

	// Iterate over all base relocation blocks inside the PE image
	while (pImgBaseRelocation->VirtualAddress) {

		// pointer to first entry of the current block
		pBaseRelocEntry = (PBASE_RELOCATION_ENTRY)(pImgBaseRelocation + 1);

		while ((PBYTE)pBaseRelocEntry != (PBYTE)pImgBaseRelocation + pImgBaseRelocation->SizeOfBlock) {
			// adjust relocation entry based on its relocation type
			switch (pBaseRelocEntry->Type) {
			case IMAGE_REL_BASED_DIR64:
				// Adjust a 64-bit field by the delta offset.
				*((ULONG_PTR*)(pPeBaseAddr + pImgBaseRelocation->VirtualAddress + pBaseRelocEntry->Offset)) += uDeltaOffset;
				break;
			case IMAGE_REL_BASED_HIGHLOW:
				// Adjust a 32-bit field by the delta offset.
				*((DWORD*)(pPeBaseAddr + pImgBaseRelocation->VirtualAddress + pBaseRelocEntry->Offset)) += (DWORD)uDeltaOffset;
				break;
			case IMAGE_REL_BASED_HIGH:
				// Adjust the high 16 bits of a 32-bit field.
				*((WORD*)(pPeBaseAddr + pImgBaseRelocation->VirtualAddress + pBaseRelocEntry->Offset)) += HIWORD(uDeltaOffset);
				break;
			case IMAGE_REL_BASED_LOW:
				// Adjust the low 16 bits of a 32-bit field.
				*((WORD*)(pPeBaseAddr + pImgBaseRelocation->VirtualAddress + pBaseRelocEntry->Offset)) += LOWORD(uDeltaOffset);
				break;
			case IMAGE_REL_BASED_ABSOLUTE:
				// No relocation is required.
				break;
			default:
				// Handle unknown relocation types.
				printf("[!] Unknown relocation type: %d | Offset: 0x%08X \n", pBaseRelocEntry->Type, pBaseRelocEntry->Offset);
				return FALSE;
			}
			pBaseRelocEntry++;

		}

		// Iterate over next realocation block
		pImgBaseRelocation = (PIMAGE_BASE_RELOCATION)pBaseRelocEntry;
	}

	return TRUE;
}

/*
	Applies correct memory protection to each section of a PE image loaded into memory
    The image is initially mapped with blanket RW permissions -- this function inspects IMAGE_SECTION_HEADER.Characteristics for each section and calls VirtualProtect to set the appropriate R/W/X combination

	pPeBaseAddr -> Base address of PE Image
	pImgNtHdr -> Pointer to NT Header struct
	pImgSecHdr -> Pointer to Section header
*/

BOOL FixMemPermissions(_In_ ULONG_PTR pPeBaseAddr, _In_ PIMAGE_NT_HEADERS pImgNtHdr, _In_ PIMAGE_SECTION_HEADER pImgSecHdr) {

	if (!pPeBaseAddr || !pImgNtHdr || !pImgSecHdr)
		return FALSE;

	// Loop through each section of the PE image.
	for (DWORD i = 0; i < pImgNtHdr->FileHeader.NumberOfSections; i++) {

		// Variables to store the new and old memory protections.
		DWORD	dwProtection = PAGE_NOACCESS, dwOldProtection = 0x00;

		// Skip sections with no raw data or no mapped address 
		if (!pImgSecHdr[i].SizeOfRawData || !pImgSecHdr[i].VirtualAddress)
			continue;

		// Determine memory protection based on section characteristics
		// These characteristics dictate whether the section is readable, writable, executable etc
		if (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_WRITE)
			dwProtection = PAGE_WRITECOPY;

		if (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_READ)
			dwProtection = PAGE_READONLY;

		if ((pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_WRITE) && (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_READ))
			dwProtection = PAGE_READWRITE;

		if (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_EXECUTE)
			dwProtection = PAGE_EXECUTE;

		if ((pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_EXECUTE) && (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_WRITE))
			dwProtection = PAGE_EXECUTE_WRITECOPY;

		if ((pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_EXECUTE) && (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_READ))
			dwProtection = PAGE_EXECUTE_READ;

		if ((pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_EXECUTE) && (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_WRITE) && (pImgSecHdr[i].Characteristics & IMAGE_SCN_MEM_READ))
			dwProtection = PAGE_EXECUTE_READWRITE;

		// Apply the determined memory protection to the section
		if (!VirtualProtect((PVOID)(pPeBaseAddr + pImgSecHdr[i].VirtualAddress), pImgSecHdr[i].SizeOfRawData, dwProtection, &dwOldProtection)) {
			printf("[!] VirtualProtect Failed: %d \n", GetLastError());
			return FALSE;
		}
	}

	return TRUE;
}


/*
	Resolves a named export's address from a loaded PE image's export directory
	Walks AddressOfNames to match cFuncName, then uses the parallel ordinal array to index into AddressOfFunctions for the resolved address 

	pEntryExportDataDir -> Pointer to the export data directory entry
	pPeBaseAddr      -> Base address of the PE image in memory
	cFuncName           -> Name of the exported function to resolve
*/
PVOID GetDllFunctionAddrs(_In_ PIMAGE_DATA_DIRECTORY pEntryExportDataDir, _In_ ULONG_PTR pPeBaseAddr, _In_ LPCSTR cFuncName) {

	if (!pEntryExportDataDir || !pPeBaseAddr || !cFuncName)
		return NULL;

	// Nothing to search if the export directory is empty
	if (!pEntryExportDataDir->VirtualAddress || !pEntryExportDataDir->Size)
		return NULL;

	PIMAGE_EXPORT_DIRECTORY pImgExportDir = (PIMAGE_EXPORT_DIRECTORY)(pPeBaseAddr + pEntryExportDataDir->VirtualAddress);
	
	// Parallel arrays: name[i] and ordinal[i] map to the same export; ordinal[i] indexes into the function address array
	PDWORD                  pFunctionNameArray = (PDWORD)(pPeBaseAddr + pImgExportDir->AddressOfNames);
	PDWORD                  pFunctionAddrArray = (PDWORD)(pPeBaseAddr + pImgExportDir->AddressOfFunctions);
	PWORD                   pFunctionOrdinalArray = (PWORD)(pPeBaseAddr + pImgExportDir->AddressOfNameOrdinals);

	// Export directory bounds -- used to detect forwarded exports
	ULONG_PTR uExportDirStart = pPeBaseAddr + pEntryExportDataDir->VirtualAddress;
	ULONG_PTR uExportDirEnd = uExportDirStart + pEntryExportDataDir->Size;

	// Iterate named exports only 
	for (DWORD i = 0; i < pImgExportDir->NumberOfNames; i++) {
		CHAR* pFunctionName = (CHAR*)(pPeBaseAddr + pFunctionNameArray[i]);
		PVOID  pFunctionAddr = (PVOID)(pPeBaseAddr + pFunctionAddrArray[pFunctionOrdinalArray[i]]);

		if (strcmp(cFuncName, pFunctionName) != 0)
			continue;

		// If the address points inside the export directory, it's a forwarded export
		// (e.g. "NTDLL.RtlAllocateHeap") -- not resolving forwarders here
		if ((ULONG_PTR)pFunctionAddr >= uExportDirStart &&
			(ULONG_PTR)pFunctionAddr < uExportDirEnd) {
			printf("[!] Forwarded export skipped: %s -> %s\n", cFuncName, (CHAR*)pFunctionAddr);
			return NULL;
		}

		return pFunctionAddr;
	}

	return NULL;
}

// Executing DLL
typedef BOOL(WINAPI* DLLMAIN)(HINSTANCE, DWORD, LPVOID);

// Executing Exe
typedef BOOL(WINAPI* MAIN)();

// Forward Declaration
SIZE_T CharStringToWCharString(_Inout_ PWCHAR Destination, _In_ PCHAR Source, _In_ SIZE_T MaximumAllowed);


VOID FixArgs(_In_ OPTIONAL LPCSTR cArgs) {
	PRTL_USER_PROCESS_PARAMETERS pParam = ((PPEB)__readgsqword(0x60))->ProcessParameters;

	// Zero out the existing command line buffer before rebuilding it
	RtlSecureZeroMemory(pParam->CommandLine.Buffer, pParam->CommandLine.MaximumLength);

	if (cArgs) {
		WCHAR* wArgumentsToPass = NULL;
		WCHAR* wNewCommand = NULL;

		// Allocate buffer for wide-char conversion of cArgs (+ null terminator)
		SIZE_T cArgsLen = strlen(cArgs);
		SIZE_T wArgsByteLen = (cArgsLen + 1) * sizeof(WCHAR);

		wArgumentsToPass = (WCHAR*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, wArgsByteLen);
		if (!wArgumentsToPass) {
			printf("HeapAlloc Failed: %d \n", GetLastError());
			return;
		}

		// Convert cArgs from char to wchar
		CharStringToWCharString(wArgumentsToPass, cArgs, wArgsByteLen);

		// Allocate buffer for full command line: "<ImagePath>" <args> + null terminator
		// ImagePathName.Length is in bytes -- divide by sizeof(WCHAR) for char count
		// +4 for the two quote chars, one space, and one null terminator
		SIZE_T wNewCmdCharCount = (pParam->ImagePathName.Length / sizeof(WCHAR)) + wcslen(wArgumentsToPass) + 4;

		wNewCommand = (WCHAR*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, wNewCmdCharCount * sizeof(WCHAR));
		if (!wNewCommand) {
			printf("HeapAlloc Failed: %d \n", GetLastError());
			HeapFree(GetProcessHeap(), 0, wArgumentsToPass);  // don't leak on failure
			return;
		}

		// Build the full command line string
		wsprintfW(wNewCommand, L"\"%s\" %s", pParam->ImagePathName.Buffer, wArgumentsToPass);

		// Overwrite PEB command line with the new value
		lstrcpyW(pParam->CommandLine.Buffer, wNewCommand);

		// Length excludes null terminator, MaximumLength includes it
		pParam->CommandLine.Length = (USHORT)(wcslen(pParam->CommandLine.Buffer) * sizeof(WCHAR));
		pParam->CommandLine.MaximumLength = pParam->CommandLine.Length + sizeof(WCHAR);

		HeapFree(GetProcessHeap(), 0, wArgumentsToPass);
		HeapFree(GetProcessHeap(), 0, wNewCommand);
		return;
	}

	// No arguments provided -- overwrite command line with image path only
	lstrcpyW(pParam->CommandLine.Buffer, pParam->ImagePathName.Buffer);

	// Length excludes null terminator, MaximumLength includes it
	pParam->CommandLine.Length = (USHORT)(wcslen(pParam->CommandLine.Buffer) * sizeof(WCHAR));
	pParam->CommandLine.MaximumLength = pParam->CommandLine.Length + sizeof(WCHAR);
}

/*
	Local PE
	pPeHdr -> Pointer to PE_HDRS structure populated by ParsePEStruct()
	cExportedFuncName -> Name of exported function to be executed
	cArgs -> The command line arguments that are passed to PE image
*/
BOOL LocalPE_Injection(_In_ PPE_HDR pPeHdr, _In_ OPTIONAL LPCSTR cExportedFuncName, _In_ OPTIONAL LPCSTR cArgs) {
	// Check if parameter is not zero
	if (!pPeHdr)
		return FALSE;

	BOOL	bState = FALSE;
	PBYTE	pPeBaseAddr = NULL;
	PVOID	pEntryPoint = NULL;
	PVOID	pExportedFuncAddress = NULL;

	// Allocate memory for PE image
	pPeBaseAddr = VirtualAlloc(NULL, pPeHdr->pImgNtHdr->OptionalHeader.SizeOfImage, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
	if (!pPeBaseAddr) {
		printf("[!] VirtualAlloc Failed: %d\n", GetLastError());
		goto _CleanUp;
	}
	// Copy PE headers into the allocated region
	memcpy(pPeBaseAddr, pPeHdr->pFileBuffer, pPeHdr->pImgNtHdr->OptionalHeader.SizeOfHeaders);
	// Copy each section from the raw file buffer into its virtual address
	for (DWORD i = 0; i < pPeHdr->pImgNtHdr->FileHeader.NumberOfSections; i++) {
		if (!pPeHdr->pImgSecHdr[i].SizeOfRawData || !pPeHdr->pImgSecHdr[i].VirtualAddress)
			continue;
		memcpy(
			pPeBaseAddr + pPeHdr->pImgSecHdr[i].VirtualAddress,
			pPeHdr->pFileBuffer + pPeHdr->pImgSecHdr[i].PointerToRawData,
			pPeHdr->pImgSecHdr[i].SizeOfRawData
		);
	}
	// Update IAT
	if (!InitIAT(pPeHdr->pEntryImportDataDir, pPeBaseAddr))
		goto _CleanUp;
	
	// Apply relocation blocks
	if (!ApplyReloc(pPeHdr->pEntryBaseRelocDataDir, pPeBaseAddr, pPeHdr->pImgNtHdr->OptionalHeader.ImageBase))
		goto _CleanUp;
	
	// Apply correct memory permissions per section
	if (!FixMemPermissions(pPeBaseAddr, pPeHdr->pImgNtHdr, pPeHdr->pImgSecHdr))
		goto _CleanUp;
	
	// Fetch exported function address (if DLL and export name was provided)
	if (pPeHdr->pEntryExportDataDir->Size && pPeHdr->pEntryExportDataDir->VirtualAddress && cExportedFuncName)
		pExportedFuncAddress = GetDllFunctionAddrs(pPeHdr->pEntryExportDataDir, (ULONG_PTR)pPeBaseAddr, cExportedFuncName);
	
	// Fix user arguments if provided
	if (cArgs)
		FixArgs(cArgs);
	
	// Register exception handlers of the loaded PE (if present)
	if (pPeHdr->pEntryExceptionDataDir->Size) {
		// Retrieve the runtime function table from the exception directory
		PIMAGE_RUNTIME_FUNCTION_ENTRY pImgRuntimeFuncEntry = (PIMAGE_RUNTIME_FUNCTION_ENTRY)(pPeBaseAddr + pPeHdr->pEntryExceptionDataDir->VirtualAddress);
		// Register the function table so SEH unwinding works correctly
		if (!RtlAddFunctionTable(pImgRuntimeFuncEntry, (pPeHdr->pEntryExceptionDataDir->Size / sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)), (DWORD64)pPeBaseAddr)) {
			printf("[!] RtlAddFunctionTable Failed: %d\n", GetLastError());
			goto _CleanUp;
		}
	}
	
	// Execute TLS callbacks (if present)
	if (pPeHdr->pEntryTLSDataDir->Size) {
		// Retrieve the TLS directory from the loaded image
		PIMAGE_TLS_DIRECTORY pImgTlsDirectory = (PIMAGE_TLS_DIRECTORY)(pPeBaseAddr + pPeHdr->pEntryTLSDataDir->VirtualAddress);
		// Get the TLS callback array -- terminated by a NULL entry
		PIMAGE_TLS_CALLBACK* pImgTlsCallback = (PIMAGE_TLS_CALLBACK*)(pImgTlsDirectory->AddressOfCallBacks);
		// Invoke each TLS callback in order until NULL terminator is reached
		for (int i = 0; pImgTlsCallback[i] != NULL; i++) {
			// Third parameter is reserved and must be NULL per the TLS callback ABI
			pImgTlsCallback[i]((LPVOID)pPeBaseAddr, DLL_PROCESS_ATTACH, NULL);
		}
	}
	
	// Resolve entry point from the optional header
	pEntryPoint = (PVOID)(pPeBaseAddr + pPeHdr->pImgNtHdr->OptionalHeader.AddressOfEntryPoint);
	
	// If DLL -- call DllMain then optionally execute the exported function on a new thread
	if (pPeHdr->bIsDLLFile) {
		DLLMAIN	pDllMain = (DLLMAIN)pEntryPoint;
		HANDLE	hThread = NULL;
		// Execute DllMain with DLL_PROCESS_ATTACH
		pDllMain((HINSTANCE)pPeBaseAddr, DLL_PROCESS_ATTACH, NULL);
		// Execute exported function on a separate thread (if resolved)
		if (pExportedFuncAddress)
			hThread = CreateThread(NULL, 0x00, pExportedFuncAddress, NULL, 0x00, NULL);
		if (hThread)
			WaitForSingleObject(hThread, INFINITE);
		bState = TRUE;
	}
	
	// If EXE -- transfer execution directly to the entry point
	else {
		MAIN pMain = (MAIN)pEntryPoint;
		bState = pMain();
	}

_CleanUp:
	if (pPeBaseAddr && !bState)
		VirtualFree(pPeBaseAddr, 0, MEM_RELEASE);
	return bState;
}



/*
	Parses command line arguments and extracts the -pe, -fptr, and -parm flags

	argc      -> Argument count from main
	argv      -> Argument vector from main
	ppe_arg   -> Receives pointer to the -pe value (path to PE file)
	pfptr_arg -> Receives pointer to the -fptr value (exported function name)
	pparm_arg -> Receives pointer to the heap-allocated -parm string, or NULL if not provided
*/
BOOL HandleCmdLineArgs(int argc, char* argv[], char** ppe_arg, char** pfptr_arg, char** pparm_arg) {
	if (!argv || !ppe_arg || !pfptr_arg || !pparm_arg)
		return FALSE;

	char* pe_arg = NULL;
	char* fptr_arg = NULL;
	char* parm_arg = NULL;

	// Heap-allocated stack buffer
	char* parm_buffer = (char*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, 1024 * 2);
	if (!parm_buffer)
		return FALSE;

	// Walk argv starting at 1 
	for (int i = 1; i < argc; i++) {
		if (strcmp(argv[i], "-pe") == 0 && i + 1 < argc) {
			pe_arg = argv[++i];
		}
		else if (strcmp(argv[i], "-fptr") == 0 && i + 1 < argc) {
			fptr_arg = argv[++i];
		}
		else if (strcmp(argv[i], "-parm") == 0 && i + 1 < argc) {
			parm_arg = parm_buffer;
			strcpy_s(parm_arg, 1024 * 2, argv[++i]);
			
			// Collect remaining argv tokens until the next flag or end of args
			while (i + 1 < argc && argv[i + 1][0] != '-') {
				strcat_s(parm_arg, 1024 * 2, " ");
				strcat_s(parm_arg, 1024 * 2, argv[++i]);
			}
		}
	}

	*ppe_arg = pe_arg;
	*pfptr_arg = fptr_arg;
	*pparm_arg = parm_arg;

	// Free the buffer if -parm was never encountered -- no point keeping it
	if (!parm_arg)
		HeapFree(GetProcessHeap(), 0, parm_buffer);

	return TRUE;
}


// Tanks to: VX-Underground
SIZE_T CharStringToWCharString(_Inout_ PWCHAR Destination, _In_ PCHAR Source, _In_ SIZE_T MaximumAllowed) {
	INT Length = (INT)MaximumAllowed;

	while (--Length >= 0)
	{
		if (!(*Destination++ = *Source++))
			return MaximumAllowed - Length - 1;
	}

	return MaximumAllowed - Length;
}

static inline LPCSTR GetFileName(_In_ LPCSTR cPath) {
	LPCSTR pSlash = strrchr(cPath, '\\');
	return pSlash ? pSlash + 1 : cPath;
}


int main(int argc, char* argv[]) {

	char* pe_arg = NULL;
	char* fptr_arg = NULL;
	char* parm_arg = NULL;

	// Parse command line flags into their respective output pointers
	HandleCmdLineArgs(argc, argv, &pe_arg, &fptr_arg, &parm_arg);

	LPCSTR cBinaryName = GetFileName(argv[0]);


	if (!pe_arg) {
		printf("[!] Usage: %s -pe <Input PE> -fptr <*optional*: Exported function> -parm <*optional*: Arguments to pass> \n", cBinaryName);
		printf("\t<i> e.g. %s -pe DllMsgBox.dll -fptr HelloWorld \n", cBinaryName);
		printf("\t<i> e.g. %s -pe DllMsgBox.dll \n", cBinaryName);
		printf("\t<i> e.g. %s -pe ExeMsgBox.exe \n", cBinaryName);
		printf("\t<i> e.g. %s -pe ExeArgs.exe -parm this is a command line arg \n", cBinaryName);
		printf("\n\n");
		
		return -1;
	}

	PBYTE		pFileBuffer = NULL;
	DWORD		dwFileSize = 0x00;
	PE_HDR		pPeHdrStruct = { 0 };

	// Read PE from disk into a heap buffer
	if (!ReadFileDisk(pe_arg, &pFileBuffer, &dwFileSize))
		return -1;

	// Parse raw buffer into PE_HDR structure
	if (!ParsePEStruct(&pPeHdrStruct, pFileBuffer, dwFileSize))
		return -1;
	
	// Map and execute the PE image in the current process
	LocalPE_Injection(&pPeHdrStruct, fptr_arg, parm_arg);

	return 0;
}