#include "manualmap.hpp"

void inject::c_map::initialize( int pid ) {
	m_handle = OpenProcess( PROCESS_ALL_ACCESS, 0, pid );
}

void inject::c_map::initialize( HANDLE process ) {
	m_handle = process;
}

void inject::c_map::write( uintptr_t address, void* data, size_t size ) {
	WriteProcessMemory( m_handle, ( void* )address, data, size, nullptr );
}

uintptr_t inject::c_map::allocate( size_t size ) {
	void* allocation = VirtualAllocEx( m_handle, 0, size, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE );
	m_allocations.push_back( allocation );

	return uintptr_t( allocation );
}

void inject::c_map::free_allocated_regions( ) {
	for( auto& it : m_allocations ) {
		VirtualFreeEx( m_handle, it, 0, MEM_FREE );
	}

	m_allocations.clear( );
}

void inject::c_map::inject( uintptr_t interfaces ) {
	HMODULE mod_data = ( HMODULE )m_inject_data.data( );

	IMAGE_DOS_HEADER dos_hdr;
	IMAGE_NT_HEADERS nt_hdrs;

	dos_hdr = *( decltype( dos_hdr )* )( mod_data );
	nt_hdrs = *( decltype( nt_hdrs )* )( uintptr_t( mod_data ) + dos_hdr.e_lfanew );

	auto allocation = allocate( nt_hdrs.OptionalHeader.SizeOfImage );

	auto size_of_headers = nt_hdrs.OptionalHeader.SizeOfHeaders;
	auto num_of_sections = nt_hdrs.FileHeader.NumberOfSections;

	m_allocation = malloc( nt_hdrs.OptionalHeader.SizeOfImage );
	memset( m_allocation, 0, nt_hdrs.OptionalHeader.SizeOfImage );
	memcpy( m_allocation, mod_data, size_of_headers );
	write( allocation, m_allocation, size_of_headers );

	auto sections = ( IMAGE_SECTION_HEADER* )( ( uintptr_t )( mod_data ) + dos_hdr.e_lfanew + sizeof( IMAGE_NT_HEADERS ) );
	for( size_t i{ }; i < num_of_sections; ++i ) {
		auto section = sections[ i ];
		uintptr_t address = ( uintptr_t )m_allocation + section.VirtualAddress;
		memcpy( ( void* )address, 
			( void* )( uintptr_t( mod_data ) + section.PointerToRawData ), 
			( size_t )section.SizeOfRawData );
	}
	
	auto base = nt_hdrs.OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_BASERELOC ].VirtualAddress;
	auto base_reloc = ( IMAGE_BASE_RELOCATION* )( ( uintptr_t )m_allocation + base );
	auto delta = allocation - nt_hdrs.OptionalHeader.ImageBase;

	while( base_reloc->VirtualAddress ) {
		if( base_reloc->SizeOfBlock >= sizeof( IMAGE_BASE_RELOCATION ) ) {
			size_t count = ( base_reloc->SizeOfBlock - sizeof( IMAGE_BASE_RELOCATION ) ) / sizeof( uint16_t );

			auto list = ( uint16_t* )( base_reloc + 1 );

			uintptr_t* ptr{ };
			for( size_t i{ }; i < count; ++i ) {
				if( list[ i ] ) {
					ptr = ( uintptr_t* )( ( uintptr_t )( m_allocation ) + ( base_reloc->VirtualAddress + ( list[ i ] & 0xfff ) ) );
					*ptr += delta;
				}
			}
		}

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

	write( allocation, m_allocation, nt_hdrs.OptionalHeader.SizeOfImage );
	free( m_allocation );

	auto shellcode_allocation = allocate( 0x1000 );

	img_data_t img_data{
		allocation,
		nt_hdrs.OptionalHeader.ImageBase,
		nt_hdrs.OptionalHeader.AddressOfEntryPoint,
		nt_hdrs.OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_BASERELOC ].VirtualAddress,
		nt_hdrs.OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_IMPORT ].VirtualAddress,
		( uintptr_t )( LoadLibraryA ),
		( uintptr_t )( GetProcAddress ),
		interfaces
	};

	write( shellcode_allocation, &img_data, sizeof( img_data_t ) );

	auto loader_code_start = shellcode_allocation + sizeof( img_data_t );
	auto loader_code_size = ( size_t )( ( uintptr_t )( &dummy_func_1 ) - ( uintptr_t )( &loader_shellcode ) );

	write( loader_code_start, &loader_shellcode, loader_code_size );
	auto thread = CreateRemoteThread( m_handle, nullptr, 0, 
		( LPTHREAD_START_ROUTINE )loader_code_start, 
		( void* )shellcode_allocation, 0, 0 );

	WaitForSingleObject( thread, INFINITE );
	ulong_t exit{ };
	GetExitCodeThread( thread, &exit );
}