#include <ntifs.h>   // ← 加这一行，IoThreadToProcess 在这里声明



// -------------------------------------------------------
// 全局变量
// -------------------------------------------------------
PVOID   g_CallbackHandle = NULL;
UNICODE_STRING g_Altitude;
WCHAR   g_AltitudeBuffer[] = L"388888";  // 海拔高度，随意取一个未被占用的值

// -------------------------------------------------------
// Pre-Operation 回调
// -------------------------------------------------------
OB_PREOP_CALLBACK_STATUS
PreOperationCallback(
    _In_ PVOID RegContext,
    _In_ POB_PRE_OPERATION_INFORMATION OpInfo
)
{
    UNREFERENCED_PARAMETER(RegContext);

    if (OpInfo == NULL) {
        return OB_PREOP_SUCCESS;
    }

    // 判断对象类型
    const BOOLEAN isProcess = (OpInfo->ObjectType == *PsProcessType);
    const BOOLEAN isThread = (OpInfo->ObjectType == *PsThreadType);

    if (!isProcess && !isThread) {
        return OB_PREOP_SUCCESS;
    }

    // 获取调用方 PID
    HANDLE callerPid = PsGetCurrentProcessId();

    // 获取目标 PID
    HANDLE targetPid = 0;
    if (isProcess && OpInfo->Object) {
        targetPid = PsGetProcessId((PEPROCESS)OpInfo->Object);
    }
    else if (isThread && OpInfo->Object) {
        PEPROCESS targetProc = IoThreadToProcess((PETHREAD)OpInfo->Object);
        if (targetProc) {
            targetPid = PsGetProcessId(targetProc);
        }
    }

    // 获取 DesiredAccess
    ACCESS_MASK desired = 0;
    if (OpInfo->Operation == OB_OPERATION_HANDLE_CREATE) {
        desired = OpInfo->Parameters->CreateHandleInformation.DesiredAccess;
    }
    else if (OpInfo->Operation == OB_OPERATION_HANDLE_DUPLICATE) {
        desired = OpInfo->Parameters->DuplicateHandleInformation.DesiredAccess;
    }

    DbgPrint("[ObDemo] %s op=%s caller=%llu target=%llu access=0x%08X\n",
        isProcess ? "PROCESS" : "THREAD",
        OpInfo->Operation == OB_OPERATION_HANDLE_CREATE ? "CREATE" : "DUPLICATE",
        (ULONG64)(ULONG_PTR)callerPid,
        (ULONG64)(ULONG_PTR)targetPid,
        desired
    );

    return OB_PREOP_SUCCESS;
}

// -------------------------------------------------------
// 注册 Object 回调
// -------------------------------------------------------
NTSTATUS RegisterCallbacks(VOID)
{
    OB_CALLBACK_REGISTRATION  cbReg = { 0 };
    OB_OPERATION_REGISTRATION opRegs[2] = { 0 };

    g_Altitude.Buffer = g_AltitudeBuffer;
    g_Altitude.Length = (USHORT)(wcslen(g_AltitudeBuffer) * sizeof(WCHAR));
    g_Altitude.MaximumLength = g_Altitude.Length + sizeof(WCHAR);

    // Process 回调
    opRegs[0].ObjectType = PsProcessType;
    opRegs[0].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
    opRegs[0].PreOperation = PreOperationCallback;
    opRegs[0].PostOperation = NULL;

    // Thread 回调
    opRegs[1].ObjectType = PsThreadType;
    opRegs[1].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
    opRegs[1].PreOperation = PreOperationCallback;
    opRegs[1].PostOperation = NULL;

    cbReg.Version = OB_FLT_REGISTRATION_VERSION;
    cbReg.OperationRegistrationCount = 2;
    cbReg.Altitude = g_Altitude;
    cbReg.RegistrationContext = NULL;
    cbReg.OperationRegistration = opRegs;

    NTSTATUS status = ObRegisterCallbacks(&cbReg, &g_CallbackHandle);
    if (NT_SUCCESS(status)) {
        DbgPrint("[ObDemo] ObRegisterCallbacks OK handle=%p\n", g_CallbackHandle);
    }
    else {
        DbgPrint("[ObDemo] ObRegisterCallbacks FAILED status=0x%08X\n", status);
    }
    return status;
}

// -------------------------------------------------------
// 驱动卸载
// -------------------------------------------------------
VOID DriverUnload(_In_ PDRIVER_OBJECT DriverObject)
{
    UNREFERENCED_PARAMETER(DriverObject);

    if (g_CallbackHandle) {
        ObUnRegisterCallbacks(g_CallbackHandle);
        g_CallbackHandle = NULL;
        DbgPrint("[ObDemo] ObUnRegisterCallbacks OK\n");
    }
    DbgPrint("[ObDemo] Driver unloaded\n");
}

// 绕过签名检查
BOOLEAN BypassCheckSign(PDRIVER_OBJECT pDriverObject)
{
#ifdef _WIN64
    typedef struct _KLDR_DATA_TABLE_ENTRY
    {
        LIST_ENTRY listEntry;
        ULONG64 __Undefined1;
        ULONG64 __Undefined2;
        ULONG64 __Undefined3;
        ULONG64 NonPagedDebugInfo;
        ULONG64 DllBase;
        ULONG64 EntryPoint;
        ULONG SizeOfImage;
        UNICODE_STRING path;
        UNICODE_STRING name;
        ULONG   Flags;
        USHORT  LoadCount;
        USHORT  __Undefined5;
        ULONG64 __Undefined6;
        ULONG   CheckSum;
        ULONG   __padding1;
        ULONG   TimeDateStamp;
        ULONG   __padding2;
    } KLDR_DATA_TABLE_ENTRY, * PKLDR_DATA_TABLE_ENTRY;
#else
    typedef struct _KLDR_DATA_TABLE_ENTRY
    {
        LIST_ENTRY listEntry;
        ULONG unknown1;
        ULONG unknown2;
        ULONG unknown3;
        ULONG unknown4;
        ULONG unknown5;
        ULONG unknown6;
        ULONG unknown7;
        UNICODE_STRING path;
        UNICODE_STRING name;
        ULONG   Flags;
    } KLDR_DATA_TABLE_ENTRY, * PKLDR_DATA_TABLE_ENTRY;
#endif

    PKLDR_DATA_TABLE_ENTRY pLdrData = (PKLDR_DATA_TABLE_ENTRY)pDriverObject->DriverSection;
    pLdrData->Flags = pLdrData->Flags | 0x20;

    return TRUE;
}

// -------------------------------------------------------
// 驱动入口
// -------------------------------------------------------
NTSTATUS DriverEntry(
    _In_ PDRIVER_OBJECT  DriverObject,
    _In_ PUNICODE_STRING RegistryPath
)
{
    UNREFERENCED_PARAMETER(RegistryPath);

    // 绕过签名检查
// LINKER_FLAGS=/INTEGRITYCHECK
    BypassCheckSign(DriverObject);

    DbgPrint("[ObDemo] DriverEntry\n");

    DriverObject->DriverUnload = DriverUnload;

    return RegisterCallbacks();
}