#include <test/intercpt.hpp>

WORD Intercept::performIntercept(PureVector<PureImport> &pureImports,DWORD baseAddress)
{
	mBaseAddress=baseAddress;
	loadImportDescriptors(pureImports);
	moduleEntryPoints();
	resolveImportNames(pureImports);
	mImportModuleNames.remove();
	size(0);
	return TRUE;
}

void Intercept::loadImportDescriptors(PureVector<PureImport> &pureImports)
{
	Block<String> moduleNameStrings;
  DWORD importCount(pureImports.size());
	loadImportModuleNames();
	for(long importIndex=0;importIndex<importCount;importIndex++)importEntryPoint(pureImports[importIndex]);
}

void Intercept::loadImportModuleNames(void)
{
	PIMAGE_DOS_HEADER npImageDosHeader;
	PIMAGE_NT_HEADERS npImageNTHeader;
	PIMAGE_IMPORT_DESCRIPTOR npImageImportDescriptor;
	String moduleString;

	mImportModuleNames.remove();
	npImageDosHeader=(PIMAGE_DOS_HEADER)baseAddress();
	if(::IsBadReadPtr((void*)baseAddress(),sizeof(PIMAGE_NT_HEADERS)))return;
	if(npImageDosHeader->e_magic!=IMAGE_DOS_SIGNATURE)return;
	npImageNTHeader=(PIMAGE_NT_HEADERS)((char*)npImageDosHeader+npImageDosHeader->e_lfanew);
	if(npImageNTHeader->Signature!=IMAGE_NT_SIGNATURE)return;
	npImageImportDescriptor=(PIMAGE_IMPORT_DESCRIPTOR)((char*)baseAddress()+npImageNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	if((char*)npImageImportDescriptor==(char*)npImageNTHeader)return;
	while(npImageImportDescriptor->Name)
	{
		moduleString=(char*)(baseAddress()+npImageImportDescriptor->Name);
		moduleString=moduleString.betweenString(0,'.');
		moduleString.upper();
		mImportModuleNames.insert(&moduleString);
		npImageImportDescriptor++;
	}
}

WORD Intercept::importEntryPoint(PureImport &pureImport)
{
	DWORD entryPoint;

	if(!pureImport.moduleName().isNull())
	{
		if(0!=(entryPoint=(DWORD)::GetProcAddress(::GetModuleHandle(pureImport.moduleName()),pureImport.importName())))
		{
			if(isWIN95Thunk((DWORD)entryPoint))pureImport.importAddress(*((DWORD*)(((char*)entryPoint)+1)));
      else pureImport.importAddress(entryPoint);
//			else pureImport.importAddress(*((DWORD*)entryPoint));
			return TRUE;
		}
	}
	for(short moduleIndex=0;moduleIndex<mImportModuleNames.size();moduleIndex++)
	{
		if(0!=(entryPoint=(DWORD)::GetProcAddress(::GetModuleHandle(mImportModuleNames[moduleIndex]),pureImport.importName())))
		{
			pureImport.moduleName(mImportModuleNames[moduleIndex]);
			if(isWIN95Thunk((DWORD)entryPoint))pureImport.importAddress(*((DWORD*)(((char*)entryPoint)+1)));
      else pureImport.importAddress(entryPoint);
//			else pureImport.importAddress(*((DWORD*)entryPoint));
			return TRUE;
		}
	}
	return FALSE;
}

void Intercept::moduleEntryPoints(void)
{
	PIMAGE_DOS_HEADER npImageDosHeader;
	PIMAGE_NT_HEADERS npImageNTHeader;
	PIMAGE_IMPORT_DESCRIPTOR npImageImportDescriptor;
	PIMAGE_THUNK_DATA pThunk;
  String moduleName;
	DWORD importIndex(0);
	DWORD importCount(0);
  QuickSort<PureImport> sortImport;

	npImageDosHeader=(PIMAGE_DOS_HEADER)baseAddress();
	if(::IsBadReadPtr((void*)baseAddress(),sizeof(PIMAGE_NT_HEADERS)))return;
	if(npImageDosHeader->e_magic!=IMAGE_DOS_SIGNATURE)return;
	npImageNTHeader=(PIMAGE_NT_HEADERS)((char*)npImageDosHeader+npImageDosHeader->e_lfanew);
	if(npImageNTHeader->Signature!=IMAGE_NT_SIGNATURE)return;
	npImageImportDescriptor=(PIMAGE_IMPORT_DESCRIPTOR)((char*)baseAddress()+npImageNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	if((char*)npImageImportDescriptor==(char*)npImageNTHeader)return;
	while(npImageImportDescriptor->Name)
	{
		pThunk=(PIMAGE_THUNK_DATA)(baseAddress()+(DWORD)npImageImportDescriptor->FirstThunk);
		while(pThunk->u1.Function){importCount++;pThunk++;}
    npImageImportDescriptor++;
	}
	npImageImportDescriptor=(PIMAGE_IMPORT_DESCRIPTOR)((char*)baseAddress()+npImageNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
  size(importCount);
	while(npImageImportDescriptor->Name)
	{
		pThunk=(PIMAGE_THUNK_DATA)(baseAddress()+(DWORD)npImageImportDescriptor->FirstThunk);
		moduleName=(char*)(baseAddress()+npImageImportDescriptor->Name);
		while(pThunk->u1.Function)
		{
			if(isWIN95Thunk((DWORD)pThunk->u1.Function))
			{
        operator[](importIndex).moduleName(moduleName);
				operator[](importIndex).importAddress(*((DWORD*)((char*)(((DWORD)pThunk->u1.Function)+1))));
				operator[](importIndex).rewriteAddress((DWORD)&(*((DWORD*)((char*)(((DWORD)pThunk->u1.Function)+1)))));
				operator[](importIndex).thunkType(PureImport::WIN95Thunk);
				importIndex++;
			}
			else
			{
				operator[](importIndex).moduleName(moduleName);
				operator[](importIndex).importAddress((DWORD)pThunk->u1.Function);
				operator[](importIndex).rewriteAddress((DWORD)&(pThunk->u1.Function));
				operator[](importIndex).thunkType(PureImport::StandardThunk);
				importIndex++;
			}
			pThunk++;
		}
		npImageImportDescriptor++;
	}
  sortImport.sortItems((PureVector<PureImport>&)*this);
}

void Intercept::resolveImportNames(PureVector<PureImport> &pureImport)
{
	PureImport moduleImport;
  DWORD importCount(pureImport.size());
	BinarySearch<PureImport> searchImport((PureVector<PureImport>&)*this);

	for(long importIndex=0;importIndex<importCount;importIndex++)
	{
		if(searchImport.searchItem(pureImport[importIndex],moduleImport))
		{
			pureImport[importIndex].importAddress(moduleImport.importAddress());
			pureImport[importIndex].rewriteAddress(moduleImport.rewriteAddress());
			pureImport[importIndex].thunkType(moduleImport.thunkType());
		}
		else pureImport[importIndex].rewriteAddress(0L);
	}
}