#include "stdafx.h"
#include "IATHook.h"
#include <Windows.h>



#ifdef HOOK64

// TODO: 64-bit hook code

#else

// Step 1 of redirecting an imported function:
// This function gets the address of the import descriptor table.
IMAGE_IMPORT_DESCRIPTOR* GetImportDescriptor(HMODULE hMod, char* pszDllName)
{
	// Get the DOS header from which we can get the "optional header":
	IMAGE_DOS_HEADER* pDOSHeader = (IMAGE_DOS_HEADER*)hMod; 
	IMAGE_OPTIONAL_HEADER* pOptionHeader = (IMAGE_OPTIONAL_HEADER*)((BYTE*)hMod + pDOSHeader->e_lfanew + 24);

	// Get the Import Descriptor data structure:
	IMAGE_IMPORT_DESCRIPTOR* pImportDesc = (IMAGE_IMPORT_DESCRIPTOR*)((BYTE*)hMod + 
									pOptionHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	char* pszImpAddr = 0;

	// Iterate through the IMPORT DESCRIPTOR until we find the address of
	// the desired function to be hooked:
	while(pImportDesc->FirstThunk)
	{
		// Get the load address of the DLL:
		pszImpAddr = (char*)(( BYTE* )hMod+ pImportDesc->Name);
		if(_stricmp(pszDllName, pszImpAddr))
		{
			pImportDesc++;
			continue;
		}
		else
		{
			return pImportDesc;   
		}
	}
	return NULL;
}

// Support routine for parsing the PE header
IMAGE_THUNK_DATA* GetOriginalFirstThunk(HMODULE hMod,IMAGE_IMPORT_DESCRIPTOR* pImportDesc)
{
	return (IMAGE_THUNK_DATA*)((BYTE*)hMod+ pImportDesc->OriginalFirstThunk );
}

// Support routine for parsing the PE header
IMAGE_THUNK_DATA* GetFirstThunk(HMODULE hMod,IMAGE_IMPORT_DESCRIPTOR* pImportDesc)
{
	return (IMAGE_THUNK_DATA*)((BYTE*)hMod+ pImportDesc->FirstThunk);
}


// This function gets the address of the imported function:
DWORD* GetCurrentFunctAddr(HMODULE hMod,
						   IMAGE_THUNK_DATA* pOriginalFirstThunk, 
						   IMAGE_THUNK_DATA* pFirstThunk,
						   char* pszFunctionName)
{
	// From the IMAGE thunk data structure, we get the address
	// of the imported function:
	char* szTest;
	while(pOriginalFirstThunk->u1.Function)
	{
		// Get the import load address of the function:
		szTest = (char*)((BYTE*)hMod + (DWORD)pOriginalFirstThunk->u1.AddressOfData+2);
		if(_stricmp(pszFunctionName,szTest)==0)
		{
			return &pFirstThunk->u1.Function;
		}
		pOriginalFirstThunk++; 
		pFirstThunk++;
	}

	return NULL;
}


// Once the import address is found, unprotect the IAT to allow us to write to it,
// then we change the import address entry to point to our hook procedure.
bool ChangeAddress(DWORD* dwOldAddress, DWORD dwNewAddress)
{
	DWORD dwOld;
	if (!(VirtualProtect(dwOldAddress, 4, PAGE_READWRITE, &dwOld))) 
	{
		return false;
	}

	// Modify the function pointer in the IAT:
	*dwOldAddress = dwNewAddress;

	// Change the memory permissions back to Excecute (?):
	if (!(VirtualProtect(dwOldAddress, 4, PAGE_EXECUTE, &dwOld))) 
	{
		return false;
	}
	else
	{
		// Flush the instruction cache so that the changes take effect:
		BOOL bResult = FlushInstructionCache(GetCurrentProcess(), NULL, NULL);

		return true;
	}
}


DWORD HookImportedFunction(HMODULE hMod, char* pszDllName, char* pszFunctionName, DWORD dwNewAddress)
{
	IMAGE_IMPORT_DESCRIPTOR* IID = GetImportDescriptor(hMod, pszDllName);
	if (IID == NULL) 
		return NULL;

	IMAGE_THUNK_DATA* OriginalFirstThunk = GetOriginalFirstThunk(hMod,IID);
	IMAGE_THUNK_DATA* FirstThunk = GetFirstThunk(hMod,IID);

	DWORD* dwOldFunctionAddress = GetCurrentFunctAddr( hMod,OriginalFirstThunk,
														FirstThunk,pszFunctionName);
	if (dwOldFunctionAddress == NULL) 
		return NULL;
	
	DWORD dwOriginalAddress = *dwOldFunctionAddress;

	if(ChangeAddress(dwOldFunctionAddress, dwNewAddress))
		return dwOriginalAddress;
	else
		return NULL;
}

#endif // end of 32-bit hoook code