#include <ntddk.h>


#define DEBUG_ALL_ACCESS        (0x1F000F)
#define DEBUG_PROCESS_ASSIGN    (0x0002)

KTIMER TimerObject;
KDPC DpcObject;
LARGE_INTEGER NextTime;
ULONGLONG   ValidAccessMaskAdr;

const ULONG kNtCreateDebugObjectIndex = 0x90;
const ULONG kTypeInfoOffset = 0x40;
const ULONG kValidAccessMaskOffset = 0x1c;
const LONG  kDbgkDebugObjectOffsetW7 = 0x7c;
const LONG  kInstructLen = 0x7;


typedef struct _KSERVICE_TABLE_DESCRIPTOR {
	unsigned long *ServiceTableBase;
	unsigned long *ServiceCounterTableBase;
	unsigned long NumberOfServices;
	unsigned char *ParamTableBase;
} ServiceDescriptorTableEntry, *pServiceDescriptorTableEntry;

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//	Description :
//		Retrieve KeServiceDescriptorTable address
//	Parameters :
//		None
//	Return value :
//		ULONGLONG : The service descriptor table address
//	Process :
//		Since KeServiceDescriptorTable isn't an exported symbol anymore, we have to retrieve it.
//		When looking at the disassembly version of nt!KiSystemServiceRepeat, we can see interesting instructions :
//			4c8d15c7202300	lea r10, [nt!KeServiceDescriptorTable (addr)]    => it's the address we are looking for (:
//			4c8d1d00212300	lea r11, [nt!KeServiceDescriptorTableShadow (addr)]
//			f7830001000080  test dword ptr[rbx+100h], 80h
//
//		Furthermore, the LSTAR MSR value (at 0xC0000082) is initialized with nt!KiSystemCall64, which is a function
//		close to nt!KiSystemServiceRepeat. We will begin to search from this address, the opcodes 0x83f7, the ones
//		after the two lea instructions, once we get here, we can finally retrieve the KeServiceDescriptorTable address
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
ULONGLONG GetKeServiceDescriptorTable64()
{
	PUCHAR      pStartSearchAddress = (PUCHAR)__readmsr(0xC0000082);
	PUCHAR      pEndSearchAddress = (PUCHAR)(((ULONG_PTR)pStartSearchAddress + PAGE_SIZE) & (~0x0FFF));
	PULONG      pFindCodeAddress = NULL;

	while (++pStartSearchAddress < pEndSearchAddress)
	{
		if ((*(PULONG)pStartSearchAddress & 0xFFFFFF00) == 0x83f70000)
		{
			pFindCodeAddress = (PULONG)(pStartSearchAddress - 12);
			return (ULONG_PTR)pFindCodeAddress + (((*(PULONG)pFindCodeAddress) >> 24) + 7) + (ULONG_PTR)(((*(PULONG)(pFindCodeAddress + 1)) & 0x0FFFF) << 8);
		}
	}
	return 0;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//	Description :
//		Retrieve the Nt* address function given its syscall number in the SSDT
//	Parameters :
//		PULONG KiServiceTable : the SSDT base address
//		ULONG  ServiceId 	  : a syscall number
//	Return value :
//		ULONGLONG : the address of the function which has the syscall number given in argument
//	Process :
//		Because the addresses contained in the SSDT have the last four bits reserved to store the number of arguments,
//		in order to retrieve only the address, we shift four bits to the right
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
ULONGLONG GetNTAddressFromSSDT(PULONG KiServiceTable, ULONG ServiceId)
{
	return (LONGLONG)(KiServiceTable[ServiceId] >> 4)
		+ (ULONGLONG)KiServiceTable;
}


/*
DbgkDebugObjectType

kd> dt nt!_OBJECT_TYPE
+0x000 TypeList         : _LIST_ENTRY
+0x010 Name             : _UNICODE_STRING
+0x020 DefaultObject    : Ptr64 Void
+0x028 Index            : UChar
+0x02c TotalNumberOfObjects : Uint4B
+0x030 TotalNumberOfHandles : Uint4B
+0x034 HighWaterNumberOfObjects : Uint4B
+0x038 HighWaterNumberOfHandles : Uint4B
+0x040 TypeInfo         : _OBJECT_TYPE_INITIALIZER
+0x0b0 TypeLock         : _EX_PUSH_LOCK
+0x0b8 Key              : Uint4B
+0x0c0 CallbackList     : _LIST_ENTRY

fffff800`02eade1c 488b151db1daff  mov     rdx,qword ptr [nt!DbgkDebugObjectType (fffff800`02c58f40)]

*/
ULONGLONG GetDbgkDebugObjectTypeAddr()
{
	LONGLONG Result = 0;

	pServiceDescriptorTableEntry KeSericeDescriptorTable = (pServiceDescriptorTableEntry)GetKeServiceDescriptorTable64();
	ULONGLONG NtCreateDebugObject = GetNTAddressFromSSDT(KeSericeDescriptorTable->ServiceTableBase, kNtCreateDebugObjectIndex);
	LONG DbgkDebugObjectTypeOffset = *(LONG*)(NtCreateDebugObject + kDbgkDebugObjectOffsetW7 + 0x3);

	Result = DbgkDebugObjectTypeOffset + (LONGLONG)NtCreateDebugObject + kDbgkDebugObjectOffsetW7 + kInstructLen;

	KdPrint(("NtCreateDebugObject:%p:%p:%x\n",NtCreateDebugObject, Result, DbgkDebugObjectTypeOffset));

	return Result;
}

//MmIsAddressValid ProbeForRead
VOID SetValidAccessMaskDpc(
	__in struct _KDPC  *Dpc,
	__in_opt PVOID  DeferredContext,
	__in_opt PVOID  SystemArgument1,
	__in_opt PVOID  SystemArgument2
)
{
	UNREFERENCED_PARAMETER(Dpc);
	UNREFERENCED_PARAMETER(SystemArgument1);
	UNREFERENCED_PARAMETER(SystemArgument2);
	__try
	{
		if (MmIsAddressValid((PVOID)ValidAccessMaskAdr))
		{
			*((ULONG*)ValidAccessMaskAdr) = (ULONG)DeferredContext;
			KdPrint(("Validaccessmask==%p\n", *(ULONGLONG*)ValidAccessMaskAdr));
		}

	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		NTSTATUS Status = GetExceptionCode();
		KeCancelTimer(&TimerObject);
		KdPrint(("SetValidAccessMaskDpc Exception:%x\n",Status));
		return;
	}
	KeSetTimer(&TimerObject, NextTime, &DpcObject);
	return;
}
/*
kd> dt _OBJECT_TYPE_INITIALIZER
ntdll!_OBJECT_TYPE_INITIALIZER
+0x000 Length           : Uint2B
+0x002 ObjectTypeFlags  : UChar
+0x002 CaseInsensitive  : Pos 0, 1 Bit
+0x002 UnnamedObjectsOnly : Pos 1, 1 Bit
+0x002 UseDefaultObject : Pos 2, 1 Bit
+0x002 SecurityRequired : Pos 3, 1 Bit
+0x002 MaintainHandleCount : Pos 4, 1 Bit
+0x002 MaintainTypeList : Pos 5, 1 Bit
+0x002 SupportsObjectCallbacks : Pos 6, 1 Bit
+0x004 ObjectTypeCode   : Uint4B
+0x008 InvalidAttributes : Uint4B
+0x00c GenericMapping   : _GENERIC_MAPPING
+0x01c ValidAccessMask  : Uint4B
+0x020 RetainAccess     : Uint4B
+0x024 PoolType         : _POOL_TYPE
+0x028 DefaultPagedPoolCharge : Uint4B
+0x02c DefaultNonPagedPoolCharge : Uint4B
+0x030 DumpProcedure    : Ptr64     void
+0x038 OpenProcedure    : Ptr64     long
+0x040 CloseProcedure   : Ptr64     void
+0x048 DeleteProcedure  : Ptr64     void
+0x050 ParseProcedure   : Ptr64     long
+0x058 SecurityProcedure : Ptr64     long
+0x060 QueryNameProcedure : Ptr64     long
+0x068 OkayToCloseProcedure : Ptr64     unsigned char
*/

VOID SetValidAccessMask(ULONG DebugAccess)
{

	ValidAccessMaskAdr = *(ULONGLONG*)GetDbgkDebugObjectTypeAddr() + kTypeInfoOffset + kValidAccessMaskOffset;

	KdPrint(("ValidAccessMask==%p\n", *(ULONGLONG*)ValidAccessMaskAdr));

	NextTime.QuadPart = -10000 * 100;
	KeInitializeTimer(&TimerObject);
	KeInitializeDpc(&DpcObject, &SetValidAccessMaskDpc, (PVOID)DebugAccess);
	KeSetTimer(&TimerObject, NextTime, &DpcObject);
}


VOID Unload(PDRIVER_OBJECT DriverObject)
{
	UNREFERENCED_PARAMETER(DriverObject);
	KeCancelTimer(&TimerObject);
}
NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegPath)
{

	UNREFERENCED_PARAMETER(DriverObject);
	UNREFERENCED_PARAMETER(RegPath);

	RTL_OSVERSIONINFOW OsVer = { 0 };
	OsVer.dwOSVersionInfoSize = sizeof(OsVer);
	RtlGetVersion(&OsVer);

	if (OsVer.dwMajorVersion != 6)
		return STATUS_UNSUCCESSFUL;

	//DEBUG_ALL_ACCESS & ~DEBUG_PROCESS_ASSIGN
	ULONG DebugAccess = DEBUG_ALL_ACCESS;
	SetValidAccessMask(DebugAccess);
	DriverObject->DriverUnload = Unload;
	return STATUS_SUCCESS;
}