#include <RemoteCode/RemoteProcess.hpp>

namespace RemoteCode
{
	// RemoteModule implementation
	RemoteModule::RemoteModule(HANDLE Module) :
		m_Module(Module) {}

	// RemoteProcess implementation
	bool RemoteProcess::Start(const char *ProcessName)
	{
		void *Toolhelp = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

		if(!Toolhelp)
			return false;

		PROCESSENTRY32 ProcessEntry{};
		ProcessEntry.dwSize = sizeof PROCESSENTRY32;

		if(!Process32First(Toolhelp, &ProcessEntry))
			return false;

		while(Process32Next(Toolhelp, &ProcessEntry))
		{
			if(strstr(ProcessName, ProcessEntry.szExeFile))
			{
				CloseHandle(Toolhelp);

				// swoo
				m_ProcessId = ProcessEntry.th32ProcessID;
				m_Process = OpenProcess(PROCESS_ALL_ACCESS, false, ProcessEntry.th32ProcessID);
				return true;
			}
		}

		CloseHandle(Toolhelp);
		return false;
	}

	void RemoteProcess::ReadMemoryWrapper_Internal(void *Address, void *Data, size_t SizeOfData)
	{
		static auto ZwReadVirtualMemory = Syscalls->Find<long(__stdcall *)(void *, void *, void *, size_t, void *)>(FNV("ZwReadVirtualMemory"));
		ZwReadVirtualMemory(m_Process, Address, Data, SizeOfData, nullptr);
	}

	void RemoteProcess::WriteMemoryWrapper_Internal(void *Address, void *Data, size_t SizeOfData)
	{
		static auto ZwWriteVirtualMemory = Syscalls->Find<long(__stdcall *)(void *, void *, void *, size_t, void *)>(FNV("ZwWriteVirtualMemory"));
		ZwWriteVirtualMemory(m_Process, Address, Data, SizeOfData, nullptr);
	}

	void *RemoteProcess::Allocate(size_t AllocationSize)
	{
		void *AllocationAddress = nullptr;
		static auto ZwAllocateVirtualMemory = Syscalls->Find<long(__stdcall *)(void *, void *, uint32_t, size_t *, uint32_t, uint32_t)>(FNV("ZwAllocateVirtualMemory"));
		
		// :b:invoke the :b:unction :b:oi
		NTSTATUS Status = ZwAllocateVirtualMemory(
			m_Process,
			&AllocationAddress,
			0,
			&AllocationSize,
			MEM_COMMIT | MEM_RESERVE,
			PAGE_EXECUTE_READWRITE
		);

		if(!NT_SUCCESS(Status))
			return nullptr;

		return AllocationAddress;
	}

	RemoteModule RemoteProcess::FindModule(const char *ModuleName)
	{
		void *Toolhelp = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32, m_ProcessId);
		
		if(!Toolhelp)
			return RemoteModule{};
		
		MODULEENTRY32 ModuleEntry{};
		ModuleEntry.dwSize = sizeof MODULEENTRY32;
		
		if(!Module32First(Toolhelp, &ModuleEntry))
			return RemoteModule{};
		
		while(Module32Next(Toolhelp, &ModuleEntry))
		{
			//printf("%s\n", ModuleEntry.szModule);
			if(strstr(ModuleEntry.szModule, ModuleName))
			{
				CloseHandle(Toolhelp);
				return RemoteModule(ModuleEntry.hModule);
			}
		}
		
		CloseHandle(Toolhelp);
		return RemoteModule{};
	}
}