diff --git a/Plugins/Filter/DriverCommon.h b/Plugins/Filter/DriverCommon.h index ec12a22..9f1be47 100644 --- a/Plugins/Filter/DriverCommon.h +++ b/Plugins/Filter/DriverCommon.h @@ -8,6 +8,7 @@ constexpr auto PYAS_PORT_NAME = L"\\PYAS_Output_Pipe"; constexpr auto MAX_PATH_LEN = 1024; constexpr auto PYAS_POOL_TAG = 'SAYP'; +constexpr auto PYAS_OB_ALTITUDE = L"320000.4101"; #ifndef PROCESS_TERMINATE #define PROCESS_TERMINATE (0x0001) @@ -60,6 +61,12 @@ constexpr auto PYAS_POOL_TAG = 'SAYP'; #ifndef THREAD_SET_THREAD_TOKEN #define THREAD_SET_THREAD_TOKEN (0x0080) #endif +#ifndef THREAD_IMPERSONATE +#define THREAD_IMPERSONATE (0x0100) +#endif +#ifndef THREAD_DIRECT_IMPERSONATION +#define THREAD_DIRECT_IMPERSONATION (0x0200) +#endif #ifndef POOL_FLAG_NON_PAGED #define POOL_FLAG_NON_PAGED 0x0000000000000040UI64 #endif @@ -94,6 +101,16 @@ typedef struct _FS_BPIO_OUTPUT { } FS_BPIO_OUTPUT, * PFS_BPIO_OUTPUT; #endif + +typedef enum _PYAS_DRIVER_STATE { + PyasDriverStateCold = 0, + PyasDriverStateStarting = 1, + PyasDriverStateRunning = 2, + PyasDriverStateStopping = 3, + PyasDriverStateStopRetry = 4, + PyasDriverStateStopped = 5 +} PYAS_DRIVER_STATE; + typedef struct _PYAS_MESSAGE { ULONG MessageCode; ULONG ProcessId; @@ -106,7 +123,8 @@ typedef enum _PYAS_COMMAND { PyasCommandLoadRuleFile = 3, PyasCommandClearRules = 4, PyasCommandAuthorizeUnload = 5, - PyasCommandRevokeUnload = 6 + PyasCommandRevokeUnload = 6, + PyasCommandQueryState = 7 } PYAS_COMMAND; typedef struct _PYAS_USER_MESSAGE { @@ -126,8 +144,12 @@ typedef struct _DRIVER_DATA { KSPIN_LOCK TrackerMutex; KSPIN_LOCK PidLock; BOOLEAN Initialized; + volatile LONG DriverState; volatile LONG UnloadAuthorized; + volatile LONG PortStopping; + volatile LONG ClientCloseActive; EX_RUNDOWN_REF PortRundown; + KEVENT CleanupCompleteEvent; PVOID ObRegistrationHandle; } DRIVER_DATA, * PDRIVER_DATA; @@ -146,40 +168,97 @@ typedef enum _RULE_CATEGORY { RuleCategoryUnknown } RULE_CATEGORY; -#define OP_WRITE 0x01 -#define OP_DELETE 0x02 -#define OP_CREATE 0x04 -#define OP_EXECUTE 0x08 -#define OP_RENAME 0x10 -#define OP_IOCTL 0x20 -#define OP_VM_READ 0x00000040 -#define OP_VM_WRITE 0x00000080 -#define OP_TERMINATE 0x00000100 -#define OP_SUSPEND_RESUME 0x00000200 -#define OP_DUP_HANDLE 0x00000400 -#define OP_SET_INFORMATION 0x00000800 +#define OP_WRITE 0x00000001 +#define OP_DELETE 0x00000002 +#define OP_CREATE 0x00000004 +#define OP_EXECUTE 0x00000008 +#define OP_RENAME 0x00000010 +#define OP_IOCTL 0x00000020 +#define OP_VM_READ 0x00000040 +#define OP_VM_WRITE 0x00000080 +#define OP_TERMINATE 0x00000100 +#define OP_SUSPEND_RESUME 0x00000200 +#define OP_DUP_HANDLE 0x00000400 +#define OP_SET_INFORMATION 0x00000800 +#define OP_VM_OPERATION 0x00001000 +#define OP_CREATE_THREAD 0x00002000 +#define OP_THREAD_SET_CONTEXT 0x00004000 +#define OP_THREAD_SET_TOKEN 0x00008000 +#define OP_CREATE_PROCESS 0x00010000 +#define OP_IMAGE_LOAD 0x00020000 +#define OP_IMPERSONATE 0x00040000 -BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode); +#define PYAS_HANDLE_CREATE 0x00000001 +#define PYAS_HANDLE_DUPLICATE 0x00000002 + +#define PYAS_OBJECT_PROCESS 0x00000001 +#define PYAS_OBJECT_THREAD 0x00000002 + +#define PYAS_MEMORY_PRIVATE 0x00000001 +#define PYAS_MEMORY_MAPPED 0x00000002 +#define PYAS_MEMORY_IMAGE 0x00000004 + +#define PYAS_PROTECT_EXECUTE 0x00000001 +#define PYAS_PROTECT_EXECUTE_WRITE 0x00000002 + +typedef enum _RULE_TRI_STATE { + RuleTriAny = 0, + RuleTriTrue = 1, + RuleTriFalse = 2 +} RULE_TRI_STATE; + +typedef enum _RULE_OPERATION_MATCH { + RuleOperationAny = 0, + RuleOperationAll = 1 +} RULE_OPERATION_MATCH; typedef struct _DYNAMIC_RULE { ULONG Code; + BOOLEAN Kill; + ULONG Priority; RULE_CATEGORY Category; ULONG Operations; + RULE_OPERATION_MATCH OperationMatch; + BOOLEAN Invalid; PRULE_NODE Initiator; PRULE_NODE InitiatorExclude; + PRULE_NODE InitiatorParent; + PRULE_NODE InitiatorParentExclude; + PRULE_NODE InitiatorProcessTree; + PRULE_NODE InitiatorProcessTreeExclude; PRULE_NODE Target; PRULE_NODE TargetExclude; PRULE_NODE TargetProcessTree; + PRULE_NODE TargetProcessTreeExclude; + PRULE_NODE Creator; + PRULE_NODE CreatorExclude; + PRULE_NODE Parent; + PRULE_NODE ParentExclude; PRULE_NODE CommandLine; + PRULE_NODE CommandLineExclude; PRULE_NODE Extensions; + ULONG HandleTypes; + ULONG ObjectTypes; + ULONG MinimumRiskScore; + ULONG MaximumRiskScore; + ULONG ThreadMemoryTypes; + ULONG ThreadMemoryProtections; + ULONGLONG MinimumRegionSize; + ULONGLONG MaximumRegionSize; + RULE_TRI_STATE ParentMismatch; + RULE_TRI_STATE FileOpenNameAvailable; + RULE_TRI_STATE SubsystemProcess; + ULONG Threshold; ULONG TimeWindow; struct _DYNAMIC_RULE* Next; } DYNAMIC_RULE, * PDYNAMIC_RULE; +BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode, PBOOLEAN OutKill); + extern DRIVER_DATA GlobalData; FLT_PREOP_CALLBACK_STATUS ProtectFile_PreCreate(PFLT_CALLBACK_DATA Data, PCFLT_RELATED_OBJECTS FltObjects, PVOID* CompletionContext); @@ -190,15 +269,19 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_SetSecurity(PFLT_CALLBACK_DATA Data, PCFLT FLT_PREOP_CALLBACK_STATUS ProtectBoot_PreDeviceControl(PFLT_CALLBACK_DATA Data, PCFLT_RELATED_OBJECTS FltObjects, PVOID* CompletionContext); NTSTATUS InitializeRegistryProtection(PDRIVER_OBJECT DriverObject); -VOID UninitializeRegistryProtection(); +NTSTATUS UninitializeRegistryProtection(); + +NTSTATUS InitializeProcessProtection(); +NTSTATUS UninitializeProcessProtection(); +BOOLEAN GetProcessRelation(HANDLE ProcessId, PHANDLE ParentProcessId, PHANDLE CreatorProcessId); NTSTATUS InitializeMemoryProtection(PDRIVER_OBJECT DriverObject); -VOID UninitializeMemoryProtection(); +NTSTATUS UninitializeMemoryProtection(BOOLEAN WaitWithoutTimeout); NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath); VOID UnloadRules(); -VOID InitializeRulesEngine(); +NTSTATUS InitializeRulesEngine(); VOID UninitializeRulesEngine(); VOID AddDynamicWhitelist(PUNICODE_STRING RuleStr); @@ -216,13 +299,45 @@ VOID SafeSetPyasPid(ULONG Pid); BOOLEAN IsProcessTrusted(HANDLE ProcessId); BOOLEAN WildcardMatch(PCWSTR Pattern, PCWSTR String, USHORT StringLengthBytes); -BOOLEAN EvaluateProcessRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, PCUNICODE_STRING CommandLine, PULONG OutCode); -BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Operation, PVOID WriteBuffer, ULONG WriteLength, PULONG OutCode); -BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG Operation, PULONG OutCode); -BOOLEAN EvaluateMemoryRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode); -BOOLEAN EvaluateProcessHandleRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode); -BOOLEAN EvaluateThreadRule(HANDLE SourcePid, HANDLE TargetPid, PVOID StartAddress, PULONG OutCode); +BOOLEAN EvaluateProcessCreateRule( + HANDLE CreatorPid, + HANDLE ParentPid, + HANDLE ProcessId, + PCUNICODE_STRING TargetPath, + PCUNICODE_STRING CommandLine, + BOOLEAN FileOpenNameAvailable, + BOOLEAN IsSubsystemProcess, + PULONG OutCode, + PBOOLEAN OutKill +); +BOOLEAN EvaluateProcessAccessRule( + HANDLE SourcePid, + HANDLE TargetPid, + ULONG RequestedOperations, + ULONG HandleType, + ULONG ObjectType, + PULONG DeniedOperations, + PULONG OutCode, + PBOOLEAN OutKill +); +BOOLEAN EvaluateThreadRule( + HANDLE SourcePid, + HANDLE TargetPid, + PVOID StartAddress, + ULONG MemoryType, + ULONG MemoryProtection, + SIZE_T RegionSize, + PULONG OutCode, + PBOOLEAN OutKill +); +BOOLEAN EvaluateImageLoadRule(HANDLE ProcessId, PCUNICODE_STRING ImagePath, PIMAGE_INFO ImageInfo, PULONG OutCode, PBOOLEAN OutKill); +BOOLEAN EvaluateProcessRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, PCUNICODE_STRING CommandLine, PULONG OutCode, PBOOLEAN OutKill); +BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Operation, PVOID WriteBuffer, ULONG WriteLength, PULONG OutCode, PBOOLEAN OutKill); +BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill); +BOOLEAN EvaluateMemoryRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill); +BOOLEAN EvaluateProcessHandleRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill); +VOID QueueRuleProcessTermination(HANDLE ProcessId, BOOLEAN Kill); NTSTATUS SendMessageToUser(ULONG Code, ULONG Pid, PWCHAR Path, USHORT PathSize); FORCEINLINE PVOID PyasAllocate(SIZE_T Size) { diff --git a/Plugins/Filter/DriverEntry.cpp b/Plugins/Filter/DriverEntry.cpp index b04c4d7..01a56e9 100644 --- a/Plugins/Filter/DriverEntry.cpp +++ b/Plugins/Filter/DriverEntry.cpp @@ -2,6 +2,7 @@ DRIVER_DATA GlobalData; static BOOLEAN g_ProcessNotifyRegistered = FALSE; +static BOOLEAN g_ProcessProtectionInitialized = FALSE; static BOOLEAN g_RegistryProtectionInitialized = FALSE; static BOOLEAN g_RulesEngineInitialized = FALSE; @@ -43,35 +44,121 @@ static VOID InstanceTeardownComplete( UNREFERENCED_PARAMETER(Flags); } -static VOID UnregisterRuntimeCallbacks() { - UninitializeMemoryProtection(); +static NTSTATUS SelectFailure(NTSTATUS CurrentStatus, NTSTATUS CandidateStatus) { + if (!NT_SUCCESS(CurrentStatus)) return CurrentStatus; + return CandidateStatus; +} + +static NTSTATUS UnregisterRuntimeCallbacksPass(BOOLEAN WaitWithoutTimeout) { + NTSTATUS Status = STATUS_SUCCESS; + + NTSTATUS MemoryStatus = UninitializeMemoryProtection(WaitWithoutTimeout); + if (!NT_SUCCESS(MemoryStatus)) { + Status = SelectFailure(Status, MemoryStatus); + } if (g_ProcessNotifyRegistered) { - PsSetCreateProcessNotifyRoutineEx(ProcessNotifyCallbackEx, TRUE); - g_ProcessNotifyRegistered = FALSE; + NTSTATUS ProcessNotifyStatus = PsSetCreateProcessNotifyRoutineEx( + ProcessNotifyCallbackEx, + TRUE + ); + + if (NT_SUCCESS(ProcessNotifyStatus) || ProcessNotifyStatus == STATUS_INVALID_PARAMETER) { + g_ProcessNotifyRegistered = FALSE; + } + else { + Status = SelectFailure(Status, ProcessNotifyStatus); + } + } + + if (!g_ProcessNotifyRegistered && g_ProcessProtectionInitialized) { + NTSTATUS ProcessStatus = UninitializeProcessProtection(); + if (NT_SUCCESS(ProcessStatus)) { + g_ProcessProtectionInitialized = FALSE; + } + else { + Status = SelectFailure(Status, ProcessStatus); + } } if (g_RegistryProtectionInitialized) { - UninitializeRegistryProtection(); - g_RegistryProtectionInitialized = FALSE; + NTSTATUS RegistryStatus = UninitializeRegistryProtection(); + if (NT_SUCCESS(RegistryStatus)) { + g_RegistryProtectionInitialized = FALSE; + } + else { + Status = SelectFailure(Status, RegistryStatus); + } } + + return Status; +} + +static NTSTATUS UnregisterRuntimeCallbacks(BOOLEAN WaitWithoutTimeout) { + for (;;) { + NTSTATUS Status = UnregisterRuntimeCallbacksPass(WaitWithoutTimeout); + if (NT_SUCCESS(Status) || !WaitWithoutTimeout) return Status; + + LARGE_INTEGER Delay; + Delay.QuadPart = -(10LL * 10 * 1000); + KeDelayExecutionThread(KernelMode, FALSE, &Delay); + } +} + +static VOID ReleaseConnectedProcessReference() { + PEPROCESS Process = NULL; + + KIRQL OldIrql; + KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql); + Process = GlobalData.PyasProcess; + GlobalData.PyasProcess = NULL; + GlobalData.ClientPort = NULL; + KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql); + + SafeSetPyasPid(0); + InterlockedExchange(&GlobalData.UnloadAuthorized, 0); + + if (Process) { + ObDereferenceObject(Process); + } +} + +static VOID CloseCurrentClientPort() { + if (InterlockedCompareExchange(&GlobalData.ClientCloseActive, 1, 0) != 0) return; + + if (GlobalData.FilterHandle && GlobalData.ClientPort) { + FltCloseClientPort(GlobalData.FilterHandle, &GlobalData.ClientPort); + } + + InterlockedExchange(&GlobalData.ClientCloseActive, 0); } static VOID ReleaseFilterResources() { + InterlockedExchange(&GlobalData.PortStopping, 1); + if (GlobalData.ServerPort) { - FltCloseCommunicationPort(GlobalData.ServerPort); + PFLT_PORT ServerPort = GlobalData.ServerPort; GlobalData.ServerPort = NULL; + FltCloseCommunicationPort(ServerPort); } + CloseCurrentClientPort(); + ExWaitForRundownProtectionRelease(&GlobalData.PortRundown); + if (GlobalData.FilterHandle) { - FltUnregisterFilter(GlobalData.FilterHandle); + PFLT_FILTER FilterHandle = GlobalData.FilterHandle; + FltUnregisterFilter(FilterHandle); GlobalData.FilterHandle = NULL; } + + ReleaseConnectedProcessReference(); } -static VOID CleanupDriverState() { +static NTSTATUS CleanupDriverState(BOOLEAN WaitWithoutTimeout) { + NTSTATUS Status = UnregisterRuntimeCallbacks(WaitWithoutTimeout); + if (!NT_SUCCESS(Status)) return Status; + GlobalData.Initialized = FALSE; - UnregisterRuntimeCallbacks(); ReleaseFilterResources(); if (g_RulesEngineInitialized) { @@ -79,17 +166,94 @@ static VOID CleanupDriverState() { UninitializeRulesEngine(); g_RulesEngineInitialized = FALSE; } + + return STATUS_SUCCESS; +} + +static NTSTATUS BeginUnload(BOOLEAN MandatoryUnload) { + for (;;) { + LONG State = InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + if (State == PyasDriverStateStopped) return STATUS_SUCCESS; + + if (State == PyasDriverStateStopping) { + LARGE_INTEGER Timeout; + PLARGE_INTEGER TimeoutPointer = NULL; + + if (!MandatoryUnload) { + Timeout.QuadPart = -(2LL * 10 * 1000 * 1000); + TimeoutPointer = &Timeout; + } + + NTSTATUS WaitStatus = KeWaitForSingleObject( + &GlobalData.CleanupCompleteEvent, + Executive, + KernelMode, + FALSE, + TimeoutPointer + ); + + if (!MandatoryUnload && WaitStatus == STATUS_TIMEOUT) { + return STATUS_DEVICE_BUSY; + } + + if (!NT_SUCCESS(WaitStatus)) return WaitStatus; + continue; + } + + if (State != PyasDriverStateRunning && State != PyasDriverStateStopRetry) { + return STATUS_DEVICE_NOT_READY; + } + + if (!MandatoryUnload && InterlockedCompareExchange(&GlobalData.UnloadAuthorized, 0, 1) != 1) { + return STATUS_ACCESS_DENIED; + } + + if (MandatoryUnload) { + InterlockedExchange(&GlobalData.UnloadAuthorized, 0); + } + + if (InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateStopping, + State + ) != State) { + if (!MandatoryUnload) return STATUS_DEVICE_BUSY; + continue; + } + + KeClearEvent(&GlobalData.CleanupCompleteEvent); + return STATUS_SUCCESS; + } } static NTSTATUS DriverUnload(FLT_FILTER_UNLOAD_FLAGS Flags) { - UNREFERENCED_PARAMETER(Flags); + BOOLEAN MandatoryUnload = FlagOn(Flags, FLTFL_FILTER_UNLOAD_MANDATORY) ? TRUE : FALSE; - if (InterlockedCompareExchange(&GlobalData.UnloadAuthorized, 0, 1) != 1) { - return STATUS_ACCESS_DENIED; - } + NTSTATUS Status = BeginUnload(MandatoryUnload); + if (!NT_SUCCESS(Status)) return Status; - CleanupDriverState(); - return STATUS_SUCCESS; + LONG State = InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + if (State == PyasDriverStateStopped) return STATUS_SUCCESS; + + Status = CleanupDriverState(MandatoryUnload); + + InterlockedExchange( + &GlobalData.DriverState, + NT_SUCCESS(Status) ? PyasDriverStateStopped : PyasDriverStateStopRetry + ); + + KeSetEvent(&GlobalData.CleanupCompleteEvent, IO_NO_INCREMENT, FALSE); + return Status; } static NTSTATUS PortMessage( @@ -100,52 +264,138 @@ static NTSTATUS PortMessage( ULONG OutputBufferLength, PULONG ReturnOutputBufferLength ) { - UNREFERENCED_PARAMETER(OutputBuffer); - UNREFERENCED_PARAMETER(OutputBufferLength); + if (ReturnOutputBufferLength) *ReturnOutputBufferLength = 0; + + if (InterlockedCompareExchange(&GlobalData.PortStopping, 0, 0) != 0) { + return STATUS_DELETE_PENDING; + } + + if (!ExAcquireRundownProtection(&GlobalData.PortRundown)) { + return STATUS_DELETE_PENDING; + } + + NTSTATUS Status = STATUS_SUCCESS; + PYAS_USER_MESSAGE Message = { 0 }; + UNICODE_STRING Path = { 0 }; + LONG DriverState = PyasDriverStateCold; + ULONG StateValue = PyasDriverStateCold; if (!PortCookie || PortCookie != GlobalData.PyasProcess) { - return STATUS_ACCESS_DENIED; + Status = STATUS_ACCESS_DENIED; + goto Exit; } if (!InputBuffer || InputBufferLength < sizeof(PYAS_USER_MESSAGE)) { - return STATUS_INVALID_PARAMETER; + Status = STATUS_INVALID_PARAMETER; + goto Exit; } - PPYAS_USER_MESSAGE Message = (PPYAS_USER_MESSAGE)InputBuffer; - Message->Path[MAX_PATH_LEN - 1] = L'\0'; + __try { + RtlCopyMemory(&Message, InputBuffer, sizeof(Message)); + } + __except (EXCEPTION_EXECUTE_HANDLER) { + Status = GetExceptionCode(); + goto Exit; + } + + Message.Path[MAX_PATH_LEN - 1] = L'\0'; + + if (Message.Command == PyasCommandQueryState) { + if (!OutputBuffer || OutputBufferLength < sizeof(ULONG) || !ReturnOutputBufferLength) { + Status = STATUS_BUFFER_TOO_SMALL; + goto Exit; + } + + StateValue = (ULONG)InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + __try { + *(PULONG)OutputBuffer = StateValue; + *ReturnOutputBufferLength = sizeof(ULONG); + } + __except (EXCEPTION_EXECUTE_HANDLER) { + Status = GetExceptionCode(); + } + goto Exit; + } + + if (Message.Command == PyasCommandAuthorizeUnload) { + DriverState = InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + if (DriverState != PyasDriverStateRunning && DriverState != PyasDriverStateStopRetry) { + Status = STATUS_DEVICE_NOT_READY; + goto Exit; + } - if (Message->Command == PyasCommandAuthorizeUnload) { InterlockedExchange(&GlobalData.UnloadAuthorized, 1); + goto Exit; } - else if (Message->Command == PyasCommandRevokeUnload) { + + if (Message.Command == PyasCommandRevokeUnload) { InterlockedExchange(&GlobalData.UnloadAuthorized, 0); - } - else if (Message->Command >= PyasCommandAddWhitelist && Message->Command <= PyasCommandClearRules) { - UNICODE_STRING Path; - RtlInitUnicodeString(&Path, Message->Path); - - if (Message->Command == PyasCommandAddWhitelist && Path.Length > 0) { - AddDynamicWhitelist(&Path); - } - else if (Message->Command == PyasCommandRemoveWhitelist && Path.Length > 0) { - RemoveDynamicWhitelist(&Path); - } - else if (Message->Command == PyasCommandLoadRuleFile && Path.Length > 0) { - LoadRuleFile(&Path); - } - else if (Message->Command == PyasCommandClearRules) { - ClearDynamicRules(); - } - } - else { - return STATUS_INVALID_DEVICE_REQUEST; + goto Exit; } - if (ReturnOutputBufferLength) { - *ReturnOutputBufferLength = 0; + if (Message.Command < PyasCommandAddWhitelist || Message.Command > PyasCommandClearRules) { + Status = STATUS_INVALID_DEVICE_REQUEST; + goto Exit; } - return STATUS_SUCCESS; + DriverState = InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + if (DriverState != PyasDriverStateRunning) { + Status = STATUS_DEVICE_NOT_READY; + goto Exit; + } + + RtlInitUnicodeString(&Path, Message.Path); + + if (Message.Command == PyasCommandAddWhitelist) { + if (Path.Length == 0) { + Status = STATUS_INVALID_PARAMETER; + goto Exit; + } + + AddDynamicWhitelist(&Path); + goto Exit; + } + + if (Message.Command == PyasCommandRemoveWhitelist) { + if (Path.Length == 0) { + Status = STATUS_INVALID_PARAMETER; + goto Exit; + } + + RemoveDynamicWhitelist(&Path); + goto Exit; + } + + if (Message.Command == PyasCommandLoadRuleFile) { + if (Path.Length == 0) { + Status = STATUS_INVALID_PARAMETER; + goto Exit; + } + + Status = LoadRuleFile(&Path); + goto Exit; + } + + ClearDynamicRules(); + +Exit: + ExReleaseRundownProtection(&GlobalData.PortRundown); + return Status; } static NTSTATUS PortConnect( @@ -159,40 +409,88 @@ static NTSTATUS PortConnect( UNREFERENCED_PARAMETER(ConnectionContext); UNREFERENCED_PARAMETER(SizeOfContext); - if (!ConnectionPortCookie) { - return STATUS_INVALID_PARAMETER; + if (!ConnectionPortCookie) return STATUS_INVALID_PARAMETER; + *ConnectionPortCookie = NULL; + + if (InterlockedCompareExchange(&GlobalData.PortStopping, 0, 0) != 0) { + return STATUS_DELETE_PENDING; + } + + LONG DriverState = InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + if (DriverState != PyasDriverStateRunning && DriverState != PyasDriverStateStopRetry) { + return STATUS_DEVICE_NOT_READY; + } + + if (!ExAcquireRundownProtection(&GlobalData.PortRundown)) { + return STATUS_DELETE_PENDING; } PEPROCESS ConnectingProcess = PsGetCurrentProcess(); ObReferenceObject(ConnectingProcess); - *ConnectionPortCookie = ConnectingProcess; - - SafeSetPyasPid((ULONG)(ULONG_PTR)PsGetCurrentProcessId()); - InterlockedExchange(&GlobalData.UnloadAuthorized, 0); + NTSTATUS Status = STATUS_SUCCESS; KIRQL OldIrql; KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql); - GlobalData.ClientPort = ClientPort; - GlobalData.PyasProcess = ConnectingProcess; - ExReInitializeRundownProtection(&GlobalData.PortRundown); + DriverState = InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + if (InterlockedCompareExchange(&GlobalData.PortStopping, 0, 0) != 0) { + Status = STATUS_DELETE_PENDING; + } + else if (DriverState != PyasDriverStateRunning && DriverState != PyasDriverStateStopRetry) { + Status = STATUS_DEVICE_NOT_READY; + } + else if (GlobalData.ClientPort || GlobalData.PyasProcess) { + Status = STATUS_DEVICE_BUSY; + } + else { + GlobalData.ClientPort = ClientPort; + GlobalData.PyasProcess = ConnectingProcess; + *ConnectionPortCookie = ConnectingProcess; + } KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql); + + if (!NT_SUCCESS(Status)) { + ObDereferenceObject(ConnectingProcess); + ExReleaseRundownProtection(&GlobalData.PortRundown); + return Status; + } + + SafeSetPyasPid((ULONG)(ULONG_PTR)PsGetCurrentProcessId()); + InterlockedExchange(&GlobalData.UnloadAuthorized, 0); + ExReleaseRundownProtection(&GlobalData.PortRundown); return STATUS_SUCCESS; } static VOID PortDisconnect(PVOID ConnectionCookie) { + CloseCurrentClientPort(); + + BOOLEAN DisconnectedCurrentClient = FALSE; + KIRQL OldIrql; KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql); - GlobalData.ClientPort = NULL; - GlobalData.PyasProcess = NULL; + if (ConnectionCookie && GlobalData.PyasProcess == (PEPROCESS)ConnectionCookie) { + GlobalData.PyasProcess = NULL; + DisconnectedCurrentClient = TRUE; + } KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql); - InterlockedExchange(&GlobalData.UnloadAuthorized, 0); - SafeSetPyasPid(0); - ExWaitForRundownProtectionRelease(&GlobalData.PortRundown); + if (DisconnectedCurrentClient) { + InterlockedExchange(&GlobalData.UnloadAuthorized, 0); + SafeSetPyasPid(0); + } if (ConnectionCookie) { ObDereferenceObject((PEPROCESS)ConnectionCookie); @@ -235,85 +533,84 @@ extern "C" NTSTATUS DriverEntry( KeInitializeSpinLock(&GlobalData.TrackerMutex); KeInitializeSpinLock(&GlobalData.PidLock); ExInitializeRundownProtection(&GlobalData.PortRundown); + KeInitializeEvent(&GlobalData.CleanupCompleteEvent, NotificationEvent, FALSE); + InterlockedExchange(&GlobalData.PortStopping, 0); + InterlockedExchange(&GlobalData.ClientCloseActive, 0); + InterlockedExchange(&GlobalData.DriverState, PyasDriverStateStarting); - InitializeRulesEngine(); + NTSTATUS Status = STATUS_SUCCESS; + PSECURITY_DESCRIPTOR SecurityDescriptor = NULL; + + Status = InitializeRulesEngine(); + if (!NT_SUCCESS(Status)) { + InterlockedExchange(&GlobalData.DriverState, PyasDriverStateStopped); + KeSetEvent(&GlobalData.CleanupCompleteEvent, IO_NO_INCREMENT, FALSE); + return Status; + } g_RulesEngineInitialized = TRUE; - NTSTATUS status = LoadRulesFromDisk(RegistryPath); - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; + Status = LoadRulesFromDisk(RegistryPath); + if (!NT_SUCCESS(Status)) goto Failure; + + Status = FltRegisterFilter(DriverObject, &FilterRegistration, &GlobalData.FilterHandle); + if (!NT_SUCCESS(Status)) goto Failure; + + Status = FltBuildDefaultSecurityDescriptor(&SecurityDescriptor, FLT_PORT_ALL_ACCESS); + if (!NT_SUCCESS(Status)) goto Failure; + + { + UNICODE_STRING PortName; + RtlInitUnicodeString(&PortName, PYAS_PORT_NAME); + + OBJECT_ATTRIBUTES ObjectAttributes; + InitializeObjectAttributes( + &ObjectAttributes, + &PortName, + OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, + NULL, + SecurityDescriptor + ); + + Status = FltCreateCommunicationPort( + GlobalData.FilterHandle, + &GlobalData.ServerPort, + &ObjectAttributes, + NULL, + PortConnect, + PortDisconnect, + PortMessage, + 1 + ); } - status = FltRegisterFilter(DriverObject, &FilterRegistration, &GlobalData.FilterHandle); - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; - } - - PSECURITY_DESCRIPTOR SecurityDescriptor = NULL; - status = FltBuildDefaultSecurityDescriptor(&SecurityDescriptor, FLT_PORT_ALL_ACCESS); - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; - } - - UNICODE_STRING PortName; - RtlInitUnicodeString(&PortName, PYAS_PORT_NAME); - - OBJECT_ATTRIBUTES ObjectAttributes; - InitializeObjectAttributes( - &ObjectAttributes, - &PortName, - OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, - NULL, - SecurityDescriptor - ); - - status = FltCreateCommunicationPort( - GlobalData.FilterHandle, - &GlobalData.ServerPort, - &ObjectAttributes, - NULL, - PortConnect, - PortDisconnect, - PortMessage, - 1 - ); - FltFreeSecurityDescriptor(SecurityDescriptor); + if (!NT_SUCCESS(Status)) goto Failure; - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; - } - - status = InitializeRegistryProtection(DriverObject); - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; - } + Status = InitializeRegistryProtection(DriverObject); + if (!NT_SUCCESS(Status)) goto Failure; g_RegistryProtectionInitialized = TRUE; - status = PsSetCreateProcessNotifyRoutineEx(ProcessNotifyCallbackEx, FALSE); - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; - } + Status = InitializeProcessProtection(); + if (!NT_SUCCESS(Status)) goto Failure; + g_ProcessProtectionInitialized = TRUE; + + Status = PsSetCreateProcessNotifyRoutineEx(ProcessNotifyCallbackEx, FALSE); + if (!NT_SUCCESS(Status)) goto Failure; g_ProcessNotifyRegistered = TRUE; - status = InitializeMemoryProtection(DriverObject); - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; - } + Status = InitializeMemoryProtection(DriverObject); + if (!NT_SUCCESS(Status)) goto Failure; - status = FltStartFiltering(GlobalData.FilterHandle); - if (!NT_SUCCESS(status)) { - CleanupDriverState(); - return status; - } + Status = FltStartFiltering(GlobalData.FilterHandle); + if (!NT_SUCCESS(Status)) goto Failure; GlobalData.Initialized = TRUE; + InterlockedExchange(&GlobalData.DriverState, PyasDriverStateRunning); return STATUS_SUCCESS; + +Failure: + CleanupDriverState(TRUE); + InterlockedExchange(&GlobalData.DriverState, PyasDriverStateStopped); + KeSetEvent(&GlobalData.CleanupCompleteEvent, IO_NO_INCREMENT, FALSE); + return Status; } diff --git a/Plugins/Filter/ProtectBoot.cpp b/Plugins/Filter/ProtectBoot.cpp index 7980061..4623fa1 100644 --- a/Plugins/Filter/ProtectBoot.cpp +++ b/Plugins/Filter/ProtectBoot.cpp @@ -23,10 +23,12 @@ FLT_PREOP_CALLBACK_STATUS ProtectBoot_PreDeviceControl(PFLT_CALLBACK_DATA Data, HANDLE Pid = PsGetCurrentProcessId(); ULONG RuleCode = 0; + BOOLEAN Kill = FALSE; - if (EvaluateDeviceRule(Pid, &RuleCode)) { + if (EvaluateDeviceRule(Pid, &RuleCode, &Kill)) { UNICODE_STRING MsgStr = RTL_CONSTANT_STRING(L"Disk_Wiper_Attempt"); SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, MsgStr.Buffer, MsgStr.Length); + QueueRuleProcessTermination(Pid, Kill); Data->IoStatus.Status = STATUS_ACCESS_DENIED; return FLT_PREOP_COMPLETE; } diff --git a/Plugins/Filter/ProtectFile.cpp b/Plugins/Filter/ProtectFile.cpp index ed27ec4..760650b 100644 --- a/Plugins/Filter/ProtectFile.cpp +++ b/Plugins/Filter/ProtectFile.cpp @@ -31,8 +31,10 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreCreate(PFLT_CALLBACK_DATA Data, PCFLT_R if (Operation != 0) { ULONG RuleCode = 0; - if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode, &Kill)) { SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length); + QueueRuleProcessTermination(Pid, Kill); Data->IoStatus.Status = STATUS_ACCESS_DENIED; callbackStatus = FLT_PREOP_COMPLETE; } @@ -73,8 +75,10 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreSetInfo(PFLT_CALLBACK_DATA Data, PCFLT_ } ULONG RuleCode = 0; - if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode, &Kill)) { SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length); + QueueRuleProcessTermination(Pid, Kill); Data->IoStatus.Status = STATUS_ACCESS_DENIED; callbackStatus = FLT_PREOP_COMPLETE; } @@ -112,8 +116,10 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreWrite(PFLT_CALLBACK_DATA Data, PCFLT_RE if (Mdl) SystemBuffer = MmGetSystemAddressForMdlSafe(Mdl, NormalPagePriority | MdlMappingNoExecute); ULONG RuleCode = 0; - if (EvaluateFileRule(Pid, &nameInfo->Name, OP_WRITE, SystemBuffer, WriteLength, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateFileRule(Pid, &nameInfo->Name, OP_WRITE, SystemBuffer, WriteLength, &RuleCode, &Kill)) { SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length); + QueueRuleProcessTermination(Pid, Kill); Data->IoStatus.Status = STATUS_ACCESS_DENIED; callbackStatus = FLT_PREOP_COMPLETE; } diff --git a/Plugins/Filter/ProtectMemory.cpp b/Plugins/Filter/ProtectMemory.cpp index 59c4cf8..9fccef7 100644 --- a/Plugins/Filter/ProtectMemory.cpp +++ b/Plugins/Filter/ProtectMemory.cpp @@ -23,8 +23,46 @@ extern "C" { #define THREAD_QUERY_SET_WIN32_START_ADDRESS 9 +constexpr ULONG PYAS_MEM_IMAGE_TYPE = 0x01000000UL; + +typedef enum _MEMORY_WORK_TYPE { + MemoryWorkThreadAnalysis = 1, + MemoryWorkViolationReport = 2 +} MEMORY_WORK_TYPE; + +typedef struct _MEMORY_WORK_ITEM { + LIST_ENTRY Entry; + MEMORY_WORK_TYPE Type; + PEPROCESS SourceProcess; + PEPROCESS TargetProcess; + HANDLE SourcePid; + HANDLE TargetPid; + HANDLE ThreadId; + ULONG RuleCode; + BOOLEAN Kill; + WCHAR FallbackMessage[64]; +} MEMORY_WORK_ITEM, * PMEMORY_WORK_ITEM; + +typedef struct _THREAD_MEMORY_REGION { + PVOID StartAddress; + ULONG MemoryType; + ULONG MemoryProtection; + SIZE_T RegionSize; +} THREAD_MEMORY_REGION, * PTHREAD_MEMORY_REGION; + static BOOLEAN g_ThreadNotifyRegistered = FALSE; static BOOLEAN g_LoadImageNotifyRegistered = FALSE; +static BOOLEAN g_MemoryProtectionInitialized = FALSE; +static BOOLEAN g_MemoryRundownCompleted = FALSE; +static EX_RUNDOWN_REF g_MemoryRundown; +static volatile LONG g_MemoryStopping = 1; +static volatile LONG g_PendingMemoryWork = 0; +static LIST_ENTRY g_MemoryWorkQueue; +static KSPIN_LOCK g_MemoryWorkLock; +static KSEMAPHORE g_MemoryWorkSemaphore; +static HANDLE g_MemoryWorkerHandle = NULL; +constexpr auto MAX_PENDING_MEMORY_WORK = 256; +constexpr auto MEMORY_WORKER_STOP_TIMEOUT_MS = 2000; static PACCESS_MASK GetDesiredAccessPointer(POB_PRE_OPERATION_INFORMATION OperationInformation) { if (OperationInformation->Operation == OB_OPERATION_HANDLE_CREATE) { @@ -38,98 +76,244 @@ static PACCESS_MASK GetDesiredAccessPointer(POB_PRE_OPERATION_INFORMATION Operat return NULL; } -static ACCESS_MASK GetProcessMemoryDeniedAccess(ULONG Operation) { +static ULONG GetHandleType(POB_PRE_OPERATION_INFORMATION OperationInformation) { + if (OperationInformation->Operation == OB_OPERATION_HANDLE_CREATE) { + return PYAS_HANDLE_CREATE; + } + + if (OperationInformation->Operation == OB_OPERATION_HANDLE_DUPLICATE) { + return PYAS_HANDLE_DUPLICATE; + } + + return 0; +} + +static ULONG GetProcessOperations(ACCESS_MASK Access) { + ULONG Operations = 0; + + if (Access & PROCESS_VM_READ) Operations |= OP_VM_READ; + if (Access & PROCESS_VM_WRITE) Operations |= OP_VM_WRITE; + if (Access & PROCESS_VM_OPERATION) Operations |= OP_VM_OPERATION; + if (Access & PROCESS_CREATE_THREAD) Operations |= OP_CREATE_THREAD; + if (Access & PROCESS_TERMINATE) Operations |= OP_TERMINATE; + if (Access & PROCESS_SUSPEND_RESUME) Operations |= OP_SUSPEND_RESUME; + if (Access & PROCESS_DUP_HANDLE) Operations |= OP_DUP_HANDLE; + if (Access & (PROCESS_SET_INFORMATION | PROCESS_SET_QUOTA)) Operations |= OP_SET_INFORMATION; + if (Access & PROCESS_CREATE_PROCESS) Operations |= OP_CREATE_PROCESS; + + return Operations; +} + +static ULONG GetThreadOperations(ACCESS_MASK Access) { + ULONG Operations = 0; + + if (Access & THREAD_SET_CONTEXT) Operations |= OP_THREAD_SET_CONTEXT; + if (Access & THREAD_TERMINATE) Operations |= OP_TERMINATE; + if (Access & THREAD_SUSPEND_RESUME) Operations |= OP_SUSPEND_RESUME; + if (Access & THREAD_SET_INFORMATION) Operations |= OP_SET_INFORMATION; + if (Access & THREAD_SET_THREAD_TOKEN) Operations |= OP_THREAD_SET_TOKEN; + if (Access & (THREAD_IMPERSONATE | THREAD_DIRECT_IMPERSONATION)) Operations |= OP_IMPERSONATE; + + return Operations; +} + +static ACCESS_MASK GetProcessDeniedAccess(ULONG Operations) { ACCESS_MASK DeniedAccess = 0; - if (Operation & OP_VM_READ) { - DeniedAccess |= PROCESS_VM_READ; - } - - if (Operation & OP_VM_WRITE) { - DeniedAccess |= PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_CREATE_THREAD; - } + if (Operations & OP_VM_READ) DeniedAccess |= PROCESS_VM_READ; + if (Operations & OP_VM_WRITE) DeniedAccess |= PROCESS_VM_WRITE; + if (Operations & OP_VM_OPERATION) DeniedAccess |= PROCESS_VM_OPERATION; + if (Operations & OP_CREATE_THREAD) DeniedAccess |= PROCESS_CREATE_THREAD; + if (Operations & OP_TERMINATE) DeniedAccess |= PROCESS_TERMINATE; + if (Operations & OP_SUSPEND_RESUME) DeniedAccess |= PROCESS_SUSPEND_RESUME; + if (Operations & OP_DUP_HANDLE) DeniedAccess |= PROCESS_DUP_HANDLE; + if (Operations & OP_SET_INFORMATION) DeniedAccess |= PROCESS_SET_INFORMATION | PROCESS_SET_QUOTA; + if (Operations & OP_CREATE_PROCESS) DeniedAccess |= PROCESS_CREATE_PROCESS; return DeniedAccess; } -static ACCESS_MASK GetProcessControlDeniedAccess(ULONG Operation) { +static ACCESS_MASK GetThreadDeniedAccess(ULONG Operations) { ACCESS_MASK DeniedAccess = 0; - if (Operation & OP_TERMINATE) { - DeniedAccess |= PROCESS_TERMINATE; - } - - if (Operation & OP_SUSPEND_RESUME) { - DeniedAccess |= PROCESS_SUSPEND_RESUME; - } - - if (Operation & OP_DUP_HANDLE) { - DeniedAccess |= PROCESS_DUP_HANDLE; - } - - if (Operation & OP_SET_INFORMATION) { - DeniedAccess |= PROCESS_SET_INFORMATION | PROCESS_SET_QUOTA; - } + if (Operations & OP_THREAD_SET_CONTEXT) DeniedAccess |= THREAD_SET_CONTEXT; + if (Operations & OP_TERMINATE) DeniedAccess |= THREAD_TERMINATE; + if (Operations & OP_SUSPEND_RESUME) DeniedAccess |= THREAD_SUSPEND_RESUME; + if (Operations & OP_SET_INFORMATION) DeniedAccess |= THREAD_SET_INFORMATION; + if (Operations & OP_THREAD_SET_TOKEN) DeniedAccess |= THREAD_SET_THREAD_TOKEN; + if (Operations & OP_IMPERSONATE) DeniedAccess |= THREAD_IMPERSONATE | THREAD_DIRECT_IMPERSONATION; return DeniedAccess; } -static ACCESS_MASK GetThreadMemoryDeniedAccess(ULONG Operation) { - ACCESS_MASK DeniedAccess = 0; +static BOOLEAN IsProcessIdentityCurrent(HANDLE ProcessId, PEPROCESS ExpectedProcess); - if (Operation & OP_VM_WRITE) { - DeniedAccess |= THREAD_SET_CONTEXT; - } +static VOID ReleaseMemoryWork(PMEMORY_WORK_ITEM Work) { + if (!Work) return; - return DeniedAccess; + if (Work->TargetProcess) ObDereferenceObject(Work->TargetProcess); + if (Work->SourceProcess) ObDereferenceObject(Work->SourceProcess); + InterlockedDecrement(&g_PendingMemoryWork); + PyasFree(Work); } -static ACCESS_MASK GetThreadControlDeniedAccess(ULONG Operation) { - ACCESS_MASK DeniedAccess = 0; +static PMEMORY_WORK_ITEM PopMemoryWork() { + PMEMORY_WORK_ITEM Work = NULL; + KIRQL OldIrql; + KeAcquireSpinLock(&g_MemoryWorkLock, &OldIrql); - if (Operation & OP_TERMINATE) { - DeniedAccess |= THREAD_TERMINATE; + if (!IsListEmpty(&g_MemoryWorkQueue)) { + PLIST_ENTRY Entry = RemoveHeadList(&g_MemoryWorkQueue); + Work = CONTAINING_RECORD(Entry, MEMORY_WORK_ITEM, Entry); } - if (Operation & OP_SUSPEND_RESUME) { - DeniedAccess |= THREAD_SUSPEND_RESUME; - } - - if (Operation & OP_SET_INFORMATION) { - DeniedAccess |= THREAD_SET_INFORMATION | THREAD_SET_THREAD_TOKEN; - } - - return DeniedAccess; + KeReleaseSpinLock(&g_MemoryWorkLock, OldIrql); + return Work; } -static VOID ReportViolation(ULONG RuleCode, HANDLE SourcePid, HANDLE TargetPid, PCWSTR FallbackMessage) { +static VOID CancelQueuedMemoryWork() { + for (;;) { + PMEMORY_WORK_ITEM Work = PopMemoryWork(); + if (!Work) return; + ReleaseMemoryWork(Work); + } +} + +static BOOLEAN QueueMemoryWork(PMEMORY_WORK_ITEM Work) { + if (!Work) return FALSE; + + LONG Pending = InterlockedIncrement(&g_PendingMemoryWork); + if (Pending > MAX_PENDING_MEMORY_WORK) { + InterlockedDecrement(&g_PendingMemoryWork); + return FALSE; + } + + KIRQL OldIrql; + KeAcquireSpinLock(&g_MemoryWorkLock, &OldIrql); + + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) { + KeReleaseSpinLock(&g_MemoryWorkLock, OldIrql); + InterlockedDecrement(&g_PendingMemoryWork); + return FALSE; + } + + InsertTailList(&g_MemoryWorkQueue, &Work->Entry); + KeReleaseSpinLock(&g_MemoryWorkLock, OldIrql); + KeReleaseSemaphore(&g_MemoryWorkSemaphore, IO_NO_INCREMENT, 1, FALSE); + return TRUE; +} + +static VOID DeliverViolation(PMEMORY_WORK_ITEM Work) { + if (!Work) return; + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) return; + + BOOLEAN SourceMatches = IsProcessIdentityCurrent(Work->SourcePid, Work->SourceProcess); + BOOLEAN TargetMatches = IsProcessIdentityCurrent(Work->TargetPid, Work->TargetProcess); + if (!SourceMatches || !TargetMatches) return; + PUNICODE_STRING TargetPath = NULL; - if (NT_SUCCESS(GetProcessImageName(TargetPid, &TargetPath)) && TargetPath && TargetPath->Buffer) { - SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)SourcePid, TargetPath->Buffer, TargetPath->Length); - ExFreePool(TargetPath); - return; + + if (NT_SUCCESS(GetProcessImageName(Work->TargetPid, &TargetPath)) && TargetPath && TargetPath->Buffer) { + SendMessageToUser( + Work->RuleCode, + (ULONG)(ULONG_PTR)Work->SourcePid, + TargetPath->Buffer, + TargetPath->Length + ); + } + else { + UNICODE_STRING MessageString; + RtlInitUnicodeString(&MessageString, Work->FallbackMessage); + SendMessageToUser( + Work->RuleCode, + (ULONG)(ULONG_PTR)Work->SourcePid, + MessageString.Buffer, + MessageString.Length + ); } if (TargetPath) ExFreePool(TargetPath); - - UNICODE_STRING MessageString; - RtlInitUnicodeString(&MessageString, FallbackMessage); - SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)SourcePid, MessageString.Buffer, MessageString.Length); + QueueRuleProcessTermination(Work->SourcePid, Work->Kill); } -BOOLEAN IsAddressInUnmappedMemory(HANDLE ProcessId, PVOID Address) { - if (KeGetCurrentIrql() != PASSIVE_LEVEL) return FALSE; +static VOID ReportViolation( + ULONG RuleCode, + BOOLEAN Kill, + HANDLE SourcePid, + HANDLE TargetPid, + PCWSTR FallbackMessage +) { + if (!RuleCode || !SourcePid || !TargetPid || !FallbackMessage) return; + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) return; - PEPROCESS Process = NULL; - NTSTATUS status = PsLookupProcessByProcessId(ProcessId, &Process); - if (!NT_SUCCESS(status)) return FALSE; + PEPROCESS SourceProcess = NULL; + PEPROCESS TargetProcess = NULL; + + NTSTATUS Status = PsLookupProcessByProcessId(SourcePid, &SourceProcess); + if (NT_SUCCESS(Status)) { + Status = PsLookupProcessByProcessId(TargetPid, &TargetProcess); + } + + if (!NT_SUCCESS(Status)) { + if (SourceProcess) ObDereferenceObject(SourceProcess); + return; + } + + PMEMORY_WORK_ITEM Work = (PMEMORY_WORK_ITEM)PyasAllocate(sizeof(MEMORY_WORK_ITEM)); + if (!Work) { + ObDereferenceObject(TargetProcess); + ObDereferenceObject(SourceProcess); + return; + } + + RtlZeroMemory(Work, sizeof(*Work)); + Work->Type = MemoryWorkViolationReport; + Work->SourceProcess = SourceProcess; + Work->TargetProcess = TargetProcess; + Work->SourcePid = SourcePid; + Work->TargetPid = TargetPid; + Work->RuleCode = RuleCode; + Work->Kill = Kill; + RtlStringCchCopyW(Work->FallbackMessage, RTL_NUMBER_OF(Work->FallbackMessage), FallbackMessage); + + if (!QueueMemoryWork(Work)) { + ObDereferenceObject(TargetProcess); + ObDereferenceObject(SourceProcess); + PyasFree(Work); + } +} + +static BOOLEAN IsExecuteProtection(ULONG Protect) { + ULONG BaseProtect = Protect & 0xFF; + + return BaseProtect == PAGE_EXECUTE || + BaseProtect == PAGE_EXECUTE_READ || + BaseProtect == PAGE_EXECUTE_READWRITE || + BaseProtect == PAGE_EXECUTE_WRITECOPY; +} + +static BOOLEAN IsExecuteWriteProtection(ULONG Protect) { + ULONG BaseProtect = Protect & 0xFF; + + return BaseProtect == PAGE_EXECUTE_READWRITE || + BaseProtect == PAGE_EXECUTE_WRITECOPY; +} + +static BOOLEAN QueryThreadMemoryRegionForProcess( + PEPROCESS Process, + PVOID Address, + PTHREAD_MEMORY_REGION Region +) { + if (!Process || !Region || !Address || KeGetCurrentIrql() != PASSIVE_LEVEL) return FALSE; + + RtlZeroMemory(Region, sizeof(*Region)); KAPC_STATE ApcState; KeStackAttachProcess(Process, &ApcState); MEMORY_BASIC_INFORMATION MemoryInformation = { 0 }; SIZE_T ReturnLength = 0; - status = ZwQueryVirtualMemory( + NTSTATUS Status = ZwQueryVirtualMemory( ZwCurrentProcess(), Address, MemoryBasicInformation, @@ -139,16 +323,140 @@ BOOLEAN IsAddressInUnmappedMemory(HANDLE ProcessId, PVOID Address) { ); KeUnstackDetachProcess(&ApcState); + + if (!NT_SUCCESS(Status)) return FALSE; + if (MemoryInformation.State != MEM_COMMIT) return FALSE; + if (!IsExecuteProtection(MemoryInformation.Protect)) return FALSE; + + Region->StartAddress = Address; + Region->RegionSize = MemoryInformation.RegionSize; + + if (MemoryInformation.Type == MEM_PRIVATE) Region->MemoryType = PYAS_MEMORY_PRIVATE; + else if (MemoryInformation.Type == MEM_MAPPED) Region->MemoryType = PYAS_MEMORY_MAPPED; + else if (MemoryInformation.Type == PYAS_MEM_IMAGE_TYPE) Region->MemoryType = PYAS_MEMORY_IMAGE; + else return FALSE; + + Region->MemoryProtection = PYAS_PROTECT_EXECUTE; + if (IsExecuteWriteProtection(MemoryInformation.Protect)) { + Region->MemoryProtection |= PYAS_PROTECT_EXECUTE_WRITE; + } + + return TRUE; +} + +static BOOLEAN QueryThreadMemoryRegion( + HANDLE ProcessId, + PVOID Address, + PTHREAD_MEMORY_REGION Region +) { + if (!ProcessId) return FALSE; + + PEPROCESS Process = NULL; + NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &Process); + if (!NT_SUCCESS(Status)) return FALSE; + + BOOLEAN Result = QueryThreadMemoryRegionForProcess(Process, Address, Region); ObDereferenceObject(Process); + return Result; +} - if (!NT_SUCCESS(status)) return FALSE; - if (!(MemoryInformation.State & MEM_COMMIT)) return FALSE; - if (!(MemoryInformation.Type & MEM_PRIVATE)) return FALSE; +BOOLEAN IsAddressInUnmappedMemory(HANDLE ProcessId, PVOID Address) { + THREAD_MEMORY_REGION Region = { 0 }; + if (!QueryThreadMemoryRegion(ProcessId, Address, &Region)) return FALSE; + return Region.MemoryType == PYAS_MEMORY_PRIVATE; +} - return MemoryInformation.Protect == PAGE_EXECUTE_READWRITE || - MemoryInformation.Protect == PAGE_EXECUTE_READ || - MemoryInformation.Protect == PAGE_EXECUTE || - MemoryInformation.Protect == PAGE_EXECUTE_WRITECOPY; +static BOOLEAN IsProcessIdentityCurrent(HANDLE ProcessId, PEPROCESS ExpectedProcess) { + if (!ProcessId || !ExpectedProcess) return FALSE; + + PEPROCESS CurrentProcess = NULL; + NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &CurrentProcess); + if (!NT_SUCCESS(Status)) return FALSE; + + BOOLEAN Matches = CurrentProcess == ExpectedProcess; + ObDereferenceObject(CurrentProcess); + return Matches; +} + +static OB_PREOP_CALLBACK_STATUS ObjectPreCallbackCore( + POB_PRE_OPERATION_INFORMATION OperationInformation +) { + if (!OperationInformation || OperationInformation->KernelHandle) { + return OB_PREOP_SUCCESS; + } + + PACCESS_MASK DesiredAccess = GetDesiredAccessPointer(OperationInformation); + ULONG HandleType = GetHandleType(OperationInformation); + + if (!DesiredAccess || *DesiredAccess == 0 || HandleType == 0) { + return OB_PREOP_SUCCESS; + } + + HANDLE SourcePid = PsGetCurrentProcessId(); + if (!SourcePid || IsProcessTrusted(SourcePid)) { + return OB_PREOP_SUCCESS; + } + + HANDLE TargetPid = NULL; + ULONG ObjectType = 0; + ULONG RequestedOperations = 0; + ACCESS_MASK OriginalAccess = *DesiredAccess; + + if (OperationInformation->ObjectType == *PsProcessType) { + TargetPid = PsGetProcessId((PEPROCESS)OperationInformation->Object); + ObjectType = PYAS_OBJECT_PROCESS; + RequestedOperations = GetProcessOperations(OriginalAccess); + } + else if (OperationInformation->ObjectType == *PsThreadType) { + TargetPid = PsGetThreadProcessId((PETHREAD)OperationInformation->Object); + ObjectType = PYAS_OBJECT_THREAD; + RequestedOperations = GetThreadOperations(OriginalAccess); + } + else { + return OB_PREOP_SUCCESS; + } + + if (!TargetPid || SourcePid == TargetPid || RequestedOperations == 0) { + return OB_PREOP_SUCCESS; + } + + ULONG DeniedOperations = 0; + ULONG RuleCode = 0; + BOOLEAN Kill = FALSE; + + if (!EvaluateProcessAccessRule( + SourcePid, + TargetPid, + RequestedOperations, + HandleType, + ObjectType, + &DeniedOperations, + &RuleCode, + &Kill + )) { + return OB_PREOP_SUCCESS; + } + + ACCESS_MASK DeniedAccess = ObjectType == PYAS_OBJECT_PROCESS + ? GetProcessDeniedAccess(DeniedOperations) + : GetThreadDeniedAccess(DeniedOperations); + + ACCESS_MASK EffectiveDeniedAccess = OriginalAccess & DeniedAccess; + if (EffectiveDeniedAccess == 0) return OB_PREOP_SUCCESS; + + *DesiredAccess = OriginalAccess & ~EffectiveDeniedAccess; + + ReportViolation( + RuleCode, + Kill, + SourcePid, + TargetPid, + ObjectType == PYAS_OBJECT_PROCESS + ? L"Process_Handle_Access_Denied" + : L"Thread_Handle_Access_Denied" + ); + + return OB_PREOP_SUCCESS; } static OB_PREOP_CALLBACK_STATUS ObjectPreCallback( @@ -157,139 +465,40 @@ static OB_PREOP_CALLBACK_STATUS ObjectPreCallback( ) { UNREFERENCED_PARAMETER(RegistrationContext); - if (!OperationInformation || OperationInformation->KernelHandle) { - return OB_PREOP_SUCCESS; - } + if (!g_MemoryProtectionInitialized) return OB_PREOP_SUCCESS; + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) return OB_PREOP_SUCCESS; + if (!ExAcquireRundownProtection(&g_MemoryRundown)) return OB_PREOP_SUCCESS; - PACCESS_MASK DesiredAccess = GetDesiredAccessPointer(OperationInformation); - if (!DesiredAccess || *DesiredAccess == 0) { - return OB_PREOP_SUCCESS; - } - - HANDLE SourcePid = PsGetCurrentProcessId(); - HANDLE TargetPid = NULL; - - if (OperationInformation->ObjectType == *PsProcessType) { - TargetPid = PsGetProcessId((PEPROCESS)OperationInformation->Object); - if (!TargetPid || SourcePid == TargetPid || IsProcessTrusted(SourcePid)) { - return OB_PREOP_SUCCESS; - } - - ACCESS_MASK OriginalAccess = *DesiredAccess; - ULONG MemoryOperation = 0; - ULONG ControlOperation = 0; - - if (OriginalAccess & PROCESS_VM_READ) { - MemoryOperation |= OP_VM_READ; - } - - if (OriginalAccess & (PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_CREATE_THREAD)) { - MemoryOperation |= OP_VM_WRITE; - } - - if (OriginalAccess & PROCESS_TERMINATE) { - ControlOperation |= OP_TERMINATE; - } - - if (OriginalAccess & PROCESS_SUSPEND_RESUME) { - ControlOperation |= OP_SUSPEND_RESUME; - } - - if (OriginalAccess & PROCESS_DUP_HANDLE) { - ControlOperation |= OP_DUP_HANDLE; - } - - if (OriginalAccess & (PROCESS_SET_INFORMATION | PROCESS_SET_QUOTA)) { - ControlOperation |= OP_SET_INFORMATION; - } - - if (MemoryOperation != 0) { - ULONG RuleCode = 0; - if (EvaluateMemoryRule(SourcePid, TargetPid, MemoryOperation, &RuleCode)) { - *DesiredAccess &= ~GetProcessMemoryDeniedAccess(MemoryOperation); - ReportViolation(RuleCode, SourcePid, TargetPid, L"Memory_Access_Violation"); - } - } - - if (ControlOperation != 0) { - ULONG RuleCode = 0; - if (EvaluateProcessHandleRule(SourcePid, TargetPid, ControlOperation, &RuleCode)) { - *DesiredAccess &= ~GetProcessControlDeniedAccess(ControlOperation); - ReportViolation(RuleCode, SourcePid, TargetPid, L"Process_Control_Access_Violation"); - } - } - - return OB_PREOP_SUCCESS; - } - - if (OperationInformation->ObjectType == *PsThreadType) { - TargetPid = PsGetThreadProcessId((PETHREAD)OperationInformation->Object); - if (!TargetPid || SourcePid == TargetPid || IsProcessTrusted(SourcePid)) { - return OB_PREOP_SUCCESS; - } - - ACCESS_MASK OriginalAccess = *DesiredAccess; - ULONG MemoryOperation = 0; - ULONG ControlOperation = 0; - - if (OriginalAccess & THREAD_SET_CONTEXT) { - MemoryOperation |= OP_VM_WRITE; - } - - if (OriginalAccess & THREAD_TERMINATE) { - ControlOperation |= OP_TERMINATE; - } - - if (OriginalAccess & THREAD_SUSPEND_RESUME) { - ControlOperation |= OP_SUSPEND_RESUME; - } - - if (OriginalAccess & (THREAD_SET_INFORMATION | THREAD_SET_THREAD_TOKEN)) { - ControlOperation |= OP_SET_INFORMATION; - } - - if (MemoryOperation != 0) { - ULONG RuleCode = 0; - if (EvaluateMemoryRule(SourcePid, TargetPid, MemoryOperation, &RuleCode)) { - *DesiredAccess &= ~GetThreadMemoryDeniedAccess(MemoryOperation); - ReportViolation(RuleCode, SourcePid, TargetPid, L"Thread_Context_Manipulation_Shield"); - } - } - - if (ControlOperation != 0) { - ULONG RuleCode = 0; - if (EvaluateProcessHandleRule(SourcePid, TargetPid, ControlOperation, &RuleCode)) { - *DesiredAccess &= ~GetThreadControlDeniedAccess(ControlOperation); - ReportViolation(RuleCode, SourcePid, TargetPid, L"Thread_Control_Access_Violation"); - } - } - } - - return OB_PREOP_SUCCESS; + OB_PREOP_CALLBACK_STATUS Result = ObjectPreCallbackCore(OperationInformation); + ExReleaseRundownProtection(&g_MemoryRundown); + return Result; } -static VOID ThreadNotifyCallback(HANDLE ProcessId, HANDLE ThreadId, BOOLEAN IsCreated) { - if (!IsCreated) return; - - HANDLE SourcePid = PsGetCurrentProcessId(); - if (SourcePid == ProcessId) return; - if (IsProcessTrusted(SourcePid)) return; +static VOID ProcessThreadAnalysis(PMEMORY_WORK_ITEM Work) { + if (!Work) return; + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) return; PVOID StartAddress = NULL; HANDLE ThreadHandle = NULL; - CLIENT_ID ClientId = { ProcessId, ThreadId }; + BOOLEAN IdentityValid = + IsProcessIdentityCurrent(Work->SourcePid, Work->SourceProcess) && + IsProcessIdentityCurrent(Work->TargetPid, Work->TargetProcess); + CLIENT_ID ClientId = { Work->TargetPid, Work->ThreadId }; OBJECT_ATTRIBUTES ObjectAttributes; InitializeObjectAttributes(&ObjectAttributes, NULL, OBJ_KERNEL_HANDLE, NULL, NULL); - NTSTATUS status = ZwOpenThread( - &ThreadHandle, - THREAD_QUERY_INFORMATION, - &ObjectAttributes, - &ClientId - ); + NTSTATUS Status = STATUS_PROCESS_IS_TERMINATING; + if (IdentityValid) { + Status = ZwOpenThread( + &ThreadHandle, + THREAD_QUERY_INFORMATION, + &ObjectAttributes, + &ClientId + ); + } - if (NT_SUCCESS(status)) { - status = ZwQueryInformationThread( + if (NT_SUCCESS(Status)) { + Status = ZwQueryInformationThread( ThreadHandle, (THREADINFOCLASS)THREAD_QUERY_SET_WIN32_START_ADDRESS, &StartAddress, @@ -299,11 +508,110 @@ static VOID ThreadNotifyCallback(HANDLE ProcessId, HANDLE ThreadId, BOOLEAN IsCr ZwClose(ThreadHandle); } - if (!NT_SUCCESS(status) || !StartAddress) return; + if (!NT_SUCCESS(Status) || !StartAddress) return; + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) return; + + THREAD_MEMORY_REGION Region = { 0 }; + if (!QueryThreadMemoryRegionForProcess(Work->TargetProcess, StartAddress, &Region)) return; ULONG RuleCode = 0; - if (EvaluateThreadRule(SourcePid, ProcessId, StartAddress, &RuleCode)) { - ReportViolation(RuleCode, SourcePid, ProcessId, L"Shellcode_Load_Detected"); + BOOLEAN Kill = FALSE; + if (!EvaluateThreadRule( + Work->SourcePid, + Work->TargetPid, + StartAddress, + Region.MemoryType, + Region.MemoryProtection, + Region.RegionSize, + &RuleCode, + &Kill + )) { + return; + } + + Work->RuleCode = RuleCode; + Work->Kill = Kill; + RtlStringCchCopyW( + Work->FallbackMessage, + RTL_NUMBER_OF(Work->FallbackMessage), + L"Remote_Executable_Thread_Detected" + ); + DeliverViolation(Work); +} + +static VOID MemoryWorkerThread(PVOID Parameter) { + UNREFERENCED_PARAMETER(Parameter); + + for (;;) { + KeWaitForSingleObject( + &g_MemoryWorkSemaphore, + Executive, + KernelMode, + FALSE, + NULL + ); + + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) { + CancelQueuedMemoryWork(); + break; + } + + PMEMORY_WORK_ITEM Work = PopMemoryWork(); + if (!Work) continue; + + if (Work->Type == MemoryWorkThreadAnalysis) { + ProcessThreadAnalysis(Work); + } + else if (Work->Type == MemoryWorkViolationReport) { + DeliverViolation(Work); + } + + ReleaseMemoryWork(Work); + } + + PsTerminateSystemThread(STATUS_SUCCESS); +} + +static VOID ThreadNotifyCallback( + HANDLE ProcessId, + HANDLE ThreadId, + BOOLEAN IsCreated +) { + if (!IsCreated || !g_MemoryProtectionInitialized) return; + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) return; + + HANDLE SourcePid = PsGetCurrentProcessId(); + if (!SourcePid || SourcePid == ProcessId || IsProcessTrusted(SourcePid)) return; + + PEPROCESS SourceProcess = PsGetCurrentProcess(); + PEPROCESS TargetProcess = NULL; + ObReferenceObject(SourceProcess); + + NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &TargetProcess); + if (!NT_SUCCESS(Status)) { + ObDereferenceObject(SourceProcess); + return; + } + + PMEMORY_WORK_ITEM Work = (PMEMORY_WORK_ITEM)PyasAllocate(sizeof(MEMORY_WORK_ITEM)); + if (!Work) { + ObDereferenceObject(TargetProcess); + ObDereferenceObject(SourceProcess); + return; + } + + RtlZeroMemory(Work, sizeof(*Work)); + Work->Type = MemoryWorkThreadAnalysis; + Work->SourceProcess = SourceProcess; + Work->TargetProcess = TargetProcess; + Work->SourcePid = SourcePid; + Work->TargetPid = ProcessId; + Work->ThreadId = ThreadId; + + if (!QueueMemoryWork(Work)) { + ObDereferenceObject(TargetProcess); + ObDereferenceObject(SourceProcess); + PyasFree(Work); } } @@ -312,36 +620,150 @@ static VOID LoadImageNotifyCallback( HANDLE ProcessId, PIMAGE_INFO ImageInfo ) { - if (!FullImageName || !ImageInfo || ImageInfo->SystemModeImage) return; + if (!g_MemoryProtectionInitialized) return; + if (InterlockedCompareExchange(&g_MemoryStopping, 0, 0) != 0) return; + if (!FullImageName || !FullImageName->Buffer || !ImageInfo || ImageInfo->SystemModeImage) return; + if (!ProcessId || !ExAcquireRundownProtection(&g_MemoryRundown)) return; ULONG RuleCode = 0; - if (EvaluateProcessRule(ProcessId, FullImageName, NULL, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateImageLoadRule(ProcessId, FullImageName, ImageInfo, &RuleCode, &Kill)) { SendMessageToUser( RuleCode, (ULONG)(ULONG_PTR)ProcessId, FullImageName->Buffer, FullImageName->Length ); + QueueRuleProcessTermination(ProcessId, Kill); } + + ExReleaseRundownProtection(&g_MemoryRundown); +} + +static NTSTATUS RemoveMemoryCallbacks() { + NTSTATUS FirstFailure = STATUS_SUCCESS; + + if (GlobalData.ObRegistrationHandle) { + ObUnRegisterCallbacks(GlobalData.ObRegistrationHandle); + GlobalData.ObRegistrationHandle = NULL; + } + + if (g_ThreadNotifyRegistered) { + NTSTATUS Status = PsRemoveCreateThreadNotifyRoutine(ThreadNotifyCallback); + if (NT_SUCCESS(Status) || Status == STATUS_PROCEDURE_NOT_FOUND) { + g_ThreadNotifyRegistered = FALSE; + } + else if (NT_SUCCESS(FirstFailure)) { + FirstFailure = Status; + } + } + + if (g_LoadImageNotifyRegistered) { + NTSTATUS Status = PsRemoveLoadImageNotifyRoutine(LoadImageNotifyCallback); + if (NT_SUCCESS(Status) || Status == STATUS_PROCEDURE_NOT_FOUND) { + g_LoadImageNotifyRegistered = FALSE; + } + else if (NT_SUCCESS(FirstFailure)) { + FirstFailure = Status; + } + } + + return FirstFailure; +} + +static NTSTATUS StartMemoryWorker() { + if (g_MemoryWorkerHandle) return STATUS_SUCCESS; + + OBJECT_ATTRIBUTES ObjectAttributes; + InitializeObjectAttributes( + &ObjectAttributes, + NULL, + OBJ_KERNEL_HANDLE, + NULL, + NULL + ); + + return PsCreateSystemThread( + &g_MemoryWorkerHandle, + SYNCHRONIZE, + &ObjectAttributes, + NULL, + NULL, + MemoryWorkerThread, + NULL + ); +} + +static NTSTATUS StopMemoryWorker(BOOLEAN WaitWithoutTimeout) { + CancelQueuedMemoryWork(); + + if (!g_MemoryWorkerHandle) return STATUS_SUCCESS; + + KeReleaseSemaphore(&g_MemoryWorkSemaphore, IO_NO_INCREMENT, 1, FALSE); + + LARGE_INTEGER Timeout; + PLARGE_INTEGER TimeoutPointer = NULL; + + if (!WaitWithoutTimeout) { + Timeout.QuadPart = -((LONGLONG)MEMORY_WORKER_STOP_TIMEOUT_MS * 10 * 1000); + TimeoutPointer = &Timeout; + } + + NTSTATUS Status = ZwWaitForSingleObject( + g_MemoryWorkerHandle, + FALSE, + TimeoutPointer + ); + + if (Status == STATUS_TIMEOUT) return STATUS_DEVICE_BUSY; + if (!NT_SUCCESS(Status)) return Status; + + ZwClose(g_MemoryWorkerHandle); + g_MemoryWorkerHandle = NULL; + return STATUS_SUCCESS; } NTSTATUS InitializeMemoryProtection(PDRIVER_OBJECT DriverObject) { UNREFERENCED_PARAMETER(DriverObject); - if (GlobalData.ObRegistrationHandle || g_ThreadNotifyRegistered) { - return STATUS_SUCCESS; + if (g_MemoryProtectionInitialized) return STATUS_SUCCESS; + + InitializeListHead(&g_MemoryWorkQueue); + KeInitializeSpinLock(&g_MemoryWorkLock); + KeInitializeSemaphore(&g_MemoryWorkSemaphore, 0, MAXLONG); + ExInitializeRundownProtection(&g_MemoryRundown); + g_MemoryRundownCompleted = FALSE; + InterlockedExchange(&g_PendingMemoryWork, 0); + InterlockedExchange(&g_MemoryStopping, 0); + g_MemoryProtectionInitialized = TRUE; + + NTSTATUS Status = StartMemoryWorker(); + if (!NT_SUCCESS(Status)) { + InterlockedExchange(&g_MemoryStopping, 1); + g_MemoryProtectionInitialized = FALSE; + return Status; } - NTSTATUS status = PsSetCreateThreadNotifyRoutine(ThreadNotifyCallback); - if (!NT_SUCCESS(status)) { - return status; + Status = PsSetCreateThreadNotifyRoutine(ThreadNotifyCallback); + if (!NT_SUCCESS(Status)) { + InterlockedExchange(&g_MemoryStopping, 1); + StopMemoryWorker(TRUE); + g_MemoryProtectionInitialized = FALSE; + return Status; } g_ThreadNotifyRegistered = TRUE; - status = PsSetLoadImageNotifyRoutine(LoadImageNotifyCallback); - if (NT_SUCCESS(status)) { - g_LoadImageNotifyRegistered = TRUE; + Status = PsSetLoadImageNotifyRoutine(LoadImageNotifyCallback); + if (!NT_SUCCESS(Status)) { + InterlockedExchange(&g_MemoryStopping, 1); + RemoveMemoryCallbacks(); + ExWaitForRundownProtectionRelease(&g_MemoryRundown); + g_MemoryRundownCompleted = TRUE; + StopMemoryWorker(TRUE); + g_MemoryProtectionInitialized = FALSE; + return Status; } + g_LoadImageNotifyRegistered = TRUE; OB_OPERATION_REGISTRATION OperationRegistration[2] = { 0 }; OperationRegistration[0].ObjectType = PsProcessType; @@ -356,41 +778,41 @@ NTSTATUS InitializeMemoryProtection(PDRIVER_OBJECT DriverObject) { CallbackRegistration.Version = OB_FLT_REGISTRATION_VERSION; CallbackRegistration.OperationRegistrationCount = RTL_NUMBER_OF(OperationRegistration); CallbackRegistration.OperationRegistration = OperationRegistration; - RtlInitUnicodeString(&CallbackRegistration.Altitude, L"320000.PYAS.Ob"); + RtlInitUnicodeString(&CallbackRegistration.Altitude, PYAS_OB_ALTITUDE); - status = ObRegisterCallbacks(&CallbackRegistration, &GlobalData.ObRegistrationHandle); - if (NT_SUCCESS(status)) { - return STATUS_SUCCESS; - } + Status = ObRegisterCallbacks(&CallbackRegistration, &GlobalData.ObRegistrationHandle); + if (NT_SUCCESS(Status)) return STATUS_SUCCESS; - GlobalData.ObRegistrationHandle = NULL; - - if (g_LoadImageNotifyRegistered) { - PsRemoveLoadImageNotifyRoutine(LoadImageNotifyCallback); - g_LoadImageNotifyRegistered = FALSE; - } - - if (g_ThreadNotifyRegistered) { - PsRemoveCreateThreadNotifyRoutine(ThreadNotifyCallback); - g_ThreadNotifyRegistered = FALSE; - } - - return status; + InterlockedExchange(&g_MemoryStopping, 1); + RemoveMemoryCallbacks(); + ExWaitForRundownProtectionRelease(&g_MemoryRundown); + g_MemoryRundownCompleted = TRUE; + StopMemoryWorker(TRUE); + g_MemoryProtectionInitialized = FALSE; + return Status; } -VOID UninitializeMemoryProtection() { - if (GlobalData.ObRegistrationHandle) { - ObUnRegisterCallbacks(GlobalData.ObRegistrationHandle); - GlobalData.ObRegistrationHandle = NULL; +NTSTATUS UninitializeMemoryProtection(BOOLEAN WaitWithoutTimeout) { + if (!g_MemoryProtectionInitialized) return STATUS_SUCCESS; + + InterlockedExchange(&g_MemoryStopping, 1); + + NTSTATUS CallbackStatus = RemoveMemoryCallbacks(); + + if (!g_MemoryRundownCompleted) { + ExWaitForRundownProtectionRelease(&g_MemoryRundown); + g_MemoryRundownCompleted = TRUE; } - if (g_ThreadNotifyRegistered) { - PsRemoveCreateThreadNotifyRoutine(ThreadNotifyCallback); - g_ThreadNotifyRegistered = FALSE; + NTSTATUS WorkerStatus = StopMemoryWorker(WaitWithoutTimeout); + if (!NT_SUCCESS(CallbackStatus)) return CallbackStatus; + if (!NT_SUCCESS(WorkerStatus)) return WorkerStatus; + + if (g_ThreadNotifyRegistered || g_LoadImageNotifyRegistered || GlobalData.ObRegistrationHandle) { + return STATUS_DEVICE_BUSY; } - if (g_LoadImageNotifyRegistered) { - PsRemoveLoadImageNotifyRoutine(LoadImageNotifyCallback); - g_LoadImageNotifyRegistered = FALSE; - } + g_MemoryProtectionInitialized = FALSE; + return STATUS_SUCCESS; } + diff --git a/Plugins/Filter/ProtectProcess.cpp b/Plugins/Filter/ProtectProcess.cpp index d3b630f..bf9bf01 100644 --- a/Plugins/Filter/ProtectProcess.cpp +++ b/Plugins/Filter/ProtectProcess.cpp @@ -1,21 +1,285 @@ #include "DriverCommon.h" -VOID ProcessNotifyCallbackEx(PEPROCESS Process, HANDLE ProcessId, PPS_CREATE_NOTIFY_INFO CreateInfo) { - UNREFERENCED_PARAMETER(Process); - UNREFERENCED_PARAMETER(ProcessId); +constexpr auto PROCESS_RELATION_BUCKET_COUNT = 1024; +constexpr auto PROCESS_RELATION_BUCKET_WAYS = 4; - if (CreateInfo != NULL) { - if (KeGetCurrentIrql() > APC_LEVEL) return; +typedef HANDLE(NTAPI* PPS_GET_PROCESS_INHERITED_FROM_UNIQUE_PROCESS_ID)(PEPROCESS Process); - HANDLE ParentId = CreateInfo->CreatingThreadId.UniqueProcess; - if (!ParentId) return; +typedef struct _PROCESS_RELATION_ENTRY { + HANDLE ProcessId; + HANDLE ParentProcessId; + HANDLE CreatorProcessId; + LARGE_INTEGER CreateTime; + LARGE_INTEGER ParentCreateTime; + ULONGLONG Sequence; +} PROCESS_RELATION_ENTRY, * PPROCESS_RELATION_ENTRY; - ULONG RuleCode = 0; - if (EvaluateProcessRule(ParentId, CreateInfo->ImageFileName, CreateInfo->CommandLine, &RuleCode)) { - CreateInfo->CreationStatus = STATUS_ACCESS_DENIED; - if (CreateInfo->ImageFileName) { - SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)ParentId, CreateInfo->ImageFileName->Buffer, CreateInfo->ImageFileName->Length); - } +static PROCESS_RELATION_ENTRY g_ProcessRelations[PROCESS_RELATION_BUCKET_COUNT][PROCESS_RELATION_BUCKET_WAYS]; +static KSPIN_LOCK g_ProcessRelationLock; +static EX_RUNDOWN_REF g_ProcessRundown; +static volatile LONG g_ProcessStopping = 1; +static volatile LONG64 g_ProcessRelationSequence = 0; +static BOOLEAN g_ProcessProtectionInitialized = FALSE; +static PPS_GET_PROCESS_INHERITED_FROM_UNIQUE_PROCESS_ID g_GetInheritedProcessId = NULL; + +static ULONG GetProcessRelationBucket(HANDLE ProcessId) { + return ((ULONG)(ULONG_PTR)ProcessId >> 2) & (PROCESS_RELATION_BUCKET_COUNT - 1); +} + +static BOOLEAN QueryProcessCreateTime(HANDLE ProcessId, PLARGE_INTEGER CreateTime) { + if (CreateTime) CreateTime->QuadPart = 0; + if (!ProcessId || !CreateTime || KeGetCurrentIrql() > APC_LEVEL) return FALSE; + + PEPROCESS Process = NULL; + NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &Process); + if (!NT_SUCCESS(Status)) return FALSE; + + CreateTime->QuadPart = PsGetProcessCreateTimeQuadPart(Process); + ObDereferenceObject(Process); + return TRUE; +} + +static VOID StoreProcessRelation( + PEPROCESS Process, + HANDLE ProcessId, + HANDLE ParentProcessId, + HANDLE CreatorProcessId, + PLARGE_INTEGER KnownParentCreateTime +) { + if (!Process || !ProcessId) return; + + PROCESS_RELATION_ENTRY NewEntry = { 0 }; + NewEntry.ProcessId = ProcessId; + NewEntry.ParentProcessId = ParentProcessId; + NewEntry.CreatorProcessId = CreatorProcessId; + NewEntry.CreateTime.QuadPart = PsGetProcessCreateTimeQuadPart(Process); + + if (KnownParentCreateTime) { + NewEntry.ParentCreateTime = *KnownParentCreateTime; + } + else if (ParentProcessId) { + QueryProcessCreateTime(ParentProcessId, &NewEntry.ParentCreateTime); + } + + NewEntry.Sequence = (ULONGLONG)InterlockedIncrement64(&g_ProcessRelationSequence); + + ULONG Bucket = GetProcessRelationBucket(ProcessId); + ULONG Selected = 0; + + KIRQL OldIrql; + KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql); + + for (ULONG Index = 0; Index < PROCESS_RELATION_BUCKET_WAYS; Index++) { + PPROCESS_RELATION_ENTRY Entry = &g_ProcessRelations[Bucket][Index]; + + if (Entry->ProcessId == ProcessId || Entry->ProcessId == NULL) { + Selected = Index; + break; + } + + if (Entry->Sequence < g_ProcessRelations[Bucket][Selected].Sequence) { + Selected = Index; } } -} \ No newline at end of file + + g_ProcessRelations[Bucket][Selected] = NewEntry; + KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql); +} + +static VOID RemoveProcessRelation(HANDLE ProcessId) { + if (!ProcessId) return; + + ULONG Bucket = GetProcessRelationBucket(ProcessId); + + KIRQL OldIrql; + KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql); + + for (ULONG Index = 0; Index < PROCESS_RELATION_BUCKET_WAYS; Index++) { + PPROCESS_RELATION_ENTRY Entry = &g_ProcessRelations[Bucket][Index]; + if (Entry->ProcessId == ProcessId) { + RtlZeroMemory(Entry, sizeof(*Entry)); + break; + } + } + + KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql); +} + +BOOLEAN GetProcessRelation( + HANDLE ProcessId, + PHANDLE ParentProcessId, + PHANDLE CreatorProcessId +) { + if (ParentProcessId) *ParentProcessId = NULL; + if (CreatorProcessId) *CreatorProcessId = NULL; + if (!ProcessId || !g_ProcessProtectionInitialized || KeGetCurrentIrql() > APC_LEVEL) return FALSE; + + PEPROCESS Process = NULL; + NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &Process); + if (!NT_SUCCESS(Status)) return FALSE; + + LARGE_INTEGER CurrentCreateTime = { 0 }; + CurrentCreateTime.QuadPart = PsGetProcessCreateTimeQuadPart(Process); + + PROCESS_RELATION_ENTRY Snapshot = { 0 }; + ULONG Bucket = GetProcessRelationBucket(ProcessId); + + KIRQL OldIrql; + KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql); + + for (ULONG Index = 0; Index < PROCESS_RELATION_BUCKET_WAYS; Index++) { + PPROCESS_RELATION_ENTRY Entry = &g_ProcessRelations[Bucket][Index]; + if (Entry->ProcessId == ProcessId) { + Snapshot = *Entry; + break; + } + } + + KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql); + + BOOLEAN CacheValid = Snapshot.ProcessId == ProcessId && + Snapshot.CreateTime.QuadPart == CurrentCreateTime.QuadPart; + + if (CacheValid && Snapshot.ParentProcessId && Snapshot.ParentCreateTime.QuadPart != 0) { + LARGE_INTEGER LiveParentCreateTime = { 0 }; + if (QueryProcessCreateTime(Snapshot.ParentProcessId, &LiveParentCreateTime) && + LiveParentCreateTime.QuadPart != Snapshot.ParentCreateTime.QuadPart) { + CacheValid = FALSE; + } + } + + if (CacheValid) { + if (ParentProcessId) *ParentProcessId = Snapshot.ParentProcessId; + if (CreatorProcessId) *CreatorProcessId = Snapshot.CreatorProcessId; + ObDereferenceObject(Process); + return Snapshot.ParentProcessId != NULL || Snapshot.CreatorProcessId != NULL; + } + + if (Snapshot.ProcessId == ProcessId) { + RemoveProcessRelation(ProcessId); + } + + HANDLE LiveParentProcessId = NULL; + LARGE_INTEGER LiveParentCreateTime = { 0 }; + + if (g_GetInheritedProcessId) { + LiveParentProcessId = g_GetInheritedProcessId(Process); + } + + if (!LiveParentProcessId || LiveParentProcessId == ProcessId) { + ObDereferenceObject(Process); + return FALSE; + } + + if (!QueryProcessCreateTime(LiveParentProcessId, &LiveParentCreateTime) || + LiveParentCreateTime.QuadPart > CurrentCreateTime.QuadPart) { + ObDereferenceObject(Process); + return FALSE; + } + + StoreProcessRelation( + Process, + ProcessId, + LiveParentProcessId, + NULL, + &LiveParentCreateTime + ); + + if (ParentProcessId) *ParentProcessId = LiveParentProcessId; + if (CreatorProcessId) *CreatorProcessId = NULL; + + ObDereferenceObject(Process); + return TRUE; +} + +NTSTATUS InitializeProcessProtection() { + if (g_ProcessProtectionInitialized) return STATUS_SUCCESS; + + KeInitializeSpinLock(&g_ProcessRelationLock); + RtlZeroMemory(g_ProcessRelations, sizeof(g_ProcessRelations)); + InterlockedExchange64(&g_ProcessRelationSequence, 0); + ExInitializeRundownProtection(&g_ProcessRundown); + + UNICODE_STRING RoutineName = RTL_CONSTANT_STRING(L"PsGetProcessInheritedFromUniqueProcessId"); + g_GetInheritedProcessId = reinterpret_cast( + MmGetSystemRoutineAddress(&RoutineName) + ); + + InterlockedExchange(&g_ProcessStopping, 0); + g_ProcessProtectionInitialized = TRUE; + return STATUS_SUCCESS; +} + +NTSTATUS UninitializeProcessProtection() { + if (!g_ProcessProtectionInitialized) return STATUS_SUCCESS; + + InterlockedExchange(&g_ProcessStopping, 1); + ExWaitForRundownProtectionRelease(&g_ProcessRundown); + + KIRQL OldIrql; + KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql); + RtlZeroMemory(g_ProcessRelations, sizeof(g_ProcessRelations)); + KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql); + + g_GetInheritedProcessId = NULL; + g_ProcessProtectionInitialized = FALSE; + return STATUS_SUCCESS; +} + +VOID ProcessNotifyCallbackEx( + PEPROCESS Process, + HANDLE ProcessId, + PPS_CREATE_NOTIFY_INFO CreateInfo +) { + if (!g_ProcessProtectionInitialized) return; + if (InterlockedCompareExchange(&g_ProcessStopping, 0, 0) != 0) return; + if (!ExAcquireRundownProtection(&g_ProcessRundown)) return; + + if (!CreateInfo) { + RemoveProcessRelation(ProcessId); + ExReleaseRundownProtection(&g_ProcessRundown); + return; + } + + if (!NT_SUCCESS(CreateInfo->CreationStatus)) { + ExReleaseRundownProtection(&g_ProcessRundown); + return; + } + + HANDLE CreatorPid = CreateInfo->CreatingThreadId.UniqueProcess; + HANDLE ParentPid = CreateInfo->ParentProcessId; + ULONG RuleCode = 0; + BOOLEAN Kill = FALSE; + + BOOLEAN Blocked = EvaluateProcessCreateRule( + CreatorPid, + ParentPid, + ProcessId, + CreateInfo->ImageFileName, + CreateInfo->CommandLine, + CreateInfo->FileOpenNameAvailable ? TRUE : FALSE, + CreateInfo->IsSubsystemProcess ? TRUE : FALSE, + &RuleCode, + &Kill + ); + + if (Blocked) { + CreateInfo->CreationStatus = STATUS_ACCESS_DENIED; + + if (CreateInfo->ImageFileName && CreateInfo->ImageFileName->Buffer) { + SendMessageToUser( + RuleCode, + (ULONG)(ULONG_PTR)CreatorPid, + CreateInfo->ImageFileName->Buffer, + CreateInfo->ImageFileName->Length + ); + } + + QueueRuleProcessTermination(CreatorPid, Kill); + } + else { + StoreProcessRelation(Process, ProcessId, ParentPid, CreatorPid, NULL); + } + + ExReleaseRundownProtection(&g_ProcessRundown); +} diff --git a/Plugins/Filter/ProtectRegistry.cpp b/Plugins/Filter/ProtectRegistry.cpp index bfec822..5a5927d 100644 --- a/Plugins/Filter/ProtectRegistry.cpp +++ b/Plugins/Filter/ProtectRegistry.cpp @@ -100,8 +100,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A HANDLE Pid = PsGetCurrentProcessId(); ULONG RuleCode = 0; - if (EvaluateRegistryRule(Pid, &FullPath, OP_CREATE | OP_WRITE, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateRegistryRule(Pid, &FullPath, OP_CREATE | OP_WRITE, &RuleCode, &Kill)) { SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length); + QueueRuleProcessTermination(Pid, Kill); status = STATUS_ACCESS_DENIED; } PyasFree(FullPathBuffer); @@ -117,8 +119,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A if (KeyName) { HANDLE Pid = PsGetCurrentProcessId(); ULONG RuleCode = 0; - if (EvaluateRegistryRule(Pid, KeyName, OP_DELETE, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateRegistryRule(Pid, KeyName, OP_DELETE, &RuleCode, &Kill)) { SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, KeyName->Buffer, KeyName->Length); + QueueRuleProcessTermination(Pid, Kill); status = STATUS_ACCESS_DENIED; } CmCallbackReleaseKeyObjectIDEx(KeyName); @@ -156,8 +160,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A HANDLE Pid = PsGetCurrentProcessId(); ULONG RuleCode = 0; - if (EvaluateRegistryRule(Pid, &FullPath, OP_WRITE, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateRegistryRule(Pid, &FullPath, OP_WRITE, &RuleCode, &Kill)) { SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length); + QueueRuleProcessTermination(Pid, Kill); status = STATUS_ACCESS_DENIED; } PyasFree(Buffer); @@ -198,8 +204,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A HANDLE Pid = PsGetCurrentProcessId(); ULONG RuleCode = 0; - if (EvaluateRegistryRule(Pid, &FullPath, OP_DELETE, &RuleCode)) { + BOOLEAN Kill = FALSE; + if (EvaluateRegistryRule(Pid, &FullPath, OP_DELETE, &RuleCode, &Kill)) { SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length); + QueueRuleProcessTermination(Pid, Kill); status = STATUS_ACCESS_DENIED; } PyasFree(Buffer); @@ -215,20 +223,20 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A } NTSTATUS InitializeRegistryProtection(PDRIVER_OBJECT DriverObject) { - static UNICODE_STRING Altitude = RTL_CONSTANT_STRING(L"320000.PYAS.Cm"); - NTSTATUS status = CmRegisterCallbackEx(RegistryCallback, &Altitude, DriverObject, NULL, &Cookie, NULL); + if (Cookie.QuadPart != 0) return STATUS_SUCCESS; - if (!NT_SUCCESS(status)) { - static UNICODE_STRING FallbackAltitude = RTL_CONSTANT_STRING(L"320000.PYAS.Cm.Fallback"); - status = CmRegisterCallbackEx(RegistryCallback, &FallbackAltitude, DriverObject, NULL, &Cookie, NULL); - } - - return status; + UNICODE_STRING Altitude; + RtlInitUnicodeString(&Altitude, L"320000.4201"); + return CmRegisterCallbackEx(RegistryCallback, &Altitude, DriverObject, NULL, &Cookie, NULL); } -VOID UninitializeRegistryProtection() { - if (Cookie.QuadPart != 0) { - CmUnRegisterCallback(Cookie); +NTSTATUS UninitializeRegistryProtection() { + if (Cookie.QuadPart == 0) return STATUS_SUCCESS; + + NTSTATUS Status = CmUnRegisterCallback(Cookie); + if (NT_SUCCESS(Status)) { Cookie.QuadPart = 0; } -} \ No newline at end of file + + return Status; +} diff --git a/Plugins/Filter/ProtectRules.cpp b/Plugins/Filter/ProtectRules.cpp index e899715..b2c7524 100644 --- a/Plugins/Filter/ProtectRules.cpp +++ b/Plugins/Filter/ProtectRules.cpp @@ -3,6 +3,7 @@ constexpr auto MAX_TRACKERS = 128; constexpr auto TRUST_CACHE_SIZE = 128; constexpr auto TRUST_CACHE_TTL_SEC = 300; +constexpr auto MAX_RULE_FILE_SIZE = 4 * 1024 * 1024; typedef struct _BEHAVIOR_TRACKER { PEPROCESS Process; @@ -19,6 +20,11 @@ typedef struct _TRUST_CACHE_ENTRY { LARGE_INTEGER CacheTime; } TRUST_CACHE_ENTRY, * PTRUST_CACHE_ENTRY; +typedef struct _RULE_TERMINATION_WORK_ITEM { + LIST_ENTRY Entry; + PEPROCESS Process; +} RULE_TERMINATION_WORK_ITEM, * PRULE_TERMINATION_WORK_ITEM; + static BEHAVIOR_TRACKER BehaviorTrackers[MAX_TRACKERS]; static TRUST_CACHE_ENTRY TrustCache[TRUST_CACHE_SIZE]; static KSPIN_LOCK TrustCacheLock; @@ -27,6 +33,13 @@ static BOOLEAN g_CacheInitialized = FALSE; static ERESOURCE RuleLock; static PDYNAMIC_RULE g_DynamicRules = NULL; static PRULE_NODE g_ProcessTrustedPaths = NULL; +static LIST_ENTRY g_TerminationQueue; +static KSPIN_LOCK g_TerminationLock; +static KSEMAPHORE g_TerminationSemaphore; +static HANDLE g_TerminationWorkerHandle = NULL; +static volatile LONG g_TerminationStopping = 1; +static volatile LONG g_PendingTerminations = 0; +constexpr auto MAX_PENDING_TERMINATIONS = 128; static BOOLEAN CheckRuleThreshold(PDYNAMIC_RULE Rule, HANDLE ProcessId); BOOLEAN IsAddressInUnmappedMemory(HANDLE ProcessId, PVOID Address); @@ -48,6 +61,225 @@ VOID SafeSetPyasPid(ULONG Pid) { KeReleaseSpinLock(&GlobalData.PidLock, OldIrql); } +static BOOLEAN IsProtectedTerminationTarget(PEPROCESS Process) { + if (!Process || Process == PsInitialSystemProcess) return TRUE; + + BOOLEAN IsConnectedClient = FALSE; + KIRQL OldIrql; + KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql); + IsConnectedClient = GlobalData.PyasProcess == Process; + KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql); + if (IsConnectedClient) return TRUE; + + PUNICODE_STRING ImagePath = NULL; + NTSTATUS Status = SeLocateProcessImageName(Process, &ImagePath); + if (!NT_SUCCESS(Status) || !ImagePath || !ImagePath->Buffer || ImagePath->Length == 0) { + if (ImagePath) ExFreePool(ImagePath); + return TRUE; + } + + USHORT CharacterCount = ImagePath->Length / sizeof(WCHAR); + LONG LastSeparator = -1; + LONG ParentSeparator = -1; + + for (LONG Index = CharacterCount - 1; Index >= 0; Index--) { + if (ImagePath->Buffer[Index] != L'\\') continue; + if (LastSeparator < 0) LastSeparator = Index; + else { + ParentSeparator = Index; + break; + } + } + + BOOLEAN Protected = TRUE; + if (LastSeparator > 0 && ParentSeparator >= 0 && LastSeparator + 1 < CharacterCount) { + UNICODE_STRING DirectoryName = { + (USHORT)((LastSeparator - ParentSeparator - 1) * sizeof(WCHAR)), + (USHORT)((LastSeparator - ParentSeparator - 1) * sizeof(WCHAR)), + ImagePath->Buffer + ParentSeparator + 1 + }; + UNICODE_STRING FileName = { + (USHORT)((CharacterCount - LastSeparator - 1) * sizeof(WCHAR)), + (USHORT)((CharacterCount - LastSeparator - 1) * sizeof(WCHAR)), + ImagePath->Buffer + LastSeparator + 1 + }; + UNICODE_STRING System32Name = RTL_CONSTANT_STRING(L"System32"); + + if (!RtlEqualUnicodeString(&DirectoryName, &System32Name, TRUE)) { + Protected = FALSE; + } + else { + PCWSTR CriticalNames[] = { + L"csrss.exe", + L"smss.exe", + L"wininit.exe", + L"winlogon.exe", + L"services.exe", + L"lsass.exe", + L"svchost.exe", + L"fontdrvhost.exe", + L"conhost.exe" + }; + + Protected = FALSE; + for (ULONG Index = 0; Index < RTL_NUMBER_OF(CriticalNames); Index++) { + UNICODE_STRING CriticalName; + RtlInitUnicodeString(&CriticalName, CriticalNames[Index]); + if (RtlEqualUnicodeString(&FileName, &CriticalName, TRUE)) { + Protected = TRUE; + break; + } + } + } + } + + ExFreePool(ImagePath); + return Protected; +} + +static VOID ReleaseRuleTerminationWork(PRULE_TERMINATION_WORK_ITEM Work) { + if (!Work) return; + if (Work->Process) ObDereferenceObject(Work->Process); + PyasFree(Work); + InterlockedDecrement(&g_PendingTerminations); +} + +static PRULE_TERMINATION_WORK_ITEM PopRuleTerminationWork() { + PRULE_TERMINATION_WORK_ITEM Work = NULL; + + KIRQL OldIrql; + KeAcquireSpinLock(&g_TerminationLock, &OldIrql); + if (!IsListEmpty(&g_TerminationQueue)) { + PLIST_ENTRY Entry = RemoveHeadList(&g_TerminationQueue); + Work = CONTAINING_RECORD(Entry, RULE_TERMINATION_WORK_ITEM, Entry); + } + KeReleaseSpinLock(&g_TerminationLock, OldIrql); + return Work; +} + +static VOID CancelRuleTerminationWork() { + for (;;) { + PRULE_TERMINATION_WORK_ITEM Work = PopRuleTerminationWork(); + if (!Work) break; + ReleaseRuleTerminationWork(Work); + } +} + +static VOID TerminateRuleProcess(PEPROCESS Process) { + if (IsProtectedTerminationTarget(Process)) return; + + HANDLE ProcessHandle = NULL; + NTSTATUS Status = ObOpenObjectByPointer( + Process, + OBJ_KERNEL_HANDLE, + NULL, + PROCESS_TERMINATE, + *PsProcessType, + KernelMode, + &ProcessHandle + ); + + if (!NT_SUCCESS(Status)) return; + ZwTerminateProcess(ProcessHandle, STATUS_ACCESS_DENIED); + ZwClose(ProcessHandle); +} + +static VOID RuleTerminationWorkerThread(PVOID Parameter) { + UNREFERENCED_PARAMETER(Parameter); + + for (;;) { + KeWaitForSingleObject( + &g_TerminationSemaphore, + Executive, + KernelMode, + FALSE, + NULL + ); + + if (InterlockedCompareExchange(&g_TerminationStopping, 0, 0) != 0) { + CancelRuleTerminationWork(); + break; + } + + PRULE_TERMINATION_WORK_ITEM Work = PopRuleTerminationWork(); + if (!Work) continue; + TerminateRuleProcess(Work->Process); + ReleaseRuleTerminationWork(Work); + } + + PsTerminateSystemThread(STATUS_SUCCESS); +} + +static NTSTATUS StartRuleTerminationWorker() { + if (g_TerminationWorkerHandle) return STATUS_SUCCESS; + + OBJECT_ATTRIBUTES ObjectAttributes; + InitializeObjectAttributes( + &ObjectAttributes, + NULL, + OBJ_KERNEL_HANDLE, + NULL, + NULL + ); + + return PsCreateSystemThread( + &g_TerminationWorkerHandle, + SYNCHRONIZE, + &ObjectAttributes, + NULL, + NULL, + RuleTerminationWorkerThread, + NULL + ); +} + +static VOID StopRuleTerminationWorker() { + CancelRuleTerminationWork(); + if (!g_TerminationWorkerHandle) return; + + KeReleaseSemaphore(&g_TerminationSemaphore, IO_NO_INCREMENT, 1, FALSE); + ZwWaitForSingleObject(g_TerminationWorkerHandle, FALSE, NULL); + ZwClose(g_TerminationWorkerHandle); + g_TerminationWorkerHandle = NULL; +} + +VOID QueueRuleProcessTermination(HANDLE ProcessId, BOOLEAN Kill) { + if (!Kill || !ProcessId || (ULONG_PTR)ProcessId <= 4) return; + if (KeGetCurrentIrql() > APC_LEVEL) return; + if (InterlockedCompareExchange(&g_TerminationStopping, 0, 0) != 0) return; + + PEPROCESS Process = NULL; + NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &Process); + if (!NT_SUCCESS(Status)) return; + + PRULE_TERMINATION_WORK_ITEM Work = (PRULE_TERMINATION_WORK_ITEM)PyasAllocate(sizeof(RULE_TERMINATION_WORK_ITEM)); + if (!Work) { + ObDereferenceObject(Process); + return; + } + + RtlZeroMemory(Work, sizeof(*Work)); + Work->Process = Process; + + LONG Pending = InterlockedIncrement(&g_PendingTerminations); + if (Pending > MAX_PENDING_TERMINATIONS) { + ReleaseRuleTerminationWork(Work); + return; + } + + KIRQL OldIrql; + KeAcquireSpinLock(&g_TerminationLock, &OldIrql); + if (InterlockedCompareExchange(&g_TerminationStopping, 0, 0) != 0) { + KeReleaseSpinLock(&g_TerminationLock, OldIrql); + ReleaseRuleTerminationWork(Work); + return; + } + + InsertTailList(&g_TerminationQueue, &Work->Entry); + KeReleaseSpinLock(&g_TerminationLock, OldIrql); + KeReleaseSemaphore(&g_TerminationSemaphore, IO_NO_INCREMENT, 1, FALSE); +} + static VOID FreeList(PRULE_NODE* Head) { PRULE_NODE Current = *Head; while (Current) { @@ -59,6 +291,29 @@ static VOID FreeList(PRULE_NODE* Head) { *Head = NULL; } +static VOID FreeDynamicRule(PDYNAMIC_RULE Rule) { + if (!Rule) return; + + FreeList(&Rule->Initiator); + FreeList(&Rule->InitiatorExclude); + FreeList(&Rule->InitiatorParent); + FreeList(&Rule->InitiatorParentExclude); + FreeList(&Rule->InitiatorProcessTree); + FreeList(&Rule->InitiatorProcessTreeExclude); + FreeList(&Rule->Target); + FreeList(&Rule->TargetExclude); + FreeList(&Rule->TargetProcessTree); + FreeList(&Rule->TargetProcessTreeExclude); + FreeList(&Rule->Creator); + FreeList(&Rule->CreatorExclude); + FreeList(&Rule->Parent); + FreeList(&Rule->ParentExclude); + FreeList(&Rule->CommandLine); + FreeList(&Rule->CommandLineExclude); + FreeList(&Rule->Extensions); + PyasFree(Rule); +} + static VOID AddRule(PRULE_NODE* Head, PUNICODE_STRING RuleStr) { if (!RuleStr || !RuleStr->Buffer) return; PRULE_NODE Check = *Head; @@ -241,198 +496,569 @@ static VOID SkipWhitespace(PCHAR* Ptr, PCHAR End) { } } -static ULONG ParseInt(PCHAR* Ptr, PCHAR End) { - ULONG val = 0; +static ULONGLONG ParseUnsigned64(PCHAR* Ptr, PCHAR End) { + ULONGLONG Value = 0; + ULONG Base = 10; + SkipWhitespace(Ptr, End); - while (*Ptr < End && **Ptr >= '0' && **Ptr <= '9') { - val = val * 10 + (**Ptr - '0'); + + if (*Ptr + 2 <= End && (*Ptr)[0] == '0' && ((*Ptr)[1] == 'x' || (*Ptr)[1] == 'X')) { + Base = 16; + *Ptr += 2; + } + + while (*Ptr < End) { + CHAR Character = **Ptr; + ULONG Digit = 0; + + if (Character >= '0' && Character <= '9') Digit = Character - '0'; + else if (Base == 16 && Character >= 'a' && Character <= 'f') Digit = Character - 'a' + 10; + else if (Base == 16 && Character >= 'A' && Character <= 'F') Digit = Character - 'A' + 10; + else break; + + if (Digit >= Base) break; + Value = Value * Base + Digit; + (*Ptr)++; + } + + return Value; +} + +static ULONG ParseInt(PCHAR* Ptr, PCHAR End) { + ULONGLONG Value = ParseUnsigned64(Ptr, End); + return Value > MAXULONG ? MAXULONG : (ULONG)Value; +} + +static BOOLEAN ParseBoolean(PCHAR* Ptr, PCHAR End, PBOOLEAN Value) { + SkipWhitespace(Ptr, End); + if (!Value) return FALSE; + + if (*Ptr + 4 <= End && RtlCompareMemory(*Ptr, "true", 4) == 4) { + *Value = TRUE; + *Ptr += 4; + return TRUE; + } + + if (*Ptr + 5 <= End && RtlCompareMemory(*Ptr, "false", 5) == 5) { + *Value = FALSE; + *Ptr += 5; + return TRUE; + } + + return FALSE; +} + +static BOOLEAN ParseStringValue(PCHAR* Ptr, PCHAR End, PCHAR Buffer, ULONG BufferSize) { + SkipWhitespace(Ptr, End); + if (!Buffer || BufferSize == 0 || *Ptr >= End || **Ptr != '"') return FALSE; + + (*Ptr)++; + ULONG Length = 0; + BOOLEAN Escaped = FALSE; + + while (*Ptr < End) { + CHAR Character = **Ptr; + if (Character == '"' && !Escaped) break; + + if (Length + 1 < BufferSize) { + Buffer[Length++] = Character; + } + + if (Character == '\\' && !Escaped) Escaped = TRUE; + else Escaped = FALSE; + (*Ptr)++; + } + + if (*Ptr >= End || **Ptr != '"') return FALSE; + (*Ptr)++; + Buffer[Length] = '\0'; + return TRUE; +} + +static VOID SkipJsonValue(PCHAR* Ptr, PCHAR End) { + SkipWhitespace(Ptr, End); + if (*Ptr >= End) return; + + if (**Ptr == '"') { + CHAR Buffer[2] = { 0 }; + ParseStringValue(Ptr, End, Buffer, sizeof(Buffer)); + return; + } + + if (**Ptr == '[' || **Ptr == '{') { + CHAR Open = **Ptr; + CHAR Close = Open == '[' ? ']' : '}'; + LONG Depth = 0; + BOOLEAN InString = FALSE; + BOOLEAN Escaped = FALSE; + + while (*Ptr < End) { + CHAR Character = **Ptr; + + if (InString) { + if (Character == '"' && !Escaped) InString = FALSE; + if (Character == '\\' && !Escaped) Escaped = TRUE; + else Escaped = FALSE; + (*Ptr)++; + continue; + } + + if (Character == '"') InString = TRUE; + else if (Character == Open) Depth++; + else if (Character == Close) { + Depth--; + (*Ptr)++; + if (Depth == 0) return; + continue; + } + + (*Ptr)++; + } + return; + } + + while (*Ptr < End && **Ptr != ',' && **Ptr != '}') { (*Ptr)++; } - return val; } static VOID ParseStringArray(PCHAR* Ptr, PCHAR End, PRULE_NODE* Head) { SkipWhitespace(Ptr, End); if (*Ptr >= End || **Ptr != '[') return; (*Ptr)++; + while (*Ptr < End) { SkipWhitespace(Ptr, End); + if (*Ptr >= End || **Ptr == ']') { if (*Ptr < End) (*Ptr)++; return; } - if (**Ptr == '"') { - PCHAR StartQuote = ++(*Ptr); - BOOLEAN Escaped = FALSE; - while (*Ptr < End) { - if (**Ptr == '"' && !Escaped) break; - Escaped = (**Ptr == '\\' && !Escaped); - (*Ptr)++; - } - if (*Ptr < End && **Ptr == '"') { - ULONG UTF8Len = (ULONG)(*Ptr - StartQuote); - if (UTF8Len > 0) { - ULONG WideSize = 0; - RtlUTF8ToUnicodeN(NULL, 0, &WideSize, StartQuote, UTF8Len); - if (WideSize > 0) { - PWCHAR WideBuffer = (PWCHAR)PyasAllocate(WideSize + sizeof(WCHAR)); - if (WideBuffer) { - ULONG ResultSize = 0; - RtlUTF8ToUnicodeN(WideBuffer, WideSize, &ResultSize, StartQuote, UTF8Len); - ULONG FinalSize = ProcessJsonUnescape(WideBuffer, ResultSize / sizeof(WCHAR)); - UNICODE_STRING Us; - Us.Buffer = WideBuffer; - Us.Length = (USHORT)FinalSize; - Us.MaximumLength = (USHORT)(WideSize + sizeof(WCHAR)); - AddRule(Head, &Us); - PyasFree(WideBuffer); - } - } - } - (*Ptr)++; - } + + if (**Ptr != '"') { + (*Ptr)++; + continue; } - else { + + PCHAR StartQuote = ++(*Ptr); + BOOLEAN Escaped = FALSE; + + while (*Ptr < End) { + CHAR Character = **Ptr; + if (Character == '"' && !Escaped) break; + if (Character == '\\' && !Escaped) Escaped = TRUE; + else Escaped = FALSE; (*Ptr)++; } + + if (*Ptr >= End || **Ptr != '"') return; + + ULONG Utf8Length = (ULONG)(*Ptr - StartQuote); + if (Utf8Length > 0) { + ULONG WideSize = 0; + RtlUTF8ToUnicodeN(NULL, 0, &WideSize, StartQuote, Utf8Length); + + if (WideSize > 0 && WideSize <= 0xFFFC) { + PWCHAR WideBuffer = (PWCHAR)PyasAllocate(WideSize + sizeof(WCHAR)); + + if (WideBuffer) { + ULONG ResultSize = 0; + RtlZeroMemory(WideBuffer, WideSize + sizeof(WCHAR)); + RtlUTF8ToUnicodeN(WideBuffer, WideSize, &ResultSize, StartQuote, Utf8Length); + ULONG FinalSize = ProcessJsonUnescape(WideBuffer, ResultSize / sizeof(WCHAR)); + + UNICODE_STRING Value; + Value.Buffer = WideBuffer; + Value.Length = (USHORT)FinalSize; + Value.MaximumLength = (USHORT)(WideSize + sizeof(WCHAR)); + AddRule(Head, &Value); + PyasFree(WideBuffer); + } + } + } + + (*Ptr)++; } } -static VOID ParseOperationsArray(PCHAR* Ptr, PCHAR End, PULONG Operations) { - PRULE_NODE TempHead = NULL; - ParseStringArray(Ptr, End, &TempHead); - PRULE_NODE Node = TempHead; - while (Node) { - UNICODE_STRING WriteStr = RTL_CONSTANT_STRING(L"Write"); - UNICODE_STRING DeleteStr = RTL_CONSTANT_STRING(L"Delete"); - UNICODE_STRING CreateStr = RTL_CONSTANT_STRING(L"Create"); - UNICODE_STRING ExecuteStr = RTL_CONSTANT_STRING(L"Execute"); - UNICODE_STRING RenameStr = RTL_CONSTANT_STRING(L"Rename"); - UNICODE_STRING IoctlStr = RTL_CONSTANT_STRING(L"Ioctl"); - UNICODE_STRING VmReadStr = RTL_CONSTANT_STRING(L"VmRead"); - UNICODE_STRING VmWriteStr = RTL_CONSTANT_STRING(L"VmWrite"); - UNICODE_STRING TerminateStr = RTL_CONSTANT_STRING(L"Terminate"); - UNICODE_STRING SuspendResumeStr = RTL_CONSTANT_STRING(L"SuspendResume"); - UNICODE_STRING DuplicateHandleStr = RTL_CONSTANT_STRING(L"DuplicateHandle"); - UNICODE_STRING SetInformationStr = RTL_CONSTANT_STRING(L"SetInformation"); +static BOOLEAN RuleNodeEquals(PRULE_NODE Node, PCWSTR Value) { + UNICODE_STRING Target; + RtlInitUnicodeString(&Target, Value); + return Node && RtlEqualUnicodeString(&Node->Pattern, &Target, TRUE); +} - if (RtlEqualUnicodeString(&Node->Pattern, &WriteStr, TRUE)) *Operations |= OP_WRITE; - else if (RtlEqualUnicodeString(&Node->Pattern, &DeleteStr, TRUE)) *Operations |= OP_DELETE; - else if (RtlEqualUnicodeString(&Node->Pattern, &CreateStr, TRUE)) *Operations |= OP_CREATE; - else if (RtlEqualUnicodeString(&Node->Pattern, &ExecuteStr, TRUE)) *Operations |= OP_EXECUTE; - else if (RtlEqualUnicodeString(&Node->Pattern, &RenameStr, TRUE)) *Operations |= OP_RENAME; - else if (RtlEqualUnicodeString(&Node->Pattern, &IoctlStr, TRUE)) *Operations |= OP_IOCTL; - else if (RtlEqualUnicodeString(&Node->Pattern, &VmReadStr, TRUE)) *Operations |= OP_VM_READ; - else if (RtlEqualUnicodeString(&Node->Pattern, &VmWriteStr, TRUE)) *Operations |= OP_VM_WRITE; - else if (RtlEqualUnicodeString(&Node->Pattern, &TerminateStr, TRUE)) *Operations |= OP_TERMINATE; - else if (RtlEqualUnicodeString(&Node->Pattern, &SuspendResumeStr, TRUE)) *Operations |= OP_SUSPEND_RESUME; - else if (RtlEqualUnicodeString(&Node->Pattern, &DuplicateHandleStr, TRUE)) *Operations |= OP_DUP_HANDLE; - else if (RtlEqualUnicodeString(&Node->Pattern, &SetInformationStr, TRUE)) *Operations |= OP_SET_INFORMATION; +static BOOLEAN ParseHandleTypesArray(PCHAR* Ptr, PCHAR End, PULONG HandleTypes) { + PRULE_NODE Values = NULL; + ParseStringArray(Ptr, End, &Values); - Node = Node->Next; + BOOLEAN Valid = Values != NULL; + for (PRULE_NODE Node = Values; Node; Node = Node->Next) { + if (RuleNodeEquals(Node, L"Create")) *HandleTypes |= PYAS_HANDLE_CREATE; + else if (RuleNodeEquals(Node, L"Duplicate")) *HandleTypes |= PYAS_HANDLE_DUPLICATE; + else Valid = FALSE; } - FreeList(&TempHead); + + FreeList(&Values); + return Valid; +} + +static BOOLEAN ParseObjectTypesArray(PCHAR* Ptr, PCHAR End, PULONG ObjectTypes) { + PRULE_NODE Values = NULL; + ParseStringArray(Ptr, End, &Values); + + BOOLEAN Valid = Values != NULL; + for (PRULE_NODE Node = Values; Node; Node = Node->Next) { + if (RuleNodeEquals(Node, L"Process")) *ObjectTypes |= PYAS_OBJECT_PROCESS; + else if (RuleNodeEquals(Node, L"Thread")) *ObjectTypes |= PYAS_OBJECT_THREAD; + else Valid = FALSE; + } + + FreeList(&Values); + return Valid; +} + +static BOOLEAN ParseThreadMemoryTypesArray(PCHAR* Ptr, PCHAR End, PULONG MemoryTypes) { + PRULE_NODE Values = NULL; + ParseStringArray(Ptr, End, &Values); + + BOOLEAN Valid = Values != NULL; + for (PRULE_NODE Node = Values; Node; Node = Node->Next) { + if (RuleNodeEquals(Node, L"Private")) *MemoryTypes |= PYAS_MEMORY_PRIVATE; + else if (RuleNodeEquals(Node, L"Mapped")) *MemoryTypes |= PYAS_MEMORY_MAPPED; + else if (RuleNodeEquals(Node, L"Image")) *MemoryTypes |= PYAS_MEMORY_IMAGE; + else Valid = FALSE; + } + + FreeList(&Values); + return Valid; +} + +static BOOLEAN ParseThreadMemoryProtectionsArray(PCHAR* Ptr, PCHAR End, PULONG Protections) { + PRULE_NODE Values = NULL; + ParseStringArray(Ptr, End, &Values); + + BOOLEAN Valid = Values != NULL; + for (PRULE_NODE Node = Values; Node; Node = Node->Next) { + if (RuleNodeEquals(Node, L"Execute")) *Protections |= PYAS_PROTECT_EXECUTE; + else if (RuleNodeEquals(Node, L"ExecuteWrite")) *Protections |= PYAS_PROTECT_EXECUTE_WRITE; + else Valid = FALSE; + } + + FreeList(&Values); + return Valid; +} + +static RULE_TRI_STATE ParseTriState(PCHAR* Ptr, PCHAR End, PBOOLEAN Parsed) { + if (Parsed) *Parsed = FALSE; + + BOOLEAN Value = FALSE; + if (!ParseBoolean(Ptr, End, &Value)) return RuleTriAny; + if (Parsed) *Parsed = TRUE; + return Value ? RuleTriTrue : RuleTriFalse; +} + +static RULE_OPERATION_MATCH ParseOperationMatch(PCHAR* Ptr, PCHAR End, PBOOLEAN Parsed) { + if (Parsed) *Parsed = FALSE; + + CHAR Value[16] = { 0 }; + if (!ParseStringValue(Ptr, End, Value, sizeof(Value))) return RuleOperationAny; + if (strcmp(Value, "Any") == 0) { + if (Parsed) *Parsed = TRUE; + return RuleOperationAny; + } + if (strcmp(Value, "All") == 0) { + if (Parsed) *Parsed = TRUE; + return RuleOperationAll; + } + return RuleOperationAny; +} + +static BOOLEAN ParseOperationsArray(PCHAR* Ptr, PCHAR End, PULONG Operations) { + PRULE_NODE Values = NULL; + ParseStringArray(Ptr, End, &Values); + + BOOLEAN Valid = Values != NULL; + for (PRULE_NODE Node = Values; Node; Node = Node->Next) { + if (RuleNodeEquals(Node, L"Write")) *Operations |= OP_WRITE; + else if (RuleNodeEquals(Node, L"Delete")) *Operations |= OP_DELETE; + else if (RuleNodeEquals(Node, L"Create")) *Operations |= OP_CREATE; + else if (RuleNodeEquals(Node, L"Execute")) *Operations |= OP_EXECUTE; + else if (RuleNodeEquals(Node, L"Rename")) *Operations |= OP_RENAME; + else if (RuleNodeEquals(Node, L"Ioctl")) *Operations |= OP_IOCTL; + else if (RuleNodeEquals(Node, L"VmRead")) *Operations |= OP_VM_READ; + else if (RuleNodeEquals(Node, L"VmWrite")) { + *Operations |= OP_VM_WRITE | OP_VM_OPERATION | OP_CREATE_THREAD | OP_THREAD_SET_CONTEXT; + } + else if (RuleNodeEquals(Node, L"WriteMemory")) *Operations |= OP_VM_WRITE; + else if (RuleNodeEquals(Node, L"VmOperation")) *Operations |= OP_VM_OPERATION; + else if (RuleNodeEquals(Node, L"CreateThread") || RuleNodeEquals(Node, L"CreateRemoteThread")) { + *Operations |= OP_CREATE_THREAD; + } + else if (RuleNodeEquals(Node, L"SetThreadContext")) *Operations |= OP_THREAD_SET_CONTEXT; + else if (RuleNodeEquals(Node, L"SetThreadToken")) *Operations |= OP_THREAD_SET_TOKEN; + else if (RuleNodeEquals(Node, L"Terminate")) *Operations |= OP_TERMINATE; + else if (RuleNodeEquals(Node, L"SuspendResume")) *Operations |= OP_SUSPEND_RESUME; + else if (RuleNodeEquals(Node, L"DuplicateHandle")) *Operations |= OP_DUP_HANDLE; + else if (RuleNodeEquals(Node, L"SetInformation")) *Operations |= OP_SET_INFORMATION; + else if (RuleNodeEquals(Node, L"CreateProcess")) *Operations |= OP_CREATE_PROCESS; + else if (RuleNodeEquals(Node, L"ImageLoad")) *Operations |= OP_IMAGE_LOAD; + else if (RuleNodeEquals(Node, L"Impersonate")) *Operations |= OP_IMPERSONATE; + else Valid = FALSE; + } + + FreeList(&Values); + return Valid; +} + +static BOOLEAN KeyEquals(PCHAR Key, ULONG KeyLength, PCSTR Expected) { + SIZE_T ExpectedLength = strlen(Expected); + return KeyLength == ExpectedLength && RtlCompareMemory(Key, Expected, ExpectedLength) == ExpectedLength; +} + +static RULE_CATEGORY ParseCategoryValue(PCHAR* Ptr, PCHAR End) { + CHAR Value[24] = { 0 }; + if (!ParseStringValue(Ptr, End, Value, sizeof(Value))) return RuleCategoryUnknown; + + if (strcmp(Value, "Process") == 0) return RuleCategoryProcess; + if (strcmp(Value, "File") == 0) return RuleCategoryFile; + if (strcmp(Value, "Registry") == 0) return RuleCategoryRegistry; + if (strcmp(Value, "Device") == 0) return RuleCategoryDevice; + if (strcmp(Value, "Memory") == 0) return RuleCategoryMemory; + if (strcmp(Value, "Thread") == 0) return RuleCategoryThread; + return RuleCategoryUnknown; +} + +static BOOLEAN IsValidDynamicRule(PDYNAMIC_RULE Rule) { + if (!Rule || Rule->Invalid) return FALSE; + if (Rule->Code == 0 || Rule->Category == RuleCategoryUnknown || Rule->Operations == 0) return FALSE; + + ULONG AllowedOperations = 0; + + if (Rule->Category == RuleCategoryProcess) { + AllowedOperations = + OP_EXECUTE | + OP_TERMINATE | + OP_SUSPEND_RESUME | + OP_DUP_HANDLE | + OP_SET_INFORMATION | + OP_THREAD_SET_TOKEN | + OP_CREATE_PROCESS | + OP_IMAGE_LOAD | + OP_IMPERSONATE; + } + else if (Rule->Category == RuleCategoryMemory) { + AllowedOperations = + OP_VM_READ | + OP_VM_WRITE | + OP_VM_OPERATION | + OP_CREATE_THREAD | + OP_THREAD_SET_CONTEXT; + } + else if (Rule->Category == RuleCategoryThread) { + AllowedOperations = OP_EXECUTE; + } + else if (Rule->Category == RuleCategoryFile) { + AllowedOperations = OP_WRITE | OP_DELETE | OP_CREATE | OP_EXECUTE | OP_RENAME; + } + else if (Rule->Category == RuleCategoryRegistry) { + AllowedOperations = OP_WRITE | OP_DELETE | OP_CREATE; + } + else if (Rule->Category == RuleCategoryDevice) { + AllowedOperations = OP_IOCTL; + } + + if ((Rule->Operations & ~AllowedOperations) != 0) return FALSE; + if (Rule->Threshold > 0 && Rule->TimeWindow == 0) Rule->TimeWindow = 1000; + if (Rule->MaximumRiskScore > 0 && Rule->MaximumRiskScore < Rule->MinimumRiskScore) return FALSE; + if (Rule->MaximumRegionSize > 0 && Rule->MaximumRegionSize < Rule->MinimumRegionSize) return FALSE; + return TRUE; } static VOID ParseDynamicRules(PCHAR JsonContent, ULONG ContentLength) { + if (!JsonContent || ContentLength == 0) return; + PCHAR Ptr = JsonContent; PCHAR End = JsonContent + ContentLength; + const CHAR DynamicRulesKey[] = "\"DynamicRules\""; + SIZE_T DynamicRulesKeyLength = sizeof(DynamicRulesKey) - 1; - const char* TargetKey = "\"DynamicRules\""; - SIZE_T KeyLen = strlen(TargetKey); + while (Ptr + DynamicRulesKeyLength <= End) { + if (RtlCompareMemory(Ptr, DynamicRulesKey, DynamicRulesKeyLength) != DynamicRulesKeyLength) { + Ptr++; + continue; + } - while (Ptr < End - KeyLen) { - if (RtlCompareMemory(Ptr, TargetKey, KeyLen) == KeyLen) { - Ptr += KeyLen; + Ptr += DynamicRulesKeyLength; + SkipWhitespace(&Ptr, End); + if (Ptr >= End || *Ptr != ':') return; + Ptr++; + SkipWhitespace(&Ptr, End); + if (Ptr >= End || *Ptr != '[') return; + Ptr++; + + while (Ptr < End) { SkipWhitespace(&Ptr, End); - if (Ptr >= End || *Ptr != ':') continue; - Ptr++; - SkipWhitespace(&Ptr, End); - if (Ptr >= End || *Ptr != '[') continue; + if (Ptr >= End || *Ptr == ']') return; + + if (*Ptr != '{') { + Ptr++; + continue; + } + Ptr++; + PDYNAMIC_RULE Rule = (PDYNAMIC_RULE)PyasAllocate(sizeof(DYNAMIC_RULE)); + if (!Rule) { + SkipJsonValue(&Ptr, End); + continue; + } + + RtlZeroMemory(Rule, sizeof(*Rule)); + Rule->Category = RuleCategoryUnknown; + Rule->OperationMatch = RuleOperationAny; while (Ptr < End) { SkipWhitespace(&Ptr, End); - if (Ptr >= End || *Ptr == ']') break; - if (*Ptr == '{') { - Ptr++; - PDYNAMIC_RULE Rule = (PDYNAMIC_RULE)PyasAllocate(sizeof(DYNAMIC_RULE)); - if (Rule) { - RtlZeroMemory(Rule, sizeof(DYNAMIC_RULE)); - while (Ptr < End) { - SkipWhitespace(&Ptr, End); - if (Ptr >= End || *Ptr == '}') { - if (Ptr < End) Ptr++; - break; - } - if (*Ptr == '"') { - PCHAR KeyStart = ++Ptr; - while (Ptr < End && *Ptr != '"') Ptr++; - ULONG KeyLenStr = (ULONG)(Ptr - KeyStart); - if (Ptr < End) Ptr++; - SkipWhitespace(&Ptr, End); - if (Ptr < End && *Ptr == ':') Ptr++; - SkipWhitespace(&Ptr, End); - if (KeyLenStr == 4 && RtlCompareMemory(KeyStart, "Code", 4) == 4) Rule->Code = ParseInt(&Ptr, End); - else if (KeyLenStr == 9 && RtlCompareMemory(KeyStart, "Threshold", 9) == 9) Rule->Threshold = ParseInt(&Ptr, End); - else if (KeyLenStr == 10 && RtlCompareMemory(KeyStart, "TimeWindow", 10) == 10) Rule->TimeWindow = ParseInt(&Ptr, End); - else if (KeyLenStr == 8 && RtlCompareMemory(KeyStart, "Category", 8) == 8) { - if (*Ptr == '"') { - Ptr++; - if (RtlCompareMemory(Ptr, "Process", 7) == 7) Rule->Category = RuleCategoryProcess; - else if (RtlCompareMemory(Ptr, "File", 4) == 4) Rule->Category = RuleCategoryFile; - else if (RtlCompareMemory(Ptr, "Registry", 8) == 8) Rule->Category = RuleCategoryRegistry; - else if (RtlCompareMemory(Ptr, "Device", 6) == 6) Rule->Category = RuleCategoryDevice; - else if (RtlCompareMemory(Ptr, "Memory", 6) == 6) Rule->Category = RuleCategoryMemory; - else if (RtlCompareMemory(Ptr, "Thread", 6) == 6) Rule->Category = RuleCategoryThread; - while (Ptr < End && *Ptr != '"') Ptr++; - if (Ptr < End) Ptr++; - } - } - else if (KeyLenStr == 9 && RtlCompareMemory(KeyStart, "Initiator", 9) == 9) ParseStringArray(&Ptr, End, &Rule->Initiator); - else if (KeyLenStr == 16 && RtlCompareMemory(KeyStart, "InitiatorExclude", 16) == 16) ParseStringArray(&Ptr, End, &Rule->InitiatorExclude); - else if (KeyLenStr == 6 && RtlCompareMemory(KeyStart, "Target", 6) == 6) ParseStringArray(&Ptr, End, &Rule->Target); - else if (KeyLenStr == 13 && RtlCompareMemory(KeyStart, "TargetExclude", 13) == 13) ParseStringArray(&Ptr, End, &Rule->TargetExclude); - else if (KeyLenStr == 11 && RtlCompareMemory(KeyStart, "CommandLine", 11) == 11) ParseStringArray(&Ptr, End, &Rule->CommandLine); - else if (KeyLenStr == 10 && RtlCompareMemory(KeyStart, "Extensions", 10) == 10) ParseStringArray(&Ptr, End, &Rule->Extensions); - else if (KeyLenStr == 10 && RtlCompareMemory(KeyStart, "Operations", 10) == 10) ParseOperationsArray(&Ptr, End, &Rule->Operations); - else { - BOOLEAN InQuote = FALSE, InArray = FALSE; - while (Ptr < End) { - if (*Ptr == '"' && *(Ptr - 1) != '\\') InQuote = !InQuote; - else if (!InQuote && *Ptr == '[') InArray = TRUE; - else if (!InQuote && *Ptr == ']') InArray = FALSE; - else if (!InQuote && !InArray && (*Ptr == ',' || *Ptr == '}')) break; - Ptr++; - } - } - } - else Ptr++; - } - Rule->Next = g_DynamicRules; - g_DynamicRules = Rule; - } + if (Ptr >= End) break; + if (*Ptr == '}') { + Ptr++; + break; } - else Ptr++; + + if (*Ptr != '"') { + Ptr++; + continue; + } + + PCHAR KeyStart = ++Ptr; + while (Ptr < End && *Ptr != '"') Ptr++; + if (Ptr >= End) break; + + ULONG KeyLength = (ULONG)(Ptr - KeyStart); + Ptr++; + SkipWhitespace(&Ptr, End); + + if (Ptr >= End || *Ptr != ':') { + SkipJsonValue(&Ptr, End); + continue; + } + + Ptr++; + SkipWhitespace(&Ptr, End); + + if (KeyEquals(KeyStart, KeyLength, "Code")) Rule->Code = ParseInt(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "Kill")) { + if (!ParseBoolean(&Ptr, End, &Rule->Kill)) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "Priority")) Rule->Priority = ParseInt(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "Threshold")) Rule->Threshold = ParseInt(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "TimeWindow")) Rule->TimeWindow = ParseInt(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "MinimumRiskScore")) Rule->MinimumRiskScore = ParseInt(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "MaximumRiskScore")) Rule->MaximumRiskScore = ParseInt(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "MinimumRegionSize")) Rule->MinimumRegionSize = ParseUnsigned64(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "MaximumRegionSize")) Rule->MaximumRegionSize = ParseUnsigned64(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "Category")) Rule->Category = ParseCategoryValue(&Ptr, End); + else if (KeyEquals(KeyStart, KeyLength, "OperationMatch")) { + BOOLEAN Parsed = FALSE; + Rule->OperationMatch = ParseOperationMatch(&Ptr, End, &Parsed); + if (!Parsed) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "ParentMismatch")) { + BOOLEAN Parsed = FALSE; + Rule->ParentMismatch = ParseTriState(&Ptr, End, &Parsed); + if (!Parsed) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "FileOpenNameAvailable")) { + BOOLEAN Parsed = FALSE; + Rule->FileOpenNameAvailable = ParseTriState(&Ptr, End, &Parsed); + if (!Parsed) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "SubsystemProcess")) { + BOOLEAN Parsed = FALSE; + Rule->SubsystemProcess = ParseTriState(&Ptr, End, &Parsed); + if (!Parsed) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "Initiator")) ParseStringArray(&Ptr, End, &Rule->Initiator); + else if (KeyEquals(KeyStart, KeyLength, "InitiatorExclude")) ParseStringArray(&Ptr, End, &Rule->InitiatorExclude); + else if (KeyEquals(KeyStart, KeyLength, "InitiatorParent")) ParseStringArray(&Ptr, End, &Rule->InitiatorParent); + else if (KeyEquals(KeyStart, KeyLength, "InitiatorParentExclude")) ParseStringArray(&Ptr, End, &Rule->InitiatorParentExclude); + else if (KeyEquals(KeyStart, KeyLength, "InitiatorProcessTree")) ParseStringArray(&Ptr, End, &Rule->InitiatorProcessTree); + else if (KeyEquals(KeyStart, KeyLength, "InitiatorProcessTreeExclude")) ParseStringArray(&Ptr, End, &Rule->InitiatorProcessTreeExclude); + else if (KeyEquals(KeyStart, KeyLength, "Target")) ParseStringArray(&Ptr, End, &Rule->Target); + else if (KeyEquals(KeyStart, KeyLength, "TargetExclude")) ParseStringArray(&Ptr, End, &Rule->TargetExclude); + else if (KeyEquals(KeyStart, KeyLength, "TargetProcessTree")) ParseStringArray(&Ptr, End, &Rule->TargetProcessTree); + else if (KeyEquals(KeyStart, KeyLength, "TargetProcessTreeExclude")) ParseStringArray(&Ptr, End, &Rule->TargetProcessTreeExclude); + else if (KeyEquals(KeyStart, KeyLength, "Creator")) ParseStringArray(&Ptr, End, &Rule->Creator); + else if (KeyEquals(KeyStart, KeyLength, "CreatorExclude")) ParseStringArray(&Ptr, End, &Rule->CreatorExclude); + else if (KeyEquals(KeyStart, KeyLength, "Parent")) ParseStringArray(&Ptr, End, &Rule->Parent); + else if (KeyEquals(KeyStart, KeyLength, "ParentExclude")) ParseStringArray(&Ptr, End, &Rule->ParentExclude); + else if (KeyEquals(KeyStart, KeyLength, "CommandLine")) ParseStringArray(&Ptr, End, &Rule->CommandLine); + else if (KeyEquals(KeyStart, KeyLength, "CommandLineExclude")) ParseStringArray(&Ptr, End, &Rule->CommandLineExclude); + else if (KeyEquals(KeyStart, KeyLength, "Extensions")) ParseStringArray(&Ptr, End, &Rule->Extensions); + else if (KeyEquals(KeyStart, KeyLength, "Operations")) { + if (!ParseOperationsArray(&Ptr, End, &Rule->Operations)) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "HandleTypes")) { + if (!ParseHandleTypesArray(&Ptr, End, &Rule->HandleTypes)) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "ObjectTypes")) { + if (!ParseObjectTypesArray(&Ptr, End, &Rule->ObjectTypes)) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "ThreadMemoryTypes")) { + if (!ParseThreadMemoryTypesArray(&Ptr, End, &Rule->ThreadMemoryTypes)) Rule->Invalid = TRUE; + } + else if (KeyEquals(KeyStart, KeyLength, "ThreadMemoryProtections")) { + if (!ParseThreadMemoryProtectionsArray(&Ptr, End, &Rule->ThreadMemoryProtections)) Rule->Invalid = TRUE; + } + else SkipJsonValue(&Ptr, End); + } + + if (IsValidDynamicRule(Rule)) { + Rule->Next = g_DynamicRules; + g_DynamicRules = Rule; + } + else { + FreeDynamicRule(Rule); } - break; } - Ptr++; + + return; } } -VOID InitializeRulesEngine() { - ExInitializeResourceLite(&RuleLock); +NTSTATUS InitializeRulesEngine() { + if (g_CacheInitialized) return STATUS_SUCCESS; + + NTSTATUS Status = ExInitializeResourceLite(&RuleLock); + if (!NT_SUCCESS(Status)) return Status; + KeInitializeSpinLock(&TrustCacheLock); + InitializeListHead(&g_TerminationQueue); + KeInitializeSpinLock(&g_TerminationLock); + KeInitializeSemaphore(&g_TerminationSemaphore, 0, MAXLONG); + InterlockedExchange(&g_PendingTerminations, 0); + InterlockedExchange(&g_TerminationStopping, 0); + + Status = StartRuleTerminationWorker(); + if (!NT_SUCCESS(Status)) { + InterlockedExchange(&g_TerminationStopping, 1); + ExDeleteResourceLite(&RuleLock); + return Status; + } + RtlZeroMemory(TrustCache, sizeof(TrustCache)); RtlZeroMemory(BehaviorTrackers, sizeof(BehaviorTrackers)); g_CacheInitialized = TRUE; + return STATUS_SUCCESS; } VOID UninitializeRulesEngine() { - ExDeleteResourceLite(&RuleLock); + if (!g_CacheInitialized) return; + + InterlockedExchange(&g_TerminationStopping, 1); + StopRuleTerminationWorker(); g_CacheInitialized = FALSE; + ExDeleteResourceLite(&RuleLock); } VOID AddDynamicWhitelist(PUNICODE_STRING RuleStr) { @@ -562,15 +1188,34 @@ NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath) { FILE_STANDARD_INFORMATION FileInfo = { 0 }; status = ZwQueryInformationFile(FileHandle, &IoStatus, &FileInfo, sizeof(FileInfo), FileStandardInformation); - if (NT_SUCCESS(status) && FileInfo.EndOfFile.LowPart > 0) { - PVOID FileBuffer = PyasAllocate(FileInfo.EndOfFile.LowPart + 1); - if (FileBuffer) { - status = ZwReadFile(FileHandle, NULL, NULL, NULL, &IoStatus, FileBuffer, FileInfo.EndOfFile.LowPart, NULL, NULL); + if (NT_SUCCESS(status) && FileInfo.EndOfFile.HighPart != 0) { + status = STATUS_FILE_TOO_LARGE; + } + else if (NT_SUCCESS(status) && FileInfo.EndOfFile.LowPart > MAX_RULE_FILE_SIZE) { + status = STATUS_FILE_TOO_LARGE; + } + else if (NT_SUCCESS(status) && FileInfo.EndOfFile.LowPart == 0) { + status = STATUS_END_OF_FILE; + } + else if (NT_SUCCESS(status)) { + ULONG FileSize = FileInfo.EndOfFile.LowPart; + PVOID FileBuffer = PyasAllocate((SIZE_T)FileSize + 1); + if (!FileBuffer) { + status = STATUS_INSUFFICIENT_RESOURCES; + } + else { + status = ZwReadFile(FileHandle, NULL, NULL, NULL, &IoStatus, FileBuffer, FileSize, NULL, NULL); + if (NT_SUCCESS(status) && IoStatus.Information != FileSize) { + status = STATUS_END_OF_FILE; + } + if (NT_SUCCESS(status)) { - ((PCHAR)FileBuffer)[FileInfo.EndOfFile.LowPart] = '\0'; + ((PCHAR)FileBuffer)[FileSize] = '\0'; KeEnterCriticalRegion(); ExAcquireResourceExclusiveLite(&RuleLock, TRUE); - ParseDynamicRules((PCHAR)FileBuffer, FileInfo.EndOfFile.LowPart); + PDYNAMIC_RULE PreviousHead = g_DynamicRules; + ParseDynamicRules((PCHAR)FileBuffer, FileSize); + if (g_DynamicRules == PreviousHead) status = STATUS_DATA_ERROR; ExReleaseResourceLite(&RuleLock); KeLeaveCriticalRegion(); } @@ -598,23 +1243,39 @@ NTSTATUS LoadRuleFile(PCUNICODE_STRING FilePath) { FILE_STANDARD_INFORMATION FileInfo = { 0 }; status = ZwQueryInformationFile(FileHandle, &IoStatus, &FileInfo, sizeof(FileInfo), FileStandardInformation); - if (NT_SUCCESS(status) && FileInfo.EndOfFile.LowPart > 0) { - PVOID FileBuffer = PyasAllocate(FileInfo.EndOfFile.LowPart + 1); - if (FileBuffer) { - status = ZwReadFile(FileHandle, NULL, NULL, NULL, &IoStatus, FileBuffer, FileInfo.EndOfFile.LowPart, NULL, NULL); + if (NT_SUCCESS(status) && FileInfo.EndOfFile.HighPart != 0) { + status = STATUS_FILE_TOO_LARGE; + } + else if (NT_SUCCESS(status) && FileInfo.EndOfFile.LowPart > MAX_RULE_FILE_SIZE) { + status = STATUS_FILE_TOO_LARGE; + } + else if (NT_SUCCESS(status) && FileInfo.EndOfFile.LowPart == 0) { + status = STATUS_END_OF_FILE; + } + else if (NT_SUCCESS(status)) { + ULONG FileSize = FileInfo.EndOfFile.LowPart; + PVOID FileBuffer = PyasAllocate((SIZE_T)FileSize + 1); + if (!FileBuffer) { + status = STATUS_INSUFFICIENT_RESOURCES; + } + else { + status = ZwReadFile(FileHandle, NULL, NULL, NULL, &IoStatus, FileBuffer, FileSize, NULL, NULL); + if (NT_SUCCESS(status) && IoStatus.Information != FileSize) { + status = STATUS_END_OF_FILE; + } + if (NT_SUCCESS(status)) { - ((PCHAR)FileBuffer)[FileInfo.EndOfFile.LowPart] = '\0'; + ((PCHAR)FileBuffer)[FileSize] = '\0'; KeEnterCriticalRegion(); ExAcquireResourceExclusiveLite(&RuleLock, TRUE); - ParseDynamicRules((PCHAR)FileBuffer, FileInfo.EndOfFile.LowPart); + PDYNAMIC_RULE PreviousHead = g_DynamicRules; + ParseDynamicRules((PCHAR)FileBuffer, FileSize); + if (g_DynamicRules == PreviousHead) status = STATUS_DATA_ERROR; ExReleaseResourceLite(&RuleLock); KeLeaveCriticalRegion(); } PyasFree(FileBuffer); } - else { - status = STATUS_INSUFFICIENT_RESOURCES; - } } ZwClose(FileHandle); } @@ -624,19 +1285,20 @@ NTSTATUS LoadRuleFile(PCUNICODE_STRING FilePath) { VOID ClearDynamicRules() { KeEnterCriticalRegion(); ExAcquireResourceExclusiveLite(&RuleLock, TRUE); + PDYNAMIC_RULE CurrentRule = g_DynamicRules; while (CurrentRule) { PDYNAMIC_RULE Next = CurrentRule->Next; - FreeList(&CurrentRule->Initiator); - FreeList(&CurrentRule->InitiatorExclude); - FreeList(&CurrentRule->Target); - FreeList(&CurrentRule->TargetExclude); - FreeList(&CurrentRule->CommandLine); - FreeList(&CurrentRule->Extensions); - PyasFree(CurrentRule); + FreeDynamicRule(CurrentRule); CurrentRule = Next; } g_DynamicRules = NULL; + + KIRQL OldIrql; + KeAcquireSpinLock(&GlobalData.TrackerMutex, &OldIrql); + RtlZeroMemory(BehaviorTrackers, sizeof(BehaviorTrackers)); + KeReleaseSpinLock(&GlobalData.TrackerMutex, OldIrql); + ExReleaseResourceLite(&RuleLock); KeLeaveCriticalRegion(); } @@ -662,13 +1324,22 @@ NTSTATUS GetProcessImageName(HANDLE ProcessId, PUNICODE_STRING* ImageName) { } BOOLEAN IsProcessTrusted(HANDLE ProcessId) { - ULONG PyasPid = SafeGetPyasPid(); - if ((ULONG)(ULONG_PTR)ProcessId == PyasPid) return TRUE; if (ProcessId == (HANDLE)4) return TRUE; PEPROCESS Process = NULL; if (!NT_SUCCESS(PsLookupProcessByProcessId(ProcessId, &Process))) return FALSE; + BOOLEAN IsConnectedClient = FALSE; + KIRQL PortIrql; + KeAcquireSpinLock(&GlobalData.PortMutex, &PortIrql); + IsConnectedClient = GlobalData.PyasProcess == Process; + KeReleaseSpinLock(&GlobalData.PortMutex, PortIrql); + + if (IsConnectedClient) { + ObDereferenceObject(Process); + return TRUE; + } + LARGE_INTEGER createTime; createTime.QuadPart = PsGetProcessCreateTimeQuadPart(Process); @@ -695,7 +1366,7 @@ BOOLEAN IsProcessTrusted(HANDLE ProcessId) { } PUNICODE_STRING imageFileName = NULL; - NTSTATUS status = GetProcessImageName(ProcessId, &imageFileName); + NTSTATUS status = SeLocateProcessImageName(Process, &imageFileName); BOOLEAN isTrusted = FALSE; if (NT_SUCCESS(status) && imageFileName && imageFileName->Buffer) { @@ -789,43 +1460,485 @@ static BOOLEAN CheckRuleThreshold(PDYNAMIC_RULE Rule, HANDLE ProcessId) { return Triggered; } -BOOLEAN EvaluateProcessRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, PCUNICODE_STRING CommandLine, PULONG OutCode) { - if (IsProcessTrusted(ProcessId)) return FALSE; - PUNICODE_STRING InitiatorPath = NULL; - GetProcessImageName(ProcessId, &InitiatorPath); +static BOOLEAN TriStateMatches(RULE_TRI_STATE Condition, BOOLEAN Value) { + if (Condition == RuleTriAny) return TRUE; + if (Condition == RuleTriTrue) return Value ? TRUE : FALSE; + return Value ? FALSE : TRUE; +} - BOOLEAN Blocked = FALSE; - KeEnterCriticalRegion(); - ExAcquireResourceSharedLite(&RuleLock, TRUE); +static BOOLEAN RiskScoreMatches(PDYNAMIC_RULE Rule, ULONG RiskScore) { + if (Rule->MinimumRiskScore > 0 && RiskScore < Rule->MinimumRiskScore) return FALSE; + if (Rule->MaximumRiskScore > 0 && RiskScore > Rule->MaximumRiskScore) return FALSE; + return TRUE; +} - PDYNAMIC_RULE Rule = g_DynamicRules; - while (Rule) { - if (Rule->Category == RuleCategoryProcess && (Rule->Operations & OP_EXECUTE)) { - BOOLEAN Match = TRUE; +static BOOLEAN GetProcessPath(HANDLE ProcessId, PUNICODE_STRING* Path) { + if (Path) *Path = NULL; + if (!Path || !ProcessId || KeGetCurrentIrql() != PASSIVE_LEVEL) return FALSE; + return NT_SUCCESS(GetProcessImageName(ProcessId, Path)) && *Path && (*Path)->Buffer; +} - if (!EvaluateNodeConflict(Rule->Initiator, Rule->InitiatorExclude, InitiatorPath)) Match = FALSE; - if (Match && !EvaluateNodeConflict(Rule->Target, Rule->TargetExclude, TargetPath)) Match = FALSE; - if (Match && Rule->CommandLine && !MatchNodeListAny(Rule->CommandLine, CommandLine)) Match = FALSE; +static BOOLEAN MatchDirectParent( + HANDLE ProcessId, + PRULE_NODE Include, + PRULE_NODE Exclude +) { + if (!Include && !Exclude) return TRUE; - if (Match && CheckRuleThreshold(Rule, ProcessId)) { - *OutCode = Rule->Code; - Blocked = TRUE; + HANDLE ParentProcessId = NULL; + if (!GetProcessRelation(ProcessId, &ParentProcessId, NULL) || !ParentProcessId) return FALSE; + + PUNICODE_STRING ParentPath = NULL; + if (!GetProcessPath(ParentProcessId, &ParentPath)) return FALSE; + + BOOLEAN Match = EvaluateNodeConflict(Include, Exclude, ParentPath); + ExFreePool(ParentPath); + return Match; +} + +static BOOLEAN MatchProcessTree( + HANDLE ProcessId, + PRULE_NODE Include, + PRULE_NODE Exclude +) { + if (!Include && !Exclude) return TRUE; + if (!ProcessId) return FALSE; + + HANDLE CurrentProcessId = ProcessId; + HANDLE Visited[12] = { 0 }; + ULONG VisitedCount = 0; + BOOLEAN IncludeMatched = Include ? FALSE : TRUE; + + for (ULONG Depth = 0; Depth < RTL_NUMBER_OF(Visited); Depth++) { + BOOLEAN Seen = FALSE; + for (ULONG Index = 0; Index < VisitedCount; Index++) { + if (Visited[Index] == CurrentProcessId) { + Seen = TRUE; break; } } - Rule = Rule->Next; + if (Seen) break; + + Visited[VisitedCount++] = CurrentProcessId; + + PUNICODE_STRING CurrentPath = NULL; + if (GetProcessPath(CurrentProcessId, &CurrentPath)) { + if (Exclude && MatchNodeListAny(Exclude, CurrentPath)) { + ExFreePool(CurrentPath); + return FALSE; + } + + if (Include && MatchNodeListAny(Include, CurrentPath)) { + IncludeMatched = TRUE; + } + + ExFreePool(CurrentPath); + } + + HANDLE ParentProcessId = NULL; + if (!GetProcessRelation(CurrentProcessId, &ParentProcessId, NULL)) break; + if (!ParentProcessId || ParentProcessId == CurrentProcessId) break; + CurrentProcessId = ParentProcessId; + } + + return IncludeMatched; +} + +static ULONG CalculateAccessRisk(ULONG Operations, ULONG HandleType, ULONG ObjectType) { + ULONG Score = 0; + + if (Operations & OP_VM_READ) Score += 10; + if (Operations & OP_VM_WRITE) Score += 35; + if (Operations & OP_VM_OPERATION) Score += 20; + if (Operations & OP_CREATE_THREAD) Score += 45; + if (Operations & OP_THREAD_SET_CONTEXT) Score += 50; + if (Operations & OP_THREAD_SET_TOKEN) Score += 55; + if (Operations & OP_TERMINATE) Score += 40; + if (Operations & OP_SUSPEND_RESUME) Score += 25; + if (Operations & OP_DUP_HANDLE) Score += 25; + if (Operations & OP_SET_INFORMATION) Score += 20; + if (Operations & OP_CREATE_PROCESS) Score += 35; + if (Operations & OP_IMPERSONATE) Score += 55; + if (HandleType == PYAS_HANDLE_DUPLICATE) Score += 10; + if (ObjectType == PYAS_OBJECT_THREAD && (Operations & (OP_THREAD_SET_CONTEXT | OP_THREAD_SET_TOKEN | OP_IMPERSONATE))) { + Score += 10; + } + + return Score > 100 ? 100 : Score; +} + +static ULONG CalculateThreadRisk(ULONG MemoryType, ULONG MemoryProtection, SIZE_T RegionSize) { + ULONG Score = 20; + + if (MemoryType & PYAS_MEMORY_PRIVATE) Score += 40; + else if (MemoryType & PYAS_MEMORY_MAPPED) Score += 20; + else if (MemoryType & PYAS_MEMORY_IMAGE) Score += 5; + + if (MemoryProtection & PYAS_PROTECT_EXECUTE_WRITE) Score += 30; + else if (MemoryProtection & PYAS_PROTECT_EXECUTE) Score += 10; + + if (RegionSize > 0 && RegionSize <= 1024 * 1024) Score += 5; + return Score > 100 ? 100 : Score; +} + +static ULONG CalculateProcessCreateRisk(HANDLE CreatorPid, HANDLE ParentPid, BOOLEAN IsSubsystemProcess) { + ULONG Score = 0; + if (CreatorPid && ParentPid && CreatorPid != ParentPid) Score += 35; + if (IsSubsystemProcess) Score += 10; + return Score; +} + +static ULONG GetMatchedOperations(PDYNAMIC_RULE Rule, ULONG RequestedOperations) { + ULONG RuleOperations = Rule->Operations & RequestedOperations; + if (RuleOperations == 0) return 0; + + if (Rule->OperationMatch == RuleOperationAll) { + ULONG RelevantOperations = Rule->Operations & ( + OP_VM_READ | + OP_VM_WRITE | + OP_VM_OPERATION | + OP_CREATE_THREAD | + OP_THREAD_SET_CONTEXT | + OP_THREAD_SET_TOKEN | + OP_TERMINATE | + OP_SUSPEND_RESUME | + OP_DUP_HANDLE | + OP_SET_INFORMATION | + OP_CREATE_PROCESS | + OP_IMPERSONATE + ); + + if (RelevantOperations == 0 || (RequestedOperations & RelevantOperations) != RelevantOperations) { + return 0; + } + + return RelevantOperations; + } + + return RuleOperations; +} + +static BOOLEAN MatchSourceTargetRule( + PDYNAMIC_RULE Rule, + HANDLE SourcePid, + HANDLE TargetPid, + PCUNICODE_STRING SourcePath, + PCUNICODE_STRING TargetPath, + ULONG HandleType, + ULONG ObjectType, + ULONG RiskScore +) { + if (Rule->HandleTypes != 0 && (Rule->HandleTypes & HandleType) == 0) return FALSE; + if (Rule->ObjectTypes != 0 && (Rule->ObjectTypes & ObjectType) == 0) return FALSE; + if (!RiskScoreMatches(Rule, RiskScore)) return FALSE; + if (!EvaluateNodeConflict(Rule->Initiator, Rule->InitiatorExclude, SourcePath)) return FALSE; + if (!EvaluateNodeConflict(Rule->Target, Rule->TargetExclude, TargetPath)) return FALSE; + if (!MatchDirectParent(SourcePid, Rule->InitiatorParent, Rule->InitiatorParentExclude)) return FALSE; + if (!MatchProcessTree(SourcePid, Rule->InitiatorProcessTree, Rule->InitiatorProcessTreeExclude)) return FALSE; + if (!MatchProcessTree(TargetPid, Rule->TargetProcessTree, Rule->TargetProcessTreeExclude)) return FALSE; + return TRUE; +} + +BOOLEAN EvaluateProcessCreateRule( + HANDLE CreatorPid, + HANDLE ParentPid, + HANDLE ProcessId, + PCUNICODE_STRING TargetPath, + PCUNICODE_STRING CommandLine, + BOOLEAN FileOpenNameAvailable, + BOOLEAN IsSubsystemProcess, + PULONG OutCode, + PBOOLEAN OutKill +) { + UNREFERENCED_PARAMETER(ProcessId); + + if (!OutCode || !OutKill || !CreatorPid) return FALSE; + *OutCode = 0; + *OutKill = FALSE; + if (IsProcessTrusted(CreatorPid)) return FALSE; + + PUNICODE_STRING CreatorPath = NULL; + PUNICODE_STRING ParentPath = NULL; + GetProcessPath(CreatorPid, &CreatorPath); + if (ParentPid) GetProcessPath(ParentPid, &ParentPath); + + BOOLEAN ParentMismatch = CreatorPid && ParentPid && CreatorPid != ParentPid; + ULONG RiskScore = CalculateProcessCreateRisk(CreatorPid, ParentPid, IsSubsystemProcess); + BOOLEAN Blocked = FALSE; + ULONG SelectedPriority = 0; + + KeEnterCriticalRegion(); + ExAcquireResourceSharedLite(&RuleLock, TRUE); + + for (PDYNAMIC_RULE Rule = g_DynamicRules; Rule; Rule = Rule->Next) { + if (Rule->Category != RuleCategoryProcess || (Rule->Operations & OP_EXECUTE) == 0) continue; + if (!TriStateMatches(Rule->ParentMismatch, ParentMismatch)) continue; + if (!TriStateMatches(Rule->FileOpenNameAvailable, FileOpenNameAvailable)) continue; + if (!TriStateMatches(Rule->SubsystemProcess, IsSubsystemProcess)) continue; + if (!RiskScoreMatches(Rule, RiskScore)) continue; + if (!EvaluateNodeConflict(Rule->Initiator, Rule->InitiatorExclude, CreatorPath)) continue; + if (!EvaluateNodeConflict(Rule->Creator, Rule->CreatorExclude, CreatorPath)) continue; + if (!EvaluateNodeConflict(Rule->Parent, Rule->ParentExclude, ParentPath)) continue; + if (!EvaluateNodeConflict(Rule->Target, Rule->TargetExclude, TargetPath)) continue; + if (!MatchDirectParent(CreatorPid, Rule->InitiatorParent, Rule->InitiatorParentExclude)) continue; + if (!MatchProcessTree(CreatorPid, Rule->InitiatorProcessTree, Rule->InitiatorProcessTreeExclude)) continue; + if (!MatchProcessTree(ParentPid, Rule->TargetProcessTree, Rule->TargetProcessTreeExclude)) continue; + if (Rule->CommandLine && !MatchNodeListAny(Rule->CommandLine, CommandLine)) continue; + if (Rule->CommandLineExclude && MatchNodeListAny(Rule->CommandLineExclude, CommandLine)) continue; + if (!CheckRuleThreshold(Rule, CreatorPid)) continue; + + if (!Blocked || Rule->Priority >= SelectedPriority) { + *OutCode = Rule->Code; + *OutKill = Rule->Kill; + SelectedPriority = Rule->Priority; + } + Blocked = TRUE; } ExReleaseResourceLite(&RuleLock); KeLeaveCriticalRegion(); - if (InitiatorPath) ExFreePool(InitiatorPath); + + if (CreatorPath) ExFreePool(CreatorPath); + if (ParentPath) ExFreePool(ParentPath); return Blocked; } -BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Operation, PVOID WriteBuffer, ULONG WriteLength, PULONG OutCode) { +BOOLEAN EvaluateProcessAccessRule( + HANDLE SourcePid, + HANDLE TargetPid, + ULONG RequestedOperations, + ULONG HandleType, + ULONG ObjectType, + PULONG DeniedOperations, + PULONG OutCode, + PBOOLEAN OutKill +) { + if (!DeniedOperations || !OutCode || !OutKill || !SourcePid || !TargetPid || RequestedOperations == 0) return FALSE; + + *DeniedOperations = 0; + *OutCode = 0; + *OutKill = FALSE; + if (IsProcessTrusted(SourcePid)) return FALSE; + + PUNICODE_STRING SourcePath = NULL; + PUNICODE_STRING TargetPath = NULL; + GetProcessPath(SourcePid, &SourcePath); + GetProcessPath(TargetPid, &TargetPath); + + ULONG RiskScore = CalculateAccessRisk(RequestedOperations, HandleType, ObjectType); + ULONG SelectedPriority = 0; + BOOLEAN Matched = FALSE; + + const ULONG MemoryOperations = + OP_VM_READ | + OP_VM_WRITE | + OP_VM_OPERATION | + OP_CREATE_THREAD | + OP_THREAD_SET_CONTEXT; + + const ULONG ControlOperations = + OP_THREAD_SET_TOKEN | + OP_TERMINATE | + OP_SUSPEND_RESUME | + OP_DUP_HANDLE | + OP_SET_INFORMATION | + OP_CREATE_PROCESS | + OP_IMPERSONATE; + + KeEnterCriticalRegion(); + ExAcquireResourceSharedLite(&RuleLock, TRUE); + + for (PDYNAMIC_RULE Rule = g_DynamicRules; Rule; Rule = Rule->Next) { + ULONG CategoryOperations = 0; + + if (Rule->Category == RuleCategoryMemory) { + CategoryOperations = RequestedOperations & MemoryOperations; + } + else if (Rule->Category == RuleCategoryProcess) { + CategoryOperations = RequestedOperations & ControlOperations; + } + else { + continue; + } + + if (CategoryOperations == 0) continue; + + ULONG RuleMatchedOperations = GetMatchedOperations(Rule, CategoryOperations); + if (RuleMatchedOperations == 0) continue; + if (!MatchSourceTargetRule( + Rule, + SourcePid, + TargetPid, + SourcePath, + TargetPath, + HandleType, + ObjectType, + RiskScore + )) continue; + if (!CheckRuleThreshold(Rule, SourcePid)) continue; + + *DeniedOperations |= RuleMatchedOperations; + + if (!Matched || Rule->Priority >= SelectedPriority) { + *OutCode = Rule->Code; + *OutKill = Rule->Kill; + SelectedPriority = Rule->Priority; + } + Matched = TRUE; + } + + ExReleaseResourceLite(&RuleLock); + KeLeaveCriticalRegion(); + + if (SourcePath) ExFreePool(SourcePath); + if (TargetPath) ExFreePool(TargetPath); + return Matched && *DeniedOperations != 0; +} + +BOOLEAN EvaluateThreadRule( + HANDLE SourcePid, + HANDLE TargetPid, + PVOID StartAddress, + ULONG MemoryType, + ULONG MemoryProtection, + SIZE_T RegionSize, + PULONG OutCode, + PBOOLEAN OutKill +) { + UNREFERENCED_PARAMETER(StartAddress); + + if (!OutCode || !OutKill || !SourcePid || !TargetPid) return FALSE; + *OutCode = 0; + *OutKill = FALSE; + if (IsProcessTrusted(SourcePid)) return FALSE; + + PUNICODE_STRING SourcePath = NULL; + PUNICODE_STRING TargetPath = NULL; + GetProcessPath(SourcePid, &SourcePath); + GetProcessPath(TargetPid, &TargetPath); + + ULONG RiskScore = CalculateThreadRisk(MemoryType, MemoryProtection, RegionSize); + ULONG SelectedPriority = 0; + BOOLEAN Matched = FALSE; + + KeEnterCriticalRegion(); + ExAcquireResourceSharedLite(&RuleLock, TRUE); + + for (PDYNAMIC_RULE Rule = g_DynamicRules; Rule; Rule = Rule->Next) { + if (Rule->Category != RuleCategoryThread || (Rule->Operations & OP_EXECUTE) == 0) continue; + + ULONG RequiredMemoryTypes = Rule->ThreadMemoryTypes != 0 + ? Rule->ThreadMemoryTypes + : PYAS_MEMORY_PRIVATE; + ULONG RequiredProtections = Rule->ThreadMemoryProtections != 0 + ? Rule->ThreadMemoryProtections + : PYAS_PROTECT_EXECUTE; + + if ((RequiredMemoryTypes & MemoryType) == 0) continue; + if ((RequiredProtections & MemoryProtection) == 0) continue; + if (Rule->MinimumRegionSize > 0 && RegionSize < Rule->MinimumRegionSize) continue; + if (Rule->MaximumRegionSize > 0 && RegionSize > Rule->MaximumRegionSize) continue; + if (!MatchSourceTargetRule( + Rule, + SourcePid, + TargetPid, + SourcePath, + TargetPath, + PYAS_HANDLE_CREATE, + PYAS_OBJECT_THREAD, + RiskScore + )) continue; + if (!CheckRuleThreshold(Rule, SourcePid)) continue; + + if (!Matched || Rule->Priority >= SelectedPriority) { + *OutCode = Rule->Code; + *OutKill = Rule->Kill; + SelectedPriority = Rule->Priority; + } + Matched = TRUE; + } + + ExReleaseResourceLite(&RuleLock); + KeLeaveCriticalRegion(); + + if (SourcePath) ExFreePool(SourcePath); + if (TargetPath) ExFreePool(TargetPath); + return Matched; +} + +BOOLEAN EvaluateImageLoadRule( + HANDLE ProcessId, + PCUNICODE_STRING ImagePath, + PIMAGE_INFO ImageInfo, + PULONG OutCode, + PBOOLEAN OutKill +) { + if (!OutCode || !OutKill || !ProcessId || !ImagePath || !ImagePath->Buffer || !ImageInfo) return FALSE; + *OutCode = 0; + *OutKill = FALSE; + + PUNICODE_STRING ProcessPath = NULL; + GetProcessPath(ProcessId, &ProcessPath); + + BOOLEAN Matched = FALSE; + ULONG SelectedPriority = 0; + + KeEnterCriticalRegion(); + ExAcquireResourceSharedLite(&RuleLock, TRUE); + + for (PDYNAMIC_RULE Rule = g_DynamicRules; Rule; Rule = Rule->Next) { + if (Rule->Category != RuleCategoryProcess || (Rule->Operations & OP_IMAGE_LOAD) == 0) continue; + if (!EvaluateNodeConflict(Rule->Initiator, Rule->InitiatorExclude, ProcessPath)) continue; + if (!EvaluateNodeConflict(Rule->Target, Rule->TargetExclude, ImagePath)) continue; + if (!MatchDirectParent(ProcessId, Rule->InitiatorParent, Rule->InitiatorParentExclude)) continue; + if (!MatchProcessTree(ProcessId, Rule->InitiatorProcessTree, Rule->InitiatorProcessTreeExclude)) continue; + if (!MatchProcessTree(ProcessId, Rule->TargetProcessTree, Rule->TargetProcessTreeExclude)) continue; + if (Rule->MinimumRegionSize > 0 && ImageInfo->ImageSize < Rule->MinimumRegionSize) continue; + if (Rule->MaximumRegionSize > 0 && ImageInfo->ImageSize > Rule->MaximumRegionSize) continue; + if (!CheckRuleThreshold(Rule, ProcessId)) continue; + + if (!Matched || Rule->Priority >= SelectedPriority) { + *OutCode = Rule->Code; + *OutKill = Rule->Kill; + SelectedPriority = Rule->Priority; + } + Matched = TRUE; + } + + ExReleaseResourceLite(&RuleLock); + KeLeaveCriticalRegion(); + + if (ProcessPath) ExFreePool(ProcessPath); + return Matched; +} + +BOOLEAN EvaluateProcessRule( + HANDLE ProcessId, + PCUNICODE_STRING TargetPath, + PCUNICODE_STRING CommandLine, + PULONG OutCode, + PBOOLEAN OutKill +) { + return EvaluateProcessCreateRule( + ProcessId, + NULL, + NULL, + TargetPath, + CommandLine, + TRUE, + FALSE, + OutCode, + OutKill + ); +} + +BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Operation, PVOID WriteBuffer, ULONG WriteLength, PULONG OutCode, PBOOLEAN OutKill) { UNREFERENCED_PARAMETER(WriteBuffer); UNREFERENCED_PARAMETER(WriteLength); + if (!OutCode || !OutKill) return FALSE; + *OutCode = 0; + *OutKill = FALSE; if (IsProcessTrusted(ProcessId)) return FALSE; PUNICODE_STRING InitiatorPath = NULL; GetProcessImageName(ProcessId, &InitiatorPath); @@ -845,6 +1958,7 @@ BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Op if (Match && CheckRuleThreshold(Rule, ProcessId)) { *OutCode = Rule->Code; + *OutKill = Rule->Kill; Blocked = TRUE; break; } @@ -858,7 +1972,10 @@ BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Op return Blocked; } -BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG Operation, PULONG OutCode) { +BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill) { + if (!OutCode || !OutKill) return FALSE; + *OutCode = 0; + *OutKill = FALSE; if (IsProcessTrusted(ProcessId)) return FALSE; PUNICODE_STRING InitiatorPath = NULL; GetProcessImageName(ProcessId, &InitiatorPath); @@ -877,6 +1994,7 @@ BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG O if (Match && CheckRuleThreshold(Rule, ProcessId)) { *OutCode = Rule->Code; + *OutKill = Rule->Kill; Blocked = TRUE; break; } @@ -890,7 +2008,10 @@ BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG O return Blocked; } -BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode) { +BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode, PBOOLEAN OutKill) { + if (!OutCode || !OutKill) return FALSE; + *OutCode = 0; + *OutKill = FALSE; if (IsProcessTrusted(ProcessId)) return FALSE; PUNICODE_STRING InitiatorPath = NULL; GetProcessImageName(ProcessId, &InitiatorPath); @@ -908,6 +2029,7 @@ BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode) { if (Match && CheckRuleThreshold(Rule, ProcessId)) { *OutCode = Rule->Code; + *OutKill = Rule->Kill; Blocked = TRUE; break; } @@ -921,97 +2043,48 @@ BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode) { return Blocked; } -static BOOLEAN EvaluateSourceTargetRule(HANDLE SourcePid, HANDLE TargetPid, RULE_CATEGORY Category, ULONG Operation, PULONG OutCode) { - if (!OutCode || Operation == 0) return FALSE; - if (IsProcessTrusted(SourcePid)) return FALSE; - - PUNICODE_STRING InitiatorPath = NULL; - PUNICODE_STRING TargetPath = NULL; - - GetProcessImageName(SourcePid, &InitiatorPath); - GetProcessImageName(TargetPid, &TargetPath); - - BOOLEAN Blocked = FALSE; - - KeEnterCriticalRegion(); - ExAcquireResourceSharedLite(&RuleLock, TRUE); - - PDYNAMIC_RULE Rule = g_DynamicRules; - while (Rule) { - if (Rule->Category == Category && (Rule->Operations & Operation)) { - BOOLEAN Match = TRUE; - - if (!EvaluateNodeConflict(Rule->Initiator, Rule->InitiatorExclude, InitiatorPath)) Match = FALSE; - if (Match && !EvaluateNodeConflict(Rule->Target, Rule->TargetExclude, TargetPath)) Match = FALSE; - - if (Match && CheckRuleThreshold(Rule, SourcePid)) { - *OutCode = Rule->Code; - Blocked = TRUE; - break; - } - } - Rule = Rule->Next; - } - - ExReleaseResourceLite(&RuleLock); - KeLeaveCriticalRegion(); - - if (InitiatorPath) ExFreePool(InitiatorPath); - if (TargetPath) ExFreePool(TargetPath); - - return Blocked; +BOOLEAN EvaluateMemoryRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill) { + ULONG DeniedOperations = 0; + return EvaluateProcessAccessRule( + SourcePid, + TargetPid, + Operation, + PYAS_HANDLE_CREATE, + PYAS_OBJECT_PROCESS, + &DeniedOperations, + OutCode, + OutKill + ); } -BOOLEAN EvaluateMemoryRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode) { - return EvaluateSourceTargetRule(SourcePid, TargetPid, RuleCategoryMemory, Operation, OutCode); -} - -BOOLEAN EvaluateProcessHandleRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode) { - return EvaluateSourceTargetRule(SourcePid, TargetPid, RuleCategoryProcess, Operation, OutCode); -} - -BOOLEAN EvaluateThreadRule(HANDLE SourcePid, HANDLE TargetPid, PVOID StartAddress, PULONG OutCode) { - if (!IsAddressInUnmappedMemory(TargetPid, StartAddress)) return FALSE; - - PUNICODE_STRING InitiatorPath = NULL; - PUNICODE_STRING TargetPath = NULL; - - GetProcessImageName(SourcePid, &InitiatorPath); - GetProcessImageName(TargetPid, &TargetPath); - - BOOLEAN Blocked = FALSE; - - KeEnterCriticalRegion(); - ExAcquireResourceSharedLite(&RuleLock, TRUE); - - PDYNAMIC_RULE Rule = g_DynamicRules; - while (Rule) { - if (Rule->Category == RuleCategoryThread && (Rule->Operations & OP_EXECUTE)) { - BOOLEAN Match = TRUE; - - if (!EvaluateNodeConflict(Rule->Initiator, Rule->InitiatorExclude, InitiatorPath)) Match = FALSE; - if (Match && !EvaluateNodeConflict(Rule->Target, Rule->TargetExclude, TargetPath)) Match = FALSE; - - if (Match && CheckRuleThreshold(Rule, SourcePid)) { - *OutCode = Rule->Code; - Blocked = TRUE; - break; - } - } - Rule = Rule->Next; - } - - ExReleaseResourceLite(&RuleLock); - KeLeaveCriticalRegion(); - - if (InitiatorPath) ExFreePool(InitiatorPath); - if (TargetPath) ExFreePool(TargetPath); - - return Blocked; +BOOLEAN EvaluateProcessHandleRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill) { + ULONG DeniedOperations = 0; + return EvaluateProcessAccessRule( + SourcePid, + TargetPid, + Operation, + PYAS_HANDLE_CREATE, + PYAS_OBJECT_PROCESS, + &DeniedOperations, + OutCode, + OutKill + ); } NTSTATUS SendMessageToUser(ULONG Code, ULONG Pid, PWCHAR Path, USHORT PathSize) { if (KeGetCurrentIrql() > APC_LEVEL) return STATUS_UNSUCCESSFUL; + + LONG DriverState = InterlockedCompareExchange( + &GlobalData.DriverState, + PyasDriverStateCold, + PyasDriverStateCold + ); + + if (DriverState != PyasDriverStateRunning && DriverState != PyasDriverStateStopRetry) { + return STATUS_PORT_DISCONNECTED; + } + + if (InterlockedCompareExchange(&GlobalData.PortStopping, 0, 0) != 0) return STATUS_PORT_DISCONNECTED; if (!ExAcquireRundownProtection(&GlobalData.PortRundown)) return STATUS_PORT_DISCONNECTED; NTSTATUS status = STATUS_PORT_DISCONNECTED; @@ -1019,6 +2092,7 @@ NTSTATUS SendMessageToUser(ULONG Code, ULONG Pid, PWCHAR Path, USHORT PathSize) if (GlobalData.ClientPort) { PPYAS_MESSAGE msg = (PPYAS_MESSAGE)PyasAllocate(sizeof(PYAS_MESSAGE)); if (msg) { + RtlZeroMemory(msg, sizeof(*msg)); msg->MessageCode = Code; msg->ProcessId = Pid;