#include "comm_win10.h"
#include "comm_dispatch.h"
#include "utils.h"
#include "kernel_function.h"
#pragma warning(disable : 4309 4838)

typedef NTSTATUS(__fastcall* fnHalpTimerConvertAuxiliaryCounterToPerformanceCounter)(PVOID, PVOID, PVOID);
fnHalpTimerConvertAuxiliaryCounterToPerformanceCounter _HalpTimerConvertAuxiliaryCounterToPerformanceCounter = 0;

uintptr_t beep_trampoline_ptr;
uintptr_t func_call_address;

namespace comm_win10
{
	auto comm_callback(PVOID a1, PVOID a2, PVOID a3) -> NTSTATUS
	{
		comm_dispatch::dispatch((CMD_COMM*)a1);
		return _HalpTimerConvertAuxiliaryCounterToPerformanceCounter(a1, a2, a3);
	}

	auto comm_init() -> BOOL
	{
		beep_trampoline_ptr = utils::get_beep_trampoline_ptr();
		return !beep_trampoline_ptr ? FALSE : TRUE;
	}

	auto comm_install_hook() -> BOOL
	{
		auto ntoskrnl = kernel_function::GetKernelModule("ntoskrnl.exe");
		if (!ntoskrnl)
			return FALSE;

		// mov rax,comm_callback
		// jmp rax

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

		auto code_addr = utils::FindSectionsCode(ntoskrnl, "\x48\x8B\x05\x00\x00\x00\x00\x00\x00\x48\x8B\x05\x00\x00\x00\x00\xE8\x00\x00\x00\x00\x8B\xC8",
			"xxx??????xxx????x????xx", "PAGE");

		func_call_address = (*(long*)(code_addr + 3)) + (code_addr + 7);
		_HalpTimerConvertAuxiliaryCounterToPerformanceCounter =
			*(fnHalpTimerConvertAuxiliaryCounterToPerformanceCounter*)(func_call_address);

		*(uintptr_t*)func_call_address = beep_trampoline_ptr;

		return TRUE;
	}
	auto comm_remov_hook() -> BOOL
	{
		if (!func_call_address)
			return FALSE;

		*(uintptr_t*)func_call_address = 
			(uintptr_t)_HalpTimerConvertAuxiliaryCounterToPerformanceCounter;

		if (!beep_trampoline_ptr)
			return FALSE;

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

		return TRUE;
	}
}