#include <remoteps/intercpt.hpp>
#include <bsptree/bintree.hpp>
#include <bsptree/iteratei.hpp>

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

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

void Intercept::loadImportModuleNames(void)
{
  loadImportModuleNamesEx();
}

void Intercept::loadImportModuleNamesEx(void)
{
  mImportModuleNames.remove();
	loadImportModuleNames(mImportModuleNames,baseAddress());
	for(int index=0;index<mImportModuleNames.size();index++)
		loadImportModuleNames(mImportModuleNames,(DWORD)::GetModuleHandle(mImportModuleNames[index]));
}

void Intercept::loadImportModuleNames(Block<String> &importModuleNames,DWORD baseAddress)
{
	PIMAGE_DOS_HEADER npImageDosHeader;
	PIMAGE_NT_HEADERS npImageNTHeader;
	PIMAGE_IMPORT_DESCRIPTOR npImageImportDescriptor;
	String strModuleName;

  if(!baseAddress)return;
  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)
	{
		strModuleName=(char*)(baseAddress+npImageImportDescriptor->Name);
		strModuleName=strModuleName.betweenString(0,'.');
		strModuleName.upper();
		if(!strModuleName.isNull()&&isascii(*((char*)strModuleName))&&!isInModuleNames(strModuleName,mImportModuleNames))mImportModuleNames.insert(&strModuleName);
		npImageImportDescriptor++;
	}
}

BOOL Intercept::isInModuleNames(const String &strModuleName,Block<String> &strModuleNames)
{
  for(int index=0;index<strModuleNames.size();index++)if(strModuleName==strModuleNames[index])return TRUE;
	return FALSE;
}

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)));
				pureImport.thunkType(PureImport::WIN95Thunk);
			}
      else
			{
			  pureImport.importAddress(entryPoint);
				pureImport.thunkType(PureImport::StandardThunk);
			}
			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)));
				pureImport.thunkType(PureImport::WIN95Thunk);
			}
			else 
			{
			  pureImport.importAddress(entryPoint);
				pureImport.thunkType(PureImport::StandardThunk);
			}

			return TRUE;
		}
	}
	return FALSE;
}

void Intercept::moduleEntryPoints(void)
{
  BinaryTree<PureImport> importEntries;
  QuickSort<PureImport> sortImport;

	entryPoints(importEntries,mBaseAddress);
	for(int index=0;index<mImportModuleNames.size();index++)entryPoints(importEntries,(DWORD)::GetModuleHandle(mImportModuleNames[index]));
	TreeIteratorInorder<PureImport> impIterator(importEntries);
	PureImport *pImport;
	index=0;
	size(importEntries.leaves());
	while(0!=(pImport=impIterator++))operator[](index++)=*pImport;
  sortImport.sortItems((Array<PureImport>&)*this);
}

// ApiSpy is watching for protection faults generated by this function by the
// while(pThunk->u1.Function) loop.  ApiSpy knows the code sequence of the first
// line of code to dereference pThunk->u1.Function.  If the dereference generates
// a fault, ApiSpy will advance the instruction pointer to the first nop instruction.
// see "apidebug.cpp" for the counterpart.
// modified:06/29/1999 for windows NT
// Added the isWIN95Thunk() code because the WINNT entry points are pure and do not need to be 
// incremented.  I have not tested the if statement for protection faults generated by this
// code change.

void Intercept::entryPoints(BinaryTree<PureImport> &importEntries,DWORD baseAddress)
{
	PIMAGE_DOS_HEADER npImageDosHeader;
	PIMAGE_NT_HEADERS npImageNTHeader;
	PIMAGE_IMPORT_DESCRIPTOR npImageImportDescriptor;
	PIMAGE_THUNK_DATA pThunk;
	PureImport pureImport;
  String moduleName;

	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);
		moduleName=(char*)(baseAddress+npImageImportDescriptor->Name);
		moduleName=moduleName.betweenString(0,'.');
		moduleName.upper();
		if(moduleName.isNull()||!isascii(*((char*)moduleName)))break;
		while(pThunk->u1.Function)
		{
		  if((int)pThunk->u1.Function==0xCCCCCCCC)break;
			if(0xCC==HIBYTE(HIWORD((int)pThunk->u1.Function)))break;
		  pureImport.moduleName(moduleName);
			if(isWIN95Thunk((DWORD)pThunk->u1.Function)||mVersionInfo.isWin95())
			{
			  pureImport.importAddress(*((DWORD*)((char*)(((DWORD)pThunk->u1.Function)+1))));
			  pureImport.rewriteAddress((DWORD)&(*((DWORD*)((char*)(((DWORD)pThunk->u1.Function)+1)))));
			  pureImport.thunkType(PureImport::WIN95Thunk);
			}
			else
			{
			  pureImport.importAddress((DWORD)pThunk->u1.Function);
			  pureImport.rewriteAddress((DWORD)&pThunk->u1.Function);
			  pureImport.thunkType(PureImport::StandardThunk);
			}
			_asm nop;
			_asm nop;
			importEntries.insert(pureImport);
			pThunk++;
		}  
    npImageImportDescriptor++;
	}
}

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

	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());
			resolved++;
		}
		else pureImport[importIndex].rewriteAddress(0L);
	}
}