#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;

// 비교 함수 (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;
}

// 메모리 할당 함수 (new와 유사)
PVOID
NTAPI
ProcessTableAllocate(
	_In_ PRTL_AVL_TABLE Table,
	_In_ CLONG ByteSize
)
{
	UNREFERENCED_PARAMETER(Table);
	// NonPagedPool에서 메모리 할당 (Tag: 'proc')
	return ExAllocatePoolWithTag(NonPagedPool, ByteSize, 'proc');
}

// 메모리 해제 함수 (delete와 유사)
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, PCUNICODE_STRING ProcessPath)
{
	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);
	
	SetLog(NULL, NULL, LOG_PROCESS, 1, 0, 0, newItem.ProcessName, newItem.ProcessPath);
	KeReleaseSpinLock(&g_TableLock, oldIrql);
}

// 프로세스 정보 삭제 (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);
		SetLog(NULL, NULL, LOG_PROCESS, 0, 0, 0, newItem.ProcessName, newItem.ProcessPath);
	}
	
	//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))


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

	UNICODE_STRING parameter = { 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(pCreateInfo->CommandLine->Buffer, L"EXCEL.EXE\" /automation -Embedding"))
			{
				PgAddPid(HandleToULong(hProcessId), PG_PID_ALLOW);
			}

			UStrFree(&parameter);
		}

		if (!g_bs1Flt.IsProcessCreate)
			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);
		//}
		if (pCreateInfo->ImageFileName && pCreateInfo->ImageFileName->Buffer)
		{
			if (!NT_SUCCESS(UStrNew(&parameter, pCreateInfo->ImageFileName->MaximumLength + 10)))
				return;

			UStrCopy(&parameter, (PUNICODE_STRING)pCreateInfo->ImageFileName);
			KLogEx(DEBUG_TRACE_INFO, "CREATE pid(%d), (%S)\n", HandleToULong(hProcessId), parameter.Buffer);
			AddProcessInfo(hProcessId, &parameter);
			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();

	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();
}

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