#include 


/*
	encapsulate all the code and data to inject into its own section
	then just try to inject that section

	MSVC only?

	For a PE w/ IMAGE_NT_SIGNATURE 
*/ 

// compiler won't mix data and code... so /merge them later
#define my_code ".my_code" 
#define my_data ".my_data" 
#define my_section_name ".Inject"
#pragma section( my_data, read, write )
#pragma section( my_code, read, execute )
#pragma code_seg  ( my_code ) // executable code
#pragma const_seg ( my_data ) // constant data
#pragma data_seg  ( my_data ) // initialized data
#pragma bss_seg   ( my_data ) // uninitialized data
//#pragma init_seg

// some random data to inject
// IIRC for the game: disciples 2
enum race_id
{ 
	HUMAN = 0,
	UNDEAD,
	HERETIC,
	DWARF,
	NEUTRAL,
	ELF,
};

struct SomeClass
{
	void* p;			// 0x00
	DWORD UnknownA;		// 0x04
	DWORD UnknownB;		// 0x08 
	DWORD UnknownC;		// 0x0C
	DWORD UnknownD;		// 0x10
};

struct __declspec( align( 16 ) ) SomeOtherClass 
{
	void* p;           // 0x00
	SomeClass* sc;     // 0x04
	race_id id;        // 0x08
};

#define DATA( name, type, address ) static type& name = (*((type*)address));
#define FUNC( name, address, ret_type, call_convention, args ) ret_type(call_convention * const name)args = (ret_type(call_convention *)args) address;
#define PFNC( name, address, ret_type, call_convention, args ) static ret_type (call_convention *& name)args = *( (ret_type(call_convention **)args)address);

namespace Game
{
	DATA( undead,  SomeOtherClass, 0x00837290 );
	DATA( elf,     SomeOtherClass, 0x008372A0 );
	DATA( neutral, SomeOtherClass, 0x008372B0 );
	DATA( human,   SomeOtherClass, 0x008372C0 );
	DATA( heretic, SomeOtherClass, 0x008372D0 );
	DATA( dwarf,   SomeOtherClass, 0x008372E0 );

	FUNC( sub_0057ED2E, 0x0057ED2E, DWORD, __stdcall, ( SomeClass*, void*, char*, char* ) );
	FUNC( sub_0057EDA6, 0x0057EDA6, DWORD, __stdcall, ( SomeOtherClass*, SomeClass*, char*, char* ) );
}

SomeClass* __fastcall my_0057EB99( SomeClass* thisptr, DWORD unused, char* szDir, void* some_pointer )
{
	MessageBox(0,"weee","splat",MB_OK);
	char szLRacedbf[] = { "LRace.dbf" };
	thisptr->p = (void*) 0x006E7234;
	thisptr->UnknownA = 0;
	thisptr->UnknownB = 0;
	thisptr->UnknownC = 0;
	thisptr->UnknownD = 0;
	Game::sub_0057ED2E( thisptr, some_pointer, szDir, szLRacedbf );
	Game::sub_0057EDA6( &Game::undead,  thisptr, "L_UNDEAD",  szLRacedbf ); 
	Game::sub_0057EDA6( &Game::elf,     thisptr, "L_ELF",     szLRacedbf ); 
	Game::sub_0057EDA6( &Game::neutral, thisptr, "L_NEUTRAL", szLRacedbf ); 
	Game::sub_0057EDA6( &Game::human,   thisptr, "L_HUMAN",   szLRacedbf ); 
	Game::sub_0057EDA6( &Game::heretic, thisptr, "L_HERETIC", szLRacedbf ); 
	Game::sub_0057EDA6( &Game::dwarf,   thisptr, "L_DWARF",   szLRacedbf ); 
	return  thisptr;
}

#pragma bss_seg   ( ) // resets bss_seg   to ".bss"
#pragma data_seg  ( ) // resets data_seg  to ".data"
#pragma const_seg ( ) // resets const_seg to ".rdata"
#pragma code_seg  ( ) // resets code_seg  to ".text"
#pragma comment( linker, "/merge:" my_data "=" my_code ) // append my_date to my_code
#pragma comment( linker, "/merge:" my_code "=" my_section_name ) // Rename the merged section
#pragma comment( linker, "/section:" my_section_name ",RWE" )// Read, Write, and Execute Access


PIMAGE_SECTION_HEADER GetSectionHeaderByName( uintptr_t module_base, const char (§ion_name)[ IMAGE_SIZEOF_SHORT_NAME ] );
unsigned long FixUpSection( uintptr_t module_base, PIMAGE_SECTION_HEADER pSectionHeader, uintptr_t new_section_address );

// ... since the bottom 16 bits of a proper HINSTANCE are always zero, different components have "borrowed" those bits for different purposes.
uintptr_t inline GetModuleBase( HINSTANCE hModule )
{
	return ( ( ( uintptr_t ) hModule ) & -65536 );
}


void __cdecl MainEntryPoint( void )
{
	PROCESS_INFORMATION   pi;
	STARTUPINFOA          si; 
	VOID*  pvInjectionAddress;
	PIMAGE_SECTION_HEADER pshInject; 

	uintptr_t module_base = GetModuleBase( GetModuleHandle( NULL ) );

	RtlSecureZeroMemory( &pi, sizeof(pi) );
	RtlSecureZeroMemory( &si, sizeof(si) );

	// browse for executable to target
	char szFileName[MAX_PATH];
	OPENFILENAME ofn;
	*szFileName = 0;
	RtlSecureZeroMemory( &ofn, sizeof( ofn ) );
	ofn.lStructSize = sizeof( ofn );
	ofn.lpstrFile = szFileName;
	ofn.nMaxFile = MAX_PATH;
	if( ! GetOpenFileName( &ofn ) )
	{
		// user canceled
		goto the_end;
	}

	if( !CreateProcessA( NULL, szFileName, NULL, NULL, FALSE, CREATE_SUSPENDED, NULL, NULL, &si, &pi ) )
	{
		MessageBox( NULL, "CreateProcess() Failed", "Error", MB_OK);
		goto the_end;
	}

	// find info about the section to inject
	pshInject = GetSectionHeaderByName( module_base, my_section_name );
	if( pshInject == NULL )
	{
		MessageBox( NULL, "Failed to GetMySectionHeader", "Error", MB_OK);
		goto the_end;
	}

	// allocate memory in target process to hold our injected section 
	pvInjectionAddress = VirtualAllocEx( pi.hProcess, NULL, min( pshInject->SizeOfRawData, pshInject->Misc.VirtualSize ), MEM_COMMIT, PAGE_EXECUTE_READWRITE );
	if( pvInjectionAddress == NULL )
	{
		MessageBox( NULL, "VirtualAllocEx() Failed", "Error", MB_OK);
		goto the_end;
	}

#if DEBUG
	char str[0x80];
	wsprintf( str, "\nInjection address: 0x%08X\n", pvInjectionAddress );
	OutputDebugString( str );
	wsprintf( str, "section address: 0x%08X\n", module_base + pshInject->VirtualAddress );
	OutputDebugString( str );
	wsprintf( str, "virtual size: 0x%08X\n", pshInject->Misc.VirtualSize );
	OutputDebugString( str );
	wsprintf( str, "raw size: 0x%08X\n\n", pshInject->SizeOfRawData );
	OutputDebugString( str );
#endif

	
	// perform in-place relocations for new base
	if( ! FixUpSection( module_base, pshInject, (uintptr_t) pvInjectionAddress ) )
	{ 
		MessageBox( NULL, "No relocation data for section to inject", "Warning", MB_OK);
	}
	

	// inject section
	if( ! WriteProcessMemory( pi.hProcess, pvInjectionAddress, 
		(void*) ( module_base + pshInject->VirtualAddress ), min( pshInject->SizeOfRawData, pshInject->Misc.VirtualSize ), NULL ) )
	{
		MessageBox( NULL, "WriteProcessMemory  Failed", "Error", MB_OK);
		goto the_end;
	}
	

	// install hook(s) in target pointing at the now injected code
	unsigned long rel = ( (unsigned long)pvInjectionAddress - 0x0058013E ) + ((unsigned long)&my_0057EB99 - module_base - pshInject->VirtualAddress );

//	wsprintf( str, "rel: 0x%08X\n", rel );
//	OutputDebugString( str );


	if( !WriteProcessMemory( pi.hProcess, (void*)0x0058013A, &rel, 4, NULL ) )
	{
		MessageBox( NULL, "WriteProcessMemory  Failed", "Error", MB_OK);
	}


the_end:
	if( pi.hThread )
	{
		ResumeThread(pi.hThread); 
		CloseHandle(pi.hThread);
	}
	if( pi.hProcess ) 
	{
		CloseHandle(pi.hProcess);
	}

	ExitProcess(0);
}


// no verification used; examining our own PE structure in memory
// section name is 8 bytes and need not be null terminated
PIMAGE_SECTION_HEADER GetSectionHeaderByName( uintptr_t module_base, const char ( §ion_name )[ IMAGE_SIZEOF_SHORT_NAME ] )
{
	PIMAGE_NT_HEADERS pe = ( PIMAGE_NT_HEADERS )( module_base + ( ( PIMAGE_DOS_HEADER ) module_base )->e_lfanew );
	PIMAGE_SECTION_HEADER section_header = IMAGE_FIRST_SECTION( pe );
	PIMAGE_SECTION_HEADER end = §ion_header[ pe->FileHeader.NumberOfSections ];
	unsigned long a = *( ( unsigned long* ) section_name ); 
	unsigned long b = *( ( unsigned long* ) §ion_name[4] ); 
	do
	{ 
		unsigned long* p = (unsigned long*)section_header->Name;
		if( ( *p == a ) && ( *++p == b ) )
		{
			return section_header;
		}
	} while( ++section_header != end );
	return NULL;
}

// dangerous
char* GetLibName( HPANDLE hProcess, uintptr_t module_base )
{
	PIMAGE_NT_HEADERS pe = ( PIMAGE_NT_HEADERS ) ( module_base + ( ( PIMAGE_DOS_HEADER ) module_base )->e_lfanew );
	PIMAGE_DATA_DIRECTORY expdir = ( PIMAGE_DATA_DIRECTORY ) ( pe->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_EXPORT );
	PIMAGE_EXPORT_DIRECTORY exports = ( PIMAGE_EXPORT_DIRECTORY ) ( module_base + expdir->VirtualAddress );
	return (char*) ( exports->Name + module_base );
}

// this isn't complete yet...
// no verification used; examining our own PE structure in memory
BOOL FixupImport( uintptr_t module_base, unsigned long address )
{
	PIMAGE_NT_HEADERS pe = ( PIMAGE_NT_HEADERS ) ( module_base + ( ( PIMAGE_DOS_HEADER ) module_base )->e_lfanew );
	PIMAGE_DATA_DIRECTORY impdir = (PIMAGE_DATA_DIRECTORY)( pe->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT );
	for( PIMAGE_IMPORT_DESCRIPTOR idata = ( PIMAGE_IMPORT_DESCRIPTOR ) ( module_base + impdir->VirtualAddress ); idata->FirstThunk != NULL; idata++ ) 
	{
		for( unsigned long* iat_entry = ( unsigned long* ) ( module_base + idata->FirstThunk ); *iat_entry != NULL; iat_entry++ )
		{
			if( iat_entry == ( unsigned long* ) address )
			{
				// found import
				char* szModule = (char*) ( idata->Name + module_base );
				char* szProc = (char*) (module_base + 2 + *((unsigned long*)((((unsigned long) iat_entry) - idata->FirstThunk) + idata->OriginalFirstThunk)));	
				char* szLib = GetLibName( GetModuleBase( GetModuleHandle( szModule ) ) );
				
				char str[0x80];
				wsprintf(str,"\nLib: %s -> module: %s -> proc: %s\n", szLib, szModule, szProc );
				OutputDebugString(str);	


			}
		}
	}
	return NULL; // not an import
}


// no verification used; examining our own PE structure in memory
// find the reloc chunk(s) that correspond to our section to inject
// use the relocation data in the PE header to patch our code so it will work when loaded at any memory address
// *section must have read/write access
// *relocations must not be stripped
unsigned long FixUpSection( uintptr_t module_base, PIMAGE_SECTION_HEADER pSectionHeader, uintptr_t new_section_address )
{
	unsigned long num_of_fixups = 0;
	PIMAGE_NT_HEADERS pe = ( PIMAGE_NT_HEADERS ) ( module_base + ( ( PIMAGE_DOS_HEADER ) module_base )->e_lfanew );
	PIMAGE_DATA_DIRECTORY relocdir = ( PIMAGE_DATA_DIRECTORY )( pe->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_BASERELOC );
	uintptr_t delta = new_section_address - ( module_base + pSectionHeader->VirtualAddress );
		
	if( relocdir->Size && relocdir->VirtualAddress ) // if module has relocations
	{
		for( PIMAGE_BASE_RELOCATION reloc = ( PIMAGE_BASE_RELOCATION ) ( module_base + relocdir->VirtualAddress );
			( reloc->SizeOfBlock != 0 ) && ( reloc < ( PIMAGE_BASE_RELOCATION )( module_base + relocdir->VirtualAddress + relocdir->Size ) );
			reloc =  ( PIMAGE_BASE_RELOCATION ) ( ( ( uintptr_t ) reloc ) + reloc->SizeOfBlock ) )
		{
			for( unsigned short* pwRel = (unsigned short*)( ( ( uintptr_t ) reloc ) + sizeof( IMAGE_BASE_RELOCATION ) );
				pwRel < ( unsigned short* ) ( ( ( uintptr_t ) reloc ) + reloc->SizeOfBlock );
				pwRel++ )
			{
				if( ( *pwRel >> 12 ) == IMAGE_REL_BASED_HIGHLOW )
				{
					uintptr_t* fixup_loc = ( uintptr_t* )( module_base + reloc->VirtualAddress + ( *pwRel & 0x0FFF ) );
					if( fixup_loc >= ( uintptr_t* )( module_base + pSectionHeader->VirtualAddress )&& 
						fixup_loc < ( uintptr_t* )( module_base + pSectionHeader->VirtualAddress + pSectionHeader->SizeOfRawData ) )
					{
						uintptr_t address = *fixup_loc;
						if( address >= module_base + pSectionHeader->VirtualAddress && 
							address < module_base + pSectionHeader->VirtualAddress + pSectionHeader->SizeOfRawData )
						{
							*fixup_loc = address + delta;
						} 
						else // pointer to address outside of our section... IAT maybe?
						{
							FixupImport( module_base, address );
							char str[64];
							wsprintf(str,"bad fixup: 0x%08X\n",address );
							OutputDebugString(str);			
						}
						num_of_fixups++;
					}
				}
			}
		}
	}
	return num_of_fixups;
}

#pragma comment(linker, "/NODEFAULTLIB") // specifically no C runtime lib 
#pragma comment( lib, "kernel32" )  
#pragma comment( lib, "user32" )  
#pragma comment(linker, "/ENTRY:\"MainEntryPoint\"") // define entry point cause no C Lib