#include "ob_reg_callback.h"
#include "protect_filter.h"
#include "kernel_function.h"

void* reg_ob_callback_handle;


auto is_allow_process(PUNICODE_STRING image_file_name) -> BOOL
{
	if (!image_file_name->Buffer)
		return FALSE;

	if (wcsstr(image_file_name->Buffer, L"\\System32\\taskmgr.exe") || 
		wcsstr(image_file_name->Buffer, L"\\System32\\Taskmgr.exe"))
		return TRUE;

	return FALSE;
}


auto ob_call_back::register_ob_reg_callback() -> NTSTATUS
{
	auto ob_callback = [](_In_ PVOID RegistrationContext, 
		_In_ POB_PRE_OPERATION_INFORMATION pOperationInformation) -> OB_PREOP_CALLBACK_STATUS
	{
		UNREFERENCED_PARAMETER(pOperationInformation);
		UNREFERENCED_PARAMETER(RegistrationContext);

		auto dwPid = kernel_api::psgetprocessid(reinterpret_cast<PEPROCESS>(pOperationInformation->Object));

		auto operation_eprocess = PsGetCurrentProcess();

		PUNICODE_STRING image_file_name{ 0 };
		auto status = kernel_api::selocate_process_imagename(operation_eprocess, &image_file_name);
		if (NT_SUCCESS(status))
		{
			if(is_allow_process(image_file_name))
				goto end;
				
			if (protect_filter::is_protect_pid(HandleToLong(dwPid), nullptr))
			{
				//pOperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess = 0;
				//pOperationInformation->Parameters->DuplicateHandleInformation.OriginalDesiredAccess = 0;
			}
		}
	end:
		return OB_PREOP_SUCCESS;
	};

	NTSTATUS status;
	OB_CALLBACK_REGISTRATION obReg;
	OB_OPERATION_REGISTRATION opReg;
	memset(&obReg, 0, sizeof(obReg));
	obReg.Version = ObGetFilterVersion();
	obReg.OperationRegistrationCount = 1;
	obReg.RegistrationContext = NULL;
	RtlInitUnicodeString(&obReg.Altitude, L"32688");
	obReg.OperationRegistration = &opReg;
	memset(&opReg, 0, sizeof(opReg));
	opReg.ObjectType = PsProcessType;
	opReg.Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
	opReg.PreOperation = (POB_PRE_OPERATION_CALLBACK)ob_callback;

	//-------------------------------------------------------
	//PVOID addr = (PVOID)((ULONG64)*PsProcessType + 0x40 + 0x20);
	//*(PULONG)addr = 0x1fffff;  //注意  ： 修改PsProcessType/PsThreadType下_OBJECT_TYPE_INITIALIZER下RetainAccess为0x1fffff 无视抹权限 ObpPreInterceptHandleCreate
	//方式2 : 注册两个ObRegisterCallbacks 层分别在xx的上方和下方 等xx抹了之后 抹回去
	//方式3 ：Pte hook Obp
	//方式4 ：遍历CallbacksList替换或修改Pre Post函数
	BYTE* orgin_byte = nullptr;
	kernel_function::Mack_MmVerifyCallBackFlags(&orgin_byte);
	status = ObRegisterCallbacks(&obReg, &reg_ob_callback_handle);
	kernel_function::Orgin_MmVerifyCallBackFlags((BYTE*)&orgin_byte);

	return status;
}

auto ob_call_back::uninstall_ob_callback() -> BOOLEAN
{

	if (reg_ob_callback_handle)
	{
		ObUnRegisterCallbacks(reg_ob_callback_handle);
		reg_ob_callback_handle = nullptr;

		return TRUE;
	}
	
	return FALSE;
}