//---------------------------------------------------------------------------
//
// ApiHook.cpp
//
// SUBSYSTEM: 
//				API Hooking system
// MODULE:    
//				ApiHook is the main module of the hooking system
//
// DESCRIPTION:
//
// AUTHOR:		Ivo Ivanov (ivopi@hotmail.com) 
//              Some part of ReplaceInOneModule() implementation is based 
//              on code published by Jeffrey Richter and John Robbins.
//
// DATE:		2000 May 05,  version 1.0 
//              2001 June 29, version 2.0
//              2002 November 30. version 2.1
//
// FIXES:
//              - 2002 May 08
//                Fixed bug that used to cause an access violation 
//                if a hooked application calls GetProcAddress() API with an 
//                ordinal value rather than string function name. For more 
//                details see the implementation of 
//                CHookedFunction* CHookedFunctions::GetHookedFunction( 
//                    PCSTR pszCalleeModName, PCSTR pszFuncName )
//              - 2002 July 25
//                In the implementation of CApiHookMgr::AddHook() method
//                a newly created object CHookedFunction must be always inserted 
//                in the collection, even the HookImport() method returns FALSE. 
//                This solve the problem with dynamically imported functions that 
//                haven't been found in the IAT during the initial call to AddHook() 
//              - 2002 August 27
//                Fixed bug related to storing incorrect values of the full path 
//                name of the DLL in the map. In fact we need to get only the name 
//                and extension parts of the full file name.
//              - 2002 November 30
//                Added two methods of class CHookedFunctions 
//                (GetFunctionNameFromExportSection() and GetFunctionNameByOrdinal()) 
//                for handling scenarios where the function for spying is imported by 
//                number. This should allow a user of the library to easily intercept 
//                Windows socket APIs that are usually imported by ordinal value.
//                Added critical section in CApiHookMgr::MyGetProcAddress() method 
//                                                                         
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
//
// Includes
//  
//---------------------------------------------------------------------------

#include "..\Common\Common.h"
#include "ApiHook.h"
#include "..\Common\SysUtils.h"
#include "..\Common\LockMgr.h"
#include "..\Common\CustomMessages.h"
#include <imagehlp.h>
#pragma comment(lib, "imagehlp.lib")
#include "ModuleScope.h"

//---------------------------------------------------------------------------
//
// class CApiHookMgr
//  
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
//
// File scope constants and typedefs
//
//---------------------------------------------------------------------------
typedef struct 
{
	char szCalleeModName[MAX_PATH]; 
	char szFuncName[MAX_PATH];	
} API_FUNC_ID;

const API_FUNC_ID MANDATORY_API_FUNCS[] =
{
	{"Kernel32.dll", "LoadLibraryA"},
	{"Kernel32.dll", "LoadLibraryW"},
	{"Kernel32.dll", "LoadLibraryExA"},
	{"Kernel32.dll", "LoadLibraryExW"},
	{"Kernel32.dll", "GetProcAddress"}
};

// This macro evaluates to the number of elements in MANDATORY_API_FUNCS
#define NUMBER_OF_MANDATORY_API_FUNCS (sizeof(MANDATORY_API_FUNCS) / sizeof(MANDATORY_API_FUNCS[0])) 

//---------------------------------------------------------------------------
//
// Static members
//
//---------------------------------------------------------------------------

CModuleScope*     CApiHookMgr::sm_pModuleScope     = NULL;
CHookedFunctions* CApiHookMgr::sm_pHookedFunctions = NULL;
CCSWrapper        CApiHookMgr::sm_CritSec;

//---------------------------------------------------------------------------
// CApiHookMgr
// 
// Ctor
//---------------------------------------------------------------------------
CApiHookMgr::CApiHookMgr(
	CModuleScope* pModuleScope
	)
{
	sm_pModuleScope = pModuleScope;
	//
	// Obtain the handle to the DLL which code executes
	//
	m_hmodThisInstance   = ModuleFromAddress(CApiHookMgr::MyGetProcAddress);
	//
	// No system functions have been hooked up yet
	//
	m_bSystemFuncsHooked = FALSE;
	//
	// Create an instance of the map container
	//
	sm_pHookedFunctions  = new CHookedFunctions(this); 
}

//---------------------------------------------------------------------------
// ~CApiHookMgr
// 
// Dtor
//---------------------------------------------------------------------------

CApiHookMgr::~CApiHookMgr()
{
	UnHookAllFuncs();
	delete sm_pHookedFunctions;
}

//---------------------------------------------------------------------------
// HookSystemFuncs
// 
// Hook all needed system functions in order to trap loading libraries
//---------------------------------------------------------------------------
BOOL CApiHookMgr::HookSystemFuncs()
{
	BOOL bResult;

	if (TRUE != m_bSystemFuncsHooked)
	{
		bResult = HookImport(
			"Kernel32.dll", 
			"LoadLibraryA", 
			(PROC) CApiHookMgr::MyLoadLibraryA
			);
		bResult = HookImport(
			"Kernel32.dll", 
			"LoadLibraryW",
			(PROC) CApiHookMgr::MyLoadLibraryW
			) || bResult;
		bResult = HookImport(
			"Kernel32.dll", 
			"LoadLibraryExA",
			(PROC) CApiHookMgr::MyLoadLibraryExA
			) || bResult;
		bResult = HookImport(
			"Kernel32.dll", 
			"LoadLibraryExW",
			(PROC) CApiHookMgr::MyLoadLibraryExW
			) || bResult;
		bResult = HookImport(
			"Kernel32.dll", 
			"GetProcAddress",
			(PROC) CApiHookMgr::MyGetProcAddress
			) || bResult;
		m_bSystemFuncsHooked = bResult;
	} // if

	return m_bSystemFuncsHooked;
}

//---------------------------------------------------------------------------
// UnHookAllFuncs
// 
// Unhook all functions and restore original ones
//---------------------------------------------------------------------------
void CApiHookMgr::UnHookAllFuncs()
{
	if (TRUE == m_bSystemFuncsHooked)
	{
		CHookedFunction* pHook;
		CHookedFunctions::const_iterator itr;
		for (itr = sm_pHookedFunctions->begin(); 
			itr != sm_pHookedFunctions->end(); 
			++itr)
		{
			pHook = itr->second;
			pHook->UnHookImport();
			delete pHook;
		} // for
		sm_pHookedFunctions->clear();
		m_bSystemFuncsHooked = FALSE;
	} // if
}

//
// Indicates whether there is hooked function
//
BOOL CApiHookMgr::AreThereHookedFunctions()
{
	return (sm_pHookedFunctions->size() > 0);
}


//---------------------------------------------------------------------------
// HookImport
//
// Hook up an API function
//---------------------------------------------------------------------------
BOOL CApiHookMgr::HookImport(
	PCSTR pszCalleeModName, 
	PCSTR pszFuncName, 
	PROC  pfnHook
	)
{
	CLockMgr<CCSWrapper>  lockMgr(sm_CritSec, TRUE);

	BOOL                  bResult = FALSE;
	PROC                  pfnOrig = NULL;
	try
	{
		if (!sm_pHookedFunctions->GetHookedFunction(
				pszCalleeModName, 
				pszFuncName
				))
		{
			pfnOrig = GetProcAddressWindows(
				::GetModuleHandleA(pszCalleeModName),
				pszFuncName
				);
			//
			// It's possible that the requested module is not loaded yet
			// so lets try to load it.
			//
			if (NULL == pfnOrig)
			{
				HMODULE hmod = ::LoadLibraryA(pszCalleeModName);
				if (NULL != hmod)
					pfnOrig = GetProcAddressWindows(
						::GetModuleHandleA(pszCalleeModName),
						pszFuncName
						);
			} // if
			if (NULL != pfnOrig)
				bResult = AddHook(
					pszCalleeModName, 
					pszFuncName, 
					pfnOrig,
					pfnHook
					);
		} // if
	}
	catch(...)
	{

	} // try..catch

	return bResult;
}

//---------------------------------------------------------------------------
// UnHookImport
//
// Restores original API function address in IAT
//---------------------------------------------------------------------------
BOOL CApiHookMgr::UnHookImport(
	PCSTR pszCalleeModName, 
	PCSTR pszFuncName
	)
{
	CLockMgr<CCSWrapper>  lockMgr(sm_CritSec, TRUE);

	BOOL bResult = TRUE;
	try
	{
		bResult = RemoveHook(pszCalleeModName, pszFuncName);
	}
	catch (...)
	{
	}
	return bResult;
}

//---------------------------------------------------------------------------
// AddHook
//
// Add a hook to the internally supported container
//---------------------------------------------------------------------------
BOOL CApiHookMgr::AddHook(
	PCSTR pszCalleeModName, 
	PCSTR pszFuncName, 
	PROC  pfnOrig,
	PROC  pfnHook
	)
{
	BOOL             bResult = FALSE;
	CHookedFunction* pHook   = NULL;

	if (!sm_pHookedFunctions->GetHookedFunction(
			pszCalleeModName, 
			pszFuncName
			))
	{
		pHook = new CHookedFunction(
			sm_pHookedFunctions,
			pszCalleeModName, 
			pszFuncName, 
			pfnOrig,
			pfnHook
			);
		// We must create the hook and insert it in the container
		pHook->HookImport();
		bResult = sm_pHookedFunctions->AddHook(pHook);
	} // if

	return bResult;
}

//---------------------------------------------------------------------------
// RemoveHook
//
// Remove a hook from the internally supported container
//---------------------------------------------------------------------------
BOOL CApiHookMgr::RemoveHook(
	PCSTR pszCalleeModName, 
	PCSTR pszFuncName
	)
{
	BOOL             bResult = FALSE;
	CHookedFunction *pHook   = NULL;

	pHook = sm_pHookedFunctions->GetHookedFunction(
		pszCalleeModName, 
		pszFuncName
		);
	if ( NULL != pHook )
	{
		bResult = pHook->UnHookImport();
		if ( bResult )
		{
			bResult = sm_pHookedFunctions->RemoveHook( pHook );
			if ( bResult )
				delete pHook;
		} // if
	} // if

	return bResult;
}


//---------------------------------------------------------------------------
// HackModuleOnLoad
//
// Used when a DLL is newly loaded after hooking a function
//---------------------------------------------------------------------------
void WINAPI CApiHookMgr::HackModuleOnLoad(HMODULE hmod, DWORD dwFlags)
{
	//
	// If a new module is loaded, just hook it
	//
	if ((hmod != NULL) && ((dwFlags & LOAD_LIBRARY_AS_DATAFILE) == 0)) 
	{
		CLockMgr<CCSWrapper>  lockMgr(sm_CritSec, TRUE);
		
		CHookedFunction* pHook;
		CHookedFunctions::const_iterator itr;
		for (itr = sm_pHookedFunctions->begin(); 
			itr != sm_pHookedFunctions->end(); 
			++itr)
		{
			pHook = itr->second;
			pHook->ReplaceInOneModule(
				pHook->Get_CalleeModName(), 
				pHook->Get_pfnOrig(), 
				pHook->Get_pfnHook(), 
				hmod
				);
		} // for
	} // if
}

//---------------------------------------------------------------------------
//
// System API hooks prototypes
//
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
// MyLoadLibraryA
//
// 
//---------------------------------------------------------------------------
HMODULE WINAPI CApiHookMgr::MyLoadLibraryA(PCSTR pszModuleName)
{
	HMODULE hmod = ::LoadLibraryA(pszModuleName);
	HackModuleOnLoad(hmod, 0);

	return hmod;
}

//---------------------------------------------------------------------------
// MyLoadLibraryW
//
// 
//---------------------------------------------------------------------------
HMODULE WINAPI CApiHookMgr::MyLoadLibraryW(PCWSTR pszModuleName)
{
	HMODULE hmod = ::LoadLibraryW(pszModuleName);
	HackModuleOnLoad(hmod, 0);

	return hmod;
}

//---------------------------------------------------------------------------
// MyLoadLibraryExA
//
// 
//---------------------------------------------------------------------------
HMODULE WINAPI CApiHookMgr::MyLoadLibraryExA(
	PCSTR  pszModuleName, 
	HANDLE hFile, 
	DWORD dwFlags)
{
	HMODULE hmod = ::LoadLibraryExA(pszModuleName, hFile, dwFlags);
	HackModuleOnLoad(hmod, 0);

	return hmod;
}

//---------------------------------------------------------------------------
// MyLoadLibraryExW
//
// 
//---------------------------------------------------------------------------
HMODULE WINAPI CApiHookMgr::MyLoadLibraryExW(
	PCWSTR pszModuleName, 
	HANDLE hFile, 
	DWORD dwFlags)
{
	HMODULE hmod = ::LoadLibraryExW(pszModuleName, hFile, dwFlags);
	HackModuleOnLoad(hmod, 0);

	return hmod;
}

//---------------------------------------------------------------------------
// MyGetProcAddress
//
// 
//---------------------------------------------------------------------------
FARPROC WINAPI CApiHookMgr::MyGetProcAddress(HMODULE hmod, PCSTR pszProcName)
{
	// It is possible that multiple threads will call hooked GetProcAddress() 
	// API, therefore we should make it thread safe because it accesses sm_pHookedFunctions 
	// shared container.
	CLockMgr<CCSWrapper>  lockMgr(sm_CritSec, TRUE);
	//
	// Get the original address of the function
	//
	FARPROC pfn = GetProcAddressWindows(hmod, pszProcName);
	//
	// Attempt to locate if the function has been hijacked
	//
	CHookedFunction* pFuncHook = 
		sm_pHookedFunctions->GetHookedFunction(
			hmod, 
			pszProcName
			);

	if (NULL != pFuncHook)
		//
		// The address to return matches an address we want to hook
		// Return the hook function address instead
		//
		pfn = pFuncHook->Get_pfnHook();

	return pfn;
}

//---------------------------------------------------------------------------
// GetProcAddressWindows
//
// Returns original address of the API function
//---------------------------------------------------------------------------
FARPROC WINAPI CApiHookMgr::GetProcAddressWindows(HMODULE hmod, PCSTR pszProcName)
{
	return ::GetProcAddress(hmod, pszProcName);
}

//---------------------------------------------------------------------------
//
// class CHookedFunction
//  
//---------------------------------------------------------------------------

//
// The highest private memory address (used for Windows 9x only)
//
PVOID CHookedFunction::sm_pvMaxAppAddr = NULL;
//
// The PUSH opcode on x86 platforms
//
const BYTE cPushOpCode = 0x68;   

//---------------------------------------------------------------------------
// CHookedFunction
//  
//
//---------------------------------------------------------------------------
CHookedFunction::CHookedFunction(
	CHookedFunctions* pHookedFunctions,
	PCSTR             pszCalleeModName, 
	PCSTR             pszFuncName, 
	PROC              pfnOrig,
	PROC              pfnHook
	):
	m_pHookedFunctions(pHookedFunctions),
	m_bHooked(FALSE),
	m_pfnOrig(pfnOrig),
	m_pfnHook(pfnHook)
{
	strcpy(m_szCalleeModName, pszCalleeModName); 
	strcpy(m_szFuncName, pszFuncName);	

	if (sm_pvMaxAppAddr == NULL) 
	{
		//