detect dll hijacking attempts

detect_dll_hijack.cpp

#include <Windows.h>
#include <stdio.h>
#include <iostream>
LPVOID  original_func = NULL;
PIMAGE_IMPORT_DESCRIPTOR import_descriptor = NULL;
const wchar_t* back = L"\\";
wchar_t current_directory[MAX_PATH];
typedef HMODULE(WINAPI  * loadlibrary_def)(
	LPCWSTR lpLibFileName
	);
PVOID original_function = GetProcAddress(LoadLibrary(L"kernel32.dll"), "LoadLibraryW");
HMODULE WINAPI hooked_loadlibrary(
	LPCWSTR lpLibFileName
) {
	HMODULE hmod1, hmod2 = NULL;
	DWORD path_len = (lstrlenW(current_directory) + lstrlenW(back) + lstrlenW(lpLibFileName)+ 4 )  * sizeof(wchar_t) + 1;
	wchar_t * path = (wchar_t *)malloc(path_len);
	ZeroMemory(path, path_len);
	lstrcpyW(path, current_directory);
	lstrcat(path, back);
	lstrcat(path, lpLibFileName);
	hmod1 = ((loadlibrary_def)original_function)(path);
	if (hmod1) {
		hmod2 = LoadLibraryEx(lpLibFileName, NULL, LOAD_LIBRARY_SEARCH_SYSTEM32);
		if (hmod2) {
			wprintf(L"suspected dll hijacking attempt => %ws", lpLibFileName);
		}
	}
	if(hmod2)
		FreeLibrary(hmod2);
	free(path);

		return hmod1;
}
inline void chk_hdrs(DWORD_PTR base, PULONG result) {
	PIMAGE_DOS_HEADER dos_header = (PIMAGE_DOS_HEADER)base;

	if (dos_header->e_magic != IMAGE_DOS_SIGNATURE) {
		*result = -1;
		return;
	}
	PIMAGE_NT_HEADERS nt_header = (PIMAGE_NT_HEADERS)(base + dos_header->e_lfanew);
	if (nt_header->Signature != IMAGE_NT_SIGNATURE) {
		*result = -1;
		return;
	}
	import_descriptor = (PIMAGE_IMPORT_DESCRIPTOR)(base + nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	if (!import_descriptor) {
		*result = -1;
		return;
	}

}
inline void hook(PIMAGE_IMPORT_DESCRIPTOR import_descriptor,char* func, DWORD_PTR base, PULONG result) {
	PIMAGE_IMPORT_BY_NAME function_name = NULL;
	HMODULE library = NULL;
	char* libname = NULL;
	while (import_descriptor->Name != NULL) {
		libname = (char*)(import_descriptor->Name + base);
		library = LoadLibraryA(libname);
		if (library) {
			PIMAGE_THUNK_DATA original_thunk, first_thunk = NULL;
			original_thunk = (PIMAGE_THUNK_DATA)(base + import_descriptor->OriginalFirstThunk);
			first_thunk = (PIMAGE_THUNK_DATA)(base + import_descriptor->FirstThunk);
			while (original_thunk->u1.AddressOfData != NULL) {
				function_name = (PIMAGE_IMPORT_BY_NAME)((DWORD_PTR)base + original_thunk->u1.AddressOfData);
				if (strcmp(function_name->Name, func) == 0) {
					SIZE_T bytes_written = 0;
					DWORD old_protect = 0;
					VirtualProtect((LPVOID)(&first_thunk->u1.Function), 8, PAGE_READWRITE, &old_protect);
					if (original_func == NULL) {
						original_func = (LPVOID)first_thunk->u1.Function;
						first_thunk->u1.Function = (DWORD_PTR)hooked_loadlibrary;

					}
					else {
						first_thunk->u1.Function = (DWORD_PTR)original_func;

					}

					*result = 0;
					goto done;
				}
				++original_thunk;
				++first_thunk;
			}
		}
	}
done:
	return;

}
inline void hook_install(char* func, PULONG result) {
	ULONG results = -1;
	DWORD_PTR base = (DWORD_PTR)GetModuleHandle(NULL);
	chk_hdrs(base, &results);
	if (result) {
		hook(import_descriptor,func, base, &results);
		if (results) {
			*result = 0;
			return;
		}
	}

}

int main() {
	ZeroMemory(current_directory, MAX_PATH);
	GetCurrentDirectory(MAX_PATH, current_directory);
	ULONG result = -1;
	hook_install((char*)"LoadLibraryW", &result);
	if (result) {
		LoadLibraryW(L"dbghelp.dll");

	}
	hook_install((char*)"LoadLibraryW", &result);
	

}