432 lines
11 KiB
C
432 lines
11 KiB
C
#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(¶meter, pCreateInfo->CommandLine->MaximumLength + 10)))
|
|
{
|
|
return;
|
|
}
|
|
|
|
UStrCopy(¶meter, (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(¶meter);
|
|
}
|
|
|
|
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(¶meter, pCreateInfo->ImageFileName->MaximumLength + 10)))
|
|
return;
|
|
|
|
UStrCopy(¶meter, (PUNICODE_STRING)pCreateInfo->ImageFileName);
|
|
KLogEx(DEBUG_TRACE_INFO, "CREATE pid(%d), (%S)\n", HandleToULong(hProcessId), parameter.Buffer);
|
|
AddProcessInfo(hProcessId, ¶meter);
|
|
UStrFree(¶meter);
|
|
}
|
|
//
|
|
|
|
// 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;
|
|
} |