#include "struct.h"

void RtlAvlRemoveNode(PVOID tree, PVOID node);

PEPROCESS LookUpProcess(HANDLE Pid)
{
	PEPROCESS eprocess = NULL;
	NTSTATUS Status = STATUS_UNSUCCESSFUL;
	Status = PsLookupProcessByProcessId(Pid, &eprocess);
	if (NT_SUCCESS(Status))
		return eprocess;
	return NULL;
}

ULONG64 GetVadOffset()
{
	UNICODE_STRING func;
	RtlInitUnicodeString(&func, L"PsGetProcessExitStatus");
	ULONG64 funcAddr = MmGetSystemRoutineAddress(&func);
	if (funcAddr == NULL)
	{
		return FALSE;
	}

	//8B 81 D4 07 00 00                       mov     eax, [rcx+7D4h]
	PUCHAR funcBytes = (PUCHAR)funcAddr;

	//DbgPrint("funcAddr : %p\n", funcAddr);
	//DbgPrint("funcBytes : %p\n", funcBytes);

	for (size_t i = 0; i < 0x10; i++)
	{
		if (funcBytes[0] == 0x8B && funcBytes[1] == 0x81)
		{
			UINT32 offset = *(UINT32*)(funcBytes + 2);
			return offset + 0x4;
		}
	}
	return FALSE;
}



PMMVAD MiLocateAddress(PEPROCESS eProcess, ULONG64 virtualAddress)
{
	PRTL_AVL_TREE vadRoot = (PRTL_AVL_TREE)((PUCHAR)eProcess + GetVadOffset());
	if (!vadRoot || !vadRoot->Root) {
		return NULL; // VAD tree is empty or invalid
	}
	PMMVAD vad = (PMMVAD)vadRoot->Root;
	if (!MmIsAddressValid(vad)) { return NULL; }
	ULONG64 address = virtualAddress >> 12;
	if (address < (vad->Core.StartingVpn | (vad->Core.StartingVpnHigh << 32)) || address >(vad->Core.EndingVpn | (vad->Core.EndingVpn << 32)))
	{
		while (vad)
		{
			if (address > (vad->Core.EndingVpn | (vad->Core.EndingVpnHigh << 32)))
			{
				vad = (PMMVAD)vad->Core.VadNode.Right;
			}
			else
			{
				if (address >= (vad->Core.StartingVpn | (vad->Core.StartingVpnHigh << 32)))
				{
					return vad;
				}
				vad = (PMMVAD)vad->Core.VadNode.Left;
			}
		}
	}

	return NULL;
}

PVOID EnumVad(PMMVAD vad)
{
	if (!vad) {
		return NULL;
	}

	PVOID stack[256];
	int stackIndex = 0;

	stack[stackIndex++] = vad;

	while (stackIndex > 0) {

		PMMVAD currentVad = (PMMVAD)stack[--stackIndex];

		DbgPrintEx(77, 0, "Vad: %p | start : %x | end : %x | vadflags : %x \n",
			currentVad,
			currentVad->Core.StartingVpn,
			currentVad->Core.EndingVpn,
			currentVad->Core.u.VadFlags);


		if (currentVad->Core.VadNode.Right) {
			stack[stackIndex++] = (PVOID)currentVad->Core.VadNode.Right;
		}

		if (currentVad->Core.VadNode.Left) {
			stack[stackIndex++] = (PVOID)currentVad->Core.VadNode.Left;
		}
	}

	return NULL;
}

PVOID ParserProcessVad(PEPROCESS eProcess)
{
	PRTL_AVL_TREE vadRoot = (PRTL_AVL_TREE)((PUCHAR)eProcess + GetVadOffset());
	if (!vadRoot || !vadRoot->Root) {
		return NULL; // VAD tree is empty or invalid
	}
	ULONG64 vadCount = *(PULONG64)((PUCHAR)eProcess + GetVadOffset() + 0x10);

	PMMVAD vad = (PMMVAD)vadRoot->Root;
	if (!MmIsAddressValid(vad)) { return NULL; }

	EnumVad(vad);
}

PVOID LockPageProction(PEPROCESS eProcess, ULONG64 virtualaddress, PMMVAD vad)
{

}

PVOID RemoveVad(PEPROCESS eProcess, ULONG64 virtualaddress)
{
	PMMVAD vad = MiLocateAddress(eProcess, virtualaddress);
	if (!MmIsAddressValid(vad)) { return NULL; }

	PRTL_AVL_TREE vadRoot = (PRTL_AVL_TREE)((PUCHAR)eProcess + GetVadOffset());
	if (!vadRoot || !vadRoot->Root)
	{
		return NULL;
	}

	RtlAvlRemoveNode(vadRoot, vad);

	ObReferenceObject(eProcess);
}

PVOID UnloadDriver(
	_In_ PDRIVER_OBJECT DriverObject
)
{
	UNREFERENCED_PARAMETER(DriverObject);
	DbgPrintEx(77, 0, "Driver unloaded successfully.\n");
	return STATUS_SUCCESS;
}

NTSTATUS DriverEntry(
	_In_ PDRIVER_OBJECT DriverObject,
	_In_ PUNICODE_STRING RegistryPath
)
{
	UNREFERENCED_PARAMETER(RegistryPath);
	DbgPrintEx(77, 0, "Driver loaded successfully.\n");

	//1. 分析内核 MiLocateAddress 这个函数，然后照抄此判断是: 给一个R3地址，返回一个VAD 节点
	//2. 给一个进程ID，解析出此进程的所有VAD节点
	//3. 代码实现WIN10 WIN11 私有内存锁页
	//4. 移除内存所在vad节点的位置

	//DbgBreakPoint();

	PEPROCESS eproc = NULL;
	eproc = LookUpProcess(0x16c8);

	if (!eproc) { return STATUS_UNSUCCESSFUL; }

	UINT32 vadOffset = GetVadOffset();
	if (!vadOffset) { return STATUS_UNSUCCESSFUL; }

	//1: over
	PMMVAD vad = MiLocateAddress(eproc, 0xDCE0000);
	DbgPrintEx(77, 0, "Find Vad: %p\n", vad);

	/*
	+0x658 VadRoot          : _RTL_AVL_TREE
	+0x660 VadHint          : 0xffffce01`8de7d9a0 Void
	+0x668 VadCount         : 0xaa
	+0x670 VadPhysicalPages : 0
	+0x678 VadPhysicalPagesLimit : 0
	*/

	//vadRoot是根节点, vadHint是上一次MiLocateAddress找到的指针, vadCount是vad的数量

	//2: over
	ParserProcessVad(eproc);

	//3. over
	RemoveVad(eproc, 0xDCE0000);

	ObDereferenceObject(eproc);
	DriverObject->DriverUnload = UnloadDriver;
	return STATUS_SUCCESS;
}

PVOID LockVadProction(PEPROCESS eProcess, ULONG64 virtualaddress)
{
	PMMVAD_SHORT vad = MiLocateAddress(eProcess, virtualaddress);
	if (!MmIsAddressValid(vad)) { return NULL; }

	vad->u.PrivateVadFlags.Graphics = 0x1;
	vad->u.PrivateVadFlags.PrivateMemoryAlwaysSet = 0x1;
}

#include <ntifs.h>

ULONG64 g_PteAddress = 0;
ULONG64 g_PdeAddress = 0;
ULONG64 g_PpeAddress = 0;
ULONG64 g_PxeAddress = 0;

ULONG64 GetPteAddress(ULONG64 Address)
{
	return (ULONG64)(((Address >> 9) & 0x7FFFFFFFF8) + g_PteAddress);
}

ULONG64 GetPdeAddress(ULONG64 Address)
{
	return (ULONG64)(((Address >> 9) & 0x7FFFFFFFF8) + g_PteAddress);
}

ULONG64 GetPpeAddress(ULONG64 Address)
{
	return (ULONG64)(((Address >> 9) & 0x7FFFFFFFF8) + g_PteAddress);
}

ULONG64 GetPxeAddress(ULONG64 Address)
{
	return (ULONG64)(((Address >> 9) & 0x7FFFFFFFF8) + g_PteAddress);
}

VOID GetFourTableBASE()
{
	ULONG64 DirectoryTableBase = *(PULONG64)((PUCHAR)PsInitialSystemProcess + 0x28);
	//后12位属性
	DirectoryTableBase &= ~0xFFF;
	PHYSICAL_ADDRESS address;
	address.QuadPart = DirectoryTableBase;
	PULONG64 VirtualDirectorBase = MmGetVirtualForPhysical(address);
	ULONG64 index = -1;

	for (size_t i = 0; i < 512; i++)
	{
		ULONG64 data = VirtualDirectorBase[i];
		//8000 0000001aa063
		//48 0xFFFFFFFFFFFF
		//ULONG64 testData = ((data & ~0xFFF) & 0xFFFFFFFFFFFF);

		if (DirectoryTableBase == ((data & ~0xFFF) & 0xFFFFFFFFFFFF))
		{
			index = i;
			break;
		}
	}

	if (index != -1)
	{
		g_PteAddress = (index + 0x1FFFE00) << 39;
		g_PdeAddress = (index << 30) + g_PteAddress;
		g_PpeAddress = (index << 21) + g_PdeAddress;
		g_PxeAddress = (index << 12) + g_PpeAddress;
	}
}


VOID DriverUnload(PDRIVER_OBJECT pDriver)
{

}

NTSTATUS DriverEntry(PDRIVER_OBJECT pDriver, PUNICODE_STRING pSeg)
{

	//1. 获取pte的基址 （ MmGetVirtualForPhysical ）
	//2. 解析 返回PTE PDE PDPTE PML4的地址 ( MmIsAddressValidEx ) 

	//1. over

	DbgBreakPoint();

	GetFourTableBASE();
	
	if (g_PteAddress != NULL)
	{
		DbgPrintEx(77, 0, "PTEBASE: %p\n", g_PteAddress);
		DbgPrintEx(77, 0, "PDEBASE: %p\n", g_PdeAddress);
		DbgPrintEx(77, 0, "PPEBASE: %p\n", g_PpeAddress);
		DbgPrintEx(77, 0, "PXEBASE: %p\n", g_PxeAddress);
	}
	else
	{
		DbgPrintEx(77, 0, "ERROR\n");
	}

	//2. over
	
	ULONG64 gdtr = 0xfffff80258bc4fb0;

	ULONG64 gdtrPte = GetPteAddress(gdtr);
	ULONG64 gdtrPde = GetPdeAddress(gdtrPte);
	ULONG64 gdtrPpe = GetPpeAddress(gdtrPde);
	ULONG64 gdtrPxe = GetPxeAddress(gdtrPpe);

	DbgPrintEx(77, 0, "gdtrPte: %p\n", gdtrPte);
	DbgPrintEx(77, 0, "gdtrPde: %p\n", gdtrPde);
	DbgPrintEx(77, 0, "gdtrPpe: %p\n", gdtrPpe);
	DbgPrintEx(77, 0, "gdtrPxe: %p\n", gdtrPxe);


	pDriver->DriverUnload = DriverUnload;

	return STATUS_SUCCESS;
}