#pragma once
#include "drv.h"
#include "ntstruct.h"
#include "..\XorStr.h"
using namespace std;

#define IOCTL_MMCOPY          CTL_CODE(FILE_DEVICE_UNKNOWN, 0x9512, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_MMCOPYEX        CTL_CODE(FILE_DEVICE_UNKNOWN, 0x9513, METHOD_BUFFERED, FILE_ANY_ACCESS)




namespace kd_driver
{

	extern HANDLE iqvw64e_device_handle;
	extern uint64_t ntoskrnlbase;
	constexpr DWORD iqvw64e_timestamp = 0x5284EAC3;

	std::wstring GetDriverPath();
	std::wstring GetDriverNameW();
	bool AcquireDebugPrivilege();

	bool Load(PDWORD status);
	bool Unload();

	bool MmCopy(HANDLE device_handle, uint64_t Base, uint64_t Buffer, SIZE_T Size);
	bool read(HANDLE device_handle, uint64_t Base, uint64_t* Buffer, SIZE_T Size);
	bool write(HANDLE device_handle, uint64_t Base, uint64_t* Buffer, SIZE_T Size);

	bool writeEx(HANDLE device_handle, uint64_t Base, uint64_t* Buffer, SIZE_T Size);


	bool MdlIoSpace(HANDLE device_handle, uint64_t address, uint64_t* Buffer, SIZE_T Size);

	uint64_t get_kernel_module_export(HANDLE device_handle, uint64_t ker_base, const char* func_name);
	uint64_t ExAllocatePool(HANDLE device_handle, BYTE pool_type, SIZE_T size);
	bool FreePool(HANDLE device_handle, uint64_t address);

	nt::PiDDBCacheEntry* LookupEntry(HANDLE device_handle, nt::PRTL_AVL_TABLE PiDDBCacheTable, ULONG timestamp, const wchar_t* name);
	PVOID RtlLookupElementGenericTableAvl(HANDLE device_handle, nt::PRTL_AVL_TABLE Table, PVOID Buffer);
	BOOLEAN RtlDeleteElementGenericTableAvl(HANDLE device_handle, PVOID Table, PVOID Buffer);

	uintptr_t FindPatternAtKernel(HANDLE device_handle, uintptr_t dwAddress, uintptr_t dwLen, BYTE* bMask, const char* szMask);
	uintptr_t FindSectionAtKernel(HANDLE device_handle, const char* sectionName, uintptr_t modulePtr, PULONG size);
	uintptr_t FindPatternInSectionAtKernel(HANDLE device_handle, const char* sectionName, uintptr_t modulePtr, BYTE* bMask, const char* szMask);

	PVOID ResolveRelativeAddress(HANDLE device_handle, _In_ PVOID Instruction, _In_ ULONG OffsetOffset, _In_ ULONG InstructionSize);
	bool ExAcquireResourceExclusiveLite(HANDLE device_handle, PVOID Resource, BOOLEAN wait);
	bool ExReleaseResourceLite(HANDLE device_handle, PVOID Resource);


	bool ClearPiDDBCacheTable(HANDLE device_handle);
	bool ClearMmUnloadedDrivers(HANDLE device_handle);
	bool ClearKernelHashBucketList(HANDLE device_handle);
	bool ClearWdFilterDriverList(HANDLE device_handle);



	template<typename T, typename ...A>
	bool func_call(HANDLE device_handle, void* func, T* out, const A...arg)
	{
		constexpr auto call_void = std::is_same_v<T, void>;

		if constexpr (!call_void) {
			if (!out)
				return false;
		}
		else {
			UNREFERENCED_PARAMETER(out);
		}



		const auto ntdll = GetModuleHandleA("ntdll.dll");
		const auto ntaddrom = reinterpret_cast<void*>(GetProcAddress(ntdll, ("NtAddAtom")));

		char jmp_shellcode[] = { (char)0x48, (char)0xb8, (char)0x00, (char)0x00, (char)0x00, (char)0x00, (char)0x00, (char)0x00, (char)0x00,(char)0x00, (char)0xff, (char)0xe0 };
		char original_func[sizeof(jmp_shellcode)];
		*(uint64_t*)&jmp_shellcode[2] = (uint64_t)func;

		uint64_t ker_NtAddatom = get_kernel_module_export(device_handle, kd_driver::ntoskrnlbase, ("NtAddAtom"));

		
		read(device_handle, ker_NtAddatom, (uint64_t*)&original_func, sizeof(jmp_shellcode));
		


		if (!writeEx(device_handle, ker_NtAddatom, (uint64_t*)&jmp_shellcode, sizeof(jmp_shellcode)))
		{
			//printf("write fail\n");
			return 0;
		}

		if constexpr (!call_void) {
			using FunctionFn = T(__stdcall*)(A...);
			const auto Function = reinterpret_cast<FunctionFn>(ntaddrom);

			*out = Function(arg...);
		}
		else {
			using FunctionFn = void(__stdcall*)(A...);
			const auto Function = reinterpret_cast<FunctionFn>(ntaddrom);

			Function(arg...);
		}

		return writeEx(device_handle, ker_NtAddatom, (uint64_t*)&original_func, sizeof(jmp_shellcode));

	}
}