#pragma once
#include <ntddk.h>
#include <ntddkbd.h>

#define MOUDEV_NAME L"\\Device\\mouflt"
#define  MOUSYM_NAME L"\\DosDevices\\mouflt"

#define CTL_CODE_BASE 0x8000
#define CTL_CMD(i) CTL_CODE(FILE_DEVICE_UNKNOWN,CTL_CODE_BASE+i,METHOD_BUFFERED,FILE_ANY_ACCESS)
#define  CTL_START_MOUFLT CTL_CMD(68)
#define  CTL_STOP_MOUFLT CTL_CMD(69)
/** ǰ */
NTKERNELAPI
UCHAR* PsGetProcessImageFileName(__in PEPROCESS Process);
/*
*

StartFltWpd(pDriver, RegPath);
رչ
unloadWpdFlt(pDriver);
*
*/
PDRIVER_OBJECT g_pdriver;
PUNICODE_STRING g_regpath;
INT g_FltFlag;

// 豸չṹ
typedef struct _Dev_exten
{
	ULONG Size;						// ýṹС
	PDEVICE_OBJECT FilterDevice;	// 豸
	PDEVICE_OBJECT TargeDevice;		// һ豸
	PDEVICE_OBJECT LowDevice;		// ײ豸
	KSPIN_LOCK IoRequestSpinLock;	// 
	KEVENT IoInProgressEvent;		// ¼
	PIRP pIrp;						// IRP
} DEV_EXTENSION, * PDEV_EXTENSION;


// ΢δObReferenceObjectByName()
NTSTATUS ObReferenceObjectByName(
	PUNICODE_STRING ObjectName,
	ULONG Attributes,
	PACCESS_STATE AccessState,
	ACCESS_MASK DesiredAccess,
	POBJECT_TYPE ObjectType,
	KPROCESSOR_MODE AccessMode,
	PVOID ParseContest,
	PVOID* Object
);

extern POBJECT_TYPE* IoDriverObjectType;

RTL_OSVERSIONINFOW osinfo;



NTSTATUS GetWindowsVersion(OUT PRTL_OSVERSIONINFOW pOsVersionInfo);
NTSTATUS DeAttach(PDEVICE_OBJECT pdevice);
NTSTATUS unloadWpdFlt(PDRIVER_OBJECT pDriver);
NTSTATUS GeneralDispatch(PDEVICE_OBJECT pDevice, PIRP pIrp);
NTSTATUS PowerDispatch(PDEVICE_OBJECT pDevice, PIRP pIrp);
NTSTATUS PnPDispatch(PDEVICE_OBJECT pDevice, PIRP pIrp);
NTSTATUS DevExtInit(PDEV_EXTENSION devExt, PDEVICE_OBJECT filterDevice, PDEVICE_OBJECT targetDevice, PDEVICE_OBJECT lowDevice);
NTSTATUS AttachDevice(PDRIVER_OBJECT pDriver, PUNICODE_STRING RegPatch);
NTSTATUS StartFltWpd(PDRIVER_OBJECT pDriver, PUNICODE_STRING RegPath);



NTSTATUS unloadWpdFlt(PDRIVER_OBJECT pDriver)
{
	PDEVICE_OBJECT pDevice;
	PDEV_EXTENSION devExt;

	DbgPrint("[db]DriverEntry Unloading...\n");

	pDevice = pDriver->DeviceObject;
	while (pDevice)
	{
		DeAttach(pDevice);
		pDevice = pDevice->NextDevice;
	}
}
NTSTATUS DeAttach(PDEVICE_OBJECT pdevice)
{
	PDEV_EXTENSION devExt;
	devExt = (PDEV_EXTENSION)pdevice->DeviceExtension;

	if (devExt == NULL)
	{
		return STATUS_SUCCESS;
	}

	IoDetachDevice(devExt->TargeDevice);
	devExt->TargeDevice = NULL;
	IoDeleteDevice(pdevice);
	devExt->FilterDevice = NULL;

	return STATUS_SUCCESS;
}
NTSTATUS GetWindowsVersion(OUT PRTL_OSVERSIONINFOW pOsVersionInfo)
{
	if (pOsVersionInfo == NULL)
	{
		return STATUS_INVALID_PARAMETER;
	}

	RTL_OSVERSIONINFOW osVersionInfo;
	RtlZeroMemory(&osVersionInfo, sizeof(RTL_OSVERSIONINFOW));
	osVersionInfo.dwOSVersionInfoSize = sizeof(RTL_OSVERSIONINFOW);

	NTSTATUS status = RtlGetVersion(&osVersionInfo);
	if (!NT_SUCCESS(status))
	{
		return status;
	}

	// Copy the version information to the output parameter
	RtlCopyMemory(pOsVersionInfo, &osVersionInfo, sizeof(RTL_OSVERSIONINFOW));

	return STATUS_SUCCESS;
}
NTSTATUS GeneralDispatch(PDEVICE_OBJECT pDevice, PIRP pIrp)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;

	PEPROCESS pEprocess = PsGetCurrentProcess();

	if (strstr(PsGetProcessImageFileName(pEprocess), "HPTS"))
	{
		pIrp->IoStatus.Information = 0;
		pIrp->IoStatus.Status = STATUS_SUCCESS;
		ObDereferenceObject(pEprocess);
		IoCompleteRequest(pIrp, IO_NO_INCREMENT);
		return STATUS_SUCCESS;
	}

	DbgPrint("[db]General Diapatch\n");

	/*PDEV_EXTENSION devExt = (PDEV_EXTENSION)pDevice->DeviceExtension;
	PDEVICE_OBJECT lowDevice = devExt->LowDevice;


	if (MmIsAddressValid(pIrp->AssociatedIrp.SystemBuffer) && g_FltFlag)
	{
		RtlZeroMemory(pIrp->AssociatedIrp.SystemBuffer, 2);
	}



	IoSkipCurrentIrpStackLocation(pIrp);
	status = IoCallDriver(lowDevice, pIrp);*/

	pIrp->IoStatus.Information = 0;
	pIrp->IoStatus.Status = STATUS_SUCCESS;
	ObDereferenceObject(pEprocess);
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS PowerDispatch(PDEVICE_OBJECT pDevice, PIRP pIrp)
{
	PDEV_EXTENSION devExt;
	devExt = (PDEV_EXTENSION)pDevice->DeviceExtension;

	PoStartNextPowerIrp(pIrp);
	IoSkipCurrentIrpStackLocation(pIrp);
	return PoCallDriver(devExt->TargeDevice, pIrp);
}
NTSTATUS PnPDispatch(PDEVICE_OBJECT pDevice, PIRP pIrp)
{
	PDEV_EXTENSION devExt;
	PIO_STACK_LOCATION stack;
	NTSTATUS status = STATUS_SUCCESS;

	devExt = (PDEV_EXTENSION)pDevice->DeviceExtension;
	stack = IoGetCurrentIrpStackLocation(pIrp);

	switch (stack->MinorFunction)
	{
	case IRP_MN_REMOVE_DEVICE:
		// Ȱȥ
		IoSkipCurrentIrpStackLocation(pIrp);
		IoCallDriver(devExt->LowDevice, pIrp);
		// Ȼ󶨡
		IoDetachDevice(devExt->LowDevice);
		// ɾԼɵ豸
		IoDeleteDevice(pDevice);
		status = STATUS_SUCCESS;
		break;

	default:
		// ͵IRPȫֱ·ɡ 
		IoSkipCurrentIrpStackLocation(pIrp);
		status = IoCallDriver(devExt->LowDevice, pIrp);
	}
	return status;
}
NTSTATUS DevExtInit(PDEV_EXTENSION devExt, PDEVICE_OBJECT filterDevice, PDEVICE_OBJECT targetDevice, PDEVICE_OBJECT lowDevice)
{
	memset(devExt, 0, sizeof(DEV_EXTENSION));
	devExt->FilterDevice = filterDevice;
	devExt->TargeDevice = targetDevice;
	devExt->LowDevice = lowDevice;
	devExt->Size = sizeof(DEV_EXTENSION);
	KeInitializeSpinLock(&devExt->IoRequestSpinLock);
	KeInitializeEvent(&devExt->IoInProgressEvent, NotificationEvent, FALSE);
	return STATUS_SUCCESS;
}



NTSTATUS IoControlDispatch(PDEVICE_OBJECT pDevice, PIRP pIrp)
{

	PIO_STACK_LOCATION pStack = IoGetCurrentIrpStackLocation(pIrp);
	ULONG CtlCode = pStack->Parameters.DeviceIoControl.IoControlCode;
	switch (CtlCode)
	{
	case CTL_START_MOUFLT:
	{
		DbgPrint("[db][DB]OVER\n");
		AttachDevice(g_pdriver, g_regpath);
		break;
	}
	case CTL_STOP_MOUFLT:
	{

		
		unloadWpdFlt(g_pdriver);
		break;
	}
	default:
		break;
	}
	pIrp->IoStatus.Information = 0;
	pIrp->IoStatus.Status = STATUS_SUCCESS;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS AttachDevice(PDRIVER_OBJECT pDriver, PUNICODE_STRING RegPatch)
{
	UNICODE_STRING kbdName;

	/*if (osinfo.dwMajorVersion == 6)
	{
		RtlInitUnicodeString(&kbdName, L"\\Driver\\WUDFRd");
	}

	if (osinfo.dwMajorVersion == 10)
	{
		RtlInitUnicodeString(&kbdName, L"\\Driver\\WpdUpFltr");
	}*/
	RtlInitUnicodeString(&kbdName, L"\\Driver\\mouclass");

	NTSTATUS status = 0;
	PDEV_EXTENSION devExt;			// 豸չ豸
	PDEVICE_OBJECT filterDevice;	// 豸 
	PDEVICE_OBJECT targetDevice;		// Ŀ豸豸
	PDEVICE_OBJECT lowDevice;		// ײ豸ĳһ豸ϼһ豸ʱһǼӵ豸ϣ豸ջջ
	PDRIVER_OBJECT kbdDriver;		// ڽմ򿪵豸

	// ȡĶ󣬱kbdDriver
	status = ObReferenceObjectByName(&kbdName, OBJ_CASE_INSENSITIVE, NULL, 0, *IoDriverObjectType, KernelMode, NULL, &kbdDriver);
	if (!NT_SUCCESS(status))
	{
		//DbgPrint("[db]Open KeyBoard Driver Failed\n");
		return status;
	}
	else
	{
		// 
		ObDereferenceObject(kbdDriver);
	}

	// ȡ豸еĵһ豸
	targetDevice = kbdDriver->DeviceObject;
	// һ̼豸е豸
	while (targetDevice)
	{
		// һ豸
		status = IoCreateDevice(pDriver, sizeof(DEV_EXTENSION), NULL, targetDevice->DeviceType, targetDevice->Characteristics, FALSE, &filterDevice);
		if (!NT_SUCCESS(status))
		{
			//DbgPrint("[db]Create New FilterDevice Failed\n");
			filterDevice = targetDevice = NULL;
			return status;
		}
		// 󶨣lowDeviceǰ֮õһ豸
		lowDevice = IoAttachDeviceToDeviceStack(filterDevice, targetDevice);
		//DbgPrint("[db]tar:%p\n low:%p\n");
		if (!lowDevice)
		{
			DbgPrint("[db]Attach Faided!\n");
			IoDeleteDevice(filterDevice);
			filterDevice = NULL;
			return status;
		}
		// ʼ豸չ
		devExt = (PDEV_EXTENSION)filterDevice->DeviceExtension;
		DevExtInit(devExt, filterDevice, targetDevice, lowDevice);

		filterDevice->DeviceType = lowDevice->DeviceType;
		filterDevice->Characteristics = lowDevice->Characteristics;
		filterDevice->StackSize = lowDevice->StackSize + 1;
		filterDevice->Flags |= lowDevice->Flags & (DO_BUFFERED_IO | DO_DIRECT_IO | DO_POWER_PAGABLE);
		// һ豸
		targetDevice = targetDevice->NextDevice;
	}
	DbgPrint("[db]Create And Attach Finshed...\n");
	return status;
}
NTSTATUS StartFltWpd(PDRIVER_OBJECT pDriver, PUNICODE_STRING RegPath)
{
	ULONG i;
	NTSTATUS status = STATUS_SUCCESS;

	g_FltFlag = 0;

	GetWindowsVersion(&osinfo);

	for (i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; i++)
	{
		if (i != IRP_MJ_DEVICE_CONTROL)
		{
			pDriver->MajorFunction[i] = GeneralDispatch;		// עͨõIRPַ
		}


	}


	pDriver->MajorFunction[IRP_MJ_DEVICE_CONTROL] = IoControlDispatch;


	pDriver->MajorFunction[IRP_MJ_POWER] = PowerDispatch;	// עԴIRPַ
	pDriver->MajorFunction[IRP_MJ_PNP] = PnPDispatch;		// עἴ弴IRPַ


}

VOID Unload(PDRIVER_OBJECT pDriver)
{
	unloadWpdFlt(pDriver);
}


NTSTATUS DriverEntry(PDRIVER_OBJECT pDriveObject, PUNICODE_STRING Regpath)
{

	NTSTATUS ntStatu = NULL;
	g_pdriver = pDriveObject;
	g_regpath = Regpath;







	PDEVICE_OBJECT pDevice = NULL;
	UNICODE_STRING uDeciceName = RTL_CONSTANT_STRING(MOUDEV_NAME);
	UNICODE_STRING uSyb_Name = RTL_CONSTANT_STRING(MOUSYM_NAME);

	ntStatu = IoCreateDevice(pDriveObject, 0, &uDeciceName, FILE_DEVICE_UNKNOWN, 0, FALSE, &pDevice);
	if (!NT_SUCCESS(ntStatu))
	{
		return STATUS_UNSUCCESSFUL;
		//DbgPrintEx(77, 0, "CreateDeviceFaild\n");
	}

	ntStatu = IoCreateSymbolicLink(&uSyb_Name, &uDeciceName);
	if (!NT_SUCCESS(ntStatu))
	{
		return STATUS_UNSUCCESSFUL;
		//DbgPrintEx(77, 0, "CreateSysbolinkFaild\n");
	}








	
	StartFltWpd(pDriveObject, Regpath);









	pDriveObject->DriverUnload = Unload;
	return STATUS_SUCCESS;
}