#include "precomp.h"

#ifndef IOCTL_WPD_MESSAGE_READ
#define IOCTL_WPD_MESSAGE_READ       CTL_CODE(FILE_DEVICE_UNKNOWN, 0x33C, METHOD_BUFFERED, FILE_READ_DATA)
#define IOCTL_WPD_MESSAGE_WRITE      CTL_CODE(FILE_DEVICE_UNKNOWN, 0x33D, METHOD_BUFFERED, FILE_WRITE_DATA)
#endif

// 버퍼 오프셋 정의 (가독성을 위해)
#define WPD_MTP_PAYLOAD_OFFSET       48
#define MTP_OPCODE_DELETE_OR_WRITE   0xDD  //0xffffffdd (signed char) -> 0xDD (unsigned char)

enum
{
	mtp_wudfRd = 0,
	mtp_maximum
};

static WCHAR* s_mptname[] =
{
	L"\\Driver\\WudfRd",
	NULL
};


static BOOLEAN enable_mtphook = FALSE;

NTSTATUS MtpDeviceIoControl(PDRIVER_DISPATCH dispatch, PDEVICE_OBJECT DeviceObject, PIRP irp);
NTSTATUS MtpHookDispatch_0(PDEVICE_OBJECT DeviceObject, PIRP Irp);

static PDRIVER_DISPATCH s_ProxyDispatchers[mtp_maximum] =
{
	MtpHookDispatch_0
};


#define MTP_COMMON_HOOK_HANDLERS \
    [IRP_MJ_DEVICE_CONTROL] = { NULL, IRP_MJ_DEVICE_CONTROL, TRUE, MtpDeviceIoControl }, \

static HOOK_CONTEXT g_MtpHookContexts[mtp_maximum] =
{
	{ NULL, FALSE, 0, { MTP_COMMON_HOOK_HANDLERS } }
};

NTSTATUS MtpDeviceIoControl(PDRIVER_DISPATCH dispatch, PDEVICE_OBJECT deviceObject, PIRP irp)
{
	PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(irp);
	ULONG controlCode = irpstack->Parameters.DeviceIoControl.IoControlCode;
	ULONG inputLen = irpstack->Parameters.DeviceIoControl.InputBufferLength;
	ULONG ProcessId = 0;
	ULONG state = 0;
	ULONG policyLog = 0;

	char szProcessName[20] = { 0, };
	wchar_t processName[50] = { 0, };
	wchar_t notice[50] = { 0, };

	if (!g_bs1Flt.IsAttached || !enable_mtphook)
		goto $MtpCleanup;

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

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

	if (state == DISABLE)
	{		
		if(policyLog)
		{
			RtlStringCbPrintfW(processName, sizeof(processName), L"%S", szProcessName);
			RtlStringCbPrintfW(notice, sizeof(notice), L"MTP Blocked control(%x)", controlCode);
			SetLog(NULL, NULL, LOG_POLICY, BDC_MTP, state, 0, processName, notice);
		}

		KLogEx(DEBUG_TRACE_INFO, " block!!");
		irp->IoStatus.Status = STATUS_UNSUCCESSFUL;
		irp->IoStatus.Information = 0;
		IoCompleteRequest(irp, IO_NO_INCREMENT);

		return STATUS_UNSUCCESSFUL;
	}

	if (state == READONLY && irpstack->FileObject)
	{
		KLogEx(DEBUG_TRACE_INFO, " controlCode(%x)", controlCode);
		switch (controlCode)
		{
			case IOCTL_WPD_MESSAGE_WRITE:
			case 0x22833C: // WudfRd에서 자주 보이는 IOCTL (IOCTL_WPD_MESSAGE_READACCESS 등)
			case 0x228340: // WudfRd에서 자주 보이는 IOCTL (IOCTL_WPD_MESSAGE_WRITEACCESS 등)
			{
				PUCHAR pBuf = NULL;

				// 버퍼 가져오기 (WPD는 주로 METHOD_BUFFERED 사용 -> SystemBuffer)
				if (irp->AssociatedIrp.SystemBuffer)
				{
					KLogEx(DEBUG_TRACE_INFO, " SystemBuffer");
					pBuf = (PUCHAR)irp->AssociatedIrp.SystemBuffer;
					PrintHexData(pBuf, inputLen);
				}
				else if (irpstack->Parameters.DeviceIoControl.Type3InputBuffer) // METHOD_NEITHER 경우
				{
					KLogEx(DEBUG_TRACE_INFO, " Type3InputBuffer");
					pBuf = (PUCHAR)irpstack->Parameters.DeviceIoControl.Type3InputBuffer;
					PrintHexData(pBuf, inputLen);
				}
				else
				{
					KLogEx(DEBUG_TRACE_INFO, " buffer invaild");
					//KLogEx(DEBUG_TRACE_INFO, "MmGetSystemAddressForMdlSafe (%p)", irp->MdlAddress);

					//if (irp->MdlAddress)
					//{
					//	ULONG	Length = 0;
					//	ULONG	Offset = 0;

					//	pBuf = MmGetSystemAddressForMdlSafe(irp->MdlAddress, HighPagePriority);
					//	Length = MmGetMdlByteCount(irp->MdlAddress);
					//	Offset = MmGetMdlByteOffset(irp->MdlAddress);

					//	if (pBuf)
					//	{
					//		PrintHexData(pBuf, Length);
					//	}
					//}

					////////////////////////////////////////////////////////////////////////////
					////	Irp->UserBuffer
					//KLogEx(DEBUG_TRACE_INFO, "UserBuffer (%p)", irp->UserBuffer);
					//if (irp->UserBuffer)
					//{
					//	pBuf = irp->UserBuffer;
					//	PrintData(pBuf, inputLen);
					//}
				}

				// MdlAddress는 보통 Read/Write IRP에서 사용되나 IOCTL에서는 드묾. 필요시 추가.

				// 5. 버퍼 내용 검사
				if (pBuf != NULL && inputLen > WPD_MTP_PAYLOAD_OFFSET)
				{
					__try
					{
						// [질문 2 답변] memcpy 대신 직접 접근 및 UCHAR 비교
						UCHAR mptOpCodeByte = pBuf[WPD_MTP_PAYLOAD_OFFSET];

						KLogEx(DEBUG_TRACE_INFO, " mptOpCodeByte(%x)", mptOpCodeByte);
						// 0xffffffdd (signed char) == 0xDD (unsigned char)
						if (mptOpCodeByte == MTP_OPCODE_DELETE_OR_WRITE)
						{
							KLogEx(DEBUG_TRACE_INFO, "MTP Write/Delete Detected at Offset %d: 0x%02X", WPD_MTP_PAYLOAD_OFFSET, mptOpCodeByte);

							if (policyLog)
							{
								RtlStringCbPrintfW(processName, sizeof(processName), L"%S", szProcessName);
								RtlStringCbPrintfW(notice, sizeof(notice), L"MTP Write Action Blocked (Op: 0x%02X)", mptOpCodeByte);
								SetLog(NULL, NULL, LOG_POLICY, BDC_MTP, state, 0, processName, notice);
							}

							irp->IoStatus.Status = STATUS_UNSUCCESSFUL;
							irp->IoStatus.Information = 0;
							IoCompleteRequest(irp, IO_NO_INCREMENT);
							return STATUS_UNSUCCESSFUL;
						}
					}
					__except (EXCEPTION_EXECUTE_HANDLER)
					{
						NTSTATUS exceptStatus = GetExceptionCode();
						KLogEx(DEBUG_TRACE_ERROR, "Exception accessing buffer: 0x%X", exceptStatus);
					}
				}
				else
				{
					KLogEx(DEBUG_TRACE_INFO, " pBuf[48] is 0 or null");
				}
				break;
			}
			default:
				break;
		}
	}

$MtpCleanup:
	return dispatch(deviceObject, irp);

}

NTSTATUS MtpHookDispatch_Common(ULONG ContextIndex, PDEVICE_OBJECT deviceObject, PIRP irp)
{
	NTSTATUS NtStatus = STATUS_SUCCESS;
	PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(irp);
	ULONG id = irpstack->MajorFunction;
	PHOOK_CONTEXT pCtx = NULL;
	PDRIVER_DISPATCH pOrgHandler = NULL;

	if (ContextIndex >= mtp_maximum)
	{
		return STATUS_UNSUCCESSFUL;
	}

	InterlockedIncrement((volatile LONG*)&g_MtpHookContexts[ContextIndex].IrpEnterCount);

	pCtx = &g_MtpHookContexts[ContextIndex];

	if (pCtx->IsHooked &&
		pCtx->HookHandlers[id].IsHook &&
		pCtx->HookHandlers[id].Work)
	{
		pOrgHandler = pCtx->HookHandlers[id].pOrgHandler;

		NtStatus = pCtx->HookHandlers[id].Work(pOrgHandler, deviceObject, irp);
	}
	else
	{
		NtStatus = STATUS_UNSUCCESSFUL;
	}

	InterlockedDecrement((volatile LONG*)&g_MtpHookContexts[ContextIndex].IrpEnterCount);

	return NtStatus;
}

NTSTATUS MtpHookDispatch_0(PDEVICE_OBJECT deviceObject, PIRP irp)
{
	return MtpHookDispatch_Common(mtp_wudfRd, deviceObject, irp);
}

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

	for (i = 0; i < mtp_maximum; i++)
	{
		hook = &g_MtpHookContexts[i];

		if (hook->IsHooked)
			continue;

		name = s_mptname[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++)
		{
			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_mtphook = TRUE;
	}

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


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

	enable_mtphook = FALSE;

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

	for (i = 0; i < mtp_maximum; i++)
	{
		hook = &g_MtpHookContexts[i];

		if (!hook->IsHooked || !hook->DriverObject)
			continue;

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

		for (mi = 0; mi <= IRP_MJ_MAXIMUM_FUNCTION; mi++)
		{
			if (!hook->HookHandlers[mi].IsHook)
				continue;
		
			if (hook->HookHandlers[mi].pOrgHandler == NULL)
				continue;

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

		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;
		//RtlZeroMemory(pCtx->HookHandlers, sizeof(pCtx->HookHandlers));
	}

	WaitTime.QuadPart = -500 * 1000 * 10;
	KeDelayExecutionThread(KernelMode, FALSE, &WaitTime);

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

	return STATUS_SUCCESS;
}