#include "precomp.h"
#include "ntddscsi.h"
#include "srb.h"

enum
{
	odd_imapi = 0,
	odd_atapi,
	odd_cdrom,
};


static WCHAR* s_oddname[] =
{
	L"\\driver\\imapi",
	L"\\driver\\atapi",
	L"\\driver\\Cdrom",
	NULL
};

static ULONG s_oddtype = 0;

NTSTATUS ODDControl(PDRIVER_DISPATCH dispath, PDEVICE_OBJECT DeviceObject, PIRP Irp);


IRP_CORE ODDIrpCore = { 0, };

#define ODDLock() InterlockedIncrement((volatile LONG *)&ODDIrpCore.IrpEnterCount);
#define ODDUnLock() InterlockedDecrement((volatile LONG *)&ODDIrpCore.IrpEnterCount);


static IRP_HOOK_HANDLER ODDIrpHook[] =
{
	{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, TRUE, ODDControl},
	{NULL, IRP_MJ_INTERNAL_DEVICE_CONTROL, TRUE, ODDControl},
	{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, FALSE}
};

static PROCESS_MANAGER s_oddallowprocess[] =
{
	{PG_PID_ALLOW, L"fph.exe",		7},
	{PG_PID_ALLOW, L"fclient.exe",	11},
	{PG_PID_ALLOW, L"iexplore.exe",13},
	{PG_PID_ALLOW, L"kcdsaclt.exe",13},
	{PG_PID_ALLOW, L"services.exe",13},
	{PG_PID_ALLOW, L"AYAgent.aye",	11}
};


static ULONG GetIrpHookHandlerCount()
{
	return sizeof(ODDIrpHook) / 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;
}


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

BOOLEAN IsODDReadOnly(PIRP irp)
{
	PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(irp);
	ULONG id = irpstack->MajorFunction;
	ULONG code = irpstack->Parameters.DeviceIoControl.IoControlCode;
	PSCSI_REQUEST_BLOCK scsi = NULL;
	PSCSI_PASS_THROUGH_DIRECT scsi_through = NULL;

	try
	{
		if (id == IRP_MJ_INTERNAL_DEVICE_CONTROL && s_oddtype == odd_atapi)
		{
			scsi = (PSCSI_REQUEST_BLOCK)irpstack->Parameters.Scsi.Srb;
			if (scsi->Function == SRB_FUNCTION_EXECUTE_SCSI)
			{
				if (scsi->DataBuffer == NULL)
					return FALSE;

				if (scsi->SrbFlags & SRB_FLAGS_DATA_OUT)
					return TRUE;
			}
		}
		else if (id == IRP_MJ_DEVICE_CONTROL && s_oddtype == odd_cdrom)
		{
			if (code == IOCTL_SCSI_PASS_THROUGH_DIRECT)
			{
				scsi_through = (PSCSI_PASS_THROUGH_DIRECT)irp->AssociatedIrp.SystemBuffer;
				if (scsi_through->DataIn == SCSI_IOCTL_DATA_OUT)
					return TRUE;
			}
		}
	}
	finally
	{

	}

	return FALSE;
}

NTSTATUS ODDControl(PDRIVER_DISPATCH dispath, PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(Irp);
	//ULONG id = irpstack->MajorFunction;
	ULONG ControlCode = irpstack->Parameters.DeviceIoControl.IoControlCode;
	ULONG state = 0;
	UNICODE_STRING processImagePath = { 0, };
	ULONG ProcessId;
	WCHAR* process_name = NULL;

	if (!g_bs1Flt.IsAttached)
		goto $ODDCleanup;

	if (IsOwnDeviceObject(DeviceObject))
		goto $ODDCleanup;

	if (g_bs1Flt.UserProcess == PsGetCurrentProcess())
		goto $ODDCleanup;

	if (IOCTL_SCSI_PASS_THROUGH != ControlCode && IOCTL_SCSI_PASS_THROUGH_DIRECT != ControlCode)
	{
		goto $ODDCleanup;
	}

	ProcessId = HandleToULong(PsGetCurrentProcessId());
	status = UGetCurrentStackProcessImageName(ProcessId, &processImagePath);
	if (!NT_SUCCESS(status))
		goto $ODDCleanup;

	// 	if( IOCTL_SCSI_PASS_THROUGH != ControlCode && IOCTL_SCSI_PASS_THROUGH_DIRECT != ControlCode)
	// 	{
	// 		KLog("ODDControl, ioctrl except s_oddtype(%d), id(%x) ControlCode(%x)\n", s_oddtype, id, ControlCode);
	// 		goto $ODDCleanup;
	// 	}


	process_name = wcsrchr(processImagePath.Buffer, '\\');
	if (!process_name)
		process_name = processImagePath.Buffer;
	else
		++process_name;

	if (IsExceptionProcess(s_oddallowprocess, ARRAYSIZE(s_oddallowprocess), process_name) == PG_PID_ALLOW)
		goto $ODDCleanup;

	state = GetPolicyState(BDC_CDROM);
	if (state == READONLY)
	{
		if (!IsODDReadOnly(Irp))
			goto $ODDCleanup;
	}
	//IOCTL_STORAGE_CHECK_VERIFY2 
	if (state != ENABLE)
	{
		/*
				/// ¿¹¿Ü ó¸®
				PSCSI_REQUEST_BLOCK scsi = NULL;
				PSCSI_PASS_THROUGH_DIRECT scsi_through = NULL;

				if(id == IRP_MJ_INTERNAL_DEVICE_CONTROL && s_oddtype == odd_atapi)
				{
					scsi = (PSCSI_REQUEST_BLOCK)irpstack->Parameters.Scsi.Srb;
					if(scsi)
					{
						KLog("ODDControl, inter	Function(%x) DataTransferLength(%d) SrbFlags(%x)\n", scsi->Function, scsi->DataTransferLength, scsi->SrbFlags);
					}

		// 			if(scsi->Function == SRB_FUNCTION_EXECUTE_SCSI)
		// 			{
		// 				if(scsi->DataBuffer == NULL)
		// 					return FALSE;
		//
		// 				if(scsi->SrbFlags & SRB_FLAGS_DATA_OUT)
		// 					return TRUE;
				}
				else if(id == IRP_MJ_DEVICE_CONTROL && s_oddtype == odd_cdrom)
				{
					if(ControlCode == IOCTL_SCSI_PASS_THROUGH_DIRECT)
					{
						scsi_through = (PSCSI_PASS_THROUGH_DIRECT)Irp->AssociatedIrp.SystemBuffer;
		// 				if(scsi_through->DataIn == SCSI_IOCTL_DATA_OUT)
		// 					return TRUE;
						if(scsi_through)
						{
							KLog("ODDControl, device	DataIn(%x) DataTransferLength(%d) TargetId(%x)\n", scsi_through->DataIn, scsi_through->DataTransferLength, scsi_through->TargetId);
						}
					}
				}

				KLog("ODDControl, processname(%S) ControlCode(%x)\n", processImagePath.Buffer, ControlCode);
		*/

		SetLog(NULL, NULL, LOG_POLICY, BDC_CDROM, state, 0, (PWCHAR)processImagePath.Buffer, NULL);


		UStrFree(&processImagePath);

		//UserNotifyMessage(DEVICE_CDROM, NULL, NULL);
		Irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
		IoCompleteRequest(Irp, IO_NO_INCREMENT);
		return STATUS_INVALID_DEVICE_REQUEST;
	}

$ODDCleanup:
	UStrFree(&processImagePath);
	return dispath(DeviceObject, Irp);

}




NTSTATUS ODDHookDispatch(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{

	NTSTATUS NtStatus = STATUS_SUCCESS;
	PIO_STACK_LOCATION pCurrentIoStack = IoGetCurrentIrpStackLocation(Irp);
	ULONG Id = pCurrentIoStack->MajorFunction;

	ODDLock();

	if (IsAdrIrpHookHandler(ODDIrpHook, Id))
	{
		NtStatus = ODDIrpHook[Id].Work(
			ODDIrpHook[Id].pOrgHandler,
			DeviceObject,
			Irp
		);
	}
	ODDUnLock();

	return NtStatus;

}


static PVOID GetHookDispatch()
{
	return (PVOID)ODDHookDispatch;
}

static PIRP_HOOK_HANDLER GetIrpHookHandler()
{
	return ODDIrpHook;
}

static PIRP_CORE GetIrpCore()
{

	ODDIrpCore.HookDispatch = GetHookDispatch();
	ODDIrpCore.IrpHookHandler = GetIrpHookHandler();
	ODDIrpCore.IrpHookHandlerCount = GetIrpHookHandlerCount();
	ODDIrpCore.IrpEnterCount = 0;
	return &ODDIrpCore;
}



NTSTATUS ODDIrpHookCleanup()
{
	ULONG cnt = 0;
	//ULONG OrSaveWP;
	LARGE_INTEGER WaitTime;
	PDRIVER_OBJECT DriverObject = ODDIrpCore.HookDriverObject;
	ULONG IrpHookHandlerCount = ODDIrpCore.IrpHookHandlerCount;
	PIRP_HOOK_HANDLER pIhh = ODDIrpCore.IrpHookHandler;

	if (!DriverObject || !IrpHookHandlerCount || !pIhh)
		return STATUS_UNSUCCESSFUL;

	if (!pIhh)
		return STATUS_UNSUCCESSFUL;

	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 (ODDIrpCore.IrpEnterCount > 0)
		KeDelayExecutionThread(KernelMode, FALSE, &WaitTime);

	KeDelayExecutionThread(KernelMode, FALSE, &WaitTime);
	KLogEx(DEBUG_TRACE_INFO, "complete\n");

	return STATUS_SUCCESS;
}

NTSTATUS ODDIrpHookInit()
{
	// 	WCHAR * xpname = L"\\driver\\imapi";
	// 	WCHAR * win7name = L"\\driver\\atapi";
	// 	WCHAR * etcname = L"\\driver\\Cdrom";
	WCHAR* name = NULL;
	//ULONG type = 0;
	ULONG cnt = 0;
	//PVOID TargetPoint = NULL;

	ULONG IrpHookHandlerCount = 0;
	PIRP_HOOK_HANDLER pIhh = NULL;

	PIRP_CORE irpcore = NULL;
	PDRIVER_OBJECT obj = NULL;

	if (g_bs1Flt.OsVersion >= WinVista)
	{
		s_oddtype = odd_atapi;
	}
	else
	{
		s_oddtype = odd_imapi;
	}

	name = s_oddname[s_oddtype];
	KLogEx(DEBUG_TRACE_INFO, "driver(%S)\n", name);
	obj = SearchDriverObject(name);
	if (!obj)
	{

		s_oddtype = odd_cdrom;
		name = s_oddname[s_oddtype];
		KLogEx(DEBUG_TRACE_INFO, "driver(%S)\n", name);
		obj = SearchDriverObject(name);
		if (!obj)
			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(%08p)\n", irpcore->HookDriverObject->MajorFunction[cnt]);
		}

		InterlockedExchangePointer((PVOID)&irpcore->HookDriverObject->MajorFunction[cnt], (PVOID)irpcore->HookDispatch);
		KLogEx(DEBUG_TRACE_INFO, "hook c (%S), (%d), (%08x), (%08x)\n", name, cnt, pIhh[cnt].pOrgHandler, irpcore->HookDriverObject->MajorFunction[cnt]);
	}

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