#include "process_notify_callback.h"
#include "kernel_function.h"
#include "protect_filter.h"
#include "utils.h"
#pragma warning(disable : 4309 4838)
namespace process_notify_callback
{
	auto process_notify_routine(
		_In_ HANDLE ParentId,
		_In_ HANDLE ProcessId,
		_In_ BOOLEAN Create) -> VOID
	{
		UNREFERENCED_PARAMETER(ParentId);
		if (!Create)	
		{

			//进程退出时进行 对插链的清理操作
			if (protect_filter::is_protect_pid(HandleToLong(ProcessId), nullptr))
				protect_filter::remove_protect_list(g_protect_list, HandleToLong(ProcessId), nullptr);

			if (protect_filter::is_white_pid(HandleToLong(ProcessId), nullptr))
			{

				DbgBreakPoint();
				protect_filter::remove_protect_list(g_white_list, HandleToLong(ProcessId), nullptr);
			}
		}

		/*DbgPrintEx(77, 0, "ParentId:%d | ProcessId:%d | Create:%d\n",
			ParentId, ProcessId, Create);*/
	}

	uintptr_t beep_trampoline_ptr;

	auto install_process_notify() -> NTSTATUS
	{
		beep_trampoline_ptr = utils::get_beep_trampoline_ptr();
		if (!beep_trampoline_ptr)
			return STATUS_UNSUCCESSFUL;

		char jmprax[] = { 0x48, 0xB8, 0x13, 0x09, 0xFC, 0xD6, 0xFC, 0x7F, 0x00, 0x00, 0xFF, 0xE0 };
		*(uintptr_t*)&jmprax[2] = (uintptr_t)process_notify_routine;
		auto status = kernel_function::MmMDLPagesCopy((PVOID)beep_trampoline_ptr, &jmprax, sizeof(jmprax), 0x0000002, KernelMode);
		if (!NT_SUCCESS(status))
			return status;

		return PsSetCreateProcessNotifyRoutine((PCREATE_PROCESS_NOTIFY_ROUTINE)beep_trampoline_ptr, FALSE);
	}

	auto remov_process_notify() -> NTSTATUS
	{
		if (!beep_trampoline_ptr)
			return STATUS_UNSUCCESSFUL;

		auto status = PsSetCreateProcessNotifyRoutine((PCREATE_PROCESS_NOTIFY_ROUTINE)beep_trampoline_ptr, TRUE);
		if (!NT_SUCCESS(status))
			return status;

		char nullcode[12]{ 0 };
		status = kernel_function::MmMDLPagesCopy((PVOID)beep_trampoline_ptr, &nullcode, sizeof(nullcode), 0x0000002, KernelMode);
		if (!NT_SUCCESS(status))
			return status;

		return status;
	}

}