proc.c

/*
 * proc.c — Process/thread/image monitoring and control.
 *
 * Stealth & security features:
 *   - ObRegisterCallbacks to protect Agent process (strip dangerous handles)
 *   - Multiple termination methods (ZwTerminate / APC-based)
 *   - Remote thread injection blocking for Agent process
 *   - DLL injection detection for Agent process
 *   - Hook detection on ZwQuerySystemInformation
 */

#include <ntifs.h>
#include <fltKernel.h>
#include "edr_shared.h"
#include "core.h"
#include "utils/inject/inject.h"
#include "utils/inject/inject_hijack.h"
#include "utils/inject/inject_instcb.h"
#include "utils/inject/mdl_alloc.h"
#include "utils/inject/mem_copy.h"
#include "utils/inject/pte_ops.h"

#ifndef PROCESS_TERMINATE
#define PROCESS_TERMINATE 0x0001
#endif

#ifndef PROCESS_CREATE_THREAD
#define PROCESS_CREATE_THREAD 0x0002
#endif

#ifndef PROCESS_VM_OPERATION
#define PROCESS_VM_OPERATION 0x0008
#endif

#ifndef PROCESS_VM_WRITE
#define PROCESS_VM_WRITE 0x0020
#endif

#ifndef PROCESS_SUSPEND_RESUME
#define PROCESS_SUSPEND_RESUME 0x0800
#endif

#ifndef THREAD_TERMINATE
#define THREAD_TERMINATE 0x0001
#endif

#ifndef THREAD_SUSPEND_RESUME
#define THREAD_SUSPEND_RESUME 0x0002
#endif

#ifndef SystemProcessInformation
#define SystemProcessInformation 5
#endif

typedef struct _SYSTEM_PROCESS_INFORMATION {
    ULONG           NextEntryOffset;
    ULONG           NumberOfThreads;
    LARGE_INTEGER   WorkingSetPrivateSize;
    ULONG           HardFaultCount;
    ULONG           NumberOfThreadsHighWatermark;
    ULONGLONG       CycleTime;
    LARGE_INTEGER   CreateTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   KernelTime;
    UNICODE_STRING  ImageName;
    LONG            BasePriority;
    HANDLE          UniqueProcessId;
    HANDLE          InheritedFromUniqueProcessId;
    ULONG           HandleCount;
    ULONG           SessionId;
    ULONG_PTR       UniqueProcessKey;
    SIZE_T          PeakVirtualSize;
    SIZE_T          VirtualSize;
    ULONG           PageFaultCount;
    SIZE_T          PeakWorkingSetSize;
    SIZE_T          WorkingSetSize;
    SIZE_T          QuotaPeakPagedPoolUsage;
    SIZE_T          QuotaPagedPoolUsage;
    SIZE_T          QuotaPeakNonPagedPoolUsage;
    SIZE_T          QuotaNonPagedPoolUsage;
    SIZE_T          PagefileUsage;
    SIZE_T          PeakPagefileUsage;
    SIZE_T          PrivatePageCount;
    LARGE_INTEGER   ReadOperationCount;
    LARGE_INTEGER   WriteOperationCount;
    LARGE_INTEGER   OtherOperationCount;
    LARGE_INTEGER   ReadTransferCount;
    LARGE_INTEGER   WriteTransferCount;
    LARGE_INTEGER   OtherTransferCount;
} SYSTEM_PROCESS_INFORMATION, *PSYSTEM_PROCESS_INFORMATION;

NTSYSCALLAPI NTSTATUS NTAPI ZwQuerySystemInformation(
    _In_      ULONG SystemInformationClass,
    _Inout_   PVOID SystemInformation,
    _In_      ULONG SystemInformationLength,
    _Out_opt_ PULONG ReturnLength
);

NTKERNELAPI BOOLEAN PsIsProtectedProcess(
    _In_ PEPROCESS Process
);

#define PROC_MAX_RULES  32
#define PROC_RULE_NAME_LEN 260

typedef struct _PROC_RULE {
    BOOLEAN Active;
    WCHAR   ImagePath[PROC_RULE_NAME_LEN];
    USHORT  PathLength;
} PROC_RULE, *PPROC_RULE;

static PROC_RULE     g_ProcRules[PROC_MAX_RULES] = { 0 };
static EX_PUSH_LOCK  g_ProcRulesLock;

static volatile LONG g_ProcInitialized = 0;
static volatile LONG g_MonitorEnabled = 0;

static PVOID g_ObCallbackHandle = NULL;

#define PROC_MAX_PROTECTED  32
static volatile ULONG g_ProtectedPids[PROC_MAX_PROTECTED] = { 0 };
static KSPIN_LOCK     g_ProtectedLock;

#define PROC_DANGEROUS_ACCESS \
    (PROCESS_TERMINATE | PROCESS_CREATE_THREAD | \
     PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_SUSPEND_RESUME)

#define THREAD_DANGEROUS_ACCESS \
    (THREAD_TERMINATE | THREAD_SUSPEND_RESUME)

static VOID ProcCreateNotify(
    _Inout_  PEPROCESS Process,
    _In_     HANDLE ProcessId,
    _Inout_opt_ PPS_CREATE_NOTIFY_INFO CreateInfo
);

static VOID ProcThreadNotify(
    _In_ HANDLE ProcessId,
    _In_ HANDLE ThreadId,
    _In_ BOOLEAN Create
);

static VOID ProcImageNotify(
    _In_opt_ PUNICODE_STRING FullImageName,
    _In_ HANDLE ProcessId,
    _In_ PIMAGE_INFO ImageInfo
);

static NTSTATUS ProcFindByName(
    _In_ PUCHAR NameBuf,
    _In_ ULONG NameLen,
    _Out_ PULONG FoundPid
);

static NTSTATUS ProcTerminate(_In_ ULONG Pid);
static BOOLEAN ProcIsAgentProcess(_In_ PEPROCESS Process);
static BOOLEAN ProcIsProtectedPid(_In_ ULONG Pid);
static NTSTATUS ProcEnableMonitor(VOID);
static VOID ProcDisableMonitor(VOID);
static NTSTATUS ProcInject(_Inout_ PRING_MESSAGE Msg);
static NTSTATUS ProcDetectInjection(_Inout_ PRING_MESSAGE Msg);

static OB_PREOP_CALLBACK_STATUS ProcObPreOpenProcess(
    _In_ PVOID RegistrationContext,
    _Inout_ POB_PRE_OPERATION_INFORMATION OperationInfo
);
static OB_PREOP_CALLBACK_STATUS ProcObPreOpenThread(
    _In_ PVOID RegistrationContext,
    _Inout_ POB_PRE_OPERATION_INFORMATION OperationInfo
);

NTSTATUS ProcModuleInitialize(VOID)
{
    ExInitializePushLock(&g_ProcRulesLock);
    RtlZeroMemory(g_ProcRules, sizeof(g_ProcRules));
    KeInitializeSpinLock(&g_ProtectedLock);
    RtlZeroMemory((PVOID)g_ProtectedPids, sizeof(g_ProtectedPids));
    g_MonitorEnabled = 0;

    ULONG agentPid = CommGetAgentPid();
    if (agentPid != 0) {
        g_ProtectedPids[0] = agentPid;
    }

    OB_OPERATION_REGISTRATION opReg[2] = { 0 };
    opReg[0].ObjectType = PsProcessType;
    opReg[0].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
    opReg[0].PreOperation = ProcObPreOpenProcess;

    opReg[1].ObjectType = PsThreadType;
    opReg[1].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
    opReg[1].PreOperation = ProcObPreOpenThread;

    OB_CALLBACK_REGISTRATION cbReg = { 0 };
    cbReg.Version = OB_FLT_REGISTRATION_VERSION;
    cbReg.OperationRegistrationCount = 2;
    cbReg.OperationRegistration = opReg;

    UNICODE_STRING altitude = RTL_CONSTANT_STRING(L"370020.2");
    cbReg.Altitude = altitude;
    cbReg.RegistrationContext = NULL;

    NTSTATUS status = ObRegisterCallbacks(&cbReg, &g_ObCallbackHandle);
    if (!NT_SUCCESS(status)) {
        g_ObCallbackHandle = NULL;
    }

    InterlockedExchange(&g_ProcInitialized, 1);
    return STATUS_SUCCESS;
}

VOID ProcModuleUninitialize(VOID)
{
    if (!InterlockedCompareExchange(&g_ProcInitialized, 0, 1)) {
        return;
    }

    ProcDisableMonitor();

    if (g_ObCallbackHandle) {
        ObUnRegisterCallbacks(g_ObCallbackHandle);
        g_ObCallbackHandle = NULL;
    }
}

NTSTATUS ProcHandleCommand(_Inout_ PRING_MESSAGE Msg)
{
    switch (Msg->Command) {

    case CMD_PROC_QUERY_PID: {
        ULONG foundPid = 0;
        NTSTATUS st = ProcFindByName(
            Msg->Data, Msg->DataLength, &foundPid
        );
        Msg->ProcessId = foundPid;
        return st;
    }

    case CMD_PROC_TERMINATE:
        return ProcTerminate(Msg->ProcessId);

    case CMD_PROC_LIST: {
        PVOID buffer = NULL;
        ULONG bufSize = 1024 * 256;
        NTSTATUS st;

        for (int attempt = 0; attempt < 3; attempt++) {
            buffer = ExAllocatePool2(POOL_FLAG_PAGED, bufSize, 'bOoI');
            if (!buffer) return STATUS_INSUFFICIENT_RESOURCES;

            st = ZwQuerySystemInformation(
                SystemProcessInformation, buffer, bufSize, &bufSize);
            if (NT_SUCCESS(st)) break;

            ExFreePoolWithTag(buffer, 'bOoI');
            buffer = NULL;
            if (st != STATUS_INFO_LENGTH_MISMATCH) return st;
            bufSize += 1024 * 64;
        }
        if (!buffer) return STATUS_INSUFFICIENT_RESOURCES;

        typedef struct { ULONG Pid; CHAR Name[16]; } PROC_ENTRY;
        ULONG maxEntries = RING_DATA_SIZE / sizeof(PROC_ENTRY);
        PROC_ENTRY *out = (PROC_ENTRY *)Msg->Data;
        ULONG count = 0;

        PSYSTEM_PROCESS_INFORMATION entry =
            (PSYSTEM_PROCESS_INFORMATION)buffer;
        while (count < maxEntries) {
            out[count].Pid = (ULONG)(ULONG_PTR)entry->UniqueProcessId;
            RtlZeroMemory(out[count].Name, 16);

            if (entry->ImageName.Buffer && entry->ImageName.Length > 0) {
                ANSI_STRING ansi;
                if (NT_SUCCESS(RtlUnicodeStringToAnsiString(&ansi, &entry->ImageName, TRUE))) {
                    ULONG copyLen = min((ULONG)ansi.Length, 15);
                    RtlCopyMemory(out[count].Name, ansi.Buffer, copyLen);
                    RtlFreeAnsiString(&ansi);
                }
            }
            count++;

            if (entry->NextEntryOffset == 0) break;
            entry = (PSYSTEM_PROCESS_INFORMATION)(
                (PUCHAR)entry + entry->NextEntryOffset);
        }

        Msg->DataLength = count * sizeof(PROC_ENTRY);
        ExFreePoolWithTag(buffer, 'bOoI');
        return STATUS_SUCCESS;
    }

    case CMD_PROC_ADD_RULE: {
        if (Msg->DataLength < sizeof(WCHAR)) {
            return STATUS_INVALID_PARAMETER;
        }
        ULONG pathBytes = Msg->DataLength;
        if (pathBytes > (PROC_RULE_NAME_LEN - 1) * sizeof(WCHAR)) {
            return STATUS_BUFFER_OVERFLOW;
        }

        KeEnterCriticalRegion();
        ExAcquirePushLockExclusive(&g_ProcRulesLock);

        NTSTATUS st = STATUS_INSUFFICIENT_RESOURCES;
        for (ULONG i = 0; i < PROC_MAX_RULES; i++) {
            if (!g_ProcRules[i].Active) {
                RtlZeroMemory(g_ProcRules[i].ImagePath, sizeof(g_ProcRules[i].ImagePath));
                RtlCopyMemory(g_ProcRules[i].ImagePath, Msg->Data, pathBytes);
                g_ProcRules[i].PathLength = (USHORT)pathBytes;
                g_ProcRules[i].Active = TRUE;
                st = STATUS_SUCCESS;
                break;
            }
        }

        ExReleasePushLockExclusive(&g_ProcRulesLock);
        KeLeaveCriticalRegion();
        return st;
    }

    case CMD_PROC_REMOVE_RULE: {
        if (Msg->DataLength < sizeof(WCHAR)) {
            return STATUS_INVALID_PARAMETER;
        }

        UNICODE_STRING target;
        target.Buffer = (PWCHAR)Msg->Data;
        target.Length = (USHORT)Msg->DataLength;
        target.MaximumLength = target.Length;

        KeEnterCriticalRegion();
        ExAcquirePushLockExclusive(&g_ProcRulesLock);

        NTSTATUS st = STATUS_NOT_FOUND;
        for (ULONG i = 0; i < PROC_MAX_RULES; i++) {
            if (!g_ProcRules[i].Active) continue;

            UNICODE_STRING rulePath;
            rulePath.Buffer = g_ProcRules[i].ImagePath;
            rulePath.Length = g_ProcRules[i].PathLength;
            rulePath.MaximumLength = sizeof(g_ProcRules[i].ImagePath);

            if (RtlEqualUnicodeString(&target, &rulePath, TRUE)) {
                g_ProcRules[i].Active = FALSE;
                st = STATUS_SUCCESS;
                break;
            }
        }

        ExReleasePushLockExclusive(&g_ProcRulesLock);
        KeLeaveCriticalRegion();
        return st;
    }

    case CMD_PROC_ENABLE_MONITOR:
        return ProcEnableMonitor();

    case CMD_PROC_DISABLE_MONITOR:
        ProcDisableMonitor();
        return STATUS_SUCCESS;

    case CMD_PROC_INJECT:
        return ProcInject(Msg);

    case CMD_PROC_DETECT_INJECT:
        return ProcDetectInjection(Msg);

    case CMD_PROC_PROTECT:
        return ProcAddProtectedPid(Msg->ProcessId);

    case CMD_PROC_UNPROTECT:
        return ProcRemoveProtectedPid(Msg->ProcessId);

    default:
        return STATUS_INVALID_PARAMETER;
    }
}

static VOID ProcCreateNotify(
    _Inout_  PEPROCESS Process,
    _In_     HANDLE ProcessId,
    _Inout_opt_ PPS_CREATE_NOTIFY_INFO CreateInfo
)
{
    UNREFERENCED_PARAMETER(Process);

    if (CreateInfo != NULL) {
        ULONG pid = (ULONG)(ULONG_PTR)ProcessId;

        if (CreateInfo->ImageFileName && CreateInfo->ImageFileName->Length > 0) {
            KeEnterCriticalRegion();
            ExAcquirePushLockShared(&g_ProcRulesLock);

            for (ULONG i = 0; i < PROC_MAX_RULES; i++) {
                if (!g_ProcRules[i].Active) continue;

                UNICODE_STRING pattern;
                pattern.Buffer = g_ProcRules[i].ImagePath;
                pattern.Length = g_ProcRules[i].PathLength;
                pattern.MaximumLength = sizeof(g_ProcRules[i].ImagePath);

                if (FsRtlIsNameInExpression(&pattern,
                        CreateInfo->ImageFileName, TRUE, NULL)) {
                    CreateInfo->CreationStatus = STATUS_ACCESS_DENIED;
                    break;
                }
            }

            ExReleasePushLockShared(&g_ProcRulesLock);
            KeLeaveCriticalRegion();
        }

        if (CreateInfo->ImageFileName) {
            CommPostEvent(
                EVT_PROC_CREATED,
                pid,
                CreateInfo->ImageFileName->Buffer,
                (ULONG)min(CreateInfo->ImageFileName->Length, RING_DATA_SIZE)
            );
        } else {
            CommPostEvent(EVT_PROC_CREATED, pid, NULL, 0);
        }
    } else {
        CommPostEvent(
            EVT_PROC_TERMINATED,
            (ULONG)(ULONG_PTR)ProcessId,
            NULL, 0
        );
    }
}

static VOID ProcThreadNotify(
    _In_ HANDLE ProcessId,
    _In_ HANDLE ThreadId,
    _In_ BOOLEAN Create
)
{
    if (!Create) return;

    ULONG targetPid = (ULONG)(ULONG_PTR)ProcessId;
    ULONG callerPid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();

    if (callerPid != targetPid && callerPid != 4) {
        struct {
            ULONG TargetPid;
            ULONG CallerPid;
            ULONG ThreadId;
        } info;
        info.TargetPid = targetPid;
        info.CallerPid = callerPid;
        info.ThreadId = (ULONG)(ULONG_PTR)ThreadId;

        CommPostEvent(EVT_THREAD_CREATED, targetPid,
                      &info, sizeof(info));

        ULONG agentPid = CommGetAgentPid();
        if (agentPid != 0 && targetPid == agentPid) {
            CommPostEvent(EVT_SELF_ATTACK, callerPid,
                          &info, sizeof(info));
        }
    }
}

static VOID ProcImageNotify(
    _In_opt_ PUNICODE_STRING FullImageName,
    _In_ HANDLE ProcessId,
    _In_ PIMAGE_INFO ImageInfo
)
{
    ULONG pid = (ULONG)(ULONG_PTR)ProcessId;

    if (FullImageName && FullImageName->Buffer) {
        CommPostEvent(
            EVT_IMAGE_LOADED, pid,
            FullImageName->Buffer,
            (ULONG)min(FullImageName->Length, RING_DATA_SIZE)
        );

        ULONG agentPid = CommGetAgentPid();
        if (agentPid != 0 && pid == agentPid && !ImageInfo->SystemModeImage) {
            ULONG callerPid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();
            if (callerPid != agentPid && callerPid != 4) {
                CommPostEvent(EVT_SELF_ATTACK, callerPid,
                    FullImageName->Buffer,
                    (ULONG)min(FullImageName->Length, RING_DATA_SIZE));
            }
        }
    }
}

static NTSTATUS ProcFindByName(
    _In_ PUCHAR NameBuf,
    _In_ ULONG NameLen,
    _Out_ PULONG FoundPid
)
{
    NTSTATUS status;
    PVOID buffer = NULL;
    ULONG bufSize = 1024 * 256;

    *FoundPid = 0;

    if (NameLen == 0 || NameLen > 260) {
        return STATUS_INVALID_PARAMETER;
    }

    CHAR targetName[261];
    RtlCopyMemory(targetName, NameBuf, NameLen);
    targetName[NameLen] = '\0';

    for (int attempt = 0; attempt < 3; attempt++) {
        buffer = ExAllocatePool2(POOL_FLAG_PAGED, bufSize, 'bOoI');
        if (!buffer) return STATUS_INSUFFICIENT_RESOURCES;

        status = ZwQuerySystemInformation(
            SystemProcessInformation,
            buffer, bufSize, &bufSize
        );
        if (NT_SUCCESS(status)) break;

        ExFreePoolWithTag(buffer, 'bOoI');
        buffer = NULL;

        if (status != STATUS_INFO_LENGTH_MISMATCH) {
            return status;
        }
        bufSize += 1024 * 64;
    }

    if (!buffer) return STATUS_INSUFFICIENT_RESOURCES;

    PSYSTEM_PROCESS_INFORMATION entry =
        (PSYSTEM_PROCESS_INFORMATION)buffer;

    while (TRUE) {
        if (entry->ImageName.Buffer && entry->ImageName.Length > 0) {
            ANSI_STRING ansi;
            UNICODE_STRING uname = entry->ImageName;

            status = RtlUnicodeStringToAnsiString(&ansi, &uname, TRUE);
            if (NT_SUCCESS(status)) {
                if (_stricmp(ansi.Buffer, targetName) == 0) {
                    *FoundPid = (ULONG)(ULONG_PTR)entry->UniqueProcessId;
                    RtlFreeAnsiString(&ansi);
                    break;
                }
                RtlFreeAnsiString(&ansi);
            }
        }

        if (entry->NextEntryOffset == 0) break;
        entry = (PSYSTEM_PROCESS_INFORMATION)(
            (PUCHAR)entry + entry->NextEntryOffset
        );
    }

    ExFreePoolWithTag(buffer, 'bOoI');
    return (*FoundPid != 0) ? STATUS_SUCCESS : STATUS_NOT_FOUND;
}

static NTSTATUS ProcTerminate(_In_ ULONG Pid)
{
    NTSTATUS status;
    PEPROCESS target = NULL;
    HANDLE processHandle = NULL;

    if (Pid == 0 || Pid == 4) {
        return STATUS_ACCESS_DENIED;
    }

    status = PsLookupProcessByProcessId((HANDLE)(ULONG_PTR)Pid, &target);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    if (PsIsProtectedProcess(target)) {
        ObDereferenceObject(target);
        return STATUS_ACCESS_DENIED;
    }

    status = ObOpenObjectByPointer(
        target,
        OBJ_KERNEL_HANDLE,
        NULL,
        PROCESS_TERMINATE,
        *PsProcessType,
        KernelMode,
        &processHandle
    );
    ObDereferenceObject(target);

    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = ZwTerminateProcess(processHandle, STATUS_SUCCESS);
    ZwClose(processHandle);

    if (NT_SUCCESS(status)) {
        CommPostEvent(EVT_PROC_TERMINATED, Pid, NULL, 0);
    }

    return status;
}

static NTSTATUS ProcEnableMonitor(VOID)
{
    if (InterlockedCompareExchange(&g_MonitorEnabled, 0, 0) != 0) {
        return STATUS_SUCCESS;
    }

    NTSTATUS status;

    status = PsSetCreateProcessNotifyRoutineEx(ProcCreateNotify, FALSE);
    if (!NT_SUCCESS(status)) return status;

    status = PsSetCreateThreadNotifyRoutine(ProcThreadNotify);
    if (!NT_SUCCESS(status)) {
        PsSetCreateProcessNotifyRoutineEx(ProcCreateNotify, TRUE);
        return status;
    }

    status = PsSetLoadImageNotifyRoutine(ProcImageNotify);
    if (!NT_SUCCESS(status)) {
        PsRemoveCreateThreadNotifyRoutine(ProcThreadNotify);
        PsSetCreateProcessNotifyRoutineEx(ProcCreateNotify, TRUE);
        return status;
    }

    InterlockedExchange(&g_MonitorEnabled, 1);
    return STATUS_SUCCESS;
}

static VOID ProcDisableMonitor(VOID)
{
    if (!InterlockedCompareExchange(&g_MonitorEnabled, 0, 1)) {
        return;
    }

    PsRemoveLoadImageNotifyRoutine(ProcImageNotify);
    PsRemoveCreateThreadNotifyRoutine(ProcThreadNotify);
    PsSetCreateProcessNotifyRoutineEx(ProcCreateNotify, TRUE);
}

static BOOLEAN ProcIsProtectedPid(_In_ ULONG Pid)
{
    if (Pid == 0) return FALSE;

    ULONG agentPid = CommGetAgentPid();
    if (agentPid != 0 && Pid == agentPid) return TRUE;

    BOOLEAN found = FALSE;
    KIRQL oldIrql;
    KeAcquireSpinLock(&g_ProtectedLock, &oldIrql);
    for (ULONG i = 0; i < PROC_MAX_PROTECTED; i++) {
        if (g_ProtectedPids[i] == Pid) { found = TRUE; break; }
    }
    KeReleaseSpinLock(&g_ProtectedLock, oldIrql);
    return found;
}

NTSTATUS ProcAddProtectedPid(_In_ ULONG Pid)
{
    if (Pid == 0) return STATUS_INVALID_PARAMETER;

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

    for (ULONG i = 0; i < PROC_MAX_PROTECTED; i++) {
        if (g_ProtectedPids[i] == Pid) {
            KeReleaseSpinLock(&g_ProtectedLock, oldIrql);
            return STATUS_SUCCESS;
        }
    }

    for (ULONG i = 0; i < PROC_MAX_PROTECTED; i++) {
        if (g_ProtectedPids[i] == 0) {
            g_ProtectedPids[i] = Pid;
            KeReleaseSpinLock(&g_ProtectedLock, oldIrql);
            return STATUS_SUCCESS;
        }
    }

    KeReleaseSpinLock(&g_ProtectedLock, oldIrql);
    return STATUS_INSUFFICIENT_RESOURCES;
}

NTSTATUS ProcRemoveProtectedPid(_In_ ULONG Pid)
{
    KIRQL oldIrql;
    KeAcquireSpinLock(&g_ProtectedLock, &oldIrql);

    for (ULONG i = 0; i < PROC_MAX_PROTECTED; i++) {
        if (g_ProtectedPids[i] == Pid) {
            g_ProtectedPids[i] = 0;
            KeReleaseSpinLock(&g_ProtectedLock, oldIrql);
            return STATUS_SUCCESS;
        }
    }

    KeReleaseSpinLock(&g_ProtectedLock, oldIrql);
    return STATUS_NOT_FOUND;
}

static BOOLEAN ProcIsAgentProcess(_In_ PEPROCESS Process)
{
    ULONG agentPid = CommGetAgentPid();
    if (agentPid == 0) return FALSE;
    ULONG pid = (ULONG)(ULONG_PTR)PsGetProcessId(Process);
    return (pid == agentPid);
}

static OB_PREOP_CALLBACK_STATUS ProcObPreOpenProcess(
    _In_ PVOID RegistrationContext,
    _Inout_ POB_PRE_OPERATION_INFORMATION OperationInfo
)
{
    UNREFERENCED_PARAMETER(RegistrationContext);

    if (OperationInfo->KernelHandle) return OB_PREOP_SUCCESS;

    PEPROCESS targetProcess = (PEPROCESS)OperationInfo->Object;
    ULONG targetPid = (ULONG)(ULONG_PTR)PsGetProcessId(targetProcess);

    if (!ProcIsProtectedPid(targetPid)) return OB_PREOP_SUCCESS;

    ULONG callerPid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();
    ULONG agentPid = CommGetAgentPid();
    if (callerPid == agentPid || callerPid == targetPid) return OB_PREOP_SUCCESS;

    if (OperationInfo->Operation == OB_OPERATION_HANDLE_CREATE) {
        OperationInfo->Parameters->CreateHandleInformation.DesiredAccess
            &= ~PROC_DANGEROUS_ACCESS;
    } else {
        OperationInfo->Parameters->DuplicateHandleInformation.DesiredAccess
            &= ~PROC_DANGEROUS_ACCESS;
    }

    CommPostEvent(EVT_SELF_ATTACK, callerPid, NULL, 0);
    return OB_PREOP_SUCCESS;
}

static OB_PREOP_CALLBACK_STATUS ProcObPreOpenThread(
    _In_ PVOID RegistrationContext,
    _Inout_ POB_PRE_OPERATION_INFORMATION OperationInfo
)
{
    UNREFERENCED_PARAMETER(RegistrationContext);

    if (OperationInfo->KernelHandle) return OB_PREOP_SUCCESS;

    PEPROCESS ownerProcess =
        IoThreadToProcess((PETHREAD)OperationInfo->Object);
    ULONG ownerPid = (ULONG)(ULONG_PTR)PsGetProcessId(ownerProcess);

    if (!ProcIsProtectedPid(ownerPid)) return OB_PREOP_SUCCESS;

    ULONG callerPid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();
    ULONG agentPid = CommGetAgentPid();
    if (callerPid == agentPid || callerPid == ownerPid) return OB_PREOP_SUCCESS;

    if (OperationInfo->Operation == OB_OPERATION_HANDLE_CREATE) {
        OperationInfo->Parameters->CreateHandleInformation.DesiredAccess
            &= ~THREAD_DANGEROUS_ACCESS;
    } else {
        OperationInfo->Parameters->DuplicateHandleInformation.DesiredAccess
            &= ~THREAD_DANGEROUS_ACCESS;
    }

    return OB_PREOP_SUCCESS;
}


static NTSTATUS ProcInjectShellcodeToProcess(
    _In_ ULONG Method,
    _In_ ULONG TargetPid,
    _In_ PVOID Shellcode,
    _In_ ULONG ShellcodeSize
)
{
    if (Method == PROC_INJECT_SHELLCODE_APC)
        return InjectViaShellcodeApc(TargetPid, Shellcode, ShellcodeSize);

    /* THREAD_HIJACK and INSTCB need PEPROCESS + mapped remote addr */
    PEPROCESS target = NULL;
    NTSTATUS status = PsLookupProcessByProcessId(
        (HANDLE)(ULONG_PTR)TargetPid, &target);
    if (!NT_SUCCESS(status)) return status;

    MDL_ALLOC alloc = { 0 };
    status = MdlAllocMapInProcess(&alloc, target,
                                  ShellcodeSize, PAGE_READWRITE);
    if (!NT_SUCCESS(status)) {
        ObDereferenceObject(target);
        return status;
    }

    status = MemCopyToProcess(target, alloc.MappedAddress,
                              Shellcode, ShellcodeSize);
    if (!NT_SUCCESS(status)) {
        MdlAllocFreeFromProcess(&alloc, target);
        ObDereferenceObject(target);
        return status;
    }

    PteSetProtection(target, alloc.MappedAddress,
                     alloc.Size, PAGE_EXECUTE_READ);

    switch (Method) {
    case PROC_INJECT_THREAD_HIJACK:
        status = InjectViaThreadHijack(
            target, alloc.MappedAddress, NULL);
        break;
    case PROC_INJECT_INSTCB:
        status = InjectViaInstrumentationCallback(
            target, alloc.MappedAddress);
        break;
    default:
        status = STATUS_INVALID_PARAMETER;
        break;
    }

    if (!NT_SUCCESS(status))
        MdlAllocFreeFromProcess(&alloc, target);
    /* On success, memory stays mapped — shellcode is executing */

    ObDereferenceObject(target);
    return status;
}

static NTSTATUS ProcInject(_Inout_ PRING_MESSAGE Msg)
{
    if (Msg->DataLength < sizeof(PROC_INJECT_REQUEST))
        return STATUS_INVALID_PARAMETER;

    PPROC_INJECT_REQUEST req = (PPROC_INJECT_REQUEST)Msg->Data;
    ULONG targetPid = Msg->ProcessId;

    /* Shellcode-based methods: DllPathLength = shellcode size */
    if (req->Method == PROC_INJECT_SHELLCODE_APC ||
        req->Method == PROC_INJECT_THREAD_HIJACK ||
        req->Method == PROC_INJECT_INSTCB) {
        ULONG scSize = req->DllPathLength;
        if (scSize == 0 ||
            scSize > Msg->DataLength - sizeof(PROC_INJECT_REQUEST))
            return STATUS_INVALID_PARAMETER;
        PVOID sc = (PUCHAR)req + sizeof(PROC_INJECT_REQUEST);
        return ProcInjectShellcodeToProcess(
            req->Method, targetPid, sc, scSize);
    }

    /* DLL-based methods: validate path */
    if (req->DllPathLength == 0 ||
        req->DllPathLength > 520 ||
        (req->DllPathLength % sizeof(WCHAR)) != 0)
        return STATUS_INVALID_PARAMETER;

    if (req->DllPathLength > Msg->DataLength - sizeof(PROC_INJECT_REQUEST))
        return STATUS_INVALID_PARAMETER;

    PWCHAR dllPath = (PWCHAR)((PUCHAR)req + sizeof(PROC_INJECT_REQUEST));
    ULONG dllChars = req->DllPathLength / sizeof(WCHAR);

    if (dllChars == 0 || dllPath[dllChars - 1] != L'\0')
        return STATUS_INVALID_PARAMETER;

    switch (req->Method) {
    case PROC_INJECT_APC:
        return InjectViaApc(targetPid, dllPath, req->DllPathLength);

    case PROC_INJECT_MANUAL_MAP: {
        ULONG headerSize = sizeof(PROC_INJECT_REQUEST) + req->DllPathLength;
        if (Msg->DataLength <= headerSize)
            return STATUS_INVALID_PARAMETER;
        PVOID peBuffer = (PUCHAR)req + headerSize;
        ULONG peSize = Msg->DataLength - headerSize;
        return InjectViaManualMap(targetPid, peBuffer, peSize);
    }

    default:
        return STATUS_INVALID_PARAMETER;
    }
}

#ifndef MemoryBasicInformation
#define MemoryBasicInformation 0
#endif

static NTSTATUS ProcDetectInjection(_Inout_ PRING_MESSAGE Msg)
{
    ULONG targetPid = Msg->ProcessId;
    if (targetPid == 0 || targetPid == 4)
        return STATUS_INVALID_PARAMETER;

    PEPROCESS target = NULL;
    NTSTATUS status = PsLookupProcessByProcessId(
        (HANDLE)(ULONG_PTR)targetPid, &target);
    if (!NT_SUCCESS(status)) return status;

    ULONG maxEntries = RING_DATA_SIZE / sizeof(PROC_DETECT_ENTRY);
    PPROC_DETECT_ENTRY out = (PPROC_DETECT_ENTRY)Msg->Data;
    ULONG count = 0;

    KAPC_STATE apc;
    KeStackAttachProcess(target, &apc);

    PVOID addr = NULL;
    MEMORY_BASIC_INFORMATION mbi;

    while (count < maxEntries) {
        SIZE_T retLen = 0;
        status = ZwQueryVirtualMemory(
            NtCurrentProcess(), addr, MemoryBasicInformation,
            &mbi, sizeof(mbi), &retLen);

        if (!NT_SUCCESS(status)) break;

        /* Detect executable private memory (not backed by image) */
        if (mbi.State == MEM_COMMIT &&
            mbi.Type == MEM_PRIVATE &&
            (mbi.Protect & (PAGE_EXECUTE | PAGE_EXECUTE_READ |
                PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY))) {
            out[count].Type = PROC_DETECT_UNBACKED_EXEC;
            out[count].Address = (ULONGLONG)mbi.BaseAddress;
            out[count].Size = (ULONGLONG)mbi.RegionSize;
            out[count].ThreadId = 0;
            count++;
        }

        /* Advance to next region */
        PVOID next = (PVOID)((ULONG_PTR)mbi.BaseAddress + mbi.RegionSize);
        if (next <= addr) break;  /* overflow guard */
        addr = next;

        /* Stay in user space */
        if (addr >= MmHighestUserAddress) break;
    }

    KeUnstackDetachProcess(&apc);
    ObDereferenceObject(target);

    Msg->DataLength = count * sizeof(PROC_DETECT_ENTRY);
    return STATUS_SUCCESS;
}