#include "precomp.h"
#include "usb.h"
#include "usbioctl.h"
#include "Usbdlib.h"


#pragma warning( disable : 4995 )
#pragma warning( disable : 4996 )


typedef enum {
	TUSB_CLASS_UNSPECIFIED = 0,
	TUSB_CLASS_AUDIO = 1,
	TUSB_CLASS_CDC = 2,
	TUSB_CLASS_HID = 3,
	TUSB_CLASS_RESERVED_4 = 4,
	TUSB_CLASS_PHYSICAL = 5,
	TUSB_CLASS_IMAGE = 6,
	TUSB_CLASS_PRINTER = 7,
	TUSB_CLASS_MSC = 8,
	TUSB_CLASS_HUB = 9,
	TUSB_CLASS_CDC_DATA = 10,
	TUSB_CLASS_SMART_CARD = 11,
	TUSB_CLASS_RESERVED_12 = 12,
	TUSB_CLASS_CONTENT_SECURITY = 13,
	TUSB_CLASS_VIDEO = 14,
	TUSB_CLASS_PERSONAL_HEALTHCARE = 15,
	TUSB_CLASS_AUDIO_VIDEO = 16,

	TUSB_CLASS_DIAGNOSTIC = 0xDC,
	TUSB_CLASS_WIRELESS_CONTROLLER = 0xE0,
	TUSB_CLASS_MISC = 0xEF,
	TUSB_CLASS_APPLICATION_SPECIFIC = 0xFE,
	TUSB_CLASS_VENDOR_SPECIFIC = 0xFF
} usb_device_class_code;

enum
{
	usb_usbhub = 0,
	usb_usbhub3 = 1,
	usb_maximum = 2
};

static WCHAR* s_usbname[] =
{
	L"\\Driver\\usbhub",
	L"\\Driver\\usbhub3",
	NULL
};


static BOOLEAN enable_usbhook = FALSE;

NTSTATUS USBPnpControl(PDRIVER_DISPATCH dispath, PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS USBHookDispatch_0(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS USBHookDispatch_3(PDEVICE_OBJECT DeviceObject, PIRP Irp);


static PDRIVER_DISPATCH s_ProxyDispatchers[usb_maximum] =
{
	USBHookDispatch_0, // for usbhub
	USBHookDispatch_3  // for usbhub3
};


//IRP_CORE USBIrpCore = { 0, };



#define USB_COMMON_HOOK_HANDLERS \
    [IRP_MJ_PNP] = { NULL, IRP_MJ_PNP, TRUE, USBPnpControl }, \
    /* ÇÊ¿äÇÑ °æ¿ì ¾Æ·¡Ã³·³ Ãß°¡ÇÏ¸é µË´Ï´Ù */ \
    /* [IRP_MJ_INTERNAL_DEVICE_CONTROL] = { NULL, IRP_MJ_INTERNAL_DEVICE_CONTROL, TRUE, MyInternalHandler }, */ \
    /* [IRP_MJ_POWER] = { NULL, IRP_MJ_POWER, TRUE, MyPowerHandler }, */

/*
static IRP_HOOK_HANDLER USBIrpHookTemplate[] =
{
	{NULL, IRP_MJ_CREATE, FALSE},
	{NULL, IRP_MJ_CREATE_NAMED_PIPE, FALSE},
	{NULL, IRP_MJ_CLOSE, FALSE},
	{NULL, IRP_MJ_READ, FALSE},
	{NULL, IRP_MJ_WRITE, FALSE},
	{NULL, IRP_MJ_QUERY_INFORMATION, FALSE},
	{NULL, IRP_MJ_SET_INFORMATION, FALSE},
	{NULL, IRP_MJ_QUERY_EA, FALSE},
	{NULL, IRP_MJ_SET_EA, FALSE},
	{NULL, IRP_MJ_FLUSH_BUFFERS, FALSE},
	{NULL, IRP_MJ_QUERY_VOLUME_INFORMATION, FALSE},
	{NULL, IRP_MJ_SET_VOLUME_INFORMATION, FALSE},
	{NULL, IRP_MJ_DIRECTORY_CONTROL, FALSE},
	{NULL, IRP_MJ_FILE_SYSTEM_CONTROL, FALSE},
	{NULL, IRP_MJ_DEVICE_CONTROL, FALSE},
	{NULL, IRP_MJ_INTERNAL_DEVICE_CONTROL, FALSE},
	{NULL, IRP_MJ_SHUTDOWN, FALSE},
	{NULL, IRP_MJ_LOCK_CONTROL, FALSE},
	{NULL, IRP_MJ_CLEANUP, FALSE},
	{NULL, IRP_MJ_CREATE_MAILSLOT, FALSE},
	{NULL, IRP_MJ_QUERY_SECURITY, FALSE},
	{NULL, IRP_MJ_SET_SECURITY, FALSE},
	{NULL, IRP_MJ_POWER, FALSE},
	{NULL, IRP_MJ_SYSTEM_CONTROL, FALSE},
	{NULL, IRP_MJ_DEVICE_CHANGE, FALSE},
	{NULL, IRP_MJ_QUERY_QUOTA, FALSE},
	{NULL, IRP_MJ_SET_QUOTA, FALSE},
	{NULL, IRP_MJ_PNP, TRUE, USBPnpControl}
};
*/


//static USB_HOOK_CONTEXT g_HookContexts[usb_maximum];

static HOOK_CONTEXT g_HookContexts[usb_maximum] =
{
	// Index 0: usbhub¿ë ÄÁÅؽºÆ® ÃʱâÈ­
	{
		NULL,   // DriverObject (³ªÁß¿¡ Init¿¡¼­ ä¿ò)
		FALSE,  // IsHooked
		0,
		{       // HookHandlers ¹è¿­ (³ª¸ÓÁö´Â ÀÚµ¿À¸·Î 0/NULL ÃʱâÈ­µÊ)
			USB_COMMON_HOOK_HANDLERS
		}
	},

	// Index 1: usbhub3¿ë ÄÁÅؽºÆ® ÃʱâÈ­
	{
		NULL,   // DriverObject
		FALSE,  // IsHooked
		0,
		{       // HookHandlers ¹è¿­
			USB_COMMON_HOOK_HANDLERS
		}
	}
};

#define USBLock() InterlockedIncrement((volatile LONG *)&g_HookContexts[usb_usbhub].IrpEnterCount);
#define USBUnLock() InterlockedDecrement((volatile LONG *)&g_HookContexts[usb_usbhub].IrpEnterCount);
#define USBLock3() InterlockedIncrement((volatile LONG *)&g_HookContexts[usb_usbhub3].IrpEnterCount);
#define USBUnLock3() InterlockedDecrement((volatile LONG *)&g_HookContexts[usb_usbhub3].IrpEnterCount);


//static ULONG GetIrpHookHandlerCount()
//{
//	return sizeof(USBIrpHookTemplate) / sizeof(IRP_HOOK_HANDLER);
//}


static BOOLEAN IsAdrIrpHookHandler(PIRP_HOOK_HANDLER pIrpH, ULONG id)
{
	if (!pIrpH[id].IsHook)
		return FALSE;

	if (!MmIsAddressValid((PVOID)pIrpH[id].pOrgHandler))
		return FALSE;

	if (!pIrpH[id].pOrgHandler)
		return FALSE;

	if (!pIrpH[id].Work)
		return FALSE;

	return TRUE;
}



static BOOLEAN IsOwnDeviceObject(PDEVICE_OBJECT DevcieObject)
{
	if (DevcieObject == g_bs1Flt.DriverObject->DeviceObject)
		return TRUE;

	return FALSE;
}


/*
NTSTATUS UDF_CallUSBDI(IN PDEVICE_OBJECT pDevObj, IN PVOID UrbEtc)
{
	IO_STATUS_BLOCK IoStatus;
	KEVENT event;
	NTSTATUS status;
	PIRP Irp = NULL;
	PIO_STACK_LOCATION NextIrpStack = NULL;

	//g_Temp = pDevObj;
	//g_Temp2 = UrbEtc;

	// Initialise IRP completion event
	KeInitializeEvent(&event, NotificationEvent, FALSE);

	// Build Internal IOCTL IRP
	Irp = IoBuildDeviceIoControlRequest(
		IOCTL_INTERNAL_USB_SUBMIT_URB, pDevObj,
		NULL, 0, // Input buffer
		NULL, 0, // Output buffer
		TRUE, &event, &IoStatus);

	// Get IRP stack location for next driver down (already set up)
	NextIrpStack = IoGetNextIrpStackLocation(Irp);
	// Store pointer to the URB etc
	NextIrpStack->Parameters.Others.Argument1 = UrbEtc;
	NextIrpStack->Parameters.Others.Argument2 = (PVOID)0;

	// Call the driver and wait for completion if necessary
	status = IoCallDriver(pDevObj, Irp);
	if (status == STATUS_PENDING)
	{
		KLogEx(DEBUG_TRACE_ERROR, "waiting for URB completion");
		status = KeWaitForSingleObject(&event, Suspended, KernelMode, FALSE, NULL);
		status = IoStatus.Status;
	}

	// return IRP completion status
	KLogEx(DEBUG_TRACE_ERROR, "returned %x", status);
	return status;
}
*/

BOOLEAN IsDefineDeviceClass(usb_device_class_code deviceclass)
{
	//BOOLEAN state = FALSE;

	if (deviceclass == TUSB_CLASS_HID
		|| deviceclass == TUSB_CLASS_PRINTER
		|| deviceclass == TUSB_CLASS_HUB
		|| deviceclass == TUSB_CLASS_AUDIO
		|| deviceclass == TUSB_CLASS_VIDEO
		)
	{
		return TRUE;
	}

	return FALSE;
}


NTSTATUS USBPnpControl(PDRIVER_DISPATCH dispath, PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(Irp);
	ULONG minor = irpstack->MinorFunction;
	ULONG state = 0;
	UNICODE_STRING processImagePath = { 0, };
	ULONG ProcessId = 0;
	ULONG policyLog = 0;
	//WCHAR* process_name = NULL;
	//PURB urb = NULL;
	//PUSB_DEVICE_DESCRIPTOR	deviceDescriptor = NULL;
	//USB_STRING_DESCRIPTOR	USD, * pFullUSD = NULL;
	//ULONG size = 0;
	char szProcessName[20] = { 0, };
	PUSB_PARSED_INFO pUsbInfo = NULL;

	if (!g_bs1Flt.IsAttached)
		goto $USBCleanup;

	if (!enable_usbhook)
		goto $USBCleanup;

	ProcessId = HandleToULong(PsGetCurrentProcessId());
	UGetProcessName(szProcessName);

	state = GetPolicyState(BDC_USB);
	policyLog = IsPolicyLog(BDC_USB);

	if (minor == IRP_MN_START_DEVICE)
	{
		KLogEx(DEBUG_TRACE_INFO, " IRP_MN_START_DEVICE (%d)(%s)\n", ProcessId, szProcessName);

		status = CollectUsbDeviceInfo(DeviceObject, &pUsbInfo);
		if (NT_SUCCESS(status) && pUsbInfo)
		{
			KLogEx(DEBUG_TRACE_INFO, "USB Connected: DeviceClass(%d), DeviceSubClass(%d), bDeviceProtocol(%d), VID=0x%04X, PID=0x%04X\n",
				pUsbInfo->DeviceDesc.bDeviceClass,
				pUsbInfo->DeviceDesc.bDeviceSubClass,
				pUsbInfo->DeviceDesc.bDeviceProtocol,
				pUsbInfo->DeviceDesc.idVendor, 
				pUsbInfo->DeviceDesc.idProduct
			);

			if (IsDefineDeviceClass(pUsbInfo->DeviceDesc.bDeviceClass))
			{
				KLogEx(DEBUG_TRACE_INFO, "define deviceclass (%d)\n", pUsbInfo->DeviceDesc.bDeviceClass);
				goto $USBCleanup;
			}

			if (pUsbInfo->ProductStr) 
			{
				KLogEx(DEBUG_TRACE_INFO, "  Product: %S\n", pUsbInfo->ProductStr);
			}

			if (pUsbInfo->SerialNumberStr) 
			{
				KLogEx(DEBUG_TRACE_INFO, "  Serial: %S\n", pUsbInfo->SerialNumberStr);
			}

			if (IsUsbPortExceptionList(BDC_USB, pUsbInfo->DeviceDesc.idVendor, pUsbInfo->DeviceDesc.idProduct, pUsbInfo->DeviceDesc.bcdDevice, pUsbInfo->SerialNumberStr))
			{
				KLogEx(DEBUG_TRACE_INFO, "user usb policy exception\n");
				goto $USBCleanup;
			}

			if (pUsbInfo->ConfigDesc && pUsbInfo->ConfigDescSize > 0)
			{
				PUSB_CONFIGURATION_DESCRIPTOR pConfigDesc = pUsbInfo->ConfigDesc;
				PUCHAR pCurr = (PUCHAR)pConfigDesc;
				PUCHAR pEnd = pCurr + pConfigDesc->wTotalLength;

				KLogEx(DEBUG_TRACE_INFO, "[Configuration Descriptor]");
				KLogEx(DEBUG_TRACE_INFO, "  wTotalLength: 0x%X, bNumInterfaces: 0x%X, bConfigurationValue: 0x%X, MaxPower: 0x%X\n",
						pConfigDesc->wTotalLength,
						pConfigDesc->bNumInterfaces,
						pConfigDesc->bConfigurationValue,
						pConfigDesc->MaxPower
					);
				
				// Loop¸¦ µ¹¸é¼­ Interface ¹× Endpoint ÆĽÌ
				// ù ¹ø° Descriptor´Â ConfigurationÀ̹ǷΠ°Ç³Ê¶Ü
				pCurr += pConfigDesc->bLength;

				while (pCurr < pEnd)
				{
					PUSB_COMMON_DESCRIPTOR pCommon = (PUSB_COMMON_DESCRIPTOR)pCurr;

					if (pCommon->bLength == 0) 
						break; // ¿¡·¯ ¹æÁö

					if (pCommon->bDescriptorType == USB_INTERFACE_DESCRIPTOR_TYPE)
					{
						PUSB_INTERFACE_DESCRIPTOR pIntf = (PUSB_INTERFACE_DESCRIPTOR)pCommon;
						KLogEx(DEBUG_TRACE_INFO, "[Interface Descriptor]");
						KLogEx(DEBUG_TRACE_INFO, "  bInterfaceNumber: 0x%X, bAlternateSetting: 0x%X, bNumEndpoints: 0x%X, bInterfaceClass: 0x%X, bInterfaceSubClass: 0x%X, bInterfaceProtocol: 0x%X",
							pIntf->bInterfaceNumber,
							pIntf->bAlternateSetting,
							pIntf->bNumEndpoints,
							pIntf->bInterfaceClass,
							pIntf->bInterfaceSubClass,
							pIntf->bInterfaceProtocol
						);
					}
					else if (pCommon->bDescriptorType == USB_ENDPOINT_DESCRIPTOR_TYPE)
					{
						PUSB_ENDPOINT_DESCRIPTOR pEp = (PUSB_ENDPOINT_DESCRIPTOR)pCommon;
						KLogEx(DEBUG_TRACE_INFO, "[Endpoint Descriptor]\n");
						KLogEx(DEBUG_TRACE_INFO, "  bEndpointAddress: 0x%X, bmAttributes: 0x%X, wMaxPacketSize: 0x%X, bInterval: 0x%X\n", 
							pEp->bEndpointAddress,
							pEp->bmAttributes,
							pEp->wMaxPacketSize,
							pEp->bInterval
						);
					}

					pCurr += pCommon->bLength;
				}
			}

			if (pUsbInfo->HidInterfaceCount > 0)
			{
				ULONG i;
				KLogEx(DEBUG_TRACE_INFO, "  Total HID Interfaces Found: %d\n", pUsbInfo->HidInterfaceCount);

				for (i = 0; i < pUsbInfo->HidInterfaceCount; i++)
				{
					HID_INTERFACE_ENTRY* pEntry = &pUsbInfo->HidEntries[i];
					PUSB_HID_DESCRIPTOR pHid = pEntry->HidDesc;

					KLogEx(DEBUG_TRACE_INFO, "    [HID Interface #%d]\n", pEntry->InterfaceNumber);
					if (pHid) 
					{
						KLogEx(DEBUG_TRACE_INFO, "      bcdHID: 0x%04X\n", pHid->bcdHID);
						KLogEx(DEBUG_TRACE_INFO, "      Report Descs: %d\n", pHid->bNumDescriptors);
						if (pHid->bNumDescriptors > 0) 
						{
							KLogEx(DEBUG_TRACE_INFO, "      Desc Length: %d\n", pHid->DescriptorList[0].wDescriptorLength);
						}
					}
				}

				// ¿¹½Ã: 1¹ø ÀÎÅÍÆäÀ̽º°¡ HID¶ó¸é Â÷´Ü?
				/*
				for (i = 0; i < pUsbInfo->HidInterfaceCount; i++) {
					if (pUsbInfo->HidEntries[i].InterfaceNumber == 1) {
						bBlock = TRUE;
					}
				}
				*/
			}
			else
			{
				KLogEx(DEBUG_TRACE_INFO, "  No HID Interfaces found.\n");
			}
		}

		if (policyLog)
		{
			
			WCHAR notice[MAX_PATH] = { 0, };
			WCHAR processName[50] = { 0, };
			ULONG DeviceClass = 0;

			if (pUsbInfo)
			{
				RtlStringCbPrintfW(notice, sizeof(notice), L"DeviceClass : %d, DeviceSubClass : %d, DeviceProtocol : %d, Vid: 0x%04X, Pid : 0x%04X, Serial : %s",
					pUsbInfo->DeviceDesc.bDeviceClass,
					pUsbInfo->DeviceDesc.bDeviceSubClass,
					pUsbInfo->DeviceDesc.bDeviceProtocol,
					pUsbInfo->DeviceDesc.idVendor,
					pUsbInfo->DeviceDesc.idProduct,
					pUsbInfo->SerialNumberStr
				);

				DeviceClass = pUsbInfo->DeviceDesc.bDeviceClass;
			}
			else
			{
				KLogEx(DEBUG_TRACE_INFO, "pUsbInfo invaild pointer\n");
				RtlStringCbPrintfW(notice, sizeof(notice), L"pUsbInfo invaild pointer");
				
			}

			RtlStringCbPrintfW(processName, sizeof(processName), L"%S", szProcessName);
			SetLog(NULL, NULL, LOG_POLICY, BDC_USB, state, DeviceClass, processName, notice);
			
		}

		if (state == DISABLE && pUsbInfo)
		{
			FreeUsbDeviceInfo(pUsbInfo);
			pUsbInfo = NULL;
			KLogEx(DEBUG_TRACE_INFO, "STATUS_ACCESS_DENIED!!!!!!\n");
			Irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;
			Irp->IoStatus.Information = 0;
			IoCompleteRequest(Irp, IO_NO_INCREMENT);
			return STATUS_INSUFFICIENT_RESOURCES;
		}

		FreeUsbDeviceInfo(pUsbInfo);
		pUsbInfo = NULL;
	}
	else if (minor == IRP_MN_REMOVE_DEVICE)
	{
		KLogEx(DEBUG_TRACE_INFO, " IRP_MN_REMOVE_DEVICE\n");
		goto $USBCleanup;
	}
	else
	{
		KLogEx(DEBUG_TRACE_INFO, " minor(%0x)\n", minor);
		goto $USBCleanup;
	}

$USBCleanup:
	FreeUsbDeviceInfo(pUsbInfo);
	UStrFree(&processImagePath);
	return dispath(DeviceObject, Irp);

}

NTSTATUS USBHookDispatch_Common(ULONG ContextIndex, PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS NtStatus = STATUS_SUCCESS;
	PIO_STACK_LOCATION pCurrentIoStack = IoGetCurrentIrpStackLocation(Irp);
	ULONG Id = pCurrentIoStack->MajorFunction;
	PHOOK_CONTEXT pCtx = NULL;
	PDRIVER_DISPATCH pOrgHandler = NULL;

	// À妽º À¯È¿¼º °Ë»ç (¹æ¾î ÄÚµå)
	if (ContextIndex >= usb_maximum)
	{
		return STATUS_UNSUCCESSFUL;
	}

	// 1. À妽º·Î ¹Ù·Î ÄÁÅؽºÆ® Á¢±Ù
	pCtx = &g_HookContexts[ContextIndex];

	// 2. ÈÄÅ· À¯È¿¼º °Ë»ç
	if (pCtx->IsHooked &&
		pCtx->HookHandlers[Id].IsHook &&
		pCtx->HookHandlers[Id].Work)
	{
		pOrgHandler = pCtx->HookHandlers[Id].pOrgHandler;

		// 3. Work ÇÔ¼ö ½ÇÇà (USBPnpControl µî)
		NtStatus = pCtx->HookHandlers[Id].Work(
			pOrgHandler,
			DeviceObject,
			Irp
		);
	}
	else
	{
		// ¿©±â·Î ¿À¸é ¾È µÇÁö¸¸, ¸¸¾à ¿À¸é ½ÇÆРó¸® ȤÀº ¿øº» È£Ã⠽õµ ºÒ°¡
		NtStatus = STATUS_UNSUCCESSFUL;
	}

	return NtStatus;
}

// ---------------------------------------------------------
// [º°µµ Dispatcher] °¢ µå¶óÀ̹öº° Àü¿ë ÇÔ¼ö
// ---------------------------------------------------------

// usbhub¿ë (Index 0)
NTSTATUS USBHookDispatch_0(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	USBLock();
	status = USBHookDispatch_Common(usb_usbhub, DeviceObject, Irp);
	USBUnLock();
	return status;
}

// usbhub3¿ë (Index 1)
NTSTATUS USBHookDispatch_3(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	USBLock3();
	status = USBHookDispatch_Common(usb_usbhub3, DeviceObject, Irp);
	USBUnLock3();
	return status;
}


//NTSTATUS USBHookDispatch(PDEVICE_OBJECT DeviceObject, PIRP Irp)
//{
//
//	NTSTATUS NtStatus = STATUS_SUCCESS;
//	PIO_STACK_LOCATION pCurrentIoStack = IoGetCurrentIrpStackLocation(Irp);
//	ULONG Id = pCurrentIoStack->MajorFunction;
//
//	USBLock();
//
//	if (IsAdrIrpHookHandler(USBIrpHook, Id))
//	{
//		NtStatus = USBIrpHook[Id].Work(
//			USBIrpHook[Id].pOrgHandler,
//			DeviceObject,
//			Irp
//		);
//	}
//	USBUnLock();
//
//	return NtStatus;
//
//}


//static PVOID GetHookDispatch()
//{
//	return (PVOID)USBHookDispatch;
//}
//
//static PIRP_HOOK_HANDLER GetIrpHookHandler()
//{
//	return USBIrpHook;
//}

//static PIRP_CORE GetIrpCore()
//{
//	USBIrpCore.HookDispatch = GetHookDispatch();
//	USBIrpCore.IrpHookHandler = GetIrpHookHandler();
//	USBIrpCore.IrpHookHandlerCount = GetIrpHookHandlerCount();
//	USBIrpCore.IrpEnterCount = 0;
//	return &USBIrpCore;
//}

//NTSTATUS USBIrpHookCleanup()
//{
//	ULONG cnt = 0;
//	//ULONG OrSaveWP;
//	LARGE_INTEGER WaitTime;
//	PDRIVER_OBJECT DriverObject = USBIrpCore.HookDriverObject;
//	ULONG IrpHookHandlerCount = USBIrpCore.IrpHookHandlerCount;
//	PIRP_HOOK_HANDLER pIhh = USBIrpCore.IrpHookHandler;
//
//	if (!DriverObject || !IrpHookHandlerCount || !pIhh)
//		return STATUS_UNSUCCESSFUL;
//
//	if (!pIhh)
//		return STATUS_UNSUCCESSFUL;
//
//	enable_usbhook = FALSE;
//	for (cnt = 0; cnt < IrpHookHandlerCount; ++cnt)
//	{
//		if (!pIhh[cnt].IsHook)
//			continue;
//
//		if (pIhh[cnt].pOrgHandler == NULL)
//			continue;
//
//		InterlockedExchangePointer((PVOID)&DriverObject->MajorFunction[cnt], (PVOID)pIhh[cnt].pOrgHandler);
//		//DbgPrint("unhook a %x %x\n",DriverObject->MajorFunction[cnt], pIhh[cnt].pOrgHandler);
//	}
//
//	WaitTime.QuadPart = -1 * 10 * 1000 * 1000;
//
//	while (USBIrpCore.IrpEnterCount > 0)
//		KeDelayExecutionThread(KernelMode, FALSE, &WaitTime);
//
//	KeDelayExecutionThread(KernelMode, FALSE, &WaitTime);
//	KLogEx(DEBUG_TRACE_INFO, "complete\n");
//
//	return STATUS_SUCCESS;
//}
//
//NTSTATUS USBIrpHookInit()
//{
//	WCHAR* name = NULL;
//	ULONG cnt = 0;
//	ULONG IrpHookHandlerCount = 0;
//	PIRP_HOOK_HANDLER pIhh = NULL;
//	PIRP_CORE irpcore = NULL;
//	PDRIVER_OBJECT obj = NULL;
//
//	//for (size_t i = 0; i < usb_maximum; i++)
//	//{	
//		name = s_usbname[usb_usbhub3];
//		KLogEx(DEBUG_TRACE_INFO, "driver(%S)\n", name);
//		obj = SearchDriverObject(name);
//		if (!obj)
//		{
//			KLogEx(DEBUG_TRACE_ERROR, "not find object (%S)\n", name);
//			return STATUS_UNSUCCESSFUL;	
//		}
//	
//		irpcore = (PIRP_CORE)GetIrpCore();
//		if (!irpcore)
//			return STATUS_UNSUCCESSFUL;
//
//		irpcore->HookDriverObject = obj;
//		//KLog("hook init %p %p %p\n",irpcore->HookDriverObject,irpcore->HookDriverObject->DriverStart,irpcore->HookDriverObject->DriverSize);
//
//		IrpHookHandlerCount = irpcore->IrpHookHandlerCount;
//		pIhh = irpcore->IrpHookHandler;
//		if (!pIhh)
//			return STATUS_UNSUCCESSFUL;
//
//		for (; cnt < IrpHookHandlerCount; ++cnt)
//		{
//			if (!pIhh[cnt].IsHook)
//				continue;
//
//			if (irpcore->HookDriverObject->MajorFunction[cnt] == irpcore->HookDispatch)
//				continue;
//
//			pIhh[cnt].pOrgHandler = irpcore->HookDriverObject->MajorFunction[cnt];
//			if (pIhh[cnt].pOrgHandler == NULL)
//				continue;
//
//			if (pIhh[cnt].Work == NULL)
//				continue;
//
//			if (irpcore->HookDriverObject->DriverStart > (PVOID)irpcore->HookDriverObject->MajorFunction[cnt] ||
//				GetPtr(irpcore->HookDriverObject->DriverStart, irpcore->HookDriverObject->DriverSize) < (PVOID)irpcore->HookDriverObject->MajorFunction[cnt])
//			{
//				KLogEx(DEBUG_TRACE_INFO, "exist hook c(%p)\n", irpcore->HookDriverObject->MajorFunction[cnt]);
//			}
//
//			InterlockedExchangePointer((PVOID)&irpcore->HookDriverObject->MajorFunction[cnt], (PVOID)irpcore->HookDispatch);
//			KLogEx(DEBUG_TRACE_INFO, "hook c (%S), (%d), (%p), (%p)\n", name, cnt, (PVOID)pIhh[cnt].pOrgHandler, (PVOID)irpcore->HookDriverObject->MajorFunction[cnt]);
//			enable_usbhook = TRUE;
//		}
//	//}
//	KLogEx(DEBUG_TRACE_INFO, "complete\n");
//	return STATUS_SUCCESS;
//}


NTSTATUS USBIrpHookInit()
{
	WCHAR* name = NULL;
	ULONG i, mi;
	PDRIVER_OBJECT obj = NULL;
	PHOOK_CONTEXT hook = NULL;

	// g_HookContexts´Â ÀÌ¹Ì ÄÄÆÄÀÏ Å¸ÀÓ¿¡ HookHandlers Á¤º¸°¡ ä¿öÁ® ÀÖÀ¸¹Ç·Î
	// memset ÃʱâÈ­³ª ÅÛÇø´ º¹»ç ·çÇÁ°¡ ÇÊ¿ä ¾ø½À´Ï´Ù.

	for (i = 0; i < usb_maximum; i++)
	{
		hook = &g_HookContexts[i];

		// ÀÌ¹Ì ÈÅ »óŸé Æнº
		if (hook->IsHooked) continue;

		name = s_usbname[i];
		if (name == NULL) continue;

		KLogEx(DEBUG_TRACE_INFO, "Target driver(%S)\n", name);

		obj = SearchDriverObject(name);
		if (!obj)
		{
			KLogEx(DEBUG_TRACE_ERROR, "Not found object (%S)\n", name);
			continue;
		}

		hook->DriverObject = obj;

		for (mi = 0; mi < IRP_MJ_MAXIMUM_FUNCTION + 1; mi++)
		{
			// ÄÄÆÄÀÏ Å¸ÀÓ¿¡ ¼³Á¤µÈ IsHook °ªÀ» È®ÀÎ
			if (!hook->HookHandlers[mi].IsHook)
				continue;

			// ÀÌ¹Ì ÈÄÅ·µÇ¾ú´ÂÁö È®ÀÎ
			if (obj->MajorFunction[mi] == s_ProxyDispatchers[i])
				continue;

			// ¿øº» ¹é¾÷
			hook->HookHandlers[mi].pOrgHandler = obj->MajorFunction[mi];

			if (hook->HookHandlers[mi].pOrgHandler == NULL)
				continue;

			// ÁÖ¼Ò ±³Ã¼
			InterlockedExchangePointer((PVOID)&obj->MajorFunction[mi],(PVOID)s_ProxyDispatchers[i]);

			KLogEx(DEBUG_TRACE_INFO, "[Hook] %S MJ:%d Org:%p New:%p\n",
				name, mi, hook->HookHandlers[mi].pOrgHandler, s_ProxyDispatchers[i]);
		}

		hook->IsHooked = TRUE;
		enable_usbhook = TRUE;
	}

	KLogEx(DEBUG_TRACE_INFO, "complete\n");
	return STATUS_SUCCESS;
}


NTSTATUS USBIrpHookCleanup()
{
	ULONG i, mi;
	LARGE_INTEGER WaitTime;
	PHOOK_CONTEXT hook = NULL;
	PDRIVER_DISPATCH pCurrentHandler = NULL;

	// 1. ´õ ÀÌ»ó »õ·Î¿î ÈÅ ·ÎÁ÷ÀÌ Å¸Áö ¾Êµµ·Ï Ç÷¡±× ºñÈ°¼ºÈ­
	enable_usbhook = FALSE;

	KLogEx(DEBUG_TRACE_INFO, "Started...\n");

	for (i = 0; i < usb_maximum; i++)
	{
		hook = &g_HookContexts[i];

		// ÈÄÅ·µÈ ÀûÀÌ ¾ø°Å³ª µå¶óÀ̹ö °´Ã¼°¡ ¾øÀ¸¸é ½ºÅµ
		if (!hook->IsHooked || !hook->DriverObject)
			continue;

		KLogEx(DEBUG_TRACE_INFO, "Unhooking Driver Index: %d\n", i);

		// 3. IRP Major Function Å×ÀÌºí º¹¿ø
		for (mi = 0; mi <= IRP_MJ_MAXIMUM_FUNCTION; mi++)
		{
			// ¿ì¸®°¡ ÈÄÅ·ÇÏ·Á°í Çß´ø ´ë»óÀÎÁö È®ÀÎ
			if (!hook->HookHandlers[mi].IsHook)
				continue;

			// ¿øº» ÇÔ¼ö Æ÷ÀÎÅÍ°¡ À¯È¿ÇÑÁö È®ÀÎ
			if (hook->HookHandlers[mi].pOrgHandler == NULL)
				continue;

			// [Áß¿ä] ÇöÀç µå¶óÀ̹öÀÇ Çڵ鷯°¡ ¿ì¸® ÇÁ·Ï½Ã ÇÔ¼ö(USBHookDispatch_X)ÀÎÁö È®ÀÎ
			// ¸¸¾à ´Ù¸¥ µå¶óÀ̹ö°¡ ±× À§¿¡ ¶Ç ÈÄÅ·Çß´Ù¸é, ¿ì¸®°¡ ¿øº»À¸·Î µÇµ¹¸®¸é ¾È µÉ ¼öµµ ÀÖÀ½
			// ÇÏÁö¸¸ º¸Åë ¾ð·Îµå ½ÃÁ¡¿¡´Â °­Á¦·Î¶óµµ ¿øº»À¸·Î º¹±¸ÇÏ´Â °ÍÀÌ ÀϹÝÀûÀÓ.
			// ¿©±â¼­´Â ¹«Á¶°Ç º¹±¸ÇÏ´Â ¹æ½ÄÀ¸·Î ÁøÇà (InterlockedExchangePointer »ç¿ë)

			pCurrentHandler = (PDRIVER_DISPATCH)InterlockedExchangePointer(
				(PVOID)&hook->DriverObject->MajorFunction[mi],
				(PVOID)hook->HookHandlers[mi].pOrgHandler
			);

			hook->HookHandlers[mi].pOrgHandler = NULL;
			// µð¹ö±ë¿ë: ±³Ã¼µÈ ÁÖ¼Ò°¡ ¿ì¸® ÇÔ¼ö¿´´ÂÁö È®ÀÎ
			/*
			if (pCurrentHandler != s_ProxyDispatchers[i]) {
				KLogEx(DEBUG_TRACE_WARNING, "Warning: Function pointer was modified by another driver! MJ:%d Curr:%p Our:%p\n",
					mj_idx, pCurrentHandler, s_ProxyDispatchers[i]);
			}
			*/			
		}

		// 4. [¸Å¿ì Áß¿ä] ½ÇÇà ÁßÀÎ IRP°¡ ¸ðµÎ ºüÁ®³ª°¥ ¶§±îÁö ´ë±â (Rundown Protection)
		// USBIrpCore.IrpEnterCount´Â Hook ÇÔ¼ö ÁøÀÔ ½Ã Áõ°¡, Å»Ã⠽à °¨¼ÒÇÔ.
		// ÀÌ °ªÀÌ 0ÀÌ µÉ ¶§±îÁö ±â´Ù·Á¾ß ¾ÈÀüÇÏ°Ô µå¶óÀ̹ö¸¦ ¾ð·ÎµåÇÒ ¼ö ÀÖÀ½.
		WaitTime.QuadPart = -50 * 1000 * 10; // 50ms ´ÜÀ§·Î ´ë±â
		while (hook->IrpEnterCount > 0)
		{
			KLogEx(DEBUG_TRACE_INFO, "Waiting for active IRPs... Count: %d\n", hook->IrpEnterCount);
			KeDelayExecutionThread(KernelMode, FALSE, &WaitTime);
		}

		// »óÅ ÃʱâÈ­
		hook->IsHooked = FALSE;
		hook->DriverObject = NULL; // µå¶óÀ̹ö °´Ã¼ ÂüÁ¶ ÇØÁ¦ (ÇÊ¿ä½Ã ObDereferenceObject °í·Á)
		//RtlZeroMemory(pCtx->HookHandlers, sizeof(pCtx->HookHandlers));
	}

	// Ȥ½Ã ¸ð¸¦ ·¹À̽º ÄÁµð¼ÇÀ» ´ëºñÇØ 0.5ÃÊ Á¤µµ ¾ÈÀüÇÏ°Ô Ãß°¡ ´ë±â
	WaitTime.QuadPart = -500 * 1000 * 10;
	KeDelayExecutionThread(KernelMode, FALSE, &WaitTime);

	KLogEx(DEBUG_TRACE_INFO, "Complete.\n");

	return STATUS_SUCCESS;
}