[REFACT] Refactored the template main

Author: hasherezade

project_template/main.cpp

@@ -1,48 +1,89 @@
+/**
+A demo of a basic manual PE loader - you can use it as a starting point for your own project,
+or delete it and start from scratch
+*/
+
 #include <Windows.h>
 #include <iostream>
 #include <tchar.h>
 
 #include <peconv.h> // include libPeConv header
 
-/**
-A demo of a basic manual PE loader - you can use it as a starting point for your own project,
-or delete it and start from the scratch
-*/
-int load_and_run(LPCTSTR pe_path)
+BYTE* g_Payload = nullptr;
+size_t g_PayloadSize = 0;
+
+// manually load the PE file using libPeConv
+bool load_payload(LPCTSTR pe_path)
 {
 	// manually load the PE file using libPeConv:
-	size_t v_size = 0;
+	size_t g_PayloadSize = 0;
 #ifdef LOAD_FROM_PATH
 	//if the PE is dropped on the disk, you can load it from the file:
-	BYTE* my_pe = peconv::load_pe_executable(pe_path, v_size);
-#else
+	g_Payload = peconv::load_pe_executable(pe_path, g_PayloadSize);
+#else // load from memory buffer
+	/*
+	in this example the memory buffer is first loaded from the disk,
+	but it can as well be fetch from resources, or a hardcoded buffer
+	*/
 	size_t bufsize = 0;
-	BYTE *buffer = peconv::load_file(pe_path, bufsize);
+	BYTE* buf = peconv::load_file(pe_path, bufsize);
+	if (buf) {
+		// if the file is NOT dropped on the disk, you can load it directly from a memory buffer:
+		g_Payload = peconv::load_pe_executable(buf, bufsize, g_PayloadSize);
 
-	// if the file is NOT dropped on the disk, you can load it directly from a memory buffer:
-	BYTE* my_pe = peconv::load_pe_executable(buffer, bufsize, v_size);
+		// at this point we can free the buffer with the raw payload:
+		peconv::free_file(buf); buf = nullptr;
+
+		if (!g_Payload) {
+			return false;
+		}
+	}
 #endif
-	if (!my_pe) {
+	return true;
+}
+
+int run_payload()
+{
+	if (!g_Payload) {
+		std::cerr << "[!] The payload is not loaded!\n";
 		return -1;
 	}
-
 	// if the loaded PE needs to access resources, you may need to connect it to the PEB:
-	peconv::set_main_module_in_peb((HMODULE)my_pe);
+	peconv::set_main_module_in_peb((HMODULE)g_Payload);
 
 	// if needed, you can run TLS callbacks before the Entry Point:
-	peconv::run_tls_callbacks(my_pe, v_size);
+	peconv::run_tls_callbacks(g_Payload, g_PayloadSize);
 
 	//calculate the Entry Point of the manually loaded module
-	DWORD ep_rva = peconv::get_entry_point_rva(my_pe);
+	DWORD ep_rva = peconv::get_entry_point_rva(g_Payload);
 	if (!ep_rva) {
+		std::cerr << "[!] Cannot fetch EP!\n";
 		return -2;
 	}
-	ULONG_PTR ep_va = ep_rva + (ULONG_PTR)my_pe;
-	//assuming that the payload is an EXE file (not DLL) this will be the simplest prototype of the main:
-	int(*new_main)() = (int(*)())ep_va;
+	const ULONG_PTR ep_va = ep_rva + (ULONG_PTR)g_Payload;
+
+	// run appropriate version of payload's main:
+	int ret = 0;
+	if (peconv::is_module_dll(g_Payload)) {
+		//the prototype of the DllMain fuction:
+		BOOL WINAPI _DllMain(
+			HINSTANCE hinstDLL,  // handle to DLL module
+			DWORD fdwReason,     // reason for calling function
+			LPVOID lpvReserved);  // reserved
+		auto new_main = reinterpret_cast<decltype(&_DllMain)>(ep_va);
+
+		// call the Entry Point of the manually loaded PE :
+		ret = new_main((HINSTANCE)g_Payload, DLL_PROCESS_ATTACH, 0);
+	}
+	else {
+		//the simplest prototype of the main fuction:
+		int basic_main(void);
+		auto new_main = reinterpret_cast<decltype(&basic_main)>(ep_va);
 
-	//call the Entry Point of the manually loaded PE:
-	return new_main();
+		//call the Entry Point of the manually loaded PE:
+		ret = new_main();
+	}
+	return ret;
 }
 
 int _tmain(int argc, LPTSTR argv[])
@@ -52,5 +93,9 @@ int _tmain(int argc, LPTSTR argv[])
 		return 0;
 	}
 	const LPTSTR pe_path = argv[1];
-	return load_and_run(pe_path);
+	if (!load_payload(pe_path)) {
+		return -1;
+	}
+	std::cout << "Payload loaded!\n";
+	return run_payload();
 }