#include "kernel_function.h"
#include "kernel_api.h"
#include "utils.h"
#pragma warning(disable : 4309)
#pragma warning(disable : 4838)
namespace kernel_function
{
	PEPROCESS process{ 0 };

	auto FindProcess(const CHAR* ProcessName) -> PEPROCESS
	{
		for (ULONG i = 4; i < 0x401000; i += 4)
		{
			auto status = PsLookupProcessByProcessId(ULongToHandle(i), &process);
			if (!NT_SUCCESS(status))
				continue;

			auto image_file_name = reinterpret_cast<char*>(kernel_api::psgetprocessimagefilname(process));
			if (!strcmp(image_file_name, ProcessName))
				return process;

			ObDereferenceObject(process);
		}

		return process;
	}
	auto GetKernelModule(CONST CHAR* Module) -> uintptr_t
	{
		UNICODE_STRING unFuncName{ 0 };

		void* buffer = nullptr;
		DWORD buffer_size = 0;


	retry:
		buffer = ExAllocatePool(NonPagedPool, buffer_size);
		NTSTATUS status = kernel_api::ntquerysysteminformation(static_cast <SYSTEM_INFORMATION_CLASS>(11), buffer, buffer_size, &buffer_size);
		if (status == STATUS_INFO_LENGTH_MISMATCH)
		{
			if (buffer)
				ExFreePool(buffer);

			goto retry;
		}


		uintptr_t base = 0;
		const auto modulebase = static_cast<PRTL_PROCESS_MODULES>(buffer);

		if (!modulebase)
			return 0;

		for (auto i = 0u; i < modulebase->NumberOfModules; i++)
		{
			const char* current_module_name = (reinterpret_cast<char*>(modulebase->Modules[i].FullPathName) + modulebase->Modules[i].OffsetToFileName);

			if (!_stricmp(current_module_name, Module))
			{
				const auto ImageBase = modulebase->Modules[i].ImageBase;
				base = (uintptr_t)ImageBase;
				break;
			}
		}


		ExFreePool(buffer);

		return base;
	}

	auto MmMDLPagesCopy(IN PVOID Address, PVOID Buff, SIZE_T Size, BYTE Type, CHAR AccessMode) -> NTSTATUS
	{
		PMDL pMDL = IoAllocateMdl(Address, (ULONG)Size, FALSE, FALSE, NULL);
		if (!pMDL) return STATUS_INSUFFICIENT_RESOURCES;

		__try
		{
			MmProbeAndLockPages(pMDL, AccessMode, IoReadAccess);
		}
		__except (EXCEPTION_EXECUTE_HANDLER)
		{
			IoFreeMdl(pMDL);
			return STATUS_ACCESS_DENIED;
		}

		PVOID NewAddress = MmMapLockedPagesSpecifyCache(pMDL, KernelMode, MmNonCached, NULL, 0, NormalPagePriority);
		if (!NewAddress)
		{
			MmUnlockPages(pMDL);
			IoFreeMdl(pMDL);
			return STATUS_INSUFFICIENT_RESOURCES;
		}

		if (Type == 0x0000001)
			RtlCopyMemory(Buff, NewAddress, Size);
		if (Type == 0x0000002)
			RtlCopyMemory(NewAddress, Buff, Size);


		MmUnmapLockedPages(NewAddress, pMDL);
		MmUnlockPages(pMDL);
		IoFreeMdl(pMDL);

		return STATUS_SUCCESS;
	}

	auto GetMmverifyCallBackFlags() -> uintptr_t
	{
		uintptr_t pfunc_address = NULL;
		auto version = utils::GetVersion();
		if (version.dwBuildNumber == 7600 || version.dwBuildNumber == 7601)
		{
			auto pfunc_extport_address = reinterpret_cast<unsigned char*>(utils::kernel_get_expotr_func(L"ObRegisterCallbacks"));
			for (size_t i = 0; i < 0x256; i++)
			{
				if (pfunc_extport_address[i] == 0x74 && pfunc_extport_address[i + 1] == 0x09)
				{
					pfunc_address = (uintptr_t)(*(long*)(&pfunc_extport_address[i + 3]) + &pfunc_extport_address[i + 7]);
					break;
				}
			}
		}
		else
		{
			auto pfunc_extport_address = reinterpret_cast<unsigned char*>(utils::kernel_get_expotr_func(L"PsSetCreateThreadNotifyRoutineEx"));
			for (size_t i = 0; i < 0x256; i++)
			{
				if (pfunc_extport_address[i] == 0xE8)
				{
					pfunc_address = (uintptr_t)(*(long*)(&pfunc_extport_address[i + 1]) + &pfunc_extport_address[i + 5]);
					break;
				}
			}
		}
		return pfunc_address;
	}

	//ÆÁ±Î0xC0000022LL´íÎó
	auto Mack_MmVerifyCallBackFlags(BYTE** orgin_byte) -> BOOL
	{
		uintptr_t pfunc_address = GetMmverifyCallBackFlags();
		if (!pfunc_address)
			return FALSE;

		//xor,rax,rax
		//mov rax,1
		//ret
		char shellcode[] = { 0x31, 0xC0, 0xB8, 0x01, 0x00 ,0x00, 0x00,0xC3 };

		memcpy(orgin_byte, (CONST void*)pfunc_address, sizeof(shellcode));

		auto status = MmMDLPagesCopy((PVOID)pfunc_address, &shellcode, sizeof(shellcode), 0x0000002, KernelMode);
		if (!NT_SUCCESS(status))
			return FALSE;

		return TRUE;
	}

	auto Orgin_MmVerifyCallBackFlags(BYTE* orgin_byte) -> BOOL
	{
		uintptr_t pfunc_address = GetMmverifyCallBackFlags();
		if (!pfunc_address)
			return FALSE;

		auto status = MmMDLPagesCopy((PVOID)pfunc_address, orgin_byte, sizeof(orgin_byte), 0x0000002, KernelMode);
		if (!NT_SUCCESS(status))
			return FALSE;

		return TRUE;
	}


}