#include "kd_driver.h"
#include "kd_utils.h"
//#include "..\Vmp\Vmp.h"

HANDLE kd_driver::iqvw64e_device_handle;
uint64_t kd_driver::ntoskrnlbase = 0;
std::string cachedDriverName = "";

std::wstring kd_driver::GetDriverPath() {
	std::wstring temp = utils::GetFullTempPath();
	if (temp.empty()) {
		return L"";
	}
	return temp + L"\\" + GetDriverNameW();
}

std::wstring kd_driver::GetDriverNameW()
{
	if (cachedDriverName.empty()) {
		//Create a random name
		char buffer[100]{};
		static const char alphanum[] =
			"abcdefghijklmnopqrstuvwxyz"
			"ABCDEFGHIJKLMNOPQRSTUVWXYZ";
		int len = rand() % 20 + 10;
		for (int i = 0; i < len; ++i)
			buffer[i] = alphanum[rand() % (sizeof(alphanum) - 1)];
		cachedDriverName = buffer;
	}

	std::wstring name(cachedDriverName.begin(), cachedDriverName.end());
	return name;
}


bool kd_driver::AcquireDebugPrivilege() {

	HMODULE ntdll = GetModuleHandleA("ntdll.dll");
	if (ntdll == NULL) {
		return false;
	}

	ULONG SE_DEBUG_PRIVILEGE = 20UL;
	BOOLEAN SeDebugWasEnabled;
	NTSTATUS Status = nt::RtlAdjustPrivilege(SE_DEBUG_PRIVILEGE, TRUE, FALSE, &SeDebugWasEnabled);
	if (!NT_SUCCESS(Status)) {
		Log((L"[-] Failed to acquire SE_DEBUG_PRIVILEGE") << std::endl);
		return false;
	}

	return true;
}

#include "Ld_BASE.h"

bool kd_driver::Load(PDWORD status)
{
	srand((unsigned)time(NULL) * GetCurrentThreadId());

	std::wstring driver_path = GetDriverPath();
	_wremove(driver_path.c_str());

	if (!utils::CreateFileFromMemory(driver_path, reinterpret_cast<const char*>(Ld_BASE), sizeof(Ld_BASE))) {

		//驱动文件创建失败

		*status = 102;
		

		Log((L"[-] Failed to create vulnerable driver file") << std::endl);
		return false;
	}

	if (!AcquireDebugPrivilege()) {

		//权限提升失败
		*status = 103;

		Log((L"[-] Failed to acquire SeDebugPrivilege") << std::endl);
		_wremove(driver_path.c_str());
		return false;
	}

	if (!nt::RegisterAndStart(driver_path, GetDriverNameW())) {

		*status = 104;
		//驱动加载或启动失败

		Log((L"[-] Failed to register and start service for the vulnerable driver") << std::endl);
		_wremove(driver_path.c_str());
		return false;
	}

	// mylink nuwer

	iqvw64e_device_handle = CreateFileA(XorStr("\\\\.\\esp_link"), GENERIC_READ | GENERIC_WRITE, 0, 0, OPEN_EXISTING, FILE_ATTRIBUTE_SYSTEM, 0);
	if (iqvw64e_device_handle == INVALID_HANDLE_VALUE)
	{
		*status = 105;
		//链接驱动通信失败

		kd_driver::Unload();
		Log((L"[-] failed handle !!") << std::endl);
		return false;
	}
	

	kd_driver::ntoskrnlbase = utils::GetKernelModule("ntoskrnl.exe");
	if (!kd_driver::ntoskrnlbase)
	{
		//获取内核模块基址失败

		*status = 106;

		kd_driver::Unload();
		Log((L"[-] Failed to get ntoskrnl.exe") << std::endl);
		return false;
	}
	auto os = utils::GetVersion();
	if (os.dwBuildNumber >= 14393)
	{
		if (!kd_driver::ClearPiDDBCacheTable(iqvw64e_device_handle)) {
			Log((L"[-] Failed to ClearPiDDBCacheTable") << std::endl);
			kd_driver::Unload();
			return false;
		}

		if (!kd_driver::ClearKernelHashBucketList(iqvw64e_device_handle)) {
			Log((L"[-] Failed to ClearKernelHashBucketList") << std::endl);
			kd_driver::Unload();
			return false;
		}

		if (!kd_driver::ClearMmUnloadedDrivers(iqvw64e_device_handle)) {
			Log((L"[!] Failed to ClearMmUnloadedDrivers") << std::endl);
			kd_driver::Unload();
			return false;
		}

		//if (!kd_driver::ClearWdFilterDriverList(iqvw64e_device_handle)) {
		//	Log((L"[!] Failed to ClearWdFilterDriverList") << std::endl);
		//	kd_driver::Unload();
		//	return false;
		//}
	}

	return true;
}

bool kd_driver::Unload()
{
	if (iqvw64e_device_handle != INVALID_HANDLE_VALUE)
		CloseHandle(iqvw64e_device_handle);

	if (!nt::StopAndRemove(GetDriverNameW()))
		return false;

	std::wstring driver_path = GetDriverPath();

	if (_wremove(driver_path.c_str()) != 0)
		return false;


	return true;
}
	
auto set_calc_random_val()->_CALC_VERIFY
{
	_CALC_VERIFY Verify{ 0 };

	auto a = utils::get_random_val(0, 1000000);
	auto b = utils::get_random_val(0, 1000000);
	auto c = utils::get_random_val(0, 1000000);
	auto d = utils::get_random_val(0, 1000000);

	Verify.a = a;
	Verify.b = b;
	Verify.c = c;
	Verify.d = d;


	return Verify;
}

bool kd_driver::MmCopy(HANDLE device_handle, uint64_t Base, uint64_t Buffer, SIZE_T Size)
{
	_CALC_VERIFY Verify = set_calc_random_val();
	Verify.calc_val = utils::calc_comm_verify_check(Verify.a, Verify.b, Verify.c, Verify.d);

	COMM_INFO memory = { 0 };
	memory.case_number = 0x13;
	memory.destination = (DWORD64)Buffer;
	memory.source = (DWORD64)Base;
	memory.length = Size;
	memory.Verify = &Verify;

	DWORD bytes_returned = 0;
	
	return DeviceIoControl(device_handle, IOCTL_MMCOPY, &memory, sizeof(memory), &memory, sizeof(memory), &bytes_returned, nullptr);
}

bool kd_driver::read(HANDLE device_handle, uint64_t Base, uint64_t* Buffer, SIZE_T Size)
{
	return MmCopy(device_handle, Base, (uint64_t)Buffer, Size);
}

bool kd_driver::write(HANDLE device_handle, uint64_t Base, uint64_t* Buffer, SIZE_T Size)
{
	return MmCopy(device_handle, (uint64_t)Buffer, (uint64_t)Base, Size);
}

bool kd_driver::writeEx(HANDLE device_handle, uint64_t Base, uint64_t* Buffer, SIZE_T Size)
{
	/*uint64_t physical_base = 0;
	if (!GetPhysicalAddress(device_handle, Base, &physical_base))
		return false;

	const auto map_physical_addr = MapIoSpace(device_handle, physical_base, (uint32_t)Size);
	if (map_physical_addr == NULL)
		return false;

	auto result = MmCopy(device_handle, (uint64_t)Buffer, (uint64_t)map_physical_addr, Size);

	UnmapIoSpace(device_handle, map_physical_addr, (uint32_t)Size);*/
	return MdlIoSpace(device_handle, Base, Buffer, Size);
//	return true;
}



bool kd_driver::MdlIoSpace(HANDLE device_handle, uint64_t address, uint64_t* Buffer, SIZE_T Size)
{
	_CALC_VERIFY Verify = set_calc_random_val();
	Verify.calc_val = utils::calc_comm_verify_check(Verify.a, Verify.b, Verify.c, Verify.d);

	COMM_INFO memory = { 0 };
	memory.case_number = 0x14;
	memory.destination = (DWORD64)Buffer;
	memory.source = (DWORD64)address;
	memory.length = Size;
	memory.Verify = &Verify;

	DWORD bytes_returned = 0;
	return DeviceIoControl(device_handle, IOCTL_MMCOPYEX, &memory, sizeof(memory), &memory, sizeof(memory), &bytes_returned, nullptr);
}


uint64_t kd_driver::get_kernel_module_export(HANDLE device_handle, uint64_t ker_base, const char* func_name)
{
	IMAGE_DOS_HEADER dos_header{ 0 };
	IMAGE_NT_HEADERS nt_hander{ 0 };

	kd_driver::read(device_handle, ker_base, (uint64_t*)&dos_header, sizeof(IMAGE_DOS_HEADER));
	kd_driver::read(device_handle, ker_base + dos_header.e_lfanew, (uint64_t*)&nt_hander, sizeof(IMAGE_NT_HEADERS));

	const auto export_table = nt_hander.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];

	const auto export_data = reinterpret_cast<PIMAGE_EXPORT_DIRECTORY>(VirtualAlloc(nullptr, export_table.Size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE));

	kd_driver::read(device_handle, ker_base + export_table.VirtualAddress, (uint64_t*)export_data, export_table.Size);

	const auto delta = reinterpret_cast<uint64_t>(export_data) - export_table.VirtualAddress;
	const auto func_name_tab = reinterpret_cast<uint32_t*>(export_data->AddressOfNames + delta);
	const auto func_index_tab = reinterpret_cast<uint16_t*>(export_data->AddressOfNameOrdinals + delta);
	const auto func_address_tab = reinterpret_cast<uint32_t*>(export_data->AddressOfFunctions + delta);

	uint64_t func_addr = NULL;
	for (size_t i = 0; i < export_data->NumberOfNames; i++)
	{
		const char* func_names = reinterpret_cast<const char*>((uint64_t)func_name_tab[i] + delta);
		if (!strcmp(func_names, func_name))
		{
			const auto func_index = func_index_tab[i];
			const auto func_address = reinterpret_cast<char*>(func_address_tab[func_index]) + ker_base;

			func_addr = (uint64_t)func_address;
			break;
		}
	}

	VirtualFree(export_data, 0, MEM_RELEASE);
	return func_addr;
}

uint64_t kd_driver::ExAllocatePool(HANDLE device_handle, BYTE pool_type, SIZE_T size)
{
	const auto ExAllocate = kd_driver::get_kernel_module_export(device_handle, kd_driver::ntoskrnlbase, ("ExAllocatePoolWithTag"));
	if (!ExAllocate) {
		Log((L"[!] Failed to find ExAllocatePool") << std::endl);
		return 0;
	}

	uint64_t alloc = 0;
	kd_driver::func_call(device_handle, (void*)ExAllocate, &alloc, pool_type, size, 'PwtA');

	return alloc;
}

bool kd_driver::FreePool(HANDLE device_handle, uint64_t address) {
	if (!address)
		return 0;

	static uint64_t kernel_ExFreePool = get_kernel_module_export(device_handle, kd_driver::ntoskrnlbase, ("ExFreePool"));

	if (!kernel_ExFreePool) {
		Log((L"[!] Failed to find ExAllocatePool") << std::endl);
		return 0;
	}

	return func_call<void>(device_handle, (void*)kernel_ExFreePool, nullptr, address);
}


PVOID kd_driver::RtlLookupElementGenericTableAvl(HANDLE device_handle, nt::PRTL_AVL_TABLE Table, PVOID Buffer) {
	if (!Table)
		return nullptr;

	static uint64_t kernel_RtlDeleteElementGenericTableAvl = get_kernel_module_export(device_handle, kd_driver::ntoskrnlbase, ("RtlLookupElementGenericTableAvl"));

	if (!kernel_RtlDeleteElementGenericTableAvl) {
		Log((L"[!] Failed to find RtlLookupElementGenericTableAvl") << std::endl);
		return nullptr;
	}

	PVOID out;

	if (!func_call(device_handle, (void*)kernel_RtlDeleteElementGenericTableAvl, &out, Table, Buffer))
		return 0;

	return out;
}

BOOLEAN kd_driver::RtlDeleteElementGenericTableAvl(HANDLE device_handle, PVOID Table, PVOID Buffer)
{
	if (!Table)
		return false;

	static uint64_t kernel_RtlDeleteElementGenericTableAvl = get_kernel_module_export(device_handle, kd_driver::ntoskrnlbase, ("RtlDeleteElementGenericTableAvl"));

	if (!kernel_RtlDeleteElementGenericTableAvl) {
		Log((L"[!] Failed to find RtlDeleteElementGenericTableAvl") << std::endl);
		return false;
	}

	bool out;
	return (func_call(device_handle, (void*)kernel_RtlDeleteElementGenericTableAvl, &out, Table, Buffer) && out);
}

nt::PiDDBCacheEntry* kd_driver::LookupEntry(HANDLE device_handle, nt::PRTL_AVL_TABLE PiDDBCacheTable, ULONG timestamp, const wchar_t* name)
{
	nt::PiDDBCacheEntry localentry{};
	localentry.TimeDateStamp = timestamp;
	localentry.DriverName.Buffer = (PWSTR)name;
	localentry.DriverName.Length = (USHORT)(wcslen(name) * 2);
	localentry.DriverName.MaximumLength = localentry.DriverName.Length + 2;

	return (nt::PiDDBCacheEntry*)RtlLookupElementGenericTableAvl(device_handle, PiDDBCacheTable, (PVOID)&localentry);
}

uintptr_t kd_driver::FindPatternAtKernel(HANDLE device_handle, uintptr_t dwAddress, uintptr_t dwLen, BYTE* bMask, const char* szMask) {
	if (!dwAddress) {
		Log(L"[-] No module address to find pattern" << std::endl);
		return 0;
	}

	if (dwLen > 1024 * 1024 * 1024) { //if read is > 1GB
		Log(L"[-] Can't find pattern, Too big section" << std::endl);
		return 0;
	}

	auto sectionData = std::make_unique<BYTE[]>(dwLen);
	if (!read(device_handle, dwAddress, (uint64_t*)sectionData.get(), dwLen)) {
		Log(L"[-] Read failed in FindPatternAtKernel" << std::endl);
		return 0;
	}

	auto result = utils::FindPattern((uintptr_t)sectionData.get(), dwLen, bMask, szMask);

	if (result <= 0) {
		return 0;
	}
	result = dwAddress - (uintptr_t)sectionData.get() + result;
	return result;
}

uintptr_t kd_driver::FindSectionAtKernel(HANDLE device_handle, const char* sectionName, uintptr_t modulePtr, PULONG size) {
	if (!modulePtr)
		return 0;
	BYTE headers[0x1000];
	if (!read(device_handle, modulePtr, (uint64_t*)headers, 0x1000)) {
		Log(L"[-] Can't read module headers" << std::endl);
		return 0;
	}
	ULONG sectionSize = 0;
	uintptr_t section = (uintptr_t)utils::FindSection(sectionName, (uintptr_t)headers, &sectionSize);
	if (!section || !sectionSize) {
		Log(L"[-] Can't find section" << std::endl);
		return 0;
	}
	if (size)
		*size = sectionSize;
	return section - (uintptr_t)headers + modulePtr;
}


uintptr_t kd_driver::FindPatternInSectionAtKernel(HANDLE device_handle, const char* sectionName, uintptr_t modulePtr, BYTE* bMask, const char* szMask) {
	ULONG sectionSize = 0;
	uintptr_t section = FindSectionAtKernel(device_handle, sectionName, modulePtr, &sectionSize);
	return FindPatternAtKernel(device_handle, section, sectionSize, bMask, szMask);
}

PVOID kd_driver::ResolveRelativeAddress(HANDLE device_handle, _In_ PVOID Instruction, _In_ ULONG OffsetOffset, _In_ ULONG InstructionSize)
{
	ULONG_PTR Instr = (ULONG_PTR)Instruction;
	LONG RipOffset = 0;
	if (!read(device_handle, Instr + OffsetOffset, (uint64_t*)&RipOffset, sizeof(LONG))) {
		return nullptr;
	}
	PVOID ResolvedAddr = (PVOID)(Instr + InstructionSize + RipOffset);
	return ResolvedAddr;
}

bool kd_driver::ExAcquireResourceExclusiveLite(HANDLE device_handle, PVOID Resource, BOOLEAN wait) {
	if (!Resource)
		return 0;

	static uint64_t kernel_ExAcquireResourceExclusiveLite = get_kernel_module_export(device_handle, kd_driver::ntoskrnlbase, XorStr("ExAcquireResourceExclusiveLite"));

	if (!kernel_ExAcquireResourceExclusiveLite) {
		Log((L"[!] Failed to find ExAcquireResourceExclusiveLite") << std::endl);
		return 0;
	}

	BOOLEAN out;
	return (func_call(device_handle, (void*)kernel_ExAcquireResourceExclusiveLite, &out, Resource, wait) && out);
}

bool kd_driver::ExReleaseResourceLite(HANDLE device_handle, PVOID Resource) {
	if (!Resource)
		return false;

	static uint64_t kernel_ExReleaseResourceLite = get_kernel_module_export(device_handle, kd_driver::ntoskrnlbase, XorStr( "ExReleaseResourceLite"));

	if (!kernel_ExReleaseResourceLite) {
		Log((L"[!] Failed to find ExReleaseResourceLite") << std::endl);
		return false;
	}

	return func_call<void>(device_handle, (void*)kernel_ExReleaseResourceLite, nullptr, Resource);
}




bool kd_driver::ClearPiDDBCacheTable(HANDLE device_handle) { //PiDDBCacheTable added on LoadDriver

#ifdef PDB_OFFSETS
	auto PiDDBLockOffset = KDSymbolsHandler::GetInstance()->GetOffset(L"PiDDBLock");
	if (!PiDDBLockOffset)
	{
		Log(L"[-] Warning PiDDBLock not found" << std::endl);
		return false;
	}

	auto PiDDBCacheTableOffset = KDSymbolsHandler::GetInstance()->GetOffset(L"PiDDBCacheTable");
	if (!PiDDBLockOffset)
	{
		Log(L"[-] Warning PiDDBCacheTable not found" << std::endl);
		return false;
	}

	PVOID PiDDBLock = (PVOID)(kd_driver::ntoskrnlbase + PiDDBLockOffset);
	nt::PRTL_AVL_TABLE PiDDBCacheTable = (nt::PRTL_AVL_TABLE)(kd_driver::ntoskrnlbase + PiDDBCacheTableOffset);
#else
	auto PiDDBLockPtr = FindPatternInSectionAtKernel(device_handle, "PAGE", kd_driver::ntoskrnlbase, (PUCHAR)"\x8B\xD8\x85\xC0\x0F\x88\x00\x00\x00\x00\x65\x48\x8B\x04\x25\x00\x00\x00\x00\x66\xFF\x88\x00\x00\x00\x00\xB2\x01\x48\x8D\x0D\x00\x00\x00\x00\xE8\x00\x00\x00\x00\x4C\x8B\x00\x24", "xxxxxx????xxxxx????xxx????xxxxx????x????xx?x"); // 8B D8 85 C0 0F 88 ? ? ? ? 65 48 8B 04 25 ? ? ? ? 66 FF 88 ? ? ? ? B2 01 48 8D 0D ? ? ? ? E8 ? ? ? ? 4C 8B ? 24 update for build 22000.132
	auto PiDDBCacheTablePtr = FindPatternInSectionAtKernel(device_handle, "PAGE", kd_driver::ntoskrnlbase, (PUCHAR)"\x66\x03\xD2\x48\x8D\x0D", "xxxxxx"); // 66 03 D2 48 8D 0D

	if (PiDDBLockPtr == NULL) { // PiDDBLock pattern changes a lot from version 1607 of windows and we will need a second pattern if we want to keep simple as possible
		PiDDBLockPtr = FindPatternInSectionAtKernel(device_handle, "PAGE", kd_driver::ntoskrnlbase, (PUCHAR)"\x48\x8B\x0D\x00\x00\x00\x00\x48\x85\xC9\x0F\x85\x00\x00\x00\x00\x48\x8D\x0D\x00\x00\x00\x00\xE8\x00\x00\x00\x00\xE8", "xxx????xxxxx????xxx????x????x"); // 48 8B 0D ? ? ? ? 48 85 C9 0F 85 ? ? ? ? 48 8D 0D ? ? ? ? E8 ? ? ? ? E8 build 22449+ (pattern can be improved but just fine for now)
		if (PiDDBLockPtr == NULL) {
			PiDDBLockPtr = FindPatternInSectionAtKernel(device_handle, "PAGE", kd_driver::ntoskrnlbase, (PUCHAR)"\x8B\xD8\x85\xC0\x0F\x88\x00\x00\x00\x00\x65\x48\x8B\x04\x25\x00\x00\x00\x00\x48\x8D\x0D\x00\x00\x00\x00\xB2\x01\x66\xFF\x88\x00\x00\x00\x00\x90\xE8\x00\x00\x00\x00\x4C\x8B\x00\x24", "xxxxxx????xxxxx????xxx????xxxxx????xx????xx?x"); // 8B D8 85 C0 0F 88 ? ? ? ? 65 48 8B 04 25 ? ? ? ? 48 8D 0D ? ? ? ? B2 01 66 FF 88 ? ? ? ? 90 E8 ? ? ? ? 4C 8B ? 24 update for build 26100.1000
			if (PiDDBLockPtr == NULL) {
				Log((L"[-] Warning PiDDBLock not found") << std::endl);
				return false;
			}
			else {
				Log((L"[+] PiDDBLock found with third pattern") << std::endl);
				PiDDBLockPtr += 19;//third pattern offset
			}
		}
		else {
			Log((L"[+] PiDDBLock found with second pattern") << std::endl);
			PiDDBLockPtr += 16; //second pattern offset
		}
	}
	else {
		PiDDBLockPtr += 28; //first pattern offset
	}

	if (PiDDBCacheTablePtr == NULL) {
		PiDDBCacheTablePtr = FindPatternInSectionAtKernel(device_handle, "PAGE", kd_driver::ntoskrnlbase, (PUCHAR)"\x48\x8B\xF9\x33\xC0\x48\x8D\x0D", "xxxxxxxx"); // 48 8B F9 33 C0 48 8D 0D
		if (PiDDBCacheTablePtr == NULL) {
			Log((L"[-] Warning PiDDBCacheTable not found") << std::endl);
			return false;
		}
		else {
			Log((L"[+] PiDDBCacheTable found with second pattern") << std::endl);
			PiDDBCacheTablePtr += 2;//second pattern offset
		}
	}

	Log((L"[+] PiDDBLock Ptr 0x") << std::hex << PiDDBLockPtr << std::endl);
	Log((L"[+] PiDDBCacheTable Ptr 0x") << std::hex << PiDDBCacheTablePtr << std::endl);

	PVOID PiDDBLock = ResolveRelativeAddress(device_handle, (PVOID)PiDDBLockPtr, 3, 7);
	nt::PRTL_AVL_TABLE PiDDBCacheTable = (nt::PRTL_AVL_TABLE)ResolveRelativeAddress(device_handle, (PVOID)PiDDBCacheTablePtr, 6, 10);
#endif
	//context part is not used by lookup, lock or delete why we should use it?

	if (!ExAcquireResourceExclusiveLite(device_handle, PiDDBLock, true)) {
		Log((L"[-] Can't lock PiDDBCacheTable") << std::endl);
		return false;
	}
	Log((L"[+] PiDDBLock Locked") << std::endl);


	wstring n = GetDriverNameW();
	// search our entry in the table

	nt::PiDDBCacheEntry* pFoundEntry = (nt::PiDDBCacheEntry*)LookupEntry(device_handle, PiDDBCacheTable, iqvw64e_timestamp, n.c_str());
	if (pFoundEntry == nullptr) {
		Log(L"[-] Not found in cache" << std::endl);
		ExReleaseResourceLite(device_handle, PiDDBLock);
		return false;
	}

	// first, unlink from the list
	PLIST_ENTRY prev;
	if (!read(device_handle, (uintptr_t)pFoundEntry + (offsetof(struct nt::_PiDDBCacheEntry, List.Blink)), (uint64_t*)&prev, sizeof(_LIST_ENTRY*))) {
		Log(L"[-] Can't get prev entry" << std::endl);
		ExReleaseResourceLite(device_handle, PiDDBLock);
		return false;
	}
	PLIST_ENTRY next;
	if (!read(device_handle, (uintptr_t)pFoundEntry + (offsetof(struct nt::_PiDDBCacheEntry, List.Flink)), (uint64_t*)&next, sizeof(_LIST_ENTRY*))) {
		Log(L"[-] Can't get next entry" << std::endl);
		ExReleaseResourceLite(device_handle, PiDDBLock);
		return false;
	}

	Log("[+] Found Table Entry = 0x" << std::hex << pFoundEntry << std::endl);

	if (!write(device_handle, (uintptr_t)prev + (offsetof(struct _LIST_ENTRY, Flink)), (uint64_t*)&next, sizeof(_LIST_ENTRY*))) {
		Log(L"[-] Can't set next entry" << std::endl);
		ExReleaseResourceLite(device_handle, PiDDBLock);
		return false;
	}
	if (!write(device_handle, (uintptr_t)next + (offsetof(struct _LIST_ENTRY, Blink)), (uint64_t*)&prev, sizeof(_LIST_ENTRY*))) {
		Log(L"[-] Can't set prev entry" << std::endl);
		ExReleaseResourceLite(device_handle, PiDDBLock);
		return false;
	}

	// then delete the element from the avl table
	if (!RtlDeleteElementGenericTableAvl(device_handle, PiDDBCacheTable, pFoundEntry)) {
		Log(XOR_STRING_W(L"[-] Can't delete from PiDDBCacheTable") << std::endl);
		ExReleaseResourceLite(device_handle, PiDDBLock);
		return false;
	}

	//Decrement delete count
	ULONG cacheDeleteCount = 0;
	read(device_handle, (uintptr_t)PiDDBCacheTable + (offsetof(struct nt::_RTL_AVL_TABLE, DeleteCount)), (uint64_t*)&cacheDeleteCount, sizeof(ULONG));
	if (cacheDeleteCount > 0) {
		cacheDeleteCount--;
		write(device_handle, (uintptr_t)PiDDBCacheTable + (offsetof(struct nt::_RTL_AVL_TABLE, DeleteCount)), (uint64_t*)&cacheDeleteCount, sizeof(ULONG));
	}

	// release the ddb resource lock
	ExReleaseResourceLite(device_handle, PiDDBLock);

	Log(XOR_STRING_W(L"[+] PiDDBCacheTable Cleaned") << std::endl);

	return true;
}

bool kd_driver::ClearMmUnloadedDrivers(HANDLE device_handle)
{
	ULONG buffer_size = 0;
	void* buffer = nullptr;

	NTSTATUS status = NtQuerySystemInformation(static_cast<SYSTEM_INFORMATION_CLASS>(nt::SystemExtendedHandleInformation), buffer, buffer_size, &buffer_size);

	while (status == nt::STATUS_INFO_LENGTH_MISMATCH)
	{
		VirtualFree(buffer, 0, MEM_RELEASE);

		buffer = VirtualAlloc(nullptr, buffer_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
		status = NtQuerySystemInformation(static_cast<SYSTEM_INFORMATION_CLASS>(nt::SystemExtendedHandleInformation), buffer, buffer_size, &buffer_size);
	}

	if (!NT_SUCCESS(status) || buffer == 0)
	{
		if (buffer != 0)
			VirtualFree(buffer, 0, MEM_RELEASE);
		return false;
	}

	uint64_t object = 0;

	auto system_handle_inforamtion = static_cast<nt::PSYSTEM_HANDLE_INFORMATION_EX>(buffer);

	for (auto i = 0u; i < system_handle_inforamtion->HandleCount; ++i)
	{
		const nt::SYSTEM_HANDLE current_system_handle = system_handle_inforamtion->Handles[i];

		if (current_system_handle.UniqueProcessId != reinterpret_cast<HANDLE>(static_cast<uint64_t>(GetCurrentProcessId())))
			continue;

		if (current_system_handle.HandleValue == device_handle)
		{
			object = reinterpret_cast<uint64_t>(current_system_handle.Object);
			break;
		}
	}

	VirtualFree(buffer, 0, MEM_RELEASE);

	if (!object)
		return false;

	uint64_t device_object = 0;

	if (!read(device_handle, object + 0x8, &device_object, sizeof(device_object)) || !device_object) {
		//Log(L"[!] Failed to find device_object" << std::endl);
		return false;
	}

	uint64_t driver_object = 0;

	if (!read(device_handle, device_object + 0x8, &driver_object, sizeof(driver_object)) || !driver_object) {
		//Log(L"[!] Failed to find driver_object" << std::endl);
		return false;
	}

	uint64_t driver_section = 0;

	if (!read(device_handle, driver_object + 0x28, &driver_section, sizeof(driver_section)) || !driver_section) {
		//Log(L"[!] Failed to find driver_section" << std::endl);
		return false;
	}

	UNICODE_STRING us_driver_base_dll_name = { 0 };

	if (!read(device_handle, driver_section + 0x58, (uint64_t*)&us_driver_base_dll_name, sizeof(us_driver_base_dll_name)) || us_driver_base_dll_name.Length == 0) {
		//Log(L"[!] Failed to find driver name" << std::endl);
		return false;
	}

	auto unloadedName = std::make_unique<wchar_t[]>((ULONG64)us_driver_base_dll_name.Length / 2ULL + 1ULL);
	if (!read(device_handle, (uintptr_t)us_driver_base_dll_name.Buffer, (uint64_t*)unloadedName.get(), us_driver_base_dll_name.Length)) {
		//Log(L"[!] Failed to read driver name" << std::endl);
		return false;
	}

	us_driver_base_dll_name.Length = 0; //MiRememberUnloadedDriver will check if the length > 0 to save the unloaded driver

	if (!write(device_handle, driver_section + 0x58, (uint64_t*)&us_driver_base_dll_name, sizeof(us_driver_base_dll_name))) {
		//Log(L"[!] Failed to write driver name length" << std::endl);
		return false;
	}

	//Log(L"[+] MmUnloadedDrivers Cleaned: " << unloadedName << std::endl);
	return true;
}

bool kd_driver::ClearKernelHashBucketList(HANDLE device_handle)
{
	uint64_t ci = utils::GetKernelModule("CI.dll");
	if (!ci) {
		Log(L"[-] Can't Find ci.dll module address" << std::endl);
		return false;
	}

	//Thanks @KDIo3 and @Swiftik from UnknownCheats
#ifdef PDB_OFFSETS
	auto g_KernelHashBucketListOffset = KDSymbolsHandler::GetInstance()->GetOffset(L"g_KernelHashBucketList");
	if (!g_KernelHashBucketListOffset)
	{
		Log(L"[-] Can't Find g_KernelHashBucketList Offset" << std::endl);
		return false;
	}

	auto g_HashCacheLockOffset = KDSymbolsHandler::GetInstance()->GetOffset(L"g_HashCacheLock");
	if (!g_KernelHashBucketListOffset)
	{
		Log(L"[-] Can't Find g_HashCacheLock Offset" << std::endl);
		return false;
	}

	PVOID g_KernelHashBucketList = (PVOID)(ci + g_KernelHashBucketListOffset);
	PVOID g_HashCacheLock = (PVOID)(ci + g_HashCacheLockOffset);
#else
	auto sig = FindPatternInSectionAtKernel(device_handle, "PAGE", ci, PUCHAR("\x48\x8B\x1D\x00\x00\x00\x00\xEB\x00\xF7\x43\x40\x00\x20\x00\x00"), "xxx????x?xxxxxxx");
	if (!sig) {
		Log(L"[-] Can't Find g_KernelHashBucketList" << std::endl);
		return false;
	}
	auto sig2 = FindPatternAtKernel(device_handle, (uintptr_t)sig - 50, 50, PUCHAR("\x48\x8D\x0D"), "xxx");
	if (!sig2) {
		Log(L"[-] Can't Find g_HashCacheLock" << std::endl);
		return false;
	}
	const auto g_KernelHashBucketList = ResolveRelativeAddress(device_handle, (PVOID)sig, 3, 7);
	const auto g_HashCacheLock = ResolveRelativeAddress(device_handle, (PVOID)sig2, 3, 7);
	if (!g_KernelHashBucketList || !g_HashCacheLock)
	{
		Log(L"[-] Can't Find g_HashCache relative address" << std::endl);
		return false;
	}
#endif

	Log(L"[+] g_KernelHashBucketList Found 0x" << std::hex << g_KernelHashBucketList << std::endl);

	if (!ExAcquireResourceExclusiveLite(device_handle, g_HashCacheLock, true)) {
		Log(L"[-] Can't lock g_HashCacheLock" << std::endl);
		return false;
	}
	Log(L"[+] g_HashCacheLock Locked" << std::endl);

	nt::HashBucketEntry* prev = (nt::HashBucketEntry*)g_KernelHashBucketList;
	nt::HashBucketEntry* entry = 0;
	if (!read(device_handle, (uintptr_t)prev, (uint64_t*)&entry, sizeof(entry))) {
		Log(L"[-] Failed to read first g_KernelHashBucketList entry!" << std::endl);
		if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
			Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
		}
		return false;
	}
	if (!entry) {
		Log(L"[!] g_KernelHashBucketList looks empty!" << std::endl);
		if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
			Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
		}
		return true;
	}

	std::wstring wdname = GetDriverNameW();
	std::wstring search_path = GetDriverPath();
	SIZE_T expected_len = (search_path.length() - 2) * 2;

	while (entry) {

		USHORT wsNameLen = 0;
		if (!read(device_handle, (uintptr_t)entry + offsetof(nt::HashBucketEntry, DriverName.Length), (uint64_t*)&wsNameLen, sizeof(wsNameLen)) || wsNameLen == 0) {
			Log(L"[-] Failed to read g_KernelHashBucketList entry text len!" << std::endl);
			if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
				Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
			}
			return false;
		}

		if (expected_len == wsNameLen) {
			wchar_t* wsNamePtr = 0;
			if (!read(device_handle, (uintptr_t)entry + offsetof(nt::HashBucketEntry, DriverName.Buffer), (uint64_t*)&wsNamePtr, sizeof(wsNamePtr)) || !wsNamePtr) {
				Log(L"[-] Failed to read g_KernelHashBucketList entry text ptr!" << std::endl);
				if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
					Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
				}
				return false;
			}

			auto wsName = std::make_unique<wchar_t[]>((ULONG64)wsNameLen / 2ULL + 1ULL);
			if (!read(device_handle, (uintptr_t)wsNamePtr, (uint64_t*)wsName.get(), wsNameLen)) {
				Log(L"[-] Failed to read g_KernelHashBucketList entry text!" << std::endl);
				if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
					Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
				}
				return false;
			}

			size_t find_result = std::wstring(wsName.get()).find(wdname);
			if (find_result != std::wstring::npos) {
				Log(L"[+] Found In g_KernelHashBucketList: " << std::wstring(&wsName[find_result]) << std::endl);
				nt::HashBucketEntry* Next = 0;
				if (!read(device_handle, (uintptr_t)entry, (uint64_t*)&Next, sizeof(Next))) {
					Log(L"[-] Failed to read g_KernelHashBucketList next entry ptr!" << std::endl);
					if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
						Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
					}
					return false;
				}

				if (!write(device_handle, (uintptr_t)prev, (uint64_t*)&Next, sizeof(Next))) {
					Log(L"[-] Failed to write g_KernelHashBucketList prev entry ptr!" << std::endl);
					if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
						Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
					}
					return false;
				}

				if (!FreePool(device_handle, (uintptr_t)entry)) {
					Log(L"[-] Failed to clear g_KernelHashBucketList entry pool!" << std::endl);
					if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
						Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
					}
					return false;
				}
				Log(L"[+] g_KernelHashBucketList Cleaned" << std::endl);
				if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
					Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
					if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
						Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
					}
					return false;
				}
				return true;
			}
		}
		prev = entry;
		//read next
		if (!read(device_handle, (uintptr_t)entry, (uint64_t*)&entry, sizeof(entry))) {
			Log(L"[-] Failed to read g_KernelHashBucketList next entry!" << std::endl);
			if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
				Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
			}
			return false;
		}
	}

	if (!ExReleaseResourceLite(device_handle, g_HashCacheLock)) {
		Log(L"[-] Failed to release g_KernelHashBucketList lock!" << std::endl);
	}
	return false;
}

bool kd_driver::ClearWdFilterDriverList(HANDLE device_handle)
{
	auto WdFilter = utils::GetKernelModule("WdFilter.sys");
	if (!WdFilter) {
		Log("[+] WdFilter.sys not loaded, clear skipped" << std::endl);
		return true;
	}

#ifdef PDB_OFFSETS
	uintptr_t MpBmDocOpenRules = KDSymbolsHandler::GetInstance()->GetOffset(L"MpBmDocOpenRules");
	if (!MpBmDocOpenRules)
	{
		Log("[-] Failed To Get MpBmDocOpenRules." << std::endl);
		return false;
	}
	MpBmDocOpenRules += WdFilter;

	uintptr_t RuntimeDriversList_Head = MpBmDocOpenRules + 0x70;
	uintptr_t RuntimeDriversCount = MpBmDocOpenRules + 0x60;
	uintptr_t RuntimeDriversArray = MpBmDocOpenRules + 0x68;
	read(device_handle, RuntimeDriversArray, &RuntimeDriversArray, sizeof(uintptr_t));

	uintptr_t MpFreeDriverInfoEx = KDSymbolsHandler::GetInstance()->GetOffset(L"MpFreeDriverInfoEx");
	if (!MpFreeDriverInfoEx)
	{
		Log("[-] Failed To Get MpFreeDriverInfoEx." << std::endl);
		return false;
	}
	MpFreeDriverInfoEx += WdFilter;
#else
	auto RuntimeDriversList = FindPatternInSectionAtKernel(device_handle, "PAGE", WdFilter, (PUCHAR)"\x48\x8B\x0D\x00\x00\x00\x00\xFF\x05", "xxx????xx");
	if (!RuntimeDriversList) {
		Log("[!] Failed to find WdFilter RuntimeDriversList" << std::endl);
		return false;
	}

	auto RuntimeDriversCountRef = FindPatternInSectionAtKernel(device_handle, "PAGE", WdFilter, (PUCHAR)"\xFF\x05\x00\x00\x00\x00\x48\x39\x11", "xx????xxx");
	if (!RuntimeDriversCountRef) {
		Log("[!] Failed to find WdFilter RuntimeDriversCount" << std::endl);
		return false;
	}

	// MpCleanupDriverInfo->MpFreeDriverInfoEx
	// The pattern only focus in the 0x8 offset and the possibility of the different order for the instructions
	/*
		49 8B C9                      mov     rcx, r9         ; P
		49 89 50 08                   mov     [r8+8], rdx
		E8 FB F0 FD FF                call    MpFreeDriverInfoEx
		48 8B 0D FC AA FA FF          mov     rcx, cs:qword_1C0021BF0
		E9 21 FF FF FF                jmp     loc_1C007701A
	*/
	auto MpFreeDriverInfoExRef = FindPatternInSectionAtKernel(device_handle, "PAGE", WdFilter, (PUCHAR)"\x89\x00\x08\xE8\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xE9", "x?xx???????????x");
	if (!MpFreeDriverInfoExRef) {
		/*
			48 89 4A 08                   mov     [rdx+8], rcx
			49 8B C8                      mov     rcx, r8         ; P
			E8 C3 58 FE FF                call    sub_1C0065308
			48 8B 0D 44 41 FA FF          mov     rcx, cs:qword_1C0023B90
			E9 39 FF FF FF                jmp     loc_1C007F98A */
		MpFreeDriverInfoExRef = FindPatternInSectionAtKernel(device_handle, "PAGE", WdFilter, (PUCHAR)"\x89\x00\x08\x00\x00\x00\xE8\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xE9", "x?x???x???????????x");
		if (!MpFreeDriverInfoExRef) {
			Log("[!] Failed to find WdFilter MpFreeDriverInfoEx" << std::endl);
			return false;
		}
		else {
			Log("[+] Found WdFilter MpFreeDriverInfoEx with second pattern" << std::endl);
		}
		MpFreeDriverInfoExRef += 0x3; // adjust for next sum offset
	}

	MpFreeDriverInfoExRef += 0x3; // skip until call instruction

	RuntimeDriversList = (uintptr_t)ResolveRelativeAddress(device_handle, (PVOID)RuntimeDriversList, 3, 7);
	uintptr_t RuntimeDriversList_Head = RuntimeDriversList - 0x8;
	uintptr_t RuntimeDriversCount = (uintptr_t)ResolveRelativeAddress(device_handle, (PVOID)RuntimeDriversCountRef, 2, 6);
	uintptr_t RuntimeDriversArray = RuntimeDriversCount + 0x8;
	read(device_handle, RuntimeDriversArray, (uint64_t*)&RuntimeDriversArray, sizeof(uintptr_t));
	uintptr_t MpFreeDriverInfoEx = (uintptr_t)ResolveRelativeAddress(device_handle, (PVOID)MpFreeDriverInfoExRef, 1, 5);
#endif

	auto ReadListEntry = [&](uintptr_t Address) -> LIST_ENTRY* { // Useful lambda to read LIST_ENTRY
		LIST_ENTRY* Entry;
		if (!read(device_handle, Address, (uint64_t*)&Entry, sizeof(LIST_ENTRY*))) return 0;
		return Entry;
		};




	for (LIST_ENTRY* Entry = ReadListEntry(RuntimeDriversList_Head);
		Entry != (LIST_ENTRY*)RuntimeDriversList_Head;
		Entry = ReadListEntry((uintptr_t)Entry + (offsetof(struct _LIST_ENTRY, Flink))))
	{
		UNICODE_STRING Unicode_String;
		if (read(device_handle, (uintptr_t)Entry + 0x10, (uint64_t*)&Unicode_String, sizeof(UNICODE_STRING))) {
			auto ImageName = std::make_unique<wchar_t[]>((ULONG64)Unicode_String.Length / 2ULL + 1ULL);
			if (read(device_handle, (uintptr_t)Unicode_String.Buffer, (uint64_t*)ImageName.get(), Unicode_String.Length)) {
				if (wcsstr(ImageName.get(), kd_driver::GetDriverNameW().c_str())) {

					//remove from RuntimeDriversArray
					bool removedRuntimeDriversArray = false;
					PVOID SameIndexList = (PVOID)((uintptr_t)Entry - 0x10);
					for (int k = 0; k < 256; k++) { // max RuntimeDriversArray elements
						PVOID value = 0;
						read(device_handle, RuntimeDriversArray + (k * 8), (uint64_t*)&value, sizeof(PVOID));
						if (value == SameIndexList) {
							PVOID emptyval = (PVOID)(RuntimeDriversCount + 1); // this is not count+1 is position of cout addr+1
							write(device_handle, RuntimeDriversArray + (k * 8), (uint64_t*)&emptyval, sizeof(PVOID));
							removedRuntimeDriversArray = true;
							break;
						}
					}

					if (!removedRuntimeDriversArray) {
						Log("[!] Failed to remove from RuntimeDriversArray" << std::endl);
						return false;
					}

					auto NextEntry = ReadListEntry(uintptr_t(Entry) + (offsetof(struct _LIST_ENTRY, Flink)));
					auto PrevEntry = ReadListEntry(uintptr_t(Entry) + (offsetof(struct _LIST_ENTRY, Blink)));

					write(device_handle, uintptr_t(NextEntry) + (offsetof(struct _LIST_ENTRY, Blink)), (uint64_t*)&PrevEntry, sizeof(LIST_ENTRY::Blink));
					write(device_handle, uintptr_t(PrevEntry) + (offsetof(struct _LIST_ENTRY, Flink)), (uint64_t*)&NextEntry, sizeof(LIST_ENTRY::Flink));

					// decrement RuntimeDriversCount
					ULONG current = 0;
					read(device_handle, RuntimeDriversCount, (uint64_t*)&current, sizeof(ULONG));
					current--;
					write(device_handle, RuntimeDriversCount, (uint64_t*)&current, sizeof(ULONG));

					// call MpFreeDriverInfoEx
					uintptr_t DriverInfo = (uintptr_t)Entry - 0x20;

					//verify DriverInfo Magic
					USHORT Magic = 0;
					read(device_handle, DriverInfo, (uint64_t*)&Magic, sizeof(USHORT));
					if (Magic != 0xDA18) {
						Log("[!] DriverInfo Magic is invalid, new wdfilter version?, driver info will not be released to prevent bsod" << std::endl);
					}
					else {
						func_call<void>(device_handle, (void*)MpFreeDriverInfoEx, nullptr, DriverInfo);
					}

					Log("[+] WdFilterDriverList Cleaned: " << ImageName << std::endl);
					return true;
				}
			}
		}
	}
	return false;
}