Лаба 1-6 Системное программирование для ОС Windows [Вариант отсутствует, общий] / Лаба 5 / Spy / DLL / APIHook

#include "CmnHdr.h"
#include <ImageHlp.h>
#pragma comment(lib, "ImageHlp")
#include "APIHook.h"
#include <tlhelp32.h>
#include <StrSafe.h>

APIHook::APIHook(PSTR moduleName, PSTR functionName, PROC hookFunction) 
{
	DllName = moduleName;
	FuncName = functionName;
	HookedFunc = hookFunction;
	OriginalFunc = GetProcAddress(GetModuleHandleA(moduleName), FuncName);
	HookInAllModuls(DllName, OriginalFunc, HookedFunc);
}

APIHook::~APIHook() 
{
	HookInAllModuls(DllName, HookedFunc, OriginalFunc);
}

static HMODULE ModuleFromAddress(PVOID pv) 
{
	MEMORY_BASIC_INFORMATION mbi;
	return((VirtualQuery(pv, &mbi, sizeof(mbi)) != 0) 
		? (HMODULE) mbi.AllocationBase : NULL);
}

void APIHook::HookInAllModuls(PCSTR moduleName, PROC currentFunc, PROC newFunc) 
{
	HMODULE hmodThisMod = ModuleFromAddress(HookInAllModuls);
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetCurrentProcessId());

	MODULEENTRY32 me = { sizeof(me) };
	for (BOOL bOk = Module32First(hSnapshot, &me); bOk; bOk = Module32Next(hSnapshot, &me)) 
	{
		if (me.hModule != hmodThisMod) 
		{
			HookInOneModul(moduleName, currentFunc, newFunc, me.hModule);
		}
	}
}

void APIHook::HookInOneModul(PCSTR moduleName, 
	PROC currentFunc, PROC newFunc, HMODULE hMod) 
{
	ULONG ulSize;
	PIMAGE_IMPORT_DESCRIPTOR pImportDesc = NULL;
	try 
	{
		pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR) ImageDirectoryEntryToData(
			hMod, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &ulSize);
	} 
	catch (...) { }
	if (pImportDesc == NULL)
		return;  

	for (; pImportDesc->Name; pImportDesc++) 
	{
		PSTR pszModName = (PSTR) ((PBYTE) hMod + pImportDesc->Name);
		if (lstrcmpiA(pszModName, moduleName) == 0) 
		{
			PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA)
				((PBYTE) hMod + pImportDesc->FirstThunk);

			for (; pThunk->u1.Function; pThunk++) 
			{
				PROC* ppfn = (PROC*) &pThunk->u1.Function;
				BOOL bFound = (*ppfn == currentFunc);
				if (bFound) 
				{
					if (!WriteProcessMemory(GetCurrentProcess(), ppfn, &newFunc, 
						sizeof(newFunc), NULL) && (GetLastError() == ERROR_NOACCESS)) 
					{
						DWORD dwOldProtect;
						if (VirtualProtect(ppfn, sizeof(newFunc), PAGE_WRITECOPY, 
							&dwOldProtect)) 
						{
							WriteProcessMemory(GetCurrentProcess(), ppfn, &newFunc, 
								sizeof(newFunc), NULL);
							VirtualProtect(ppfn, sizeof(newFunc), dwOldProtect, 
								&dwOldProtect);
						}
					}
					return;  
				}
			}
		}  
	}
}