Some bug fixes and small improvements

-bugs fixed: leaked handles, handle comparisons
-loging macro

mmap/kernelmode_proc_handler.cpp

@@ -1,23 +1,26 @@
 #include "kernelmode_proc_handler.hpp"
 
-bool kernelmode_proc_handler::is_attached() {
-	return handle;
-}
+kernelmode_proc_handler::kernelmode_proc_handler() 
+	:handle{ INVALID_HANDLE_VALUE }, pid{ 0 } {}
+
+kernelmode_proc_handler::~kernelmode_proc_handler() { if (is_attached()) CloseHandle(handle); }
+
+bool kernelmode_proc_handler::is_attached() { return handle != INVALID_HANDLE_VALUE; }
 
 bool kernelmode_proc_handler::attach(const char* proc_name) {
 	bool is_admin{ false };
-	HANDLE hToken{ NULL };
-	if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken)) {
+	HANDLE token_handle{ NULL };
+	if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token_handle)) {
 		TOKEN_ELEVATION token;
 		DWORD size = sizeof(TOKEN_ELEVATION);
-		if (GetTokenInformation(hToken, TokenElevation, &token, sizeof(TOKEN_ELEVATION), &size)) {
+		if (GetTokenInformation(token_handle, TokenElevation, &token, sizeof(TOKEN_ELEVATION), &size)) {
 			is_admin = true;
 		}
-		CloseHandle(hToken);
+		CloseHandle(token_handle);
 	}
 
 	if (!is_admin) {
-		logger::log_error("Launch as admin");
+		LOG_ERROR("Launch as admin");
 		return false;
 	}
 
@@ -28,7 +31,7 @@ bool kernelmode_proc_handler::attach(const char* proc_name) {
 		FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 
 	if (handle == INVALID_HANDLE_VALUE) {
-		logger::log_error("Load the driver first");
+		LOG_ERROR("Load the driver first");
 		return false;
 	}
 

mmap/kernelmode_proc_handler.hpp

@@ -39,6 +39,9 @@ class kernelmode_proc_handler final : public process_handler {
 	HANDLE handle;
 	uint32_t pid;
 public:
+	kernelmode_proc_handler();
+
+	~kernelmode_proc_handler();
 
 	virtual bool is_attached() override;
 

mmap/logger.hpp

@@ -1,8 +1,5 @@
-#include <iostream>
-#include <string>
+#include <cstdio>
 
-namespace logger {
-	void log(const std::string& message);
-	void log_error(const std::string& message);
-	void log_address(const std::string& name, uint64_t value);
-}
+#define LOG_ERROR(str, ...) fprintf(stderr,"ERROR: " str "\n", ##__VA_ARGS__)
+
+#define LOG(str, ...) fprintf(stdout, str "\n", ##__VA_ARGS__)

mmap/main.cpp

@@ -1,8 +1,4 @@
-#include <iostream>
-#include "mmap.hpp"
-#include <thread>
-#include <chrono>
-#include <array>
+#include "mmap.hpp"
 
 int main(int argc, char **argv) { 
 	mmap mapper(INJECTION_TYPE::KERNEL);
@@ -16,7 +12,7 @@ int main(int argc, char **argv) {
 	if (!mapper.inject())
 		return 1;
 
-	std::cout << "\nPress any key to close.\n";
+	LOG("\nPress any key to close.");
 	std::getchar();
 
 	return 0;

mmap/mmap.cpp

@@ -1,28 +1,28 @@
 #include "mmap.hpp"
 
-mmap::mmap(INJECTION_TYPE type) { 
-	if(type == INJECTION_TYPE::KERNEL)
-		proc = new kernelmode_proc_handler();
+mmap::mmap(INJECTION_TYPE type) {
+	if (type == INJECTION_TYPE::KERNEL)
+		proc = std::make_unique<kernelmode_proc_handler>();
 	else
-		proc = new usermode_proc_handler();
+		proc = std::make_unique<usermode_proc_handler>();
 }
 
 bool mmap::attach_to_process(const char* process_name) {
 	this->process_name = process_name;
 	if (!proc->attach(process_name)) {
-		logger::log_error("Unable to attach to process!");
+		LOG_ERROR("Unable to attach to process!");
 		return false;
 	}
 
-	std::cout << "Attached to process " << process_name << " successfully...\n";
+	LOG("Attached to process %s successfully...", process_name); 
 	return true;
 }
 
 bool mmap::load_dll(const char* file_name) {
 	std::ifstream f(file_name, std::ios::binary | std::ios::ate);
 
 	if (!f) {
-		logger::log_error("Unable to open DLL file!");
+		LOG_ERROR("Unable to open DLL file!");
 		return false;
 	}
 
@@ -44,12 +44,12 @@ bool mmap::load_dll(const char* file_name) {
 bool mmap::inject() {
 
 	if (!proc->is_attached()) {
-		logger::log_error("Not attached to process!");
+		LOG_ERROR("Not attached to process!");
 		return false;
 	}
 
 	if (!raw_data) {
-		logger::log_error("Data buffer is empty!");
+		LOG_ERROR("Data buffer is empty!");
 		return false;
 	}
 
@@ -109,14 +109,14 @@ bool mmap::inject() {
 	IMAGE_DOS_HEADER *dos_header{ (IMAGE_DOS_HEADER *)raw_data };
 
 	if (dos_header->e_magic != IMAGE_DOS_SIGNATURE) {
-		logger::log_error("Invalid DOS header signature!");
+		LOG_ERROR("Invalid DOS header signature!");
 		return false;
 	}
 
 	IMAGE_NT_HEADERS *nt_header{ (IMAGE_NT_HEADERS *)(&raw_data[dos_header->e_lfanew]) };
 
 	if (nt_header->Signature != IMAGE_NT_SIGNATURE) {
-		logger::log_error("Invalid NT header signature!");
+		LOG_ERROR("Invalid NT header signature!");
 		return false;
 	}
 
@@ -125,30 +125,30 @@ bool mmap::inject() {
 									   PAGE_EXECUTE_READWRITE) };
 
 	if (!base) {
-		logger::log_error("Unable to allocate memory for the image!");
+		LOG_ERROR("Unable to allocate memory for the image!");
 		return false;
 	}
 
-	logger::log_address("Image base", base);
+	LOG("Image base: 0x%p", base);
 
 	uint64_t stub_base{ proc->virtual_alloc(sizeof(dll_stub),
 											MEM_COMMIT | MEM_RESERVE,
 											PAGE_EXECUTE_READWRITE) };
 
 	if (!stub_base) {
-		logger::log_error("Unable to allocate memory for the stub!");
+		LOG_ERROR("Unable to allocate memory for the stub!");
 		return false;
 	}
 
-	logger::log_address("Stub base", stub_base);
+	LOG("Stub base: 0x%p", stub_base);
 
 	PIMAGE_IMPORT_DESCRIPTOR import_descriptor{ (PIMAGE_IMPORT_DESCRIPTOR)get_ptr_from_rva(
 												(uint64_t)(nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress),
 												nt_header,
 												raw_data) };
 
 	if (nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size) {
-		logger::log("Solving imports...");
+		LOG("Solving imports...");
 		solve_imports(raw_data, nt_header, import_descriptor);
 	}
 
@@ -158,7 +158,7 @@ bool mmap::inject() {
 																		raw_data)};
 
 	if (nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size) {
-		logger::log("Solving relocations..."); 
+		LOG("Solving relocations..."); 
 		solve_relocations((uint64_t) raw_data,
 						  base,
 						  nt_header,
@@ -168,18 +168,18 @@ bool mmap::inject() {
 
 
 	if (!parse_imports()) {
-		logger::log_error("Unable to parse imports!");
+		LOG_ERROR("Unable to parse imports!");
 		return false;
 	} 
 
 	uint64_t iat_function_ptr{ imports["TranslateMessage"] };
 	if (!iat_function_ptr) { 
-		logger::log_error("Cannot find import");
+		LOG_ERROR("Cannot find import");
 		return false;
 	}
 
 	uint64_t orginal_function_addr{ read_memory<uint64_t>(iat_function_ptr) };
-	logger::log_address("IAT function pointer", iat_function_ptr);
+	LOG("IAT function pointer: 0x%p", iat_function_ptr);
 
 	*(uint64_t*)(dll_stub + 0x18) = iat_function_ptr;
 	*(uint64_t*)(dll_stub + 0x22) = orginal_function_addr;
@@ -193,7 +193,7 @@ bool mmap::inject() {
 
 	proc->write_memory(base, (uintptr_t)raw_data, nt_header->FileHeader.SizeOfOptionalHeader + sizeof(nt_header->FileHeader) + sizeof(nt_header->Signature));
 
-	logger::log("Mapping PE sections...");
+	LOG("Mapping PE sections...");
 	map_pe_sections(base, nt_header);
 
 	uint64_t entry_point{ (uint64_t)base + nt_header->OptionalHeader.AddressOfEntryPoint };
@@ -208,17 +208,17 @@ bool mmap::inject() {
 	call rax
 	*/
 
-	logger::log_address("Entry point", entry_point);	
+	LOG("Entry point: 0x%p", entry_point);	
 
 	proc->write_memory(stub_base, (uintptr_t)dll_stub, sizeof(dll_stub));
 
-	auto old_protect = proc->virtual_protect(iat_function_ptr, sizeof(uint64_t), PAGE_READWRITE);
+	proc->virtual_protect(iat_function_ptr, sizeof(uint64_t), PAGE_READWRITE);
 	proc->write_memory(iat_function_ptr, (uintptr_t)&stub_base, sizeof(uint64_t));
-	
-	logger::log("Injected successfully!");
+
+	LOG("Injected successfully!");
 
 	system("Pause");
-	proc->virtual_protect(iat_function_ptr, sizeof(uint64_t), old_protect);
+	proc->virtual_protect(iat_function_ptr, sizeof(uint64_t), PAGE_READONLY);
 
 	delete [] raw_data;
 	return true;
@@ -325,7 +325,7 @@ uint64_t mmap::get_proc_address(const char* module_name, const char* func) {
 bool mmap::parse_imports() {
 	auto base{ proc->get_module_base(process_name.c_str()) };
 	if (!base) {
-		logger::log_error("Cannot get module base");
+		LOG_ERROR("Cannot get module base");
 		return false;
 	}
 

mmap/mmap.hpp

@@ -14,7 +14,7 @@ enum INJECTION_TYPE{
 };
 
 class mmap {
-	process_handler *proc;
+	std::unique_ptr<process_handler> proc;
 	std::string process_name;
 	std::map<std::string, uint64_t> imports;
 	uint8_t *raw_data;

mmap/process_handler.hpp

@@ -12,6 +12,8 @@
 
 class process_handler { 
 public:
+	virtual ~process_handler() { };
+
 	virtual bool is_attached() = 0;
 
 	virtual bool attach(const char* proc_name) = 0;

mmap/usermode_proc_handler.cpp

@@ -1,31 +1,36 @@
 #include "usermode_proc_handler.hpp"
 
-bool usermode_proc_handler::is_attached() {
-	return handle;
-}
+usermode_proc_handler::usermode_proc_handler() 
+	:handle{ NULL }, pid{ 0 } {}
+
+usermode_proc_handler::~usermode_proc_handler() { if (handle) CloseHandle(handle); }
+
+bool usermode_proc_handler::is_attached() {	return handle; }
 
 bool usermode_proc_handler::attach(const char* proc_name) {
 	while (!is_process_running(proc_name, pid))
 		std::this_thread::sleep_for(std::chrono::seconds(1));
 
 	handle = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE, FALSE, pid);
 
-	if (!handle)
-		return false;
-
-	return true;
+	return handle;
 }
 
 uint64_t usermode_proc_handler::get_module_base(const std::string &module_name) {
-	MODULEENTRY32 module_entry;
-	HANDLE snapshot{ CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid) };
+	MODULEENTRY32 module_entry{};
 	module_entry.dwSize = sizeof(MODULEENTRY32);
-	Module32First(snapshot, &module_entry);
-	do {
-		if (!_stricmp(module_entry.szModule, module_name.c_str()))
-			return (uint64_t)module_entry.hModule;
-		module_entry.dwSize = sizeof(MODULEENTRY32);
-	} while (Module32Next(snapshot, &module_entry));
+	auto snapshot{ CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid) };
+	if (snapshot == INVALID_HANDLE_VALUE)
+		return false;
+	if (Module32First(snapshot, &module_entry)) {
+		do {
+			if (!_stricmp(module_entry.szModule, module_name.c_str())) {
+				CloseHandle(snapshot);
+				return (uint64_t)module_entry.hModule;
+			}
+			module_entry.dwSize = sizeof(MODULEENTRY32);
+		} while (Module32Next(snapshot, &module_entry));
+	}
 	CloseHandle(snapshot);
 	return NULL;
 }

mmap/usermode_proc_handler.hpp

@@ -2,9 +2,12 @@
 #include "process_handler.hpp"
 
 class usermode_proc_handler final : public process_handler {
-	uint32_t pid;
 	HANDLE handle;
+	uint32_t pid;
 public:
+	usermode_proc_handler();
+
+	~usermode_proc_handler();
 
 	virtual bool is_attached() override;
 

mmap/utils.cpp

@@ -5,10 +5,8 @@ bool is_process_running(const char* process_name, uint32_t& pid) {
 	process_entry.dwSize = sizeof(PROCESSENTRY32);
 	pid = 0;
 	auto snapshot{ CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL) };
-
-	if (!snapshot)
+	if (snapshot == INVALID_HANDLE_VALUE)
 		return false;
-
 	if (Process32First(snapshot, &process_entry)) {
 		do {
 			if (!strcmp(process_name, process_entry.szExeFile)) {
@@ -17,8 +15,7 @@ bool is_process_running(const char* process_name, uint32_t& pid) {
 				return true;
 			}
 		} while (Process32Next(snapshot, &process_entry));
-	}
-
+	} 
 	CloseHandle(snapshot);
 	return false;
 }