#pragma once
#include <Windows.h>
#include <vector>
#include <memory>
#include "util.hpp"

enum DllSections_t {
	SECTION_TEXT,   //.text		| allocation + write( obviously )
	SECTION_RDATA,	//.rdata	| allocation + write
	SECTION_DATA,	//.data		| need to allocate
	SECTION_RSRC,	//.rsrc		| not needed
	SETCION_RELOC,	//.reloc	| will need to do on server
	SECTION_MAX
};

namespace inject
{
	struct img_data_t {
		uintptr_t m_base;
		uintptr_t m_image;
		uintptr_t m_entry;
		uintptr_t m_relocation;
		uintptr_t m_imports;
		uintptr_t m_loadlib;
		uintptr_t m_get_procaddr;
		uintptr_t m_interface_ptr;
	};

	using dllmain_t = int( __stdcall* )( void*, ulong_t, void* );

	static __declspec( naked ) ulong_t __stdcall loader_shellcode( void* address ) {
		__asm {
			push ebp
			mov ebp, esp
			sub esp, __LOCAL_SIZE
		}

		img_data_t* data; data = ( img_data_t* )address;

		uintptr_t base;			base = data->m_base;
		uintptr_t entry_point;  entry_point = base + data->m_entry;
		uintptr_t delta;		delta = base - data->m_image;

		IMAGE_BASE_RELOCATION* base_reloc;
		IMAGE_IMPORT_DESCRIPTOR* import_dir;

		base_reloc = ( IMAGE_BASE_RELOCATION* )( base + data->m_relocation );
		import_dir = ( IMAGE_IMPORT_DESCRIPTOR* )( base + data->m_imports );

		decltype( &LoadLibraryA ) loadlib;
		decltype( &GetProcAddress ) get_procaddr;

		loadlib = ( decltype( &LoadLibraryA ) )( data->m_loadlib );
		get_procaddr = ( decltype( &GetProcAddress ) )( data->m_get_procaddr );

		IMAGE_THUNK_DATA* orig_first_thunk;
		IMAGE_THUNK_DATA* first_thunk;

		uintptr_t name;
		HMODULE import_module;
		uintptr_t ordinal;
		uintptr_t import_fn;

		IMAGE_IMPORT_BY_NAME* import_;

		while( import_dir->Characteristics ) {
			orig_first_thunk = ( IMAGE_THUNK_DATA* )( base + import_dir->OriginalFirstThunk );
			first_thunk = ( IMAGE_THUNK_DATA* )( base + import_dir->FirstThunk );

			import_module = 0;

			name = base + import_dir->Name;
			__asm {
				push name
				call loadlib
				mov import_module, eax
			}

			if( !import_module ) {
				//return 0
				__asm mov eax, 0;
				goto END;
			}

			while( orig_first_thunk->u1.AddressOfData ) {
				if( orig_first_thunk->u1.Ordinal & IMAGE_ORDINAL_FLAG ) {
					ordinal = orig_first_thunk->u1.Ordinal & 0xffff;
					import_fn = 0;

					__asm {
						push ordinal
						push import_module
						call get_procaddr
						mov import_fn, eax
					}

					if( !import_fn ) {
						__asm mov eax, 0;
						goto END;
					}


					first_thunk->u1.Function = import_fn;
				}
				else {
					import_ = ( IMAGE_IMPORT_BY_NAME* )( base + orig_first_thunk->u1.AddressOfData );
					name = ( uintptr_t )( import_->Name );

					import_fn = 0;

					__asm {
						push name
						push import_module
						call get_procaddr
						mov import_fn, eax
					}

					if( !import_fn ) {
						__asm mov eax, 0;
						goto END;
					}

					first_thunk->u1.Function = import_fn;
				}

				++orig_first_thunk;
				++first_thunk;
			}

			++import_dir;
		}

		void* interface_ptr;
		interface_ptr = ( void* )( data->m_interface_ptr );

		dllmain_t fn; fn = reinterpret_cast< dllmain_t >( entry_point );
		fn( ( void* )base, DLL_PROCESS_ATTACH, interface_ptr );

		__asm mov eax, 1;

	END:
		__asm {
			mov esp, ebp
			pop ebp
			ret
		}
	}

	static ulong_t __stdcall dummy_func_1( ) { return 0; }

	class c_map {
		HANDLE					m_handle;
		std::vector< void* >	m_allocations;
		std::vector< uint8_t >	m_inject_data;
		void*					m_allocation;

		void write( uintptr_t address, void* data, size_t size );
		uintptr_t allocate( size_t size );

		void free_allocated_regions( );
	public:
		c_map( std::vector< uint8_t >& file ) : m_inject_data( file ) { };
		~c_map( ) { if( m_handle ) { CloseHandle( m_handle ); } }

		void initialize( int process_id );
		void initialize( HANDLE process );
		void inject( uintptr_t interfaces );
	};
}