#include <Windows.h>
#include <intrin.h>
#include <string>
#include <TlHelp32.h>
#include <psapi.h>

BOOL PatchTheRet(HMODULE realModule) {

	// Get primary module info

	PBYTE baseAddress = NULL;
	DWORD baseSize = 0;

	WCHAR fileName[MAX_PATH];
	GetProcessImageFileName((HANDLE)-1, fileName, MAX_PATH);
	std::wstring pathString = std::wstring(fileName);

	HANDLE hSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetCurrentProcessId());

	MODULEENTRY32 pEntry;
	pEntry.dwSize = sizeof(pEntry);
	BOOL hRes = Module32Next(hSnapShot, &pEntry);
	while (hRes)
	{
		if (pathString.find(pEntry.szModule) != std::wstring::npos) {
			baseAddress = pEntry.modBaseAddr;
			baseSize = pEntry.modBaseSize;
			break;
		}
		hRes = Module32Next(hSnapShot, &pEntry);
	}
	CloseHandle(hSnapShot);

	if (!baseAddress || !baseSize)
		return FALSE;

	// Hunt the stack

	PBYTE loadLibrary = (PBYTE)GetProcAddress(GetModuleHandle(L"KERNELBASE"), "LoadLibraryExW");
	PBYTE* stack = (PBYTE*)_AddressOfReturnAddress();
	BOOL foundLoadDll = FALSE;

	ULONG_PTR lowLimit, highLimit;
	GetCurrentThreadStackLimits(&lowLimit, &highLimit);

	for (; (ULONG_PTR)stack < highLimit; stack++) {
		if (*stack < (PBYTE)0x1000)
			continue;

		if (*stack > loadLibrary&&* stack < loadLibrary + 0x1000) {
			// LdrLoadDll is in the stack, let's start looking for our module
			foundLoadDll = TRUE;
		}

		if (foundLoadDll && *stack > baseAddress&&* stack < (baseAddress + baseSize)) {
			MEMORY_BASIC_INFORMATION mInfo = { 0 };
			VirtualQuery(*stack, &mInfo, sizeof(mInfo));

			if (!(mInfo.Protect & PAGE_EXECUTE_READ))
				continue;

			// Prepare a small rewrite+jmp block

			BYTE RetRewrite[22] = {
				0x48, 0xB8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // mov rax, <val>
				0x48, 0xBB, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // mov rbx, <val>
				0xFF, 0xE3 // jmp rbx
			};

			PBYTE rewrite = (PBYTE)VirtualAlloc(0, sizeof(RetRewrite), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
			if (!rewrite) return FALSE;

			RtlCopyMemory(rewrite, RetRewrite, sizeof(RetRewrite));
			RtlCopyMemory(rewrite + 2, &realModule, 8);
			RtlCopyMemory(rewrite + 12, stack, 8);

			// Move return address to our rewrite
			*(PBYTE *)stack = rewrite;

			return TRUE;
		}
	}
}

BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved)
{

	if (ul_reason_for_call != DLL_PROCESS_ATTACH)
		return TRUE;

	HMODULE kernel32 = LoadLibrary(L"kernel32.dll");
	PatchTheRet(kernel32);

	return TRUE;
}