#include "precomp.h"


static DEVICE_POLICY s_policy[BDC_MAX_DEVICE_TYPE] = { 0, };


void Initpolicy()
{
	RtlZeroMemory(s_policy, sizeof(s_policy));
}

void SetPolicy(ULONG device_type, ULONG state, ULONG islog)
{
	s_policy[device_type].device_type = device_type;
	s_policy[device_type].state = state;
	s_policy[device_type].islog = islog;
}

ULONG GetPolicyState(ULONG device_type)
{
	return s_policy[device_type].state;
}

BOOLEAN IsPolicyDisable(ULONG device_type)
{
	if (s_policy[device_type].state == DISABLE)
		return TRUE;

	return FALSE;
}

ULONG IsPolicyLog(ULONG device_type)
{
	return s_policy[device_type].islog;
}


//ULONG IsNetwork(PFLT_CALLBACK_DATA data, PCFLT_RELATED_OBJECTS fltobject)
//{
//	ULONG create_character = 0;
//	ULONG device_type = 0;
//	ULONG state = DEVICE_UNKNOWN;
//	PVOLUME_CONTEXT vctx = NULL;
//	NTSTATUS ntstatus = STATUS_SUCCESS;
//
//	create_character = fltobject->FileObject->DeviceObject->Characteristics;
//	device_type = fltobject->FileObject->DeviceObject->DeviceType;
//
//	if ((create_character & FILE_REMOTE_DEVICE) && (device_type == FILE_DEVICE_NETWORK_FILE_SYSTEM))
//	{
//		return DEVICE_NETWORKDRIVEOUT;
//	}
//	return FALSE;
//}

ULONG GetDeviceType(PFLT_CALLBACK_DATA data, PCFLT_RELATED_OBJECTS fltobject)
{
	ULONG create_character = fltobject->FileObject->DeviceObject->Characteristics;
	ULONG device_type = fltobject->FileObject->DeviceObject->DeviceType;
	ULONG state = BDC_UNKNOWN_DEV;
	PVOLUME_CONTEXT vctx = NULL;
	NTSTATUS ntstatus = STATUS_SUCCESS;


	///USB DISK 인경우
	if ((create_character & FILE_REMOVABLE_MEDIA) && (device_type == FILE_DEVICE_DISK))
	{
		if ((create_character & FILE_FLOPPY_DISKETTE))
			state = BDC_FLOOPY;
		else
			state = BDC_USB_DISK;
	}
	else if ((create_character & FILE_FLOPPY_DISKETTE) && (device_type == FILE_DEVICE_DISK))
	{
		state = BDC_FLOOPY;
	}
	///CD ROM 인 경우
	else if (device_type == FILE_DEVICE_CD_ROM)
	{
		state = BDC_CDROM;
	}
	///네트워크 일 경우
	else if ((create_character & FILE_REMOTE_DEVICE) && (device_type == FILE_DEVICE_NETWORK_FILE_SYSTEM))
	{
		state = BDC_NETWORKDRIVEOUT;
	}
	else
	{
		if (GetPolicyState(BDC_NETWORKDRIVEIN) != ENABLE)
		{
			PSECURITY_SUBJECT_CONTEXT subjectContext = NULL;
			PACCESS_TOKEN    accessToken = NULL;
			PTOKEN_SOURCE    tokenSource = NULL;
			ULONG majorfunction = data->Iopb->MajorFunction;

			if (majorfunction == IRP_MJ_CREATE)
			{
				// 1. IRP 파라미터에서 직접 보안 컨텍스트 가져오기 (SeCaptureSubjectContext 불필요)
				if (data->Iopb->Parameters.Create.SecurityContext != NULL &&
					data->Iopb->Parameters.Create.SecurityContext->AccessState != NULL)
				{
					subjectContext = &data->Iopb->Parameters.Create.SecurityContext->AccessState->SubjectSecurityContext;

					// 2. 토큰 가져오기
					accessToken = SeQuerySubjectContextToken(subjectContext);

					if (accessToken)
					{
						// 3. 토큰 정보 쿼리 (TokenSource)
						// SeQueryInformationToken은 메모리를 할당하므로 반드시 해제해야 함
						ntstatus = SeQueryInformationToken(accessToken, TokenSource, (PVOID)&tokenSource);

						if (NT_SUCCESS(ntstatus) && tokenSource != NULL)
						{
							// [중요] SourceName은 Null로 끝나지 않으므로 %.8s 사용


							// 4. 소스 이름 비교 (NtLmSsp, KSecDD)
							// TOKEN_SOURCE_LENGTH는 8입니다.
							if ((RtlCompareMemory(tokenSource->SourceName, "NtLmSsp ", TOKEN_SOURCE_LENGTH) == TOKEN_SOURCE_LENGTH) ||
								(RtlCompareMemory(tokenSource->SourceName, "KSecDD  ", TOKEN_SOURCE_LENGTH) == TOKEN_SOURCE_LENGTH))
							{
								KLogEx(DEBUG_TRACE_ERROR, "DEVICE_NETWORKDRIVEIN Check: pid(%x), Token Source=[%.8s], create_char=[%x], type=[%x]\n",
									PsGetCurrentProcessId(), tokenSource->SourceName, create_character, device_type);
								KLogEx(DEBUG_TRACE_ERROR, "DEVICE_NETWORKDRIVEIN Detected! Blocking.\n");

								// 메모리 해제 후 리턴
								ExFreePool(tokenSource);
								return BDC_NETWORKDRIVEIN;
							}

							// 매칭되지 않은 경우 메모리 해제
							ExFreePool(tokenSource);
						}
					}
				}
			}
		}
		
		state = BDC_LOCAL_DISK;
		if (GetPolicyState(BDC_EXTERNALHDD) != ENABLE || IsPolicyLog(BDC_EXTERNALHDD) == TRUE)
		{
			if (device_type == FILE_DEVICE_DISK)
			{
				///device_type 이 7일때만 아래의 함수를 실행한다..
				ntstatus = FltGetVolumeContext(fltobject->Filter, fltobject->Volume, &vctx);
				if (!NT_SUCCESS(ntstatus))
					return BDC_UNKNOWN_DEV;

				if (vctx->bustype == 7) ///BusTypeUsb
				{
					/// 외장 하드 디스크의 경우
					state = BDC_EXTERNALHDD;
				}

				FltReleaseContext(vctx);
			}
		}
	}

	return state;
}