#include "precomp.h"
#include <Ntstrsafe.h>
typedef NTKERNELAPI NTSTATUS(__stdcall* fpPsSetCreateProcessNotifyRoutineEx)(__in PCREATE_PROCESS_NOTIFY_ROUTINE_EX NotifyRoutine, __in BOOLEAN Remove);

typedef struct _PROCESS_INFO_ENTRY {
	RTL_BALANCED_LINKS Links;   // AVL 트리 노드 연결용 (필수)
	HANDLE ProcessId;           // Key (PID는 64비트에서 HANDLE로 관리하는게 안전)
	WCHAR ProcessName[260];     // Value 1
	WCHAR ProcessPath[512];     // Value 2
} PROCESS_INFO_ENTRY, * PPROCESS_INFO_ENTRY;

// 전역 변수
static RTL_AVL_TABLE    g_ProcessTable;
static KSPIN_LOCK       g_TableLock;
static BOOLEAN          g_IsTableInitialized = FALSE;

static NTSTATUS s_ntstatus = STATUS_UNSUCCESSFUL;
static BOOLEAN bIsVistaLater = FALSE;
static fpPsSetCreateProcessNotifyRoutineEx pPsSetCreateProcessNotifyRoutineEx = NULL;

// 차단 규칙을 관리할 내부 리스트 노드 구조체
typedef struct _PROCESS_CREATE_BLOCK_RULE_NODE {
	LIST_ENTRY ListEntry;
	BS1FLT_PROCESS_CREATE_BLOCK_RULE Rule;
} PROCESS_CREATE_BLOCK_RULE_NODE, * PPROCESS_CREATE_BLOCK_RULE_NODE;

static LIST_ENTRY g_BlockRuleListHead;
static KSPIN_LOCK g_BlockRuleLock;
static BOOLEAN g_IsBlockRuleInit = FALSE;

//// =============================================================

VOID InitBlockRuleList()
{
	if (g_IsBlockRuleInit)
		return;
	
	InitializeListHead(&g_BlockRuleListHead);
	KeInitializeSpinLock(&g_BlockRuleLock);
	g_IsBlockRuleInit = TRUE;
}

NTSTATUS AddBlockRule(WCHAR* pProcName, WCHAR* pCmdLine, WCHAR* ParentProcessName)
{
	PPROCESS_CREATE_BLOCK_RULE_NODE pNode = NULL;
	KIRQL oldIrql;

	if (!g_IsBlockRuleInit) 
		InitBlockRuleList();

	pNode = (PPROCESS_CREATE_BLOCK_RULE_NODE)ExAllocatePoolWithTag(NonPagedPool, sizeof(PROCESS_CREATE_BLOCK_RULE_NODE), 'blkR');
	if (!pNode) return STATUS_INSUFFICIENT_RESOURCES;

	RtlZeroMemory(pNode, sizeof(PROCESS_CREATE_BLOCK_RULE_NODE));

	if (pProcName) RtlStringCchCopyW(pNode->Rule.ProcessName, 260, pProcName);	
	if (pCmdLine) RtlStringCchCopyW(pNode->Rule.CommandLine, 512, pCmdLine);
	if (ParentProcessName) RtlStringCchCopyW(pNode->Rule.ParentProcessName, 50, ParentProcessName);
	
	KeAcquireSpinLock(&g_BlockRuleLock, &oldIrql);
	InsertTailList(&g_BlockRuleListHead, &pNode->ListEntry);
	KeReleaseSpinLock(&g_BlockRuleLock, oldIrql);

	return STATUS_SUCCESS;
}

VOID ClearBlockRules()
{
	KIRQL oldIrql;
	PLIST_ENTRY pEntry;
	PPROCESS_CREATE_BLOCK_RULE_NODE pNode;

	if (!g_IsBlockRuleInit) return;

	KeAcquireSpinLock(&g_BlockRuleLock, &oldIrql);
	while (!IsListEmpty(&g_BlockRuleListHead)) {
		pEntry = RemoveHeadList(&g_BlockRuleListHead);
		pNode = CONTAINING_RECORD(pEntry, PROCESS_CREATE_BLOCK_RULE_NODE, ListEntry);
		ExFreePoolWithTag(pNode, 'blkR');
	}
	KeReleaseSpinLock(&g_BlockRuleLock, oldIrql);
}


BOOLEAN IsProcessBlocked(PUNICODE_STRING pImageName, PUNICODE_STRING pCommandLine, PWSTR parentNameStr)
{
	KIRQL oldIrql;
	PLIST_ENTRY pEntry;
	PPROCESS_CREATE_BLOCK_RULE_NODE pNode;
	BOOLEAN bBlocked = FALSE;
	WCHAR* pFileName = NULL;
	UNICODE_STRING usUpcaseCmd = { 0 };
	BOOLEAN bMatch = TRUE;
	if (!g_IsBlockRuleInit || !pImageName || !pImageName->Buffer)
		return FALSE;

	pFileName = wcsrchr(pImageName->Buffer, L'\\');
	if (pFileName)
		pFileName++;
	else
		pFileName = pImageName->Buffer;

	if (pCommandLine && pCommandLine->Buffer && pCommandLine->Length > 0)
	{
		RtlUpcaseUnicodeString(&usUpcaseCmd, pCommandLine, TRUE);
	}

	KeAcquireSpinLock(&g_BlockRuleLock, &oldIrql);

	for (pEntry = g_BlockRuleListHead.Flink; pEntry != &g_BlockRuleListHead; pEntry = pEntry->Flink)
	{
		bMatch = TRUE;
		pNode = CONTAINING_RECORD(pEntry, PROCESS_CREATE_BLOCK_RULE_NODE, ListEntry);

		// ParentProcessName 검사 (값이 비어있지 않고 '*'가 아닐 때만 비교)
		if (pNode->Rule.ParentProcessName[0] != L'\0' && pNode->Rule.ParentProcessName[0] != L'*')
		{
			if (!parentNameStr || _wcsicmp(pNode->Rule.ParentProcessName, parentNameStr) != 0)
			{
				bMatch = FALSE;
			}
		}

		// ProcessName 검사 (값이 비어있지 않고 '*'가 아닐 때만 비교)
		if (bMatch && pNode->Rule.ProcessName[0] != L'\0' && pNode->Rule.ProcessName[0] != L'*')
		{
			if (!pFileName || _wcsicmp(pNode->Rule.ProcessName, pFileName) != 0)
			{
				bMatch = FALSE;
			}
		}

		// CommandLine 검사 (값이 비어있지 않을 때만 비교)
		if (bMatch && pNode->Rule.CommandLine[0] != L'\0')
		{
			if (usUpcaseCmd.Buffer == NULL)
			{
				bMatch = FALSE;
			}
			else
			{
				if (wcsstr(usUpcaseCmd.Buffer, pNode->Rule.CommandLine) == NULL)
				{
					bMatch = FALSE;
				}
			}
		}

		if (bMatch)
		{
			bBlocked = TRUE;
			break;
		}
	}

	KeReleaseSpinLock(&g_BlockRuleLock, oldIrql);
	
	if (usUpcaseCmd.Buffer)
	{
		RtlFreeUnicodeString(&usUpcaseCmd);
	}

	return bBlocked;
}

// 비교 함수 (std::less와 유사)
RTL_GENERIC_COMPARE_RESULTS
NTAPI
ProcessTableCompare(
	_In_ PRTL_AVL_TABLE Table,
	_In_ PVOID FirstStruct,
	_In_ PVOID SecondStruct
)
{
	UNREFERENCED_PARAMETER(Table);
	PPROCESS_INFO_ENTRY p1 = (PPROCESS_INFO_ENTRY)FirstStruct;
	PPROCESS_INFO_ENTRY p2 = (PPROCESS_INFO_ENTRY)SecondStruct;

	// 핸들(PID) 값 자체를 비교
	if (p1->ProcessId < p2->ProcessId) {
		return GenericLessThan;
	}
	else if (p1->ProcessId > p2->ProcessId) {
		return GenericGreaterThan;
	}

	return GenericEqual;
}

PVOID
NTAPI
ProcessTableAllocate(
	_In_ PRTL_AVL_TABLE Table,
	_In_ CLONG ByteSize
)
{
	UNREFERENCED_PARAMETER(Table);
	return ExAllocatePoolWithTag(NonPagedPool, ByteSize, 'proc');
}

VOID
NTAPI
ProcessTableFree(
	_In_ PRTL_AVL_TABLE Table,
	_In_ PVOID Buffer
)
{
	UNREFERENCED_PARAMETER(Table);
	ExFreePoolWithTag(Buffer, 'proc');
}

// =============================================================
// [3] Map 관리 함수 (Init, Add, Find, Remove)
// =============================================================

// 초기화
VOID InitProcessMap()
{
	if (!g_IsTableInitialized) {
		KeInitializeSpinLock(&g_TableLock);

		// AVL 테이블 초기화 (콜백 함수 등록)
		RtlInitializeGenericTableAvl(
			&g_ProcessTable,
			ProcessTableCompare,
			ProcessTableAllocate,
			ProcessTableFree,
			NULL
		);
		g_IsTableInitialized = TRUE;
	}
}

// 정리 (모든 노드 삭제)
VOID CleanupProcessMap()
{
	if (g_IsTableInitialized)
	{
		KIRQL oldIrql;
		PVOID pElement;

		// 락을 걸고 하나씩 지우거나, 더 효율적인 방법은 아래와 같습니다.
		// 드라이버 언로드 시에는 락 없이도 안전하다면 락 제거 가능
		KeAcquireSpinLock(&g_TableLock, &oldIrql);

		// 테이블이 빌 때까지 반복 삭제
		while (!RtlIsGenericTableEmptyAvl(&g_ProcessTable)) 
		{
			// 루트 노드(혹은 임의의 노드)를 가져와서 삭제
			pElement = RtlGetElementGenericTableAvl(&g_ProcessTable, 0);
			if (pElement) 
			{
				RtlDeleteElementGenericTableAvl(&g_ProcessTable, pElement);
			}
		}

		KeReleaseSpinLock(&g_TableLock, oldIrql);
		g_IsTableInitialized = FALSE;
	}
}

// 프로세스 정보 추가 (Insert)
VOID AddProcessInfo(HANDLE ProcessId, 
	_In_ PCUNICODE_STRING ProcessPath,
	_In_ PWCHAR processParameter,
	_In_ PWCHAR parentProcessName
)
{
	if (!g_IsTableInitialized) return;

	PROCESS_INFO_ENTRY newItem = { 0 };
	newItem.ProcessId = ProcessId; // Key 설정

	// 데이터 복사
	if (ProcessPath && ProcessPath->Buffer) {
		RtlStringCchCopyW(newItem.ProcessPath, 512, ProcessPath->Buffer);

		WCHAR* name = wcsrchr(newItem.ProcessPath, L'\\');
		if (name) name++;
		else name = newItem.ProcessPath;

		RtlStringCchCopyW(newItem.ProcessName, 260, name);
	}
	else {
		RtlStringCchCopyW(newItem.ProcessName, 260, L"Unknown");
	}

	KIRQL oldIrql;
	KeAcquireSpinLock(&g_TableLock, &oldIrql);

	// 테이블에 삽입 (이미 존재하면 업데이트하거나 무시됨)
	// RtlInsertElementGenericTableAvl은 내부적으로 AllocateRoutine을 호출하여 메모리를 복사합니다.
	// 따라서 지역변수 newItem을 넘겨줘도 안전합니다.
	RtlInsertElementGenericTableAvl(&g_ProcessTable, &newItem, sizeof(PROCESS_INFO_ENTRY), NULL);
	KeReleaseSpinLock(&g_TableLock, oldIrql);

	SetProcessLog(LOG_PROCESS_MONITOR, 1, 0, 0, newItem.ProcessName, processParameter, parentProcessName);
}

// 프로세스 정보 삭제 (Remove)
VOID RemoveProcessInfo(HANDLE ProcessId)
{
	if (!g_IsTableInitialized) return;

	PROCESS_INFO_ENTRY searchKey = { 0 };
	searchKey.ProcessId = ProcessId; // 검색할 Key만 설정

	KIRQL oldIrql;
	KeAcquireSpinLock(&g_TableLock, &oldIrql);

	// 삭제 (Key를 기반으로 찾아서 삭제함)
	RtlDeleteElementGenericTableAvl(&g_ProcessTable, &searchKey);

	KeReleaseSpinLock(&g_TableLock, oldIrql);
}

// PID로 정보 조회 (Find)
BOOLEAN GetProcessInfoByPid(HANDLE ProcessId, WCHAR* pOutName, ULONG NameLen, WCHAR* pOutPath, ULONG PathLen)
{
	if (!g_IsTableInitialized) return FALSE;

	PROCESS_INFO_ENTRY searchKey = { 0 };
	searchKey.ProcessId = ProcessId;
	BOOLEAN bFound = FALSE;

	KIRQL oldIrql;
	KeAcquireSpinLock(&g_TableLock, &oldIrql);

	// 검색
	PPROCESS_INFO_ENTRY pEntry = (PPROCESS_INFO_ENTRY)RtlLookupElementGenericTableAvl(&g_ProcessTable, &searchKey);

	if (pEntry) {
		if (pOutName) RtlStringCchCopyW(pOutName, NameLen, pEntry->ProcessName);
		if (pOutPath) RtlStringCchCopyW(pOutPath, PathLen, pEntry->ProcessPath);
		bFound = TRUE;
	}

	KeReleaseSpinLock(&g_TableLock, oldIrql);
	return bFound;
}


VOID UserNotifyEvent()
{
	{
		HANDLE rptEventHandle;
		PKEVENT rptEvent;
		UNICODE_STRING name;
		RtlInitUnicodeString(&name, PROCESS_TERMINATE_NOTIFY_KERNEL_EVENT_NAME);
		rptEvent = IoCreateNotificationEvent(&name, &rptEventHandle);
		if (rptEvent)
		{
			KeSetEvent(rptEvent, FALSE, FALSE);
			ZwClose(rptEventHandle);
		}
	}
}

VOID TerminateProcessNotify(HANDLE ulProcessId)
{
	//BOOLEAN Is = FALSE;
	ULONG state = 0;
	PROCESS_INFO_ENTRY newItem = { 0 };
	/// 프로세스 아이디 리스트 제거
	state = PgRemovePid(HandleToULong(ulProcessId));

	if (GetProcessInfoByPid(ulProcessId, newItem.ProcessName, ARRAYSIZE(newItem.ProcessName), newItem.ProcessPath, ARRAYSIZE(newItem.ProcessPath)))
	{
		RemoveProcessInfo(ulProcessId);
		KLogEx(DEBUG_TRACE_INFO, "TERMINATE pid(%d), (%S)\n", HandleToULong(ulProcessId), newItem.ProcessName);
		SetProcessLog(LOG_PROCESS_MONITOR, 0, 0, 0, newItem.ProcessName, newItem.ProcessPath, NULL);
	}
	
	//if (state & (PG_PID_WHITE | PG_PID_GREEN))
	//{
	//	/// 종료 프로세스 아이디 등록
	//	SetExitPid(ulProcessId);

	//	/// 유저 공유 이벤트 시그널
	//	UserNotifyEvent();
	//}
}

VOID IsWhiteProcess(ULONG hProcessId)
{
	UNICODE_STRING ProcessPath = { 0, };
	ULONG ulType = 0;

	/// 프로세스 이름을 얻는다.
	if (!NT_SUCCESS(UGetCurrentStackProcessImageName(hProcessId, &ProcessPath)))
	{
		return;
	}

	/// 정의된 허용 프로세스 인지 확인
	KLogEx(DEBUG_TRACE_INFO, "ProcessCreate, Pid(%d), %S(%d)\n", hProcessId, ProcessPath.Buffer, ProcessPath.Length);
	if (!IsAllowProcessName(ProcessPath.Buffer, ProcessPath.Length, &ulType))
	{
		/// 정의 안된 프로세스 정책에 따른 차단
		PgAddPid(hProcessId, PG_PID_BLACK);
	}
	else
	{
		/// 허용 프로세스 등록
		PgAddPid(hProcessId, ulType);
	}

	UStrFree(&ProcessPath);
	return;
}

typedef struct _PROCESS_MESSAGE {
	ULONG ProcessId;
	WCHAR ProcessName[260];
	WCHAR ProcessPath[512];
} PROCESS_MESSAGE, * PPROCESS_MESSAGE;


#define SCANNER_MESSAGE_SIZE   (sizeof(FILTER_MESSAGE_HEADER) + sizeof(PROCESS_MESSAGE))

typedef struct _KILL_CONTEXT {
	WORK_QUEUE_ITEM WorkItem;
	HANDLE ProcessId;
} KILL_CONTEXT, * PKILL_CONTEXT;

// [작업자 스레드] 실제 프로세스를 종료하는 함수
VOID KillProcessWorker(PVOID Context)
{
	PKILL_CONTEXT pContext = (PKILL_CONTEXT)Context;
	NTSTATUS status;
	HANDLE hProcess = NULL;
	OBJECT_ATTRIBUTES oa;
	CLIENT_ID clientId;

	if (!pContext) return;

	InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);

	clientId.UniqueProcess = pContext->ProcessId;
	clientId.UniqueThread = NULL;

	// 1. 프로세스 핸들 열기
	status = ZwOpenProcess(&hProcess, 0x1, &oa, &clientId);

	if (NT_SUCCESS(status)) {
		// 2. 프로세스 강제 종료 (Exit Code: 0으로 하여 오류 메시지 최소화)
		ZwTerminateProcess(hProcess, 0);
		ZwClose(hProcess);

		KLogEx(DEBUG_TRACE_INFO, "[Blocking] Silent Kill PID(%d) Success\n", HandleToULong(pContext->ProcessId));
	}
	else {
		KLogEx(DEBUG_TRACE_ERROR, "[Blocking] Silent Kill Failed PID(%d) Status(0x%x)\n", HandleToULong(pContext->ProcessId), status);
	}

	// 메모리 해제
	ExFreePoolWithTag(pContext, 'kill');
}

// [큐 삽입] 종료 작업을 시스템 큐에 등록하는 함수
VOID ScheduleSilentKill(HANDLE ProcessId)
{
	PKILL_CONTEXT pContext = (PKILL_CONTEXT)ExAllocatePoolWithTag(NonPagedPool, sizeof(KILL_CONTEXT), 'kill');

	if (pContext) {
		pContext->ProcessId = ProcessId;

		// WorkItem 초기화 및 대기열 등록 (DelayedWorkQueue 사용)
		ExInitializeWorkItem(&pContext->WorkItem, KillProcessWorker, pContext);
		ExQueueWorkItem(&pContext->WorkItem, DelayedWorkQueue);
	}
}


VOID CreateProcessNotifyEx(PEPROCESS pEprocess, HANDLE hProcessId, PPS_CREATE_NOTIFY_INFO pCreateInfo)
{

	UNICODE_STRING parameter = { 0, };
	UNICODE_STRING imageFileName = { 0, };
	//NTSTATUS status = 0;

	UNREFERENCED_PARAMETER(pEprocess);

	if (!g_bs1Flt.IsAttached)
		return;

	// 프로세스 종료
	if (!pCreateInfo)
	{
		//KLogEx("	Terminate	%d\n",  HandleToULong(hProcessId));
		TerminateProcessNotify(hProcessId);
	}
	else
	{


		if (pCreateInfo->CommandLine && pCreateInfo->CommandLine->Buffer)
		{
			if (!NT_SUCCESS(UStrNew(&parameter, pCreateInfo->CommandLine->MaximumLength + 10)))
			{
				return;
			}

			UStrCopy(&parameter, (PUNICODE_STRING)pCreateInfo->CommandLine);
			//EXCEL.EXE" /automation -Embedding
			if (wcsstr(parameter.Buffer, L"EXCEL.EXE\" /automation -Embedding"))
			{
				PgAddPid(HandleToULong(hProcessId), PG_PID_ALLOW);
			}

			
		}

		if (!g_bs1Flt.IsProcessCreate)
		{
			UStrFree(&parameter);
			return;
		}
		
		//PROCESS_MESSAGE msg = { 0, };
		//msg.ProcessId = (ULONG)(ULONG_PTR)hProcessId;

		//if (pCreateInfo->ImageFileName && pCreateInfo->ImageFileName->Buffer)
		//{
		//	RtlStringCchCopyW(msg.ProcessPath, 512, pCreateInfo->ImageFileName->Buffer);

		//	wchar_t* name = wcsrchr(msg.ProcessPath, L'\\');
		//	if (name) 
		//		name++;
		//	else 
		//		name = msg.ProcessPath;

		//	RtlStringCchCopyW(msg.ProcessName, 260, name);
		//}
		//else
		//{
		//	RtlStringCchCopyW(msg.ProcessName, 260, L"Unknown");
		//}

		//KLogEx(DEBUG_TRACE_INFO, "Create pid(%d), (%S)\n", msg.ProcessId, msg.ProcessName);
		//
		//if (g_bs1Flt.ClientPort)
		//{
		//	LARGE_INTEGER timeout;
		//	timeout.QuadPart = (LONGLONG)-10 * 1000 * 1000;  //3 seconds
		//	
		//	status = FltSendMessage(g_bs1Flt.Filter, &g_bs1Flt.ClientPort, &msg, sizeof(msg), NULL, NULL, &timeout);
		//	KLogEx(DEBUG_TRACE_ERROR, "FltSendMessage, status(%x)\n", status);
		//}
		HANDLE parentPid = pCreateInfo->ParentProcessId;		
		WCHAR parentNameStrW[100] = { 0, };		
		wchar_t curProcessPath[260] = { 0, };
		//UGetProcessNameFromPid

		if (!NT_SUCCESS(UGetProcessFullPathFromPid(parentPid, curProcessPath, sizeof(curProcessPath), parentNameStrW, sizeof(parentNameStrW))))
		{
			RtlStringCbPrintfW(parentNameStrW, sizeof(parentNameStrW), L"Unknown Parent(%p)", parentPid);
		}
		
		if (pCreateInfo->ImageFileName && pCreateInfo->ImageFileName->Buffer)
		{
			if (!NT_SUCCESS(UStrNew(&imageFileName, pCreateInfo->ImageFileName->MaximumLength + 10)))
				return;

			UStrCopy(&imageFileName, (PUNICODE_STRING)pCreateInfo->ImageFileName);
			KLogEx(DEBUG_TRACE_INFO, "CREATE pid(%d), (%S)\n", HandleToULong(hProcessId), imageFileName.Buffer);
			
		}

		if (IsProcessBlocked(&imageFileName, &parameter, parentNameStrW))
		{

			KLogEx(DEBUG_TRACE_INFO, "[BLOCK] Blocked by Rule: PID(%d) Image(%wZ) Cmd(%wZ) PraentName(%S)\n",
				HandleToULong(hProcessId),
				pCreateInfo->ImageFileName,
				pCreateInfo->CommandLine,
				parentNameStrW);

			// 프로세스 생성 차단 (접근 거부 리턴)
			pCreateInfo->CreationStatus = STATUS_UNSUCCESSFUL;//STATUS_ACCESS_DENIED;
			//ScheduleSilentKill(hProcessId);
			//pCreateInfo->CreationStatus = STATUS_SUCCESS;
			SetProcessLog(LOG_PROCESS_BLOCK, 2, 0, 0, imageFileName.Buffer, parameter.Buffer, parentNameStrW);
			
		}
		else
		{
			AddProcessInfo(hProcessId, &imageFileName, parameter.Buffer, parentNameStrW);
		}

		UStrFree(&imageFileName);
		UStrFree(&parameter);
		//

// 		IsWhiteProcess(hProcessId);
	}


}

VOID CreateProcessNotify(IN HANDLE  ParentId, IN HANDLE  hProcessId, IN BOOLEAN  Create)
{
	UNREFERENCED_PARAMETER(ParentId);

	if (!g_bs1Flt.IsAttached)
	{
		return;
	}

	if (Create == FALSE)
	{
		//KLogEx("	Terminate	%d\n",  HandleToULong(hProcessId));
		TerminateProcessNotify(hProcessId);
	}
	else
	{
		// 생성된 프로세스 체크
		//KLogEx("	Create	%d\n",  HandleToULong(hProcessId));
//		IsWhiteProcess(hProcessId);
	}
}


NTSTATUS InitProcessNotify()
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING funcname = { 0, };

	InitProcessMap();
	InitBlockRuleList();

	RtlInitUnicodeString(&funcname, L"PsSetCreateProcessNotifyRoutineEx");
	pPsSetCreateProcessNotifyRoutineEx = (fpPsSetCreateProcessNotifyRoutineEx)MmGetSystemRoutineAddress(&funcname);
	if (pPsSetCreateProcessNotifyRoutineEx)
	{
		bIsVistaLater = TRUE;
		status = pPsSetCreateProcessNotifyRoutineEx(CreateProcessNotifyEx, FALSE);
	}
	else
	{
		status = PsSetCreateProcessNotifyRoutine(CreateProcessNotify, FALSE);
	}

	KLogEx(DEBUG_TRACE_INFO, "status(%x), (%p)\n", status, pPsSetCreateProcessNotifyRoutineEx);
	s_ntstatus = status;
	return status;
}


VOID CleanupProcessNotify()
{
	if (bIsVistaLater)
	{
		if(NT_SUCCESS(s_ntstatus))
			pPsSetCreateProcessNotifyRoutineEx(CreateProcessNotifyEx, TRUE);
	}
	else
	{
		if (NT_SUCCESS(s_ntstatus))
			PsSetCreateProcessNotifyRoutine(CreateProcessNotify, TRUE);
	}

	CleanupProcessMap();
	ClearBlockRules();
}

DWORD GetProcessNotifyStatus()
{
	KLogEx(DEBUG_TRACE_INFO, "status = %x\n", s_ntstatus);
	return s_ntstatus;
}