#include <RemoteCode/RemoteMapper.hpp>

namespace RemoteCode
{
	bool RemoteMapper::Start(RemoteProcess Process)
	{
		m_Process = Process;
		return true;
	}

	bool RemoteMapper::WriteCodeToMap(ByteArray Code)
	{
		// Check if the PE file is valid.
		/*uint8_t			 *Buffer    = Code.data();
		IMAGE_DOS_HEADER *DosHeader = (IMAGE_DOS_HEADER *)Buffer;

		if(!DosHeader || DosHeader->e_magic != IMAGE_DOS_SIGNATURE)
			return false;

		IMAGE_NT_HEADERS *NtHeaders = (IMAGE_NT_HEADERS *)(Buffer + DosHeader->e_lfanew);

		if(!NtHeaders || NtHeaders->Signature != IMAGE_NT_SIGNATURE)
			return false;

		WRAP_IF_DEBUG(
			printf("[DEBUG] PE file validated!\n");
		);
		
		IMAGE_SECTION_HEADER *Sections = (IMAGE_SECTION_HEADER *)((uintptr_t)&NtHeaders->OptionalHeader + NtHeaders->FileHeader.SizeOfOptionalHeader);

		if(!Sections)
			return false;

		// Allocate map in process.
		m_Map = m_Process.Allocate(Code.size());

		if(!m_Map)
			return false;

		// Ghetto workaround so heap doesn't get corrupted (i'm lazy)
		size_t SizeOfCode = Code.size();
		m_Mapped = new uint8_t[SizeOfCode];

		uint16_t SectionCount = NtHeaders->FileHeader.NumberOfSections;
		
		// Write each section to the process.
		for(uint16_t n{}; n < SectionCount; n++)
		{
			uint32_t	VirtualAddress = Sections[n].VirtualAddress;
			uint32_t	PointerToData  = Sections[n].PointerToRawData;
			uint32_t	SizeOfData     = Sections[n].SizeOfRawData;

			WRAP_IF_DEBUG(
				printf("[DEBUG] Writing PE section #%d!\n", n);
				printf("[DEBUG] Virtual address: 0x%x\n", VirtualAddress);
				printf("[DEBUG] Pointer to data: 0x%x\n", PointerToData);
				printf("[DEBUG] Size of data   : %d bytes\n", SizeOfData);
			);

			// Write the data to map.
			uint8_t *Data     = (uint8_t *)((uintptr_t)Buffer + PointerToData);
			void    *ImageRva = (void *)((uintptr_t)&m_Mapped[0] + VirtualAddress);
			
			std::memcpy(ImageRva, Data, SizeOfData);
		}

		// Copy over the PE header (temporarily).
		uint32_t SizeOfHeaders = NtHeaders->OptionalHeader.SizeOfHeaders;
		std::memcpy(&m_Mapped[0], Buffer, SizeOfHeaders);

		WRAP_IF_DEBUG(
			printf("[DEBUG] Successfully copied over PE header!\n");
			printf("[DEBUG] Relocating image...\n");
		);

		// Relocate the image.
		IMAGE_BASE_RELOCATION *Reloc = (IMAGE_BASE_RELOCATION *)(&m_Mapped[0] + NtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);

		uintptr_t ImageDelta = (uintptr_t)m_Map - NtHeaders->OptionalHeader.ImageBase;

		while(Reloc->SizeOfBlock > 0)
		{
			// address in binary where we need to relocate an address used by code
			uintptr_t RelocationAddress     = (uintptr_t)&m_Mapped[Reloc->VirtualAddress];
			uintptr_t RelocationDataAddress = (uintptr_t)Reloc + sizeof IMAGE_BASE_RELOCATION;
			size_t    RelocationCount       = (Reloc->SizeOfBlock - sizeof IMAGE_BASE_RELOCATION) / sizeof uint16_t;

			for(size_t i = 0; i < RelocationCount; i++)
			{
				uint16_t RelocationData   = *(uint16_t *)(RelocationDataAddress + sizeof uint16_t * i);
				uint16_t RelocationType   = RelocationData >> 12;
				uint16_t RelocationOffset = RelocationData & 0x0FFF;

				WRAP_IF_DEBUG(
					printf("[DEBUG] Processing relocation at %llx!\n", RelocationAddress + RelocationOffset);
				);

				switch(RelocationType)
				{
					case IMAGE_REL_BASED_MIPS_JMPADDR:
					case IMAGE_REL_BASED_HIGH:
					case IMAGE_REL_BASED_LOW:
					case IMAGE_REL_BASED_HIGHLOW:
					{
						*(uint32_t*)(RelocationAddress + RelocationOffset) += ImageDelta;
						break;
					}

					case IMAGE_REL_BASED_ABSOLUTE:
					default:
						break;
				};
			}

			Reloc = (IMAGE_BASE_RELOCATION*)((uintptr_t)Reloc + Reloc->SizeOfBlock);
		}

		// Cripple the entire PE header.
		std::memset(&m_Mapped[0], 0, SizeOfHeaders);

		WRAP_IF_DEBUG(
			printf("[DEBUG] Successfully crippled PE header!\n");
		);

		// Write the mapped image to the process.
		m_Process.Write(m_Map, &m_Mapped[0], Code.size());*/

		return true;
	}

	bool RemoteMapper::ExecuteCodeFromMap()
	{
		WRAP_IF_DEBUG(
			printf("[DEBUG] Executing payload!\n");
		)

		return true;
	}
}