Injection and server changes.

4 files changed, 26 insertions, 34 deletions

diff --git a/client/src/client/client.h b/client/src/client/client.h
index 9d18345..b76ac9f 100644
--- a/client/src/client/client.h
+++ b/client/src/client/client.h
@@ -29,7 +29,7 @@ namespace tcp {
 	};
 
 	enum client_state {
-		idle = 0, logged_in, waiting
+		idle = 0, logged_in, waiting, injected
 	};
 
 	enum login_result {
diff --git a/client/src/injection/mapper.cpp b/client/src/injection/mapper.cpp
index 7fcb8b3..570155e 100644
--- a/client/src/injection/mapper.cpp
+++ b/client/src/injection/mapper.cpp
@@ -12,20 +12,31 @@ void mmap::thread(tcp::client& client) {
 	util::fetch_processes();
 
 	auto needle = std::find_if(util::process_list.begin(), util::process_list.end(), [&](util::process& proc) {
-		return proc.name() == "notepad++.exe";
+		return strcmp(proc.name().c_str(), "notepad++.exe") == 0;
 	});
 
 	while (needle == util::process_list.end()) {
 		std::this_thread::sleep_for(std::chrono::seconds(5));
+
 		util::fetch_processes();
+		
+		io::logger->info("size {}", util::process_list.size());
+
 		io::logger->info("waiting for process..");
+
 		needle = std::find_if(util::process_list.begin(), util::process_list.end(), [&](util::process& proc) {
-			return proc.name() == "notepad++.exe";
+			return strcmp(proc.name().c_str(), "notepad++.exe") == 0;
 		});
 	}
 
-	needle->open();
-	needle->enum_modules();
+	if (!needle->open()) {
+		return;
+	}
+
+	if (!needle->enum_modules()) {
+		io::logger->error("failed to enum {} modules", needle->name());
+		return;
+	}
 
 	auto image = needle->allocate(client.mapper_data.image_size, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
 	if (!image) {
@@ -101,5 +112,5 @@ void mmap::thread(tcp::client& client) {
 
 	io::logger->info("done");
 
-	std::cin.get();
+	client.state = tcp::client_state::injected;
 }
\ No newline at end of file
diff --git a/client/src/injection/process.cpp b/client/src/injection/process.cpp
index 9b05963..6552d3c 100644
--- a/client/src/injection/process.cpp
+++ b/client/src/injection/process.cpp
@@ -5,17 +5,10 @@
 #include "process.h"
 
 util::process::process(const SYSTEM_PROCESS_INFORMATION* info) {
-	std::wstring name;
-	name.resize(info->ImageName.Length);
-
-	std::memcpy(&name[0], &info->ImageName.Buffer[0], name.size());
-
-	name.assign(name.data());
+	std::wstring name(info->ImageName.Buffer, info->ImageName.Length / sizeof(wchar_t));
 
 	m_name = util::wide_to_multibyte(name);
 	m_id = int(info->UniqueProcessId);
-
-	m_handle = INVALID_HANDLE_VALUE;
 }
 
 util::process::~process() {
@@ -42,20 +35,11 @@ bool util::process::open() {
 
 	io::logger->info("opened handle to {}.", m_name);
 
-	if (!enum_modules()) {
-		io::logger->error("failed to enumerate process modules.");
-		return false;
-	}
-
 	return true;
 }
 
 bool util::process::read(const uintptr_t addr, void* data, size_t size) {
 	static auto nt_read = g_syscalls.get<native::NtReadVirtualMemory>("NtReadVirtualMemory");
-	if (!m_handle) {
-		io::logger->error("invalid {} handle.", m_name);
-		return false;
-	}
 
 	ULONG read;
 	auto status = nt_read(m_handle, reinterpret_cast<void*>(addr), data, size, &read);
@@ -69,10 +53,6 @@ bool util::process::read(const uintptr_t addr, void* data, size_t size) {
 
 bool util::process::write(const uintptr_t addr, void* data, size_t size) {
 	static auto nt_write = g_syscalls.get<native::NtWiteVirtualMemory>("NtWriteVirtualMemory");
-	if (!m_handle) {
-		io::logger->error("invalid {} handle.", m_name);
-		return false;
-	}
 
 	ULONG wrote;
 	auto status = nt_write(m_handle, reinterpret_cast<void*>(addr), data, size, &wrote);
@@ -246,10 +226,6 @@ uintptr_t util::process::peb() {
 
 uintptr_t util::process::allocate(size_t size, uint32_t type, uint32_t protection) {
 	static auto nt_alloc = g_syscalls.get<native::NtAllocateVirtualMemory>("NtAllocateVirtualMemory");
-	if (!m_handle) {
-		io::logger->error("invalid {} handle.", m_name);
-		return {};
-	}
 
 	void* alloc = nullptr;
 	SIZE_T win_size = size;
@@ -368,7 +344,9 @@ bool util::process::close() {
 std::vector<util::process> util::process_list;
 
 bool util::fetch_processes() {
-	auto info = g_syscalls.get<native::NtQuerySystemInformation>("NtQuerySystemInformation");
+	process_list.clear();
+
+	static auto info = g_syscalls.get<native::NtQuerySystemInformation>("NtQuerySystemInformation");
 
 	std::vector<char> buf(1);
 	ULONG size_needed = 0;
diff --git a/client/src/main.cpp b/client/src/main.cpp
index 6248460..0b6580c 100644
--- a/client/src/main.cpp
+++ b/client/src/main.cpp
@@ -21,6 +21,7 @@ int main(int argc, char* argv[]) {
 	t.detach();
 
 	std::thread t1{ mmap::thread, std::ref(client) };
+	t1.detach();
 
 	client.start("127.0.0.1", 6666);
 
@@ -170,6 +171,8 @@ int main(int argc, char* argv[]) {
 
 	}
 
-	t1.join();
-	std::cin.get();
+	while (client.state != tcp::client_state::injected) {
+		std::this_thread::sleep_for(std::chrono::seconds(1));
+	}
+
 }