#include "precomp.h"


/*************************************************************************
    Prototypes
*************************************************************************/

//
//  Assign text sections for each routine.
//


#ifdef ALLOC_PRAGMA
#pragma alloc_text(PAGE, Bs1FltPortConnect)
#pragma alloc_text(PAGE, Bs1FltPortDisconnect)
#pragma alloc_text(PAGE, Bs1FltMssageProc)
#endif

typedef struct _CLIENT_CONTEXT {
	LIST_ENTRY      ListEntry;      // 리스트 연결용
	PFLT_PORT       ClientPort;     // 클라이언트 포트 핸들
	PEPROCESS       UserProcess;    // 연결한 유저 프로세스
	HANDLE          ProcessId;      // PID (로깅/식별용)
} CLIENT_CONTEXT, * PCLIENT_CONTEXT;

// 전역 변수 선언 (DriverEntry 등에서 초기화 필수)
LIST_ENTRY      g_ClientListHead;
KSPIN_LOCK      g_ClientListLock;
LONG g_refCnt = 0;

#define Reference() InterlockedIncrement(&g_refCnt)
#define Dereference() InterlockedDecrement(&g_refCnt);


void Bs1FltPortInit()
{
	InitializeListHead(&g_ClientListHead);
	KeInitializeSpinLock(&g_ClientListLock);
}

void Bs1FltPortUnload()
{
	KIRQL oldIrql;
	PCLIENT_CONTEXT clientCtx = NULL;
	PFLT_PORT portToClose = NULL;

	while (TRUE)
	{
		KeAcquireSpinLock(&g_ClientListLock, &oldIrql);

		if (IsListEmpty(&g_ClientListHead))
		{
			KeReleaseSpinLock(&g_ClientListLock, oldIrql);
			break;
		}

		PLIST_ENTRY pEntry = g_ClientListHead.Flink;
		clientCtx = CONTAINING_RECORD(pEntry, CLIENT_CONTEXT, ListEntry);
		portToClose = (PFLT_PORT)InterlockedExchangePointer((PVOID*)&clientCtx->ClientPort, NULL);
		KeReleaseSpinLock(&g_ClientListLock, oldIrql);

		if (portToClose != NULL)
		{
			KLogEx(DEBUG_TRACE_INFO, "Force closing port: 0x%p\n", portToClose);
			FltCloseClientPort(g_bs1Flt.Filter, &portToClose);
		}
		else
		{
			KeAcquireSpinLock(&g_ClientListLock, &oldIrql);
			if (!IsListEmpty(&g_ClientListHead) && g_ClientListHead.Flink == pEntry)
			{
				RemoveEntryList(pEntry);
				ExFreePoolWithTag(clientCtx, '1SB');
			}
			KeReleaseSpinLock(&g_ClientListLock, oldIrql);
		}
	}
	return;
}

NTSTATUS
Bs1FltPortConnect(
	__in PFLT_PORT ClientPort,
	__in_opt PVOID ServerPortCookie,
	__in_bcount_opt(SizeOfContext) PVOID ConnectionContext,
	__in ULONG SizeOfContext,
	__deref_out_opt PVOID* ConnectionCookie
)
{
	PCLIENT_CONTEXT clientCtx = NULL;
	KIRQL oldIrql;

	PAGED_CODE();

	UNREFERENCED_PARAMETER(ServerPortCookie);
	UNREFERENCED_PARAMETER(ConnectionContext);
	UNREFERENCED_PARAMETER(SizeOfContext);
	UNREFERENCED_PARAMETER(ConnectionCookie);

	ASSERT(g_bs1Flt.ClientPort == NULL);
	ASSERT(g_bs1Flt.UserProcess == NULL);

	//
	//  Set the user process and port.
	//
	
	clientCtx = (PCLIENT_CONTEXT)ExAllocatePoolWithTag(NonPagedPool, sizeof(CLIENT_CONTEXT), '1SB');
	if (clientCtx == NULL)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	RtlZeroMemory(clientCtx, sizeof(CLIENT_CONTEXT));
	clientCtx->ClientPort = ClientPort;
	clientCtx->UserProcess = PsGetCurrentProcess();
	clientCtx->ProcessId = PsGetCurrentProcessId();

	KeAcquireSpinLock(&g_ClientListLock, &oldIrql);
	InsertTailList(&g_ClientListHead, &clientCtx->ListEntry);
	KeReleaseSpinLock(&g_ClientListLock, oldIrql);

	*ConnectionCookie = (PVOID)clientCtx;
	
	if( g_refCnt == 0 )
	{
		// 초기화 작업
		g_bs1Flt.UserProcess = PsGetCurrentProcess();
		g_bs1Flt.ClientPort = ClientPort;
	}
	
	KLogEx(DEBUG_TRACE_INFO, "connected, port=0x%p\n", ClientPort);
	Reference();

	return STATUS_SUCCESS;
}

VOID
Bs1FltPortDisconnect(
	__in_opt PVOID ConnectionCookie
)
{
	PCLIENT_CONTEXT clientCtx = (PCLIENT_CONTEXT)ConnectionCookie;
	KIRQL oldIrql;
	PFLT_PORT portToClose = NULL;
	//UNREFERENCED_PARAMETER(ConnectionCookie);

	PAGED_CODE();

	if (clientCtx == NULL)
	{
		return;
	}

	KeAcquireSpinLock(&g_ClientListLock, &oldIrql);
	if (!IsListEmpty(&clientCtx->ListEntry))
	{
		RemoveEntryList(&clientCtx->ListEntry);
		//InitializeListHead(&clientCtx->ListEntry); // 제거 후 초기화 (중복 제거 방지)
		clientCtx->ListEntry.Flink = NULL;
		clientCtx->ListEntry.Blink = NULL;
	}

	BOOLEAN bIsLastClient = IsListEmpty(&g_ClientListHead);
	KeReleaseSpinLock(&g_ClientListLock, oldIrql);
	
	portToClose = (PFLT_PORT)InterlockedExchangePointer((PVOID*)&clientCtx->ClientPort, NULL);
	KLogEx(DEBUG_TRACE_INFO, "disconnected, port=0x%p\n", portToClose);
	if (portToClose)
	{
		FltCloseClientPort(g_bs1Flt.Filter, &portToClose);
		clientCtx->ClientPort = NULL;
	}

	ExFreePoolWithTag(clientCtx, '1SB');
	Dereference();
	
	if (bIsLastClient)
	{
		g_bs1Flt.UserProcess = NULL;
		g_bs1Flt.ClientPort = NULL;
		g_bs1Flt.IsAttached = FALSE;
		g_bs1Flt.IsShareWatched = FALSE;
		g_bs1Flt.IsDeviceProtect = FALSE;
		// 	ResetFileObjectList();
		// 	ResetFileList();
		// 	CleanupPathlist();
		// 	CleanupProcesslist();
		// 	PgReset();
	}
	
	if (g_refCnt == 0)
	{
		/*
		NTSTATUS ntstatus = 0;
		UNICODE_STRING drivername = {0,};
		RtlInitUnicodeString(&drivername, DRIVERNAME);

// 		PBShareLockUnload(0);
//
// 		KLogEx("Unload RegPath-----");
//
		KLogEx(DEBUG_TRACE_INFO, "Unload RegPath = %S\n", drivername.Buffer);
		ntstatus = FltUnloadFilter(&drivername);
		KLogEx(DEBUG_TRACE_INFO, "Unload ntstatus(%x)\n", ntstatus);
		*/
	}
	else
	{

	}

}


NTSTATUS
Bs1FltMssageProc(
	__in PVOID ConnectionCookie,
	__in_bcount_opt(InputBufferSize) PVOID InputBuffer,
	__in ULONG InputBufferSize,
	__out_bcount_part_opt(OutputBufferSize, *ReturnOutputBufferLength) PVOID OutputBuffer,
	__in ULONG OutputBufferSize,
	__out PULONG ReturnOutputBufferLength
)
{
	NTSTATUS		ntStatus = STATUS_SUCCESS;
	PBS1FLT_MESSAGE msg = NULL;

	PAGED_CODE();

	UNREFERENCED_PARAMETER(InputBufferSize);
	UNREFERENCED_PARAMETER(ConnectionCookie);
	UNREFERENCED_PARAMETER(ReturnOutputBufferLength);
	UNREFERENCED_PARAMETER(OutputBufferSize);
	try
	{
		msg = (PBS1FLT_MESSAGE)InputBuffer;
		if (!InputBuffer)
			leave;

		switch (msg->id_)
		{
		case SET_DEBUG_LEVEL:
		{
			g_DebugLevel = msg->pid_;
			KLogEx(DEBUG_TRACE_INFO, "SET_DEBUG_LEVEL(%x)\n", g_DebugLevel);			
		}
		break;
		case START:
		{
			g_bs1Flt.IsAttached = TRUE;

			KLogEx(DEBUG_TRACE_INFO, "START\n");
		}
		break;
		case STOP:
		{
			g_bs1Flt.IsAttached = FALSE;

			KLogEx(DEBUG_TRACE_INFO, "STOP\n");
		}
		break;
		case CLEAR:
		{
			CleanupPathlist();
			CleanupProcesslist();
			CleanupFilelist();
			CleanupProcessProtectList();
			PgReset();
			Initpolicy();
			//CleanupUnloadlist();

			PgAddPid(0, PG_PID_ALLOW);
			PgAddPid(4, PG_PID_ALLOW);
			PgAddPid(8, PG_PID_ALLOW);

			KLogEx(DEBUG_TRACE_INFO, "CLEAR\n");
		}
		break;
		case SET_LOG_TYPE:
		{
			/*if (msg->type_ == 0 || msg->type_ == LOG_ALL)
				g_bs1Flt.LogType = msg->type_;
			else*/
				g_bs1Flt.LogType = msg->type_;

			KLogEx(DEBUG_TRACE_INFO, "SET_LOG_TYPE(%d)\n", g_bs1Flt.LogType);
		}
		break;
		case SET_PROCESSNAME:
		{
			if (msg->w.process_path_.size == 0)
				break;

			SetProcessPath(msg->w.process_path_.type, msg->w.process_path_.path, msg->w.process_path_.size);
			KLogEx(DEBUG_TRACE_INFO, "SET_PROCESSNAME(%d)(%S)(%d)\n", \
					msg->w.process_path_.type, \
					msg->w.process_path_.path, \
					msg->w.process_path_.size \
			);
		}
		break;
		case DEL_PROCESSNAME:
		{
			DelProcessPath(msg->w.process_path_.type, msg->w.process_path_.path, msg->w.process_path_.size);
			KLogEx(DEBUG_TRACE_INFO, "DEL_PROCESSNAME(%d)(%S)(%d)\n", \
					msg->w.process_path_.type, \
					msg->w.process_path_.path, \
					msg->w.process_path_.size	\
			);
		}
		break;
		case SET_PATH:
		{
			if (msg->w.file_path_.size == 0)
				break;

			SetProtectPath(msg->w.file_path_.type, msg->w.file_path_.path, msg->w.file_path_.size);
			KLogEx(DEBUG_TRACE_INFO, "SET_PATH(%d)(%S)(%d)\n", msg->w.file_path_.type, msg->w.file_path_.path, msg->w.file_path_.size);
		}
		break;
		case DEL_PATH:
		{
			DelPathlist(msg->w.file_path_.type, msg->w.file_path_.path, msg->w.file_path_.size);
		}
		break;
		case SET_FILENAME:
		{
			if (msg->w.file_path_.size == 0)
				break;

			SetFileName(msg->w.file_path_.type, msg->w.file_path_.path, msg->w.file_path_.size);
			KLogEx(DEBUG_TRACE_INFO, "SET_FILENAME(%d)(%S)(%d)\n", msg->w.file_path_.type, msg->w.file_path_.path, msg->w.file_path_.size);
		}
		break;
		case START_FOLDER_PROTECT:
		{
			if (msg->state_)
				g_bs1Flt.IsFolderProtect = TRUE;
			else
				g_bs1Flt.IsFolderProtect = FALSE;

			KLogEx(DEBUG_TRACE_INFO, "START_FOLDER_PROTECT(%d)\n", msg->state_);
		}
		break;
		case START_DEVICE_PROTECT:
		{
			if (msg->state_)
				g_bs1Flt.IsDeviceProtect = TRUE;
			else
				g_bs1Flt.IsDeviceProtect = FALSE;

			KLogEx(DEBUG_TRACE_INFO, "START_DEVICE_PROTECT(%d)\n", msg->state_);
		}
		break;

		case START_IS_SHARE_FOLDER_WATCHE:
		{			
			if (msg->state_)
				g_bs1Flt.IsShareWatched = TRUE;
			else
				g_bs1Flt.IsShareWatched = FALSE;

			KLogEx(DEBUG_TRACE_INFO, "START_IS_SHARE_FOLDER_WATCHE(%d)\n", msg->state_);
		}
		break;
		case GET_PROCESS_NOTIFY_STATUS:
		{
			*(DWORD*)OutputBuffer = (DWORD)GetProcessNotifyStatus();
			*ReturnOutputBufferLength = sizeof(DWORD);
			KLogEx(DEBUG_TRACE_INFO, "GET_PROCESS_NOTIFY_STATUS(%d)\n", *(DWORD*)OutputBuffer);
		}
		break;
		case SET_POLICY:
		{			
			KLogEx(DEBUG_TRACE_INFO, "SET_POLICY, devicetype(%d), state(%d), log(%d)\n", msg->device_policy_.device_type, msg->device_policy_.state, msg->device_policy_.islog);
			SetPolicy(msg->device_policy_.device_type, msg->device_policy_.state, msg->device_policy_.islog);
			break;
		}
		case GET_LOG:
		{
			ULONG req = 0;

			ntStatus = GetLog(OutputBuffer, OutputBufferSize, &req);
			if (ntStatus == STATUS_SUCCESS)
				*ReturnOutputBufferLength = OutputBufferSize;
			else
				*ReturnOutputBufferLength = req;

			break;
		}
		case SET_HOOK:
		{			
#ifdef USB_PORT_HOOK
			if(msg->type_ == BDC_USB) // USB Port Hook
			{
				if (msg->state_)
					USBIrpHookInit();
				else
					USBIrpHookCleanup();
			}
			else if (msg->type_ == BDC_BLUETOOTH)
			{
				if (msg->state_)
					BlueToothIrpHookInit();
				else
					BlueToothIrpHookCleanup();
			}
			else if (msg->type_ == BDC_MTP)
			{
				if (msg->state_)
					MtpIrpHookInit();
				else
					MtpIrpHookCleanup();
			}
			else
			{
				switch (msg->type_)
				{
				case BDC_1394:
				case BDC_SERIAL:
				case BDC_PARALLEL:
				case BDC_IRDA:
				case BDC_MODEM:
					if (msg->state_)
						EtcIrpHookInit();
					else
						EtcIrpHookCleanup();
					break;
				default:
					break;
				}				
			}
#endif
			KLogEx(DEBUG_TRACE_INFO, "SET_HOOK(%d)(%d)\n", msg->type_, msg->state_);
			break;
		}
		case SET_REG_PROTECT:
		{
			if (msg->state_)
				StartRegFlt(g_bs1Flt.DriverObject);
			else
				g_bs1Flt.IsRegProtect = FALSE;

			KLogEx(DEBUG_TRACE_INFO, "SET_REG_PROTECT(%d)\n", msg->state_);
			break;
		}
		case SET_REG_KEY:
		{
			SetRegKeylist(msg->w.regkey_.regkey);
			KLogEx(DEBUG_TRACE_INFO, "SET_REG_KEY(%S)\n", msg->w.regkey_.regkey);
			break;
		}
		case DEL_REG_KEY:
		{
			DelRegKeylist(msg->w.regkey_.regkey);
			KLogEx(DEBUG_TRACE_INFO, "DEL_REG_KEY(%S)\n", msg->w.regkey_.regkey);
			break;
		}
		case SET_PROCESS_PROTECT:
		{
			if (msg->state_)
				InstallProcessProtect();
			else
				g_bs1Flt.IsProcessProtect = FALSE;

			KLogEx(DEBUG_TRACE_INFO, "SET_PROCESS_PROTECT(%d)\n", g_bs1Flt.IsProcessProtect);
			break;
		}
		case SET_PROCESS_PROTECT_PID:
		{
			PgAddPid(msg->pid_, PG_PID_PROTECT);
			KLogEx(DEBUG_TRACE_INFO, "SET_PROCESS_PROTECT_PID(%d)\n", msg->pid_);
			break;	
		}
		case DEL_PROCESS_PROTECT_PID:
		{
			if (msg->pid_ == 0)
			{
				PgResetState(PG_PID_PROTECT);
			}
			else
			{
				PgRemovePidState(msg->pid_, PG_PID_PROTECT);
			}
			
			KLogEx(DEBUG_TRACE_INFO, "DEL_PROCESS_PROTECT_PID(%d)\n", msg->pid_);
			break;
		}
		case SET_PROCESS_PROTECT_PROCESSNAME:
		{			
			SetProcessProtectList(msg->w.process_path_.type, msg->w.process_path_.path, msg->w.process_path_.size);
			KLogEx(DEBUG_TRACE_INFO, "SET_PROCESS_PROTECT_PROCESSNAME(%S)(%d)\n", msg->w.process_path_.path, msg->w.process_path_.size);
			break;
		}
		case DEL_PROCESS_PROTECT_PROCESSNAME:
		{
			if (msg->w.process_path_.size == 0)
			{
				CleanupProcessProtectList();
			}
			else
			{
				DelProcessProtectList(msg->w.process_path_.type, msg->w.process_path_.path, msg->w.process_path_.size);
			}
			KLogEx(DEBUG_TRACE_INFO, "DEL_PROCESS_PROTECT_PROCESSNAME(%S)(%d)\n", msg->w.process_path_.path, msg->w.process_path_.size);
			break;
		}
		case SET_USB_DISK_EXCEPT:
		{
			if (msg->state_ == STATE_SET)
			{
				SetUsbDiskException(
					msg->w.usb_except_.vendorid,
					msg->w.usb_except_.productid,
					msg->w.usb_except_.productrevisionlevel,
					msg->w.usb_except_.vendorspecific
				);

				KLogEx(DEBUG_TRACE_INFO, "SET_USB_DISK_EXCEPT, (STATE_SET), vid(%s), pid(%s)(%s)\n", 
					msg->w.usb_except_.vendorid,
					msg->w.usb_except_.productid,
					msg->w.usb_except_.vendorspecific
				);
			}
			else if (msg->state_ == STATE_DEL)
			{
				KLogEx(DEBUG_TRACE_INFO, "SET_USB_DISK_EXCEPT, (STATE_DEL), vid(%s), pid(%s)(%s)\n",
					msg->w.usb_except_.vendorid,
					msg->w.usb_except_.productid,
					msg->w.usb_except_.vendorspecific
				);

				DelUsbDiskExceptionList(
					msg->w.usb_except_.vendorid,
					msg->w.usb_except_.productid,
					msg->w.usb_except_.productrevisionlevel,
					msg->w.usb_except_.vendorspecific
				);
			}
			else
			{
				CleanupUsbDiskExceptionList();
				KLogEx(DEBUG_TRACE_INFO, "SET_USB_DISK_EXCEPT, (STATE_CLEAR)\n");
			}

			
			break;
		}
		case SET_USB_PORT_EXCEPT:
		{
#ifdef USB_PORT_HOOK
			if (msg->state_ == STATE_SET)
			{
				SetUsbPortException(
					msg->w.usb_port_except_.devicetype,
					msg->w.usb_port_except_.vendorid,
					msg->w.usb_port_except_.productid,
					msg->w.usb_port_except_.bcddevice,
					msg->w.usb_port_except_.serial
				);

				KLogEx(DEBUG_TRACE_INFO, "SET_USB_PORT_EXCEPT, (STATE_SET), vid(%x), pid(%x), serial(%S)\n", msg->w.usb_port_except_.vendorid,
					msg->w.usb_port_except_.productid, msg->w.usb_port_except_.serial
				);
			}
			else if (msg->state_ == STATE_DEL)
			{
				DelUsbPortExceptionList(
					msg->w.usb_port_except_.devicetype,
					msg->w.usb_port_except_.vendorid,
					msg->w.usb_port_except_.productid,
					msg->w.usb_port_except_.bcddevice,
					msg->w.usb_port_except_.serial
				);

				KLogEx(DEBUG_TRACE_INFO, "SET_USB_PORT_EXCEPT, (STATE_DEL), vid(%x), pid(%x), serial(%S)\n", msg->w.usb_port_except_.vendorid,
					msg->w.usb_port_except_.productid, msg->w.usb_port_except_.serial
				);
			}
			else
			{
				CleanupUsbPortExceptionList();
				KLogEx(DEBUG_TRACE_INFO, "SET_USB_PORT_EXCEPT, (STATE_CLEAR)(%d)\n" , msg->state_);
			}	
#endif
			break;
		}
		default:
			break;
		}
	}
	except(EXCEPTION_EXECUTE_HANDLER)
	{
		ULONG ulEceptcode = GetExceptionCode();
		KLogEx(DEBUG_TRACE_INFO, "eception = %08x\n", ulEceptcode);
		return ulEceptcode;
	}

	return ntStatus;
}