Started security.

4 files changed, 244 insertions, 72 deletions

diff --git a/client/src/injection/mapper.cpp b/client/src/injection/mapper.cpp
index c7f771c..28b5409 100644
--- a/client/src/injection/mapper.cpp
+++ b/client/src/injection/mapper.cpp
@@ -5,46 +5,74 @@
 #include "mapper.h"
 
 void mmap::thread(tcp::client& client) {
-	while (client.state != tcp::client_state::imports_ready) {
-		std::this_thread::sleep_for(std::chrono::seconds(1));
-	}
+	while (client) {
+		if (client.state != tcp::client_state::imports_ready) {
+			std::this_thread::sleep_for(std::chrono::seconds(5));
+			continue;
+		}
 
-	if (client.selected_game.x64) {
-		map64(client);
+		if (client.selected_game.x64) {
+			map64(client);
 
-		return;
-	}
+			break;
+		}
 
-	map32(client);
+		map32(client);
+		break;
+	}
 }
 
 void mmap::map32(tcp::client& client) {
-	util::system_data_t dat;
-	util::fetch_system_data(dat);
+	std::vector<util::process_data_t> dat;
+	if (!util::fetch_processes(dat)) {
+		io::log_error("failed to fetch processes.");
+		client.shutdown();
+		return;
+	}
 
-	auto needle = std::find_if(dat.processes.begin(), dat.processes.end(), [&](util::process_data_t& dat) {
+	auto needle = std::find_if(dat.begin(), dat.end(), [&](util::process_data_t& dat) {
 		return dat.name == client.selected_game.process_name;
+	});
+
+	io::log("waiting for {}.", client.selected_game.process_name);
+	while (needle == dat.end()) {
+		std::this_thread::sleep_for(std::chrono::seconds(5));
+		if (!client) {
+			return;
+		}
+
+		if (!util::fetch_processes(dat)) {
+			io::log_error("failed to fetch processes.");
+			client.shutdown();
+			return;
+		}
+
+		needle = std::find_if(dat.begin(), dat.end(), [&](util::process_data_t& dat) {
+			return dat.name == client.selected_game.process_name;
 		});
 
-	if (needle == dat.processes.end()) {
-		io::log_error("failed to find process.");
-		return;
+		io::log(".");
 	}
 
+	io::log("found!");
+
 	util::process<uint32_t> proc(*needle);
 
 	if (!proc.open()) {
+		client.shutdown();
 		return;
 	}
 
 	if (!proc.enum_modules()) {
 		io::log_error("failed to enum {} modules", proc.name());
+		client.shutdown();
 		return;
 	}
 
 	auto image = proc.allocate(client.mapper_data.image_size, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
 	if (!image) {
 		io::log_error("failed to allocate memory for image.");
+		client.shutdown();
 		return;
 	}
 
@@ -75,11 +103,16 @@ void mmap::map32(tcp::client& client) {
 
 	io::log("please wait...");
 	while (client.state != tcp::client_state::image_ready) {
+		if (!client) {
+			return;
+		}
+
 		std::this_thread::sleep_for(std::chrono::seconds(1));
 	}
 
 	if (!proc.write(image, client.mapper_data.image.data(), client.mapper_data.image.size())) {
 		io::log_error("failed to write image.");
+		client.shutdown();
 		return;
 	}
 	client.mapper_data.image.clear();
@@ -97,6 +130,8 @@ void mmap::map32(tcp::client& client) {
 	auto code = proc.allocate(shellcode.size(), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
 	if (!proc.write(code, shellcode.data(), shellcode.size())) {
 		io::log_error("failed to write shellcode.");
+		client.shutdown();
+
 		return;
 	}
 
@@ -111,35 +146,64 @@ void mmap::map32(tcp::client& client) {
 	client.state = tcp::client_state::injected;
 
 	io::log("done");
+
+	client.shutdown();
 }
 
 void mmap::map64(tcp::client& client) {
-	util::system_data_t dat;
-	util::fetch_system_data(dat);
+	std::vector<util::process_data_t> dat;
+	if (!util::fetch_processes(dat)) {
+		io::log_error("failed to fetch processes.");
+		client.shutdown();
+		return;
+	}
 
-	auto needle = std::find_if(dat.processes.begin(), dat.processes.end(), [&](util::process_data_t& dat) {
+	auto needle = std::find_if(dat.begin(), dat.end(), [&](util::process_data_t& dat) {
 		return dat.name == client.selected_game.process_name;
 	});
 
-	if (needle == dat.processes.end()) {
-		io::log_error("failed to find process.");
-		return;
+	io::log("waiting for {}.", client.selected_game.process_name);
+	while (needle == dat.end()) {
+		std::this_thread::sleep_for(std::chrono::seconds(5));
+
+		if (!client) {
+			return;
+		}
+
+		if (!util::fetch_processes(dat)) {
+			io::log_error("failed to fetch processes.");
+			client.shutdown();
+			return;
+		}
+
+		needle = std::find_if(dat.begin(), dat.end(), [&](util::process_data_t& dat) {
+			return dat.name == client.selected_game.process_name;
+		});
+
+		io::log(".");
 	}
 
+	io::log("found!");
+
 	util::process<uint64_t> proc(*needle);
 
 	if (!proc.open()) {
+		client.shutdown();
 		return;
 	}
 
 	if (!proc.enum_modules()) {
 		io::log_error("failed to enum {} modules", proc.name());
+
+		client.shutdown();
+
 		return;
 	}
 
 	auto image = proc.allocate(client.mapper_data.image_size, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
 	if (!image) {
 		io::log_error("failed to allocate memory for image.");
+		client.shutdown();
 		return;
 	}
 
@@ -170,11 +234,16 @@ void mmap::map64(tcp::client& client) {
 
 	io::log("please wait...");
 	while (client.state != tcp::client_state::image_ready) {
+		if (!client) {
+			return;
+		}
+		
 		std::this_thread::sleep_for(std::chrono::seconds(1));
 	}
 
 	if (!proc.write(image, client.mapper_data.image.data(), client.mapper_data.image.size())) {
 		io::log_error("failed to write image.");
+		client.shutdown();
 		return;
 	}
 	client.mapper_data.image.clear();
@@ -193,6 +262,7 @@ void mmap::map64(tcp::client& client) {
 	auto code = proc.allocate(shellcode.size(), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
 	if (!proc.write(code, shellcode.data(), shellcode.size())) {
 		io::log_error("failed to write shellcode.");
+		client.shutdown();
 		return;
 	}
 
@@ -207,4 +277,6 @@ void mmap::map64(tcp::client& client) {
 	client.state = tcp::client_state::injected;
 
 	io::log("done");
+
+	client.shutdown();
 }
\ No newline at end of file
diff --git a/client/src/injection/pe.h b/client/src/injection/pe.h
index f246108..fab41bb 100644
--- a/client/src/injection/pe.h
+++ b/client/src/injection/pe.h
@@ -3,14 +3,93 @@
 #include "../util/native.h"
 
 namespace pe {
+
+#pragma pack(push, 4)
+	struct reloc_entry_t {
+		uint16_t                    offset : 12;
+		uint16_t                    type : 4;
+	};
+
+	struct reloc_block_t {
+		uint32_t                    base_rva;
+		uint32_t                    size_block;
+		reloc_entry_t               entries[1];   // Variable length array
+
+
+		inline reloc_block_t* get_next() { return (reloc_block_t*)((char*)this + this->size_block); }
+		inline uint32_t num_entries() { return (reloc_entry_t*)get_next() - &entries[0]; }
+	};
+
+	struct image_named_import_t
+	{
+		uint16_t            hint;
+		char                name[1];
+	};
+
+#pragma pack(push, 8)
+	struct image_thunk_data_x64_t
+	{
+		union
+		{
+			uint64_t        forwarder_string;
+			uint64_t        function;
+			uint64_t        address;                   // -> image_named_import_t
+			struct
+			{
+				uint64_t    ordinal : 16;
+				uint64_t    _reserved0 : 47;
+				uint64_t    is_ordinal : 1;
+			};
+		};
+	};
+#pragma pack(pop)
+
+	struct image_thunk_data_x86_t
+	{
+		union
+		{
+			uint32_t        forwarder_string;
+			uint32_t        function;
+			uint32_t        address;                   // -> image_named_import_t
+			struct
+			{
+				uint32_t    ordinal : 16;
+				uint32_t    _reserved0 : 15;
+				uint32_t    is_ordinal : 1;
+			};
+		};
+	};
+#pragma pack(pop)
+
+	template<bool x64,
+		typename base_type = typename std::conditional<x64, image_thunk_data_x64_t, image_thunk_data_x86_t>::type>
+	struct image_thunk_data_t : base_type {};
+
+	template<bool x64, typename base_type = typename std::conditional<x64, IMAGE_NT_HEADERS64, IMAGE_NT_HEADERS32>::type>
+	struct nt_headers_t : base_type {};
+
+	struct import_t {
+		std::string name;
+		uint32_t rva;
+	};
+
+	struct section_t {
+		std::string name;
+		size_t size;
+		size_t v_size;
+		uint32_t rva;
+		uint32_t va;
+	};
+
 	class virtual_image {
 		std::unordered_map<std::string, uintptr_t> m_exports;
+		std::vector<section_t> m_sections;
 
 		IMAGE_NT_HEADERS64* m_nt;
 		uintptr_t m_base;
 	public:
 		virtual_image() : m_nt{ nullptr }, m_base{ 0 } {};
-		virtual_image(const std::string_view mod) : m_base{0}, m_nt{nullptr} {
+		virtual_image(const std::string_view mod) : m_base{ 0 }, m_nt{ nullptr } {
 			auto peb = util::peb();
 			if (!peb) return;
 
@@ -30,7 +109,7 @@ namespace pe {
 					m_base = uintptr_t(entry->DllBase);
 					auto dos = reinterpret_cast<IMAGE_DOS_HEADER*>(m_base);
 
-					m_nt = reinterpret_cast<native::nt_headers_t<true>*>(m_base + dos->e_lfanew);
+					m_nt = reinterpret_cast<nt_headers_t<true>*>(m_base + dos->e_lfanew);
 
 					parse_exports();
 					break;
@@ -38,6 +117,24 @@ namespace pe {
 			}
 		}
 
+		virtual_image(const uintptr_t base) : m_base{ base }, m_nt{ nullptr } {
+			auto dos = reinterpret_cast<IMAGE_DOS_HEADER*>(m_base);
+
+			m_nt = reinterpret_cast<nt_headers_t<true>*>(m_base + dos->e_lfanew);
+		}
+
+		void parse_sections() {
+			auto secs = IMAGE_FIRST_SECTION(m_nt);
+			const size_t n = m_nt->FileHeader.NumberOfSections;
+
+			for (size_t i = 0; i < n; i++) {
+				auto sec = secs[i];
+
+				auto name = reinterpret_cast<const char*>(sec.Name);
+				m_sections.emplace_back(section_t{ name, sec.SizeOfRawData, sec.Misc.VirtualSize, sec.PointerToRawData, sec.VirtualAddress });
+			}
+		};
+
 		void parse_exports() {
 			auto dir = m_nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
 			auto exp =
@@ -60,20 +157,11 @@ namespace pe {
 		}
 
 		auto& exports() { return m_exports; }
-		operator bool() { return m_base != 0; }
-	};
-
-	struct import_t {
-		std::string name;
-		uint32_t rva;
-	};
+		auto &base() { return m_base; }
+		auto& nt() { return m_nt; }
+		auto& sections() { return m_sections; }
 
-	struct section_t {
-		std::string name;
-		size_t size;
-		size_t v_size;
-		uint32_t rva;
-		uint32_t va;
+		operator bool() { return m_base != 0; }
 	};
 
 	template <bool x64 = false>
@@ -83,9 +171,9 @@ namespace pe {
 		std::unordered_map<std::string, std::vector<import_t>> m_imports;
 		std::vector<section_t> m_sections;
 
-		native::nt_headers_t<x64>* m_nt;
+		nt_headers_t<x64>* m_nt;
 		IMAGE_DOS_HEADER* m_dos;
-		std::vector<std::pair<uint32_t, native::reloc_entry_t>> m_relocs;
+		std::vector<std::pair<uint32_t, reloc_entry_t>> m_relocs;
 
 	public:
 		image() = default;
@@ -94,7 +182,7 @@ namespace pe {
 
 
 			m_dos = reinterpret_cast<IMAGE_DOS_HEADER*>(m_buffer.data());
-			m_nt = reinterpret_cast<native::nt_headers_t<x64>*>(m_buffer.data() + m_dos->e_lfanew);
+			m_nt = reinterpret_cast<nt_headers_t<x64>*>(m_buffer.data() + m_dos->e_lfanew);
 
 			load();
 		}
@@ -134,8 +222,8 @@ namespace pe {
 			if (!reloc_dir.Size) return;
 
 			const auto ptr = rva_to_ptr(reloc_dir.VirtualAddress);
-			auto block = reinterpret_cast<native::reloc_block_t*>(ptr);
-			
+			auto block = reinterpret_cast<reloc_block_t*>(ptr);
+
 			while (block->base_rva) {
 				for (size_t i = 0; i < block->num_entries(); ++i) {
 					auto entry = block->entries[i];
@@ -156,11 +244,11 @@ namespace pe {
 			for (uint32_t i = 0; i < table->Name; i = table->Name, ++table) {
 				auto mod_name = std::string(reinterpret_cast<char*>(rva_to_ptr(table->Name)));
 
-				auto thunk = reinterpret_cast<native::image_thunk_data_t<x64>*>(rva_to_ptr(table->OriginalFirstThunk));
+				auto thunk = reinterpret_cast<image_thunk_data_t<x64>*>(rva_to_ptr(table->OriginalFirstThunk));
 
 				auto step = x64 ? sizeof(uint64_t) : sizeof(uint32_t);
 				for (uint32_t index = 0; thunk->address; index += step, ++thunk) {
-					auto named_import = reinterpret_cast<native::image_named_import_t*>(rva_to_ptr(thunk->address));
+					auto named_import = reinterpret_cast<image_named_import_t*>(rva_to_ptr(thunk->address));
 
 					if (!thunk->is_ordinal) {
 						import_t data;
@@ -178,7 +266,7 @@ namespace pe {
 		void copy(std::vector<char>& out) {
 			out.resize(m_nt->OptionalHeader.SizeOfImage);
 
-			std::memcpy(&out[0], &m_buffer[0], 4096);
+			std::memcpy(&out[0], &m_buffer[0], m_nt->OptionalHeader.SizeOfHeaders);
 
 			for (auto& sec : m_sections) {
 				std::memcpy(&out[sec.va], &m_buffer[sec.rva], sec.size);
@@ -209,4 +297,15 @@ namespace pe {
 		auto& relocs() const { return m_relocs; }
 		auto& sections() const { return m_sections; }
 	};
+
+
+	static virtual_image &ntdll() {
+		static virtual_image img{};
+		if (!img) {
+			img = virtual_image("ntdll.dll");
+		}
+		return img;
+	}
+
+	void get_all_modules(std::unordered_map<std::string, virtual_image>& modules);
 };  // namespace pe
\ No newline at end of file
diff --git a/client/src/injection/process.cpp b/client/src/injection/process.cpp
index 38a676b..3f3c96e 100644
--- a/client/src/injection/process.cpp
+++ b/client/src/injection/process.cpp
@@ -42,7 +42,7 @@ bool util::base_process::read(const uintptr_t addr, void* data, size_t size) {
 }
 
 bool util::base_process::write(const uintptr_t addr, void* data, size_t size) {
-	static auto nt_write = g_syscalls.get<native::NtWiteVirtualMemory>("NtWriteVirtualMemory");
+	static auto nt_write = g_syscalls.get<native::NtWriteVirtualMemory>("NtWriteVirtualMemory");
 
 	ULONG wrote;
 	auto status = nt_write(m_handle, reinterpret_cast<void*>(addr), data, size, &wrote);
@@ -175,7 +175,7 @@ bool util::process<T>::enum_modules() {
 
 template<typename T>
 uintptr_t util::process<T>::peb() {
-	constexpr bool is64 = sizeof(T) == sizeof(uint64_t);
+	constexpr bool is64 = std::is_same_v<T, uint64_t>;
 	if (is64) {
 		native::PROCESS_EXTENDED_BASIC_INFORMATION proc_info;
 		proc_info.Size = sizeof(proc_info);
@@ -209,8 +209,8 @@ uintptr_t util::process<T>::module_export(const uintptr_t base, const std::strin
 	if (dos.e_magic != IMAGE_DOS_SIGNATURE)
 		return {};
 
-	constexpr bool is64 = sizeof(T) == sizeof(uint64_t);
-	native::nt_headers_t<is64> nt{};
+	constexpr bool is64 = std::is_same_v<T, uint64_t>;
+	pe::nt_headers_t<is64> nt{};
 	if (!read(base + dos.e_lfanew, &nt, sizeof(nt))) {
 		io::log_error("failed to read nt header for {}", m_name);
 		return {};
@@ -291,7 +291,7 @@ uintptr_t util::process<T>::module_export(const uintptr_t base, const std::strin
 template<typename T>
 uintptr_t util::process<T>::map(const std::string_view module_name) {
 	std::string mod{ module_name };
-	if (g_apiset(mod)) {
+	if (g_apiset.find(mod)) {
 		io::log("resolved {} -> {}", module_name, mod);
 	}
 
@@ -302,7 +302,7 @@ uintptr_t util::process<T>::map(const std::string_view module_name) {
 
 	io::log("mapping {}", module_name);
 
-	constexpr bool is64 = sizeof(T) == sizeof(uint64_t);
+	constexpr bool is64 = std::is_same_v<T, uint64_t>;
 	std::string path{ is64 ? "C:\\Windows\\System32\\" : "C:\\Windows\\SysWOW64\\" };
 	path.append(mod);
 
@@ -331,7 +331,9 @@ uintptr_t util::process<T>::map(const std::string_view module_name) {
 	for (auto& [mod, funcs] : img.imports()) {
 		for (auto& func : funcs) {
 			auto addr = module_export(map(mod), func.name);
+
 			//io::log("{}:{}->{:x}", mod, func.name, addr);
+
 			*reinterpret_cast<T*>(&remote_image[func.rva]) = addr;
 		}
 	}
@@ -343,6 +345,7 @@ uintptr_t util::process<T>::map(const std::string_view module_name) {
 	}
 
 	io::log("{}->{:x}", mod, base);
+
 	m_modules[mod] = base;
 
 	return base;
@@ -352,14 +355,15 @@ uintptr_t util::process<T>::map(const std::string_view module_name) {
 template class util::process<uint64_t>;
 template class util::process<uint32_t>;
 
-bool util::fetch_system_data(system_data_t& out) {
+bool util::fetch_processes(std::vector<process_data_t>& out, bool threads /*= false*/) {
 	static auto info = g_syscalls.get<native::NtQuerySystemInformation>("NtQuerySystemInformation");
 
+	out.clear();
 	std::vector<uint8_t> buf(1);
 
 	ULONG size_needed = 0;
 	NTSTATUS status;
-	while ((status = info(native::SystemProcessInformation, buf.data(), buf.size(), &size_needed)) == STATUS_INFO_LENGTH_MISMATCH) {
+	while ((status = info(SystemProcessInformation, buf.data(), buf.size(), &size_needed)) == STATUS_INFO_LENGTH_MISMATCH) {
 		buf.resize(size_needed);
 	};
 
@@ -368,28 +372,29 @@ bool util::fetch_system_data(system_data_t& out) {
 		return false;
 	}
 
-	std::vector<thread_data_t> threads;
-	std::vector<process_data_t> processes;
 	auto pi = reinterpret_cast<SYSTEM_PROCESS_INFORMATION*>(buf.data());
 
 	while (pi->NextEntryOffset) {
 		std::wstring name(pi->ImageName.Buffer, pi->ImageName.Length / sizeof(wchar_t));
-		processes.emplace_back(process_data_t{ util::wide_to_multibyte(name), int(pi->UniqueProcessId) });
+		process_data_t data{int(pi->UniqueProcessId), util::wide_to_multibyte(name)};
+		
+		if (!threads) {
+			out.emplace_back(data);
+			
+			pi = reinterpret_cast<SYSTEM_PROCESS_INFORMATION*>(uintptr_t(pi) + pi->NextEntryOffset);
+			continue;
+		}
 
+		std::vector<thread_data_t> threads;
 		auto ti = reinterpret_cast<SYSTEM_THREAD_INFORMATION*>(uintptr_t(pi) + sizeof(SYSTEM_PROCESS_INFORMATION));
-
 		for (auto i = 0; i < pi->NumberOfThreads; ++i) {
-			auto dat = ti[i];
-			threads.emplace_back(thread_data_t{ int(dat.ClientId.UniqueProcess), uintptr_t(dat.ClientId.UniqueThread), dat.ThreadState });
+			auto thread = ti[i];
+			threads.emplace_back(thread_data_t{ thread.ClientId.UniqueThread, thread.ThreadState });
 		}
 
 		pi = reinterpret_cast<SYSTEM_PROCESS_INFORMATION*>(uintptr_t(pi) + pi->NextEntryOffset);
 	}
 
-
-	out.processes = std::move(processes);
-	out.threads = std::move(threads);
-
 	return true;
 }
 
@@ -400,7 +405,8 @@ bool util::fetch_process_handles(const int pid, std::vector<handle_info_t>& out)
 
 	ULONG size_needed = 0;
 	NTSTATUS status;
-	while ((status = info(native::SystemHandleInformation, buf.data(), buf.size(), &size_needed)) == STATUS_INFO_LENGTH_MISMATCH) {
+	/* SystemHandleInformation */
+	while ((status = info(static_cast<SYSTEM_INFORMATION_CLASS>(16), buf.data(), buf.size(), &size_needed)) == STATUS_INFO_LENGTH_MISMATCH) {
 		buf.resize(size_needed);
 	};
 
diff --git a/client/src/injection/process.h b/client/src/injection/process.h
index e5d3f18..74cfccd 100644
--- a/client/src/injection/process.h
+++ b/client/src/injection/process.h
@@ -25,19 +25,14 @@ namespace util {
 		auto& id() { return m_id; }
 	};
 
-	struct process_data_t {
-		std::string name;
-		int id;
-	};
-
 	struct thread_data_t {
-		int id;
-		uintptr_t handle;
+		HANDLE handle;
 		uint32_t state;
 	};
 
-	struct system_data_t {
-		std::vector<process_data_t> processes;
+	struct process_data_t {
+		int id;
+		std::string name;
 		std::vector<thread_data_t> threads;
 	};
 
@@ -64,6 +59,6 @@ namespace util {
 		uint32_t obj_type;
 	};
 
-	bool fetch_system_data(system_data_t& out);
+	bool fetch_processes(std::vector<process_data_t>& out, bool threads = false);
 	bool fetch_process_handles(const int pid, std::vector<handle_info_t>& out);
 };