#include "USBFilter.h"

#define CURINFO1 " >> 進(jìn)入: "
#define CURINFO2 __FILE__ 
#define CURINFO3 " USBFilter.sys "  
#define CURINFO4 __DATE__   

#define CURINFO   CURINFO4 CURINFO3 CURINFO2  CURINFO1
/*************************************************************************************/
HANDLE pWriteFile;
HANDLE pReadFile;
PSTR WriteFilePath = "\\??\\C:\\WriteLog.txt";
PSTR ReadFilePath = "\\??\\C:\\ReadLog.txt";
UCHAR IoInfo[512*0x80];
typedef struct _FILE_WORK_ITEM
{
	PVOID FileContext;
	WORK_QUEUE_ITEM WorkItem;
	HANDLE FileHandle;
	PUNICODE_STRING pUFileName;
	ULONG Length;
	char OutFromFile[20];
	ULONG OutFromFileLen;
} FILE_WORK_ITEM,*PFILE_WORK_ITEM;


#define InitializeObjectAttributes( p, n, a, r, s ) { \
	(p)->Length = sizeof( OBJECT_ATTRIBUTES ); \
	(p)->RootDirectory = r; \
	(p)->Attributes = a; \
	(p)->ObjectName = n; \
	(p)->SecurityDescriptor = s; \
	(p)->SecurityQualityOfService = NULL; \
}

NTSTATUS
CreateLogFile(
			  IN PVOID Context,
			  IN OUT HANDLE *FileHandle,
			  IN PSTR FileName
			  );
VOID MyDriverCreateFileWorkItem(
								PVOID Context);
VOID MyDriverWriteFile(
					   IN PVOID Buffer,
					   IN ULONG Length,
					   IN OUT HANDLE FileHandle
					   );
VOID MyDriverWriteFileWorkItem(
							   PVOID Context);


VOID MyDriverCloseFile(
					   IN OUT HANDLE FileHandle);
VOID MyDriverCloseFileWorkItem(
							   PVOID Context);

/******************************************************************************************/

NTSTATUS
DriverEntry(PDRIVER_OBJECT DriverObject,
			PUNICODE_STRING RegistryPath)
{
	NTSTATUS ntStatus=STATUS_SUCCESS;
	PDRIVER_DISPATCH *DriverDispatch;
	ULONG i=0;
	
	DbgPrint(CURINFO "DriverEntry...\n");
	//+++++++++++++++++++++++++++++++++++++++++++
	UNREFERENCED_PARAMETER(RegistryPath);

	for (i=0,DriverDispatch=DriverObject->MajorFunction;
					i<=IRP_MJ_MAXIMUM_FUNCTION;
					i++,DriverDispatch++)
	{
		*DriverDispatch = USBFilter;
	}
	DriverObject->MajorFunction[IRP_MJ_SCSI]=USBFilterSCSI;
	DriverObject->MajorFunction[IRP_MJ_PNP] =USBFilterPnp;
	DriverObject->DriverExtension->AddDevice=USBFilterAddDevice;
	DriverObject->DriverUnload=USBFilterUnload;

	/*******************************************
	創(chuàng)建讀寫文件
	********************************************/
	RtlZeroMemory(IoInfo,512*0x80);
	ntStatus=CreateLogFile(NULL,&pWriteFile,WriteFilePath);
	if(!NT_SUCCESS(ntStatus))
	{
		DbgPrint("Create WriteFile is falied with error code 0x%08x\n",ntStatus);
		return ntStatus;
	}
	ntStatus=CreateLogFile(NULL,&pReadFile,ReadFilePath);
	if(!NT_SUCCESS(ntStatus))
	{
		DbgPrint("Create ReadFile is falied with error code 0x%08x\n",ntStatus);
		return ntStatus;
	}
	return ntStatus;
}
VOID
USBFilterUnload(PDRIVER_OBJECT DriverObject)
{
	PAGED_CODE();

	ASSERT(DriverObject->DeviceObject==NULL);
	DbgPrint(CURINFO "USBFilterUnload...\n");

	return;	
}

NTSTATUS
USBFilterAddDevice(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT PhysicalDevice)
{
	NTSTATUS ntStatus=STATUS_SUCCESS;
	PDEVICE_OBJECT object=NULL;
	PUSB_DEVICE_EXTENSION usbDeviceExe;
	ULONG deviceType=FILE_DEVICE_DISK;
	UNICODE_STRING ntName;
	UNICODE_STRING win32Name;

	DbgPrint(CURINFO "USBFilterAddDevice...\n");

	RtlInitUnicodeString(&ntName,L"\\Device\\USBFilter");
	RtlInitUnicodeString(&win32Name,L"\\DosDevices\\USBFilter");
	
	PAGED_CODE();

	if(!IoIsWdmVersionAvailable(1,0x20))
	{
		object=IoGetAttachedDeviceReference(PhysicalDevice);
		deviceType=object->DeviceType;
		ObDereferenceObject(object);
	}
	
	ntStatus=IoCreateDevice(DriverObject,\
							sizeof(USB_DEVICE_EXTENSION),\
							&ntName,\
							FILE_DEVICE_DISK,\
							FILE_DEVICE_SECURE_OPEN,\
							FALSE,\
							&object);
	if(!NT_SUCCESS(ntStatus))
	{
		DbgPrint(CURINFO "USBFilterAddDevice is falied...\n");
		return ntStatus;
	}
	
	usbDeviceExe=(PUSB_DEVICE_EXTENSION)object->DeviceExtension;

	usbDeviceExe->NextLowerDriver=IoAttachDeviceToDeviceStack(object,PhysicalDevice);
	if(!usbDeviceExe->NextLowerDriver)
	{
		DbgPrint(CURINFO "IoAttachDeviceToDeviceStack IS FALIED!...\n");
		IoDeleteDevice(object);
		return STATUS_UNSUCCESSFUL;
	}
	
	ntStatus=IoCreateSymbolicLink(&win32Name,&ntName);
	if(!NT_SUCCESS(ntStatus))
	{
		DbgPrint(CURINFO "IoCreateSymbolicLink is falied with error code 0x%x\n",ntStatus);
		return ntStatus;
	}

  object->Flags |= usbDeviceExe->NextLowerDriver->Flags & (DO_BUFFERED_IO | DO_DIRECT_IO | DO_POWER_PAGABLE);
	object->DeviceType=FILE_DEVICE_DISK;//deviceType;
	object->Characteristics=usbDeviceExe->NextLowerDriver->Characteristics;;
	usbDeviceExe->CurrentDeviceObject=object;
	
	
	IoInitializeRemoveLock(&usbDeviceExe->RemoveLock,POOL_TAG,1,100);
	
	INITIALIZE_PNP_STATE(usbDeviceExe);
	
	object->Flags&=~DO_DEVICE_INITIALIZING;
	return STATUS_SUCCESS;	
}




NTSTATUS
USBFilter(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS ntStatus=STATUS_SUCCESS;
	PUSB_DEVICE_EXTENSION USBDeviceExe;
	PIO_STACK_LOCATION IrpStack;
	PVOID InputBuffer;
	CHAR* OutputBuffer=" USBFilter ";
	
	DbgPrint(CURINFO " USBFilter...\n");
	IrpStack=IoGetCurrentIrpStackLocation(Irp);
	
	/****這里暫時(shí)用不上就現(xiàn)屏蔽了，如果需要這部分功能可以在從新打開
	if(IrpStack->MajorFunction==IRP_MJ_DEVICE_CONTROL)
	{
		if(IrpStack->Parameters.DeviceIoControl.IoControlCode==IOCTL_800_WRITE)
		{
		//	InputBuffer=Irp->AssociatedIrp.SystemBuffer;
		//	DbgPrint(CURINFO "USBFilter GetInformation IOCTL_800_WRITE: %s...\n",InputBuffer);
		}else if(IrpStack->Parameters.DeviceIoControl.IoControlCode==IOCTL_800_READ)
		{
		//	Irp->AssociatedIrp.SystemBuffer=(VOID*)OutputBuffer;
		//	DbgPrint(CURINFO "USBFilter GetInformation IOCTL_800_READ: %s...\n",OutputBuffer);
		}
	}*/
	
		USBDeviceExe=(PUSB_DEVICE_EXTENSION) DeviceObject->DeviceExtension;
	ntStatus=IoAcquireRemoveLock(&USBDeviceExe->RemoveLock,Irp);

	if (!NT_SUCCESS(ntStatus))
	{
		Irp->IoStatus.Status=ntStatus;
		IoCompleteRequest(Irp,IO_NO_INCREMENT);
		return ntStatus;
	}
	
	IoSkipCurrentIrpStackLocation(Irp);
	ntStatus=IoCallDriver(USBDeviceExe->NextLowerDriver,Irp);
	IoReleaseRemoveLock(&USBDeviceExe->RemoveLock,Irp);
	
	return	ntStatus;
}






NTSTATUS
USBFilterPnp(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS ntStatus=STATUS_SUCCESS;
	PUSB_DEVICE_EXTENSION USBDeviceExe;
	PIO_STACK_LOCATION IrpStack;
	KEVENT event;
	
	DbgPrint(CURINFO "USBFilterPnp...\n");

	PAGED_CODE();
	
	USBDeviceExe=(PUSB_DEVICE_EXTENSION)DeviceObject->DeviceExtension;
	IrpStack=IoGetCurrentIrpStackLocation(Irp);
	ntStatus=IoAcquireRemoveLock(&USBDeviceExe->RemoveLock,Irp);
	if(!NT_SUCCESS(ntStatus))
	{
		DbgPrint(CURINFO " USBFilterPnp IoAcquireRemoveLock is falied with error code 0x%x\n",ntStatus);
		Irp->IoStatus.Status=ntStatus;
		IoCompleteRequest(Irp,IO_NO_INCREMENT);
		return ntStatus;
	}
	switch(IrpStack->MinorFunction)
	{
	case IRP_MN_START_DEVICE:
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_START_DEVICE ...\n");
		KeInitializeEvent(&event,NotificationEvent,FALSE);
		IoCopyCurrentIrpStackLocationToNext(Irp);
		IoSetCompletionRoutine(Irp,\
							   (PIO_COMPLETION_ROUTINE)USBFilterStartCompletionRoutine,\
							   &event,\
							   TRUE,\
							   TRUE,\
							   TRUE);
		ntStatus=IoCallDriver(USBDeviceExe->NextLowerDriver,Irp);
		if(ntStatus==STATUS_PENDING)
		{
			KeWaitForSingleObject(&event,Executive,KernelMode,FALSE,NULL);
			ntStatus=Irp->IoStatus.Status;
		}
		if (NT_SUCCESS(ntStatus))
		{
			SET_NEW_PNP_STATE(USBDeviceExe,Started);

			if (USBDeviceExe->NextLowerDriver->Characteristics & FILE_REMOVABLE_MEDIA)
			{
				DeviceObject->Characteristics|=FILE_REMOVABLE_MEDIA;
			}
		}
		Irp->IoStatus.Status = ntStatus;
		IoCompleteRequest(Irp,IO_NO_INCREMENT);
		IoReleaseRemoveLock(&USBDeviceExe->RemoveLock,Irp);

		return ntStatus;
	case IRP_MN_REMOVE_DEVICE:		
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_REMOVE_DEVICE ...\n");		
		IoReleaseRemoveLockAndWait(&USBDeviceExe->RemoveLock,Irp);		
		IoSkipCurrentIrpStackLocation(Irp);		
		ntStatus = IoCallDriver(USBDeviceExe->NextLowerDriver,Irp);		
		SET_NEW_PNP_STATE(USBDeviceExe,Deleted);		
		IoDetachDevice(USBDeviceExe->NextLowerDriver);
		IoDeleteDevice(DeviceObject);
		
		return	ntStatus;
	case IRP_MN_QUERY_STOP_DEVICE:
		
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_QUERY_STOP_DEVICE ...\n");		
		
		SET_NEW_PNP_STATE(USBDeviceExe,StopPending);
		ntStatus = STATUS_SUCCESS;
		break;
	case IRP_MN_CANCEL_STOP_DEVICE:
		
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_CANCEL_STOP_DEVICE ...\n");		
		
		if (StopPending == USBDeviceExe->DevicePnPState)
		{
			RESTORE_PREVIOUS_PNP_STATE(USBDeviceExe);
		}
		
		ntStatus = STATUS_SUCCESS;
		break;
		
	case IRP_MN_STOP_DEVICE:
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_STOP_DEVICE ...\n");		
        SET_NEW_PNP_STATE(USBDeviceExe, Stopped);
        ntStatus = STATUS_SUCCESS;
        break;
		
    case IRP_MN_QUERY_REMOVE_DEVICE:
		
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_QUERY_REMOVE_DEVICE ...\n");		
		
        SET_NEW_PNP_STATE(USBDeviceExe, RemovePending);
        ntStatus = STATUS_SUCCESS;
        break;
		
    case IRP_MN_SURPRISE_REMOVAL:
		
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_SURPRISE_REMOVAL ...\n");		
		
        SET_NEW_PNP_STATE(USBDeviceExe, SurpriseRemovePending);
        ntStatus = STATUS_SUCCESS;
        break;
		
    case IRP_MN_CANCEL_REMOVE_DEVICE:
		
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_CANCEL_REMOVE_DEVICE ...\n");		
		
		if(RemovePending == USBDeviceExe->DevicePnPState)
        {
            RESTORE_PREVIOUS_PNP_STATE(USBDeviceExe);
        }
		
        ntStatus = STATUS_SUCCESS; // We must not fail this IRP.
        break;
	case IRP_MN_DEVICE_USAGE_NOTIFICATION:
		
		DbgPrint(CURINFO " USBFilterPnp IRP_MN_DEVICE_USAGE_NOTIFICATION ...\n");		
		
		if ((DeviceObject->AttachedDevice == NULL) ||
            (DeviceObject->AttachedDevice->Flags & DO_POWER_PAGABLE)) {
			
            DeviceObject->Flags |= DO_POWER_PAGABLE;
        }
		
		IoCopyCurrentIrpStackLocationToNext(Irp);
		
		IoSetCompletionRoutine(
            Irp,
            FilterDeviceUsageNotificationCompletionRoutine,
            NULL,
            TRUE,
            TRUE,
            TRUE
            );
		
		return IoCallDriver(USBDeviceExe->NextLowerDriver, Irp);
		
	default:
		
		DbgPrint(CURINFO " USBFilterPnp Default ...\n");		
		
		ntStatus = Irp->IoStatus.Status;		
        break;
	}
	Irp->IoStatus.Status = ntStatus;
	IoSkipCurrentIrpStackLocation(Irp);
	ntStatus = IoCallDriver(USBDeviceExe->NextLowerDriver,Irp);
	IoReleaseRemoveLock(&USBDeviceExe->RemoveLock,Irp);

	return	ntStatus;
}


NTSTATUS
USBFilterStartCompletionRoutine(
	IN PDEVICE_OBJECT DeviceObject, 
	IN PIRP Irp, IN PVOID Context 
	)
{
	PKEVENT		event = (PKEVENT)Context;

	DbgPrint(CURINFO " FilterStartCompletionRoutine!...\n");

	UNREFERENCED_PARAMETER(DeviceObject);

	if (Irp->PendingReturned == TRUE)
	{
		KeSetEvent(event,IO_NO_INCREMENT,FALSE);
	}

	return STATUS_MORE_PROCESSING_REQUIRED;
}

NTSTATUS
FilterDeviceUsageNotificationCompletionRoutine(
    IN PDEVICE_OBJECT   DeviceObject,
    IN PIRP             Irp,
    IN PVOID            Context
    )
{
	PUSB_DEVICE_EXTENSION       USBDeviceExe;

	DbgPrint(CURINFO "FilterDeviceUsageNotificationCompletionRoutine!...\n");

    UNREFERENCED_PARAMETER(Context);

    USBDeviceExe = (PUSB_DEVICE_EXTENSION) DeviceObject->DeviceExtension;

	if (Irp->PendingReturned) 
	{
        IoMarkIrpPending(Irp);
    }