#include <windows.h>
#include <vector>
#include <cstdint>
#include <algorithm>
#include <functional>
#include <string>

#include "ReClassNET_Plugin.hpp"

bool IsValidMemoryRange(LPCVOID address, int length)
{
	const auto endAddress = static_cast<const uint8_t*>(address) + length;

	do
	{
		MEMORY_BASIC_INFORMATION info;
		if (!VirtualQuery(address, &info, sizeof(MEMORY_BASIC_INFORMATION)))
		{
			return false;
		}

		if (info.State != MEM_COMMIT)
		{
			return false;
		}

		switch (info.Protect)
		{
		case PAGE_EXECUTE_READ:
		case PAGE_EXECUTE_READWRITE:
		case PAGE_EXECUTE_WRITECOPY:
		case PAGE_READONLY:
		case PAGE_READWRITE:
		case PAGE_WRITECOPY:
			break;
		default:
			return false;
		}

		address = static_cast<uint8_t*>(info.BaseAddress) + info.RegionSize;
	} while (endAddress > address);

	return true;
}
//---------------------------------------------------------------------------
bool ReadMemory(LPCVOID address, std::vector<uint8_t>& buffer)
{
	if (!IsValidMemoryRange(address, static_cast<int>(buffer.size())))
	{
		return false;
	}

	std::memcpy(buffer.data(), address, buffer.size());

	return true;
}
//---------------------------------------------------------------------------
bool WriteMemory(LPVOID address, const std::vector<uint8_t>& buffer)
{
	if (!IsValidMemoryRange(address, static_cast<int>(buffer.size())))
	{
		return false;
	}

	DWORD oldProtect;
	if (VirtualProtect(address, buffer.size(), PAGE_EXECUTE_READWRITE, &oldProtect))
	{
		std::memcpy(address, buffer.data(), buffer.size());

		VirtualProtect(address, buffer.size(), oldProtect, nullptr);

		return true;
	}

	return false;
}
//---------------------------------------------------------------------------
void EnumerateRemoteSectionsAndModules(const std::function<void(RC_Pointer, RC_Pointer, std::wstring&&)>& moduleCallback, const std::function<void(RC_Pointer, RC_Pointer, SectionType, SectionCategory, SectionProtection, std::wstring&&, std::wstring&&)>& sectionCallback)
{
	std::vector<EnumerateRemoteSectionData> sections;

	// First enumerate all memory sections.
	MEMORY_BASIC_INFORMATION memInfo = { };
	memInfo.RegionSize = 0x1000;
	size_t address = 0;
	while (VirtualQuery(reinterpret_cast<LPCVOID>(address), &memInfo, sizeof(MEMORY_BASIC_INFORMATION)) != 0 && address + memInfo.RegionSize > address)
	{
		if (memInfo.State == MEM_COMMIT)
		{
			EnumerateRemoteSectionData section = {};
			section.BaseAddress = memInfo.BaseAddress;
			section.Size = memInfo.RegionSize;

			switch (memInfo.Protect & 0xFF)
			{
			case PAGE_EXECUTE:
				section.Protection = SectionProtection::Execute;
				break;
			case PAGE_EXECUTE_READ:
				section.Protection = SectionProtection::Execute | SectionProtection::Read;
				break;
			case PAGE_EXECUTE_READWRITE:
			case PAGE_EXECUTE_WRITECOPY:
				section.Protection = SectionProtection::Execute | SectionProtection::Read | SectionProtection::Write;
				break;
			case PAGE_NOACCESS:
				section.Protection = SectionProtection::NoAccess;
				break;
			case PAGE_READONLY:
				section.Protection = SectionProtection::Read;
				break;
			case PAGE_READWRITE:
			case PAGE_WRITECOPY:
				section.Protection = SectionProtection::Read | SectionProtection::Write;
				break;
			}
			if ((memInfo.Protect & PAGE_GUARD) == PAGE_GUARD)
			{
				section.Protection |= SectionProtection::Guard;
			}

			switch (memInfo.Type)
			{
			case MEM_IMAGE:
				section.Type = SectionType::Image;
				break;
			case MEM_MAPPED:
				section.Type = SectionType::Mapped;
				break;
			case MEM_PRIVATE:
				section.Type = SectionType::Private;
				break;
			}

			section.Category = section.Type == SectionType::Private ? SectionCategory::HEAP : SectionCategory::Unknown;

			sections.push_back(std::move(section));
		}
		address = reinterpret_cast<size_t>(memInfo.BaseAddress) + memInfo.RegionSize;
	}

	struct UNICODE_STRING
	{
		USHORT Length;
		USHORT MaximumLength;
		PWSTR  Buffer;
	};

	struct LDR_MODULE
	{
		LIST_ENTRY InLoadOrderModuleList;
		LIST_ENTRY InMemoryOrderModuleList;
		LIST_ENTRY InInitializationOrderModuleList;
		PVOID BaseAddress;
		PVOID EntryPoint;
		ULONG SizeOfImage;
		UNICODE_STRING FullDllName;
		UNICODE_STRING BaseDllName;
		ULONG Flags;
		SHORT LoadCount;
		SHORT TlsIndex;
		LIST_ENTRY HashTableEntry;
		ULONG TimeDateStamp;
	};

	struct PEB_LDR_DATA
	{
		ULONG Length;
		BOOLEAN Initialized;
		PVOID SsHandle;
		LIST_ENTRY InLoadOrderModuleList;
		LIST_ENTRY InMemoryOrderModuleList;
		LIST_ENTRY InInitializationOrderModuleList;
	};

	struct PEB
	{
		BOOLEAN InheritedAddressSpace;
		BOOLEAN ReadImageFileExecOptions;
		BOOLEAN BeingDebugged;
		BOOLEAN Spare;
		HANDLE Mutant;
		PVOID ImageBaseAddress;
		PEB_LDR_DATA *LoaderData;
	};

	// Second enumerate all modules.
#ifdef _WIN64
	const auto peb = reinterpret_cast<PEB*>(__readgsqword(0x60));
#else
	const auto peb = reinterpret_cast<PEB*>(__readfsdword(0x30));
#endif
	auto ldr = reinterpret_cast<LDR_MODULE*>(peb->LoaderData->InLoadOrderModuleList.Flink);
	while (ldr->BaseAddress != nullptr)
	{
		moduleCallback(static_cast<RC_Pointer>(ldr->BaseAddress), reinterpret_cast<RC_Pointer>(static_cast<intptr_t>(ldr->SizeOfImage)), ldr->FullDllName.Buffer);

		const auto it = std::lower_bound(std::begin(sections), std::end(sections), ldr->BaseAddress, [&sections](const auto& lhs, const LPVOID& rhs)
		{
			return lhs.BaseAddress < rhs;
		});

		const auto dosHeader = reinterpret_cast<IMAGE_DOS_HEADER*>(ldr->BaseAddress);
		const auto ntHeader = reinterpret_cast<IMAGE_NT_HEADERS*>(reinterpret_cast<intptr_t>(ldr->BaseAddress) + dosHeader->e_lfanew);

		int i = 0;
		for (auto sectionHeader = IMAGE_FIRST_SECTION(ntHeader); i < ntHeader->FileHeader.NumberOfSections; i++, sectionHeader++)
		{
			const auto sectionAddress = reinterpret_cast<intptr_t>(ldr->BaseAddress) + static_cast<intptr_t>(sectionHeader->VirtualAddress);
			for (auto j = it; j != std::end(sections); ++j)
			{
				if (sectionAddress >= reinterpret_cast<intptr_t>(j->BaseAddress) && sectionAddress < reinterpret_cast<intptr_t>(j->BaseAddress) + static_cast<intptr_t>(j->Size))
				{
					// Copy the name because it is not null padded.
					char buffer[IMAGE_SIZEOF_SHORT_NAME + 1] = { 0 };
					std::memcpy(buffer, sectionHeader->Name, IMAGE_SIZEOF_SHORT_NAME);

					if (std::strcmp(buffer, ".text") == 0 || std::strcmp(buffer, "code") == 0)
					{
						j->Category = SectionCategory::CODE;
					}
					else if (std::strcmp(buffer, ".data") == 0 || std::strcmp(buffer, "data") == 0 || std::strcmp(buffer, ".rdata") == 0 || std::strcmp(buffer, ".idata") == 0)
					{
						j->Category = SectionCategory::DATA;
					}

					size_t convertedChars = 0;
					mbstowcs_s(&convertedChars, j->Name, IMAGE_SIZEOF_SHORT_NAME, buffer, _TRUNCATE);
					std::memcpy(j->ModulePath, ldr->FullDllName.Buffer, sizeof(EnumerateRemoteSectionData::ModulePath));
					break;
				}
			}
		}

		ldr = reinterpret_cast<LDR_MODULE*>(ldr->InLoadOrderModuleList.Flink);
	}

	for (auto&& section : sections)
	{
		sectionCallback(section.BaseAddress, reinterpret_cast<RC_Pointer>(section.Size), section.Type, section.Category, section.Protection, section.Name, section.ModulePath);
	}
}
//---------------------------------------------------------------------------