#include "precomp.h"

typedef struct _TD_CALLBACK_PARAMETERS {
	ACCESS_MASK AccessBitsToClear;
	ACCESS_MASK AccessBitsToSet;
}TD_CALLBACK_PARAMETERS, * PTD_CALLBACK_PARAMETERS;

typedef struct _TD_CALLBACK_REGISTRATION {

	PVOID RegistrationHandle;
	PVOID TargetProcess;
	HANDLE TargetProcessId;
	TD_CALLBACK_PARAMETERS ProcessParams;
	TD_CALLBACK_PARAMETERS ThreadParams;
	ULONG RegistrationId;        // Index in the global TdCallbacks array.

}TD_CALLBACK_REGISTRATION, * PTD_CALLBACK_REGISTRATION;

typedef struct _TD_CALL_CONTEXT
{
	PTD_CALLBACK_REGISTRATION CallbackRegistration;

	OB_OPERATION Operation;
	PVOID Object;
	POBJECT_TYPE ObjectType;
}
TD_CALL_CONTEXT, * PTD_CALL_CONTEXT;

static OB_CALLBACK_REGISTRATION  s_obRegistration = { 0 };
static OB_OPERATION_REGISTRATION s_operationRegistrations[2] = { { 0 }, { 0 } };
static PVOID s_registrationHandle = NULL;
static TD_CALLBACK_REGISTRATION CallbackRegistration = { 0 };
static BOOLEAN s_callbacksInstalled = FALSE;

#define TD_CALL_CONTEXT_TAG           '1bCO'
#define CB_PROCESS_TERMINATE 0x0001
#define CB_THREAD_TERMINATE  0x0001


#define PROCESS_TERMINATE                  (0x0001)  
#define PROCESS_CREATE_THREAD              (0x0002)  
#define PROCESS_SET_SESSIONID              (0x0004)  
#define PROCESS_VM_OPERATION               (0x0008)  
#define PROCESS_VM_READ                    (0x0010)  
#define PROCESS_VM_WRITE                   (0x0020)  
#define PROCESS_DUP_HANDLE                 (0x0040)  
#define PROCESS_CREATE_PROCESS             (0x0080)  
#define PROCESS_SET_QUOTA                  (0x0100)  
#define PROCESS_SET_INFORMATION            (0x0200)  
#define PROCESS_QUERY_INFORMATION          (0x0400)  
#define PROCESS_SUSPEND_RESUME             (0x0800)  
#define PROCESS_QUERY_LIMITED_INFORMATION  (0x1000)  
#define PROCESS_SET_LIMITED_INFORMATION    (0x2000)  

#define THREAD_TERMINATE                 (0x0001)  
#define THREAD_SUSPEND_RESUME            (0x0002)  
#define THREAD_GET_CONTEXT               (0x0008)  
#define THREAD_SET_CONTEXT               (0x0010)  
#define THREAD_QUERY_INFORMATION         (0x0040)  
#define THREAD_SET_INFORMATION           (0x0020)  
#define THREAD_SET_THREAD_TOKEN          (0x0080)
#define THREAD_IMPERSONATE               (0x0100)
#define THREAD_DIRECT_IMPERSONATION      (0x0200)
#define THREAD_SET_LIMITED_INFORMATION   (0x0400)  // winnt
#define THREAD_QUERY_LIMITED_INFORMATION (0x0800)  // winnt
#define THREAD_RESUME                    (0x1000)  // winnt


#define TD_ASSERT(_exp) \
    ((!(_exp)) ? \
	(__annotation(L"Debug", L"AssertFail", L#_exp), \
	DbgRaiseAssertionFailure(), FALSE) : \
TRUE)



BOOLEAN IsObExceptProcess(LPWSTR name)
{
	static wchar_t* s_obExceptionprocess[] = {
		L"MsMpEng.exe",
		L"System",
		L"csrss.exe",
		L"lsm.exe",
		L"spoolsv.exe",
		L"lsass.exe",
		L"services.exe",
		L"winlogon.exe",
		L"smss.exe",
		L"alg.exe",
		L"conhost.exe",
		L"svchost.exe",
		NULL
	};

	LPWSTR* ep = s_obExceptionprocess;
	while (*ep)
	{
		if (_wcsnicmp(name, *ep, wcslen(name)) == 0)
			return TRUE;

		++ep;
	}

	return FALSE;
}




OB_PREOP_CALLBACK_STATUS
PreOperationCallback(
	IN PVOID RegistrationContext,
	IN OUT POB_PRE_OPERATION_INFORMATION PreInfo
)
{
	UNREFERENCED_PARAMETER(RegistrationContext);

	OB_PREOP_CALLBACK_STATUS state = OB_PREOP_SUCCESS;
	PACCESS_MASK DesiredAccess = NULL;
	ACCESS_MASK OriginalDesiredAccess = 0;

	HANDLE pid = 0;
	HANDLE curPid = PsGetCurrentProcessId();
	wchar_t curProcessName[260] = { 0, };
	wchar_t curProcessPath[260] = { 0, };
	wchar_t processName[260] = { 0, };
	wchar_t processPath[260] = { 0, };
	HANDLE ProcessIdOfTargetThread;
	
	LPCWSTR ObjectTypeName = NULL;
	LPCWSTR OperationName = NULL;
	ACCESS_MASK AccessBitsToClear = 0;
	ACCESS_MASK InitialDesiredAccess = 0;
	BOOLEAN isProtected = FALSE;
	//if (PreInfo->Operation != OB_OPERATION_HANDLE_CREATE)
	//	return state;

	if (!g_bs1Flt.IsProcessProtect)
		return state;

	if (KeGetCurrentIrql() > PASSIVE_LEVEL)
		return state;
	
	if (PreInfo->ObjectType == *PsProcessType) 
	{
		if (PreInfo->Object == PsGetCurrentProcess()) 
			return OB_PREOP_SUCCESS;
		
		ObjectTypeName = L"PsProcessType";
		AccessBitsToClear = PROCESS_TERMINATE | PROCESS_SUSPEND_RESUME;
		pid = PsGetProcessId(PreInfo->Object);
	}
	else if (PreInfo->ObjectType == *PsThreadType) 
	{
		ProcessIdOfTargetThread = PsGetThreadProcessId((PETHREAD)PreInfo->Object);
		if (ProcessIdOfTargetThread == curPid)
			return OB_PREOP_SUCCESS;

		ObjectTypeName = L"PsThreadType";
		AccessBitsToClear = THREAD_TERMINATE | THREAD_SUSPEND_RESUME;
		pid = ProcessIdOfTargetThread;
	}
	else
	{
		KLogEx(DEBUG_TRACE_ERROR, "unexpected object type\n");
		return OB_PREOP_SUCCESS;

	}

	switch (PreInfo->Operation) {
	case OB_OPERATION_HANDLE_CREATE:
		DesiredAccess = &PreInfo->Parameters->CreateHandleInformation.DesiredAccess;
		OriginalDesiredAccess = PreInfo->Parameters->CreateHandleInformation.OriginalDesiredAccess;

		if (!FlagOn(*DesiredAccess, (PROCESS_TERMINATE | PROCESS_SUSPEND_RESUME)))
		{
			return OB_PREOP_SUCCESS;
		}

		OperationName = L"OB_OPERATION_HANDLE_CREATE";
		break;

	case OB_OPERATION_HANDLE_DUPLICATE:
		DesiredAccess = &PreInfo->Parameters->DuplicateHandleInformation.DesiredAccess;
		OriginalDesiredAccess = PreInfo->Parameters->DuplicateHandleInformation.OriginalDesiredAccess;

		if (!FlagOn(*DesiredAccess, (THREAD_TERMINATE | THREAD_SUSPEND_RESUME)))
		{
			return OB_PREOP_SUCCESS;
		}

		OperationName = L"OB_OPERATION_HANDLE_DUPLICATE";
		break;

	default:
		KLogEx(DEBUG_TRACE_ERROR, "unexpected object type\n");
		return state;
	}

	if (PreInfo->KernelHandle == 1) {
		return OB_PREOP_SUCCESS;
	}
	
	if (curPid == pid)
		return OB_PREOP_SUCCESS;

	InitialDesiredAccess = *DesiredAccess;

	state = UGetProcessFullPathFromPid(curPid, curProcessPath, sizeof(curProcessPath), curProcessName, sizeof(curProcessName));
	if (!NT_SUCCESS(state))
	{
		KLogEx(DEBUG_TRACE_ERROR, "UGetProcessFullPathFromPid:: cur stats %x\n", state);
		return OB_PREOP_SUCCESS;
	}

	if(IsObExceptProcess(curProcessName))
		return OB_PREOP_SUCCESS;

	state = UGetProcessFullPathFromPid(pid, processPath, sizeof(processPath), processName, sizeof(processName));
	if (!NT_SUCCESS(state))
	{
		KLogEx(DEBUG_TRACE_ERROR, "UGetProcessFullPathFromPid:: tag stats %x\n", state);
		return OB_PREOP_SUCCESS;
	}

	KLogEx(DEBUG_TRACE_ERROR, "curProcessName(%S), processName(%S)\n", curProcessName, processName);

	isProtected = TRUE;
	if (_wcsicmp(processName, L"bs1fltputool.exe") == 0)
	{	
		//PreInfo->Parameters->CreateHandleInformation.DesiredAccess = 0;		
	}
	else if (PgGetProtectPidState(HandleToULong(pid)) == PG_PID_PROTECT)
	{
		//KLogEx(DEBUG_TRACE_INFO, "PgGetProtectPidState(%S)\n", processName);
	}
	else if (IsProcessProtectList(PG_PID_PROTECT, processName, (ULONG)wcslen(processName) * (ULONG)sizeof(WCHAR)) == PG_PID_PROTECT)
	{
		//KLogEx(DEBUG_TRACE_INFO, "IsProcessProtectList(%S)\n", pocessName);
	}
	else
	{
		isProtected = FALSE;
	}

	if(isProtected)
	{
		*DesiredAccess &= ~AccessBitsToClear;
		SetLog(NULL,NULL,LOG_PROCESS_PROTECT,HandleToULong(pid),HandleToULong(curPid),InitialDesiredAccess,curProcessPath,processPath);
		KLogEx(DEBUG_TRACE_ERROR, "[PROTECT] tried to kill Target\n");
	}

	KLogEx(DEBUG_TRACE_ERROR,
		"    Client Id:    %p:%p:%p\n"
		"    Object:       %p\n"
		"    Type:         %ls\n"
		"    Operation:    %ls (KernelHandle=%d)\n"
		"    OriginalDesiredAccess: 0x%x\n"
		"    DesiredAccess (in):    0x%x\n"
		"    DesiredAccess (out):   0x%x\n",
		pid,
		curPid,
		PsGetCurrentThreadId(),
		PreInfo->Object,
		ObjectTypeName,
		OperationName,
		PreInfo->KernelHandle,
		OriginalDesiredAccess,
		InitialDesiredAccess,
		*DesiredAccess
	);

/*
	if (PreInfo->Operation == OB_OPERATION_HANDLE_CREATE) 
	{
		DesiredAccess = &PreInfo->Parameters->CreateHandleInformation.DesiredAccess;
		OriginalDesiredAccess = PreInfo->Parameters->CreateHandleInformation.OriginalDesiredAccess;
	
		if (0x0001 | 0x800 & PreInfo->Parameters->CreateHandleInformation.DesiredAccess) // TERMINATE_PROCESS
		{
			if ((PreInfo->ObjectType == *PsProcessType) || (PreInfo->ObjectType == *PsThreadType))
			{
				if (PreInfo->ObjectType == *PsProcessType)
				{
					pid = PsGetProcessId(PreInfo->Object);
				}
				if (PreInfo->ObjectType == *PsThreadType)
				{
					pid = PsGetThreadProcessId((PETHREAD)PreInfo->Object);
				}

				if (curPid == pid)
					return state;

				UGetProcessNameFormPid(curProcessName, curPid);
				

				if (_stricmp(curProcessName, "MsMpEng.exe") == 0
					|| _stricmp(curProcessName, "System") == 0
					|| _stricmp(curProcessName, "csrss.exe") == 0
					|| _stricmp(curProcessName, "lsm.exe") == 0
					|| _stricmp(curProcessName, "spoolsv.exe") == 0
					|| _stricmp(curProcessName, "lsass.exe") == 0
					|| _stricmp(curProcessName, "services.exe") == 0
					|| _stricmp(curProcessName, "winlogon.exe") == 0
					|| _stricmp(curProcessName, "smss.exe") == 0
					|| _stricmp(curProcessName, "alg.exe") == 0
					|| _stricmp(curProcessName, "conhost.exe") == 0)
				{
					//KLogEx(DEBUG_TRACE_INFO, "OB_PREOP_SUCCESS\n");
					return OB_PREOP_SUCCESS;
				}

				UGetProcessNameFormPid(processName, pid);
				KLogEx(DEBUG_TRACE_INFO, "curProcessName(%s)(%s)\n", curProcessName, processName);
				if (_stricmp(processName, "bs1flt_tool.ex") == 0)
				{
					PreInfo->Parameters->CreateHandleInformation.DesiredAccess = 0;
					KLogEx(DEBUG_TRACE_INFO, "CREATE PROTECT curProcessName(%s), processName(%s), (%x)\n", curProcessName, processName,
						(ULONG)*DesiredAccess, OriginalDesiredAccess);
				}
			}
		}
	}
	else if (PreInfo->Operation == OB_OPERATION_HANDLE_DUPLICATE)
	{
		DesiredAccess = &PreInfo->Parameters->DuplicateHandleInformation.DesiredAccess;
		OriginalDesiredAccess = PreInfo->Parameters->DuplicateHandleInformation.OriginalDesiredAccess;

		if ((PreInfo->ObjectType == *PsProcessType) || (PreInfo->ObjectType == *PsThreadType))
		{
			if (PreInfo->ObjectType == *PsProcessType)
			{
				pid = PsGetProcessId(PreInfo->Object);
			}
			if (PreInfo->ObjectType == *PsThreadType)
			{
				pid = PsGetThreadProcessId((PETHREAD)PreInfo->Object);
			}

			if (curPid == pid)
				return state;

			UGetProcessNameFormPid(curProcessName, curPid);
			UGetProcessNameFormPid(processName, pid);			
			KLogEx(DEBUG_TRACE_INFO, "DUPLICATE PROTECT curProcessName(%s), processName(%s), (%x)\n", curProcessName, processName,
				(ULONG)*DesiredAccess, OriginalDesiredAccess);
			
		}
	}
*/
	return state;
}

void TdCheckAndFreeCallContext(
	IN OUT POB_POST_OPERATION_INFORMATION PostInfo,
	IN PTD_CALLBACK_REGISTRATION _CallbackRegistration
)
{
	PTD_CALL_CONTEXT CallContext = (PTD_CALL_CONTEXT)PostInfo->CallContext;

	if (CallContext != NULL)
	{
		TD_ASSERT(CallContext->CallbackRegistration == _CallbackRegistration);

		TD_ASSERT(CallContext->Operation == PostInfo->Operation);
		TD_ASSERT(CallContext->Object == PostInfo->Object);
		TD_ASSERT(CallContext->ObjectType == PostInfo->ObjectType);

		ExFreePoolWithTag(CallContext, TD_CALL_CONTEXT_TAG);
	}
}

VOID
PostOperationCallback(
	IN PVOID RegistrationContext,
	IN POB_POST_OPERATION_INFORMATION PostInfo
)
{
	PTD_CALLBACK_REGISTRATION CallbackRegistration_ = (PTD_CALLBACK_REGISTRATION)RegistrationContext;

	TdCheckAndFreeCallContext(PostInfo, CallbackRegistration_);

	if (PostInfo->ObjectType == *PsProcessType) {
		//
		// Ignore requests for processes other than our target process.
		//

		if (CallbackRegistration_->TargetProcess != NULL &&
			CallbackRegistration_->TargetProcess != PostInfo->Object
			) {
			return;
		}

		//
		// Also ignore requests that are trying to open/duplicate the current
		// process.
		//

		if (PostInfo->Object == PsGetCurrentProcess()) {
			return;
		}
	}
	else if (PostInfo->ObjectType == *PsThreadType) {
		HANDLE ProcessIdOfTargetThread = PsGetThreadProcessId((PETHREAD)PostInfo->Object);

		//
		// Ignore requests for threads belonging to processes other than our
		// target process.
		//

		if (CallbackRegistration_->TargetProcess != NULL &&
			CallbackRegistration_->TargetProcessId != ProcessIdOfTargetThread
			) {
			return;
		}

		//
		// Also ignore requests for threads belonging to the current processes.
		//

		if (ProcessIdOfTargetThread == PsGetCurrentProcessId()) {
			return;
		}
	}
	else {
		TD_ASSERT(FALSE);
	}
}

NTSTATUS InstallProcessProtect() {

	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING altitude = { 0 };

	if(s_callbacksInstalled)
	{
		KLogEx(DEBUG_TRACE_INFO, "OB callbacks already installed\n");
		return STATUS_SUCCESS;
	}

	s_operationRegistrations[0].ObjectType = PsProcessType;
	s_operationRegistrations[0].Operations |= OB_OPERATION_HANDLE_CREATE;
	s_operationRegistrations[0].Operations |= OB_OPERATION_HANDLE_DUPLICATE;
	s_operationRegistrations[0].PreOperation = PreOperationCallback;
	s_operationRegistrations[0].PostOperation = PostOperationCallback;

	s_operationRegistrations[1].ObjectType = PsThreadType;
	s_operationRegistrations[1].Operations |= OB_OPERATION_HANDLE_CREATE;
	s_operationRegistrations[1].Operations |= OB_OPERATION_HANDLE_DUPLICATE;
	s_operationRegistrations[1].PreOperation = PreOperationCallback;
	s_operationRegistrations[1].PostOperation = PostOperationCallback;

	RtlInitUnicodeString(&altitude, OBJECT_ALTITUDE);

	s_obRegistration.Version = OB_FLT_REGISTRATION_VERSION;
	s_obRegistration.OperationRegistrationCount = 2;
	s_obRegistration.Altitude = altitude;
	s_obRegistration.RegistrationContext = &CallbackRegistration;
	s_obRegistration.OperationRegistration = s_operationRegistrations;

	status = ObRegisterCallbacks(&s_obRegistration, &s_registrationHandle);

	if (!NT_SUCCESS(status)) {
		KLogEx(DEBUG_TRACE_INFO, "ObRegisterCallbacks: installing OB callbacks failed  status (%x)\n", status);
		return status;
	}

	g_bs1Flt.IsProcessProtect = TRUE;
	s_callbacksInstalled = TRUE;
	KLogEx(DEBUG_TRACE_INFO, "Sucess\n");
	return STATUS_SUCCESS;

}

VOID UnInstallProcessProtect() 
{
	if (!s_callbacksInstalled)
		return;

	if (!s_registrationHandle)
		return;
	
	g_bs1Flt.IsProcessProtect = FALSE;
	ObUnRegisterCallbacks(s_registrationHandle);
	s_registrationHandle = NULL;
	s_callbacksInstalled = FALSE;
	KLogEx(DEBUG_TRACE_ERROR, "Sucess\n");
}