#include "comm_win7.h"
#include "comm_dispatch.h"
#include "utils.h"
#include "kernel_function.h"

#pragma warning(disable : 4309 4838)

PFILE_OBJECT file_obj{ 0 };
PDEVICE_OBJECT device_obj{ 0 };

LPVOID orign_maj_function = nullptr;

namespace comm_win7
{
	auto comm_init() -> BOOL
	{
		UNICODE_STRING unName{ 0 };
		RtlInitUnicodeString(&unName, L"\\Device\\Null");	
		auto status = IoGetDeviceObjectPointer(&unName, FILE_ALL_ACCESS, &file_obj, &device_obj);
		if (!NT_SUCCESS(status))
			return FALSE;

		ObDereferenceObject(file_obj);
		return TRUE;
	}

	auto dispatch_device_io(PDEVICE_OBJECT drvice_object, PIRP irp)->NTSTATUS
	{
		UNREFERENCED_PARAMETER(drvice_object);
		
		auto comm_buf_data =  reinterpret_cast<CMD_COMM*>(irp->AssociatedIrp.SystemBuffer);

		auto status = comm_dispatch::dispatch(comm_buf_data);

		irp->IoStatus.Information = sizeof(CMD_COMM);
		irp->IoStatus.Status = status;
		IoCompleteRequest(irp, IO_NO_INCREMENT);
		return status;
	};

	uintptr_t beep_trampoline = NULL;

	auto comm_install_hook() -> BOOL
	{
		auto drv_obj = device_obj->DriverObject;
		if (!drv_obj)
			return FALSE;

		orign_maj_function = drv_obj->MajorFunction[IRP_MJ_DEVICE_CONTROL];

		beep_trampoline = utils::get_beep_trampoline_ptr();
		if (!beep_trampoline)
			return FALSE;

		//mov rax,dispatch_device_io
        //jmp rax
		char jmprax[] = { 0x48, 0xB8, 0xA0, 0x9A, 0xD5, 0x3E, 0xFE, 0x7F, 0x00, 0x00, 0xFF, 0xE0 };

		*(uintptr_t*)&jmprax[2] = (uintptr_t)dispatch_device_io;
		kernel_function::MmMDLPagesCopy((PVOID)beep_trampoline, &jmprax, sizeof(jmprax), 0x0000002, KernelMode);

		drv_obj->MajorFunction[IRP_MJ_DEVICE_CONTROL] = (PDRIVER_DISPATCH)beep_trampoline/*直接就是你的函数*/;
		
		return TRUE;
	}

	auto comm_remov_hook() -> BOOL
	{
		auto drv_obj = device_obj->DriverObject;
		if (!drv_obj)
			return FALSE;

		if (drv_obj->MajorFunction[IRP_MJ_DEVICE_CONTROL])
		{
			drv_obj->MajorFunction[IRP_MJ_DEVICE_CONTROL] = 
				(PDRIVER_DISPATCH)orign_maj_function;

			if (!beep_trampoline)
				return FALSE;

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

		}

		return TRUE;
	}

}