Kernel_Based_ProcessMemory_Editor/main.cpp at main · lastime1650/Kernel_Based_ProcessMemory_Editor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#include "helper.hpp"
#include "remote_call.hpp"
#include "MANUAL_KERNEL.hpp"
#include "ob_callbacks.hpp"
#include "ioctl_helper.hpp"

//
// Track the control device created by this driver for safe destruction
//
static PDEVICE_OBJECT g_ControlDeviceObject = nullptr;

extern "C" DRIVER_INITIALIZE DriverEntry;

DRIVER_UNLOAD DriverUnloadRoutine;

VOID DriverUnloadRoutine(
    _In_ PDRIVER_OBJECT DriverObject
)
{
    UNREFERENCED_PARAMETER(DriverObject);

    //
    // 1. Drain active requests and deferred DLL cleanup
    //
    helper::lifecycle::Shutdown();
    // Prevent new control requests, then cancel/join the registration worker and drain callbacks.
    helper::ob_callbacks::Shutdown();
    helper::remote_call::Shutdown();
    helper::process::ReapDllCleanup(TRUE);

    //
    // 2. Unregister device objects and symbolic links
    //
    PDEVICE_OBJECT device = g_ControlDeviceObject;
    g_ControlDeviceObject = nullptr;
    helper::ioctl::UnregisterIoctlRoutine(device, L"\\DosDevices\\MyDriver");
}

extern "C"
NTSTATUS DriverEntry(
    PDRIVER_OBJECT DriverObject,
    PUNICODE_STRING RegistryPath
)
{
    UNREFERENCED_PARAMETER(RegistryPath);

    // Record driver load start time for diagnostics telemetry
    helper::diagnostics::RecordDriverStart();
    helper::ob_callbacks::PrepareDiagnostics();

    // Initialize request rundown, user-function call tracking and DLL cleanup
    helper::lifecycle::Initialize();
    helper::remote_call::Initialize();
    helper::manual_map::Initialize();
    helper::process::InitializeDllCleanup();

    // A normal load uses its owning driver object. Keep the dynamic path
    // for callers that do not supply one, with the same device/IOCTL ABI.
    PDRIVER_OBJECT target_driver = DriverObject;
    if (target_driver == nullptr)
    {
        NTSTATUS status = helper::driver::CreateDriver(L"\\Driver\\MyDriver2", &target_driver);
        if (!NT_SUCCESS(status)) return status;
    }

    if (target_driver != nullptr)
    {
        target_driver->DriverUnload = DriverUnloadRoutine;
        target_driver->MajorFunction[IRP_MJ_CREATE] = helper::remote_call::FileRoutine;
        target_driver->MajorFunction[IRP_MJ_CLEANUP] = helper::remote_call::FileRoutine;
        target_driver->MajorFunction[IRP_MJ_CLOSE] = helper::remote_call::FileRoutine;
        target_driver->MajorFunction[IRP_MJ_DEVICE_CONTROL] = helper::ioctl_helper::DeviceControlRoutine;
    }
    else
    {
        return STATUS_UNSUCCESSFUL;
    }

    NTSTATUS status = helper::ioctl::RegisterIoctlRoutine(
        target_driver,
        L"\\Device\\MyDriver",
        L"\\DosDevices\\MyDriver",
        &g_ControlDeviceObject
    );
    if (!NT_SUCCESS(status))
    {
        if (DriverObject == nullptr) IoDeleteDriver(target_driver);
        g_ControlDeviceObject = nullptr;
    }
    // OB observers start unregistered. The global control IOCTL starts registration explicitly.
    return status;
}