作者： NewbieCoder

#define PROJECT_NAME "FlowWatcher"

//
// 输出网络使用情况的间隔，单位毫秒。
//
// 由于 MILLISECOND_TO_SECOND 在将毫秒转换为秒时可能会有精度损失
// 因此此处最好是整数秒。
//

#define NETWORK_USAGE_DUMP_INTERVAL 1000

//
// 毫秒转换为内核下的等待时长。
//
// 内核下的等待时长为负数时表示超时间隔为相对时长，单位是纳秒。
//
// 有关于内核下的等待时长（即 KeWaitForSingleObject 的 Timeout 参数），参考
// http://msdn.microsoft.com/en-us/library/windows/hardware/ff553350(v=vs.85).aspx。
//

#define MILLISECOND_TO_KERNEL_DELAY(Millisecond) (-1LL * (Millisecond) * 1000 * 10)

//
// 毫秒转换为秒。
//
// 此处是整除，若不是整数秒，则会有精度损失。
//

#define MILLISECOND_TO_SECOND(Millisecond) ((Millisecond) / 1000)

//
// 获取用于附加操作的设备的设备扩展。
//

#define ATTACHING_DEVICE_EXTENSION(DeviceObject) ((PATTACHING_DEVICE_EXTENSION)((DeviceObject)->DeviceExtension))

//
// 用于记录网络使用信息的通用表的节点的缓冲池标记。
//

#define NETWORK_USAGE_GENERIC_TABLE_NODE_TAG 'NUGT'

//
// TDI 的 TCP 设备名。
//

#define TCP_DEVICE_NAME L"\\Device\\Tcp"

//
// TDI 的 UDP 设备名。
//

#define UDP_DEVICE_NAME L"\\Device\\Udp"

//
// TDI 的 RAW IP 设备名。
//

#define RAWIP_DEVICE_NAME L"\\Device\\RawIp"

//
// 计算静态数组中的元素个数。
//

#define CountOf(Array) ((sizeof(Array) / sizeof((Array)[0])))

//
// 单进程的网络使用的统计信息。
//

typedef struct _NETWORK_USAGE_INFORMATION {

//
// 自从上次输出统计信息后发送的字节数。
//
// 每次输出统计信息都会清空发送字节数。
//

ULONGLONG BytesSent;

//
// 自从上次输出统计信息后接受的字节数。
//
// 每次输出统计信息都会清空接受字节数。
//

ULONGLONG BytesReceived;
} NETWORK_USAGE_INFORMATION, *PNETWORK_USAGE_INFORMATION;

//
// 储存在通用表中的节点。
//

typedef struct _NETWORK_USAGE_GENERIC_TABLE_NODE {

//
// 进程的 PCB 指针。
//
// 在通用表中存放时，用于比较结构体相对大小使用的就是 Process 成员。
//

PEPROCESS Process;

//
// 进程的网络使用信息。
//

NETWORK_USAGE_INFORMATION NetworkUsageInformation;
} NETWORK_USAGE_GENERIC_TABLE_NODE, *PNETWORK_USAGE_GENERIC_TABLE_NODE;

//
// 用于统计进程网络使用信息的表。
//

typedef struct _NETWORK_USAGE_GENERIC_TABLE {

//
// 用于记录信息的通用表。
//

RTL_GENERIC_TABLE Table;

//
// 用于同步访问的自旋锁。
//

KSPIN_LOCK TableLock;
} NETWORK_USAGE_GENERIC_TABLE, *PNETWORK_USAGE_GENERIC_TABLE;

//
// 输出网络使用情况的工作线程控制块。
//

typedef struct _NETWORK_USAGE_DUMPER_CONTROL_BLOCK {

//
// 工作线程的指针。
//

PETHREAD Worker;

//
// 为 TRUE 时要求工作线程退出。
//

BOOLEAN Exit;

//
// 输出网络使用情况的工作线程每隔 NETWORK_USAGE_DUMP_INTERVAL 毫秒
// 输出一次。如果直接使用 KeDelayExecutionThread，那么在 DriverUnload
// 中等待工作线程退出最多会需要大约 NETWORK_USAGE_DUMP_INTERVAL
// 毫秒的时间。当 NETWORK_USAGE_DUMP_INTERVAL 增大时，DriverUnload
// 可能会需要相当长的等待时间。
//
// 考虑到 DriverUnload 是由系统的工作线程调用的，因此应当尽快返回。
// 所以此处令工作线程等待 SleepingEvent。SleepingEvent 除非在 Exit
// 为 TRUE 的时候才会置为信号态。这样，在驱动退出之前。等待 SleepingEvent
// 的结果总是会超时，这样也就实现了 KeDelayExecutionThread 的效果，
// 而 DriverUnload 中将 SleepingEvent 置为信号态，这样等待就会立即
// 结束，从而让工作线程立即退出。
//
// 如果将来 WDK 提供了类似 KeAlertThread 的例程，则可以考虑使用
// KeDelayExecutionThread + KeAlertThread 的组合。
//
// 同样，通过投递 APC 也可以唤醒工作线程，但是为了唤醒线程而投递
// APC 并非一个好的选择。
//

KEVENT SleepingEvent;
} NETWORK_USAGE_DUMPER_CONTROL_BLOCK, *PNETWORK_USAGE_DUMPER_CONTROL_BLOCK;

//
// 附加到设备的操作的状态。
//

typedef struct _DEVICE_ATTACH_STATE {

//
// 用于附加操作所创建的设备。
//

PDEVICE_OBJECT AttachingDevice;

//
// 附加到的设备。
//
// 由于我们是附加到目标设备之上，因此此设备也是紧挨我们下层的设备。
//

PDEVICE_OBJECT AttachedDevice;
} DEVICE_ATTACH_STATE, *PDEVICE_ATTACH_STATE;

//
// 用于附加操作所创建的设备的设备扩展。
//

typedef struct _ATTACHING_DEVICE_EXTENSION {

//
// 用于描述附加操作的状态。
//

DEVICE_ATTACH_STATE AttachState;
} ATTACHING_DEVICE_EXTENSION, *PATTACHING_DEVICE_EXTENSION;

//
// 用于描述附加操作的数据结构。
//

typedef struct _DEVICE_ATTACHMENT {

//
// 要附加到的设备的设备名。
//

UNICODE_STRING TargetDeviceName;

//
// 为附加操作创建的设备对象。
//

PDEVICE_OBJECT AttachingDevice;
} DEVICE_ATTACHMENT, *PDEVICE_ATTACHMENT;

//
// 输出调试信息。
//

#define DbgWrite(Format, ...) DbgPrint("[" PROJECT_NAME "] " __FUNCTION__ ": " Format, __VA_ARGS__)
#define DbgWriteLine(Format, ...) DbgWrite(Format "\n", __VA_ARGS__)

//
// 输出网络使用情况的工作线程的控制块。
//

NETWORK_USAGE_DUMPER_CONTROL_BLOCK NetworkUsageDumperControlBlock;

//
// 统计网络使用状况的表。
//

NETWORK_USAGE_GENERIC_TABLE NetworkUsageGenericTable;

//
// 需要进行的附加驱动的操作。
//

DEVICE_ATTACHMENT DeviceAttachments[] = {
{RTL_CONSTANT_STRING(TCP_DEVICE_NAME), 0},
{RTL_CONSTANT_STRING(UDP_DEVICE_NAME), 0},
{RTL_CONSTANT_STRING(RAWIP_DEVICE_NAME), 0}
};

//
// 驱动对象。
//

PDRIVER_OBJECT FlowWatcherDriverObject;

//
// 函数原型声明。
//

RTL_GENERIC_COMPARE_ROUTINE NetworkUsageGenericTableComparaRoutine;
RTL_GENERIC_ALLOCATE_ROUTINE NetworkUsageGenericTableAllocateRoutine;
RTL_GENERIC_FREE_ROUTINE NetworkUsageGenericTableFreeRoutine;
DRIVER_DISPATCH GeneralPassthroughDispath;
DRIVER_INITIALIZE DriverEntry;
DRIVER_UNLOAD DriverUnload;
KSTART_ROUTINE NetworkUsageDumper;

VOID InitializeNetworkUsageGenericTable(
VOID
);

NTSTATUS AttachDevice(
__out PDEVICE_ATTACH_STATE DeviceAttachState,
__in PDEVICE_OBJECT AttachingDevice,
__in PUNICODE_STRING DeviceName
);

NTSTATUS CreateAndAttachDevice(
__out PDEVICE_OBJECT *AttachingDevice,
__in PUNICODE_STRING DeviceName
);

VOID DetachDevice(
__in PDEVICE_ATTACH_STATE DeviceAttachState
);

VOID DeleteAndDetachDevice(
__in PDEVICE_ATTACH_STATE DeviceAttachState
);

NTSTATUS AttachDevices(
VOID
);

VOID DetachDevices(
VOID
);

NTSTATUS CreateNetworkUsageDumper(
VOID
);

VOID TerminateNetworkUsageDumper(
VOID
);

//
// 将函数关联到对应的代码段。
//

#ifdef ALLOC_PRAGMA

#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(INIT, InitializeNetworkUsageGenericTable)
#pragma alloc_text(INIT, CreateNetworkUsageDumper)
#pragma alloc_text(INIT, AttachDevice)
#pragma alloc_text(INIT, CreateAndAttachDevice)
#pragma alloc_text(INIT, AttachDevices)

#pragma alloc_text(PAGE, DriverUnload)
#pragma alloc_text(PAGE, TerminateNetworkUsageDumper)
#pragma alloc_text(PAGE, NetworkUsageDumper)
#pragma alloc_text(PAGE, DetachDevice)
#pragma alloc_text(PAGE, DeleteAndDetachDevice)
#pragma alloc_text(PAGE, DetachDevices)

#endif

//
// 用于比较两个 NETWORK_USAGE_GENERIC_TABLE_NODE 的比较函数。
//

RTL_GENERIC_COMPARE_RESULTS NTAPI NetworkUsageGenericTableComparaRoutine(
__in PRTL_GENERIC_TABLE GenericTable,
__in PVOID FirstObject,
__in PVOID SecondObject
)
{
PNETWORK_USAGE_GENERIC_TABLE_NODE LeftNode, RightNode;

LeftNode = (PNETWORK_USAGE_GENERIC_TABLE_NODE)FirstObject;
RightNode = (PNETWORK_USAGE_GENERIC_TABLE_NODE)SecondObject;

//
// 此处使用 NETWORK_USAGE_GENERIC_TABLE_NODE::Process 来比较结构体的大小。
//

if (LeftNode->Process < RightNode->Process) {
return GenericLessThan;
} else if (LeftNode->Process > RightNode->Process) {
return GenericGreaterThan;
} else {
return GenericEqual;
}
}

//
// 用于分配 NETWORK_USAGE_GENERIC_TABLE_NODE 的函数。
//

PVOID NTAPI NetworkUsageGenericTableAllocateRoutine(
__in PRTL_GENERIC_TABLE GenericTable,
__in CLONG BufferSize
)
{
return ExAllocatePoolWithTag(NonPagedPool, BufferSize, NETWORK_USAGE_GENERIC_TABLE_NODE_TAG);
}

//
// 用于释放 NETWORK_USAGE_GENERIC_TABLE_NODE 的函数。
//

VOID NTAPI NetworkUsageGenericTableFreeRoutine(
__in PRTL_GENERIC_TABLE GenericTable,
__in PVOID Buffer
)
{
ExFreePoolWithTag(Buffer, NETWORK_USAGE_GENERIC_TABLE_NODE_TAG);

return;
}

//
// 初始化统计网络使用情况的通用表。
//

VOID InitializeNetworkUsageGenericTable(
VOID
)
{
//
// 初始化用于同步访问的自旋锁。
//

KeInitializeSpinLock(&NetworkUsageGenericTable.TableLock);

//
// 初始化用于记录信息的通用表。
//

RtlInitializeGenericTable(&NetworkUsageGenericTable.Table,
NetworkUsageGenericTableComparaRoutine,
NetworkUsageGenericTableAllocateRoutine,
NetworkUsageGenericTableFreeRoutine,
NULL);

return;
}

//
// 释放统计网络使用情况的通用表（不上锁）。
//

VOID DeleteNetworkUsageGenericTableUnsafe(
VOID
)
{
PVOID CurrentEntry;

//
// 逐个删除元素。
//

for (CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, TRUE);
CurrentEntry;
CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, FALSE)) {
BOOLEAN DeleteStatus;

//
// 删除元素。
//

DeleteStatus = RtlDeleteElementGenericTable(&NetworkUsageGenericTable.Table, CurrentEntry);

ASSERT(DeleteStatus == TRUE);
}

return;
}

//
// 释放统计网络使用情况的通用表。
//

VOID DeleteNetworkUsageGenericTable(
VOID
)
{
KLOCK_QUEUE_HANDLE LockHandle;
PVOID CurrentEntry;

//
// 获取自旋锁。
//

KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
&LockHandle);

//
// 清空通用表。
//

DeleteNetworkUsageGenericTableUnsafe();

//
// 释放自旋锁。
//

KeReleaseInStackQueuedSpinLock(&LockHandle);

return;
}

//
// 获取某进程的网络使用情况（不上锁）。
//

BOOLEAN QueryProcessNetworkUsageUnsafe(
__in PEPROCESS Process,
__out PNETWORK_USAGE_INFORMATION NetworkUsageInformation
)
{
NETWORK_USAGE_GENERIC_TABLE_NODE TableNode;
PVOID EntryFound;

TableNode.Process = Process;

//
// 查询使用情况。
//

EntryFound = RtlLookupElementGenericTable(&NetworkUsageGenericTable.Table,
(PVOID)&TableNode);

if (!EntryFound) {
return FALSE;
}

//
// 返回信息。
//

*NetworkUsageInformation = ((PNETWORK_USAGE_GENERIC_TABLE_NODE)EntryFound)->NetworkUsageInformation;

return TRUE;
}

//
// 获取某进程的网络使用情况。
//

BOOLEAN QueryProcessNetworkUsage(
__in PEPROCESS Process,
__out PNETWORK_USAGE_INFORMATION NetworkUsageInformation
)
{
KLOCK_QUEUE_HANDLE LockHandle;
BOOLEAN Status;

//
// 获取自旋锁。
//

KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
&LockHandle);

//
// 查询使用情况。
//

Status = QueryProcessNetworkUsageUnsafe(Process, NetworkUsageInformation);

//
// 释放自旋锁。
//

KeReleaseInStackQueuedSpinLock(&LockHandle);

if (Status != TRUE) {
return Status;
}

return Status;
}

//
// 更新某进程的网络使用情况（不上锁）。
//

VOID UpdateProcessNetworkUsageUnsafe(
__in PEPROCESS Process,
__in PNETWORK_USAGE_INFORMATION NetworkUsageInformation
)
{
NETWORK_USAGE_GENERIC_TABLE_NODE TableNode;
PVOID EntryFound;

TableNode.Process = Process;
TableNode.NetworkUsageInformation = *NetworkUsageInformation;

//
// 删除已有节点。
//

EntryFound = RtlLookupElementGenericTable(&NetworkUsageGenericTable.Table,
(PVOID)&TableNode);

if (EntryFound) {
BOOLEAN Status;

Status = RtlDeleteElementGenericTable(&NetworkUsageGenericTable.Table, EntryFound);

ASSERT(Status == TRUE);
}

//
// 插入新节点。
//

(VOID)RtlInsertElementGenericTable(&NetworkUsageGenericTable.Table,
(PVOID)&TableNode, sizeof(NETWORK_USAGE_GENERIC_TABLE_NODE), NULL);

return;
}

//
// 更新某进程的网络使用情况。
//

VOID UpdateProcessNetworkUsage(
__in PEPROCESS Process,
__in PNETWORK_USAGE_INFORMATION NetworkUsageInformation
)
{
KLOCK_QUEUE_HANDLE LockHandle;

//
// 获取自旋锁。
//

KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
&LockHandle);

//
// 更新网络使用情况。
//

UpdateProcessNetworkUsageUnsafe(Process, NetworkUsageInformation);

//
// 释放自旋锁。
//

KeReleaseInStackQueuedSpinLock(&LockHandle);

return;
}

//
// 输出网络使用情况的统计信息。
//

VOID DumpNetworkUsage(
VOID
)
{
KLOCK_QUEUE_HANDLE LockHandle;
PVOID CurrentEntry;

//
// 获取自旋锁。
//

KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
&LockHandle);

//
// 遍历统计网络使用情况的通用表。
//

for (CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, TRUE);
CurrentEntry;
CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, FALSE)) {
PNETWORK_USAGE_GENERIC_TABLE_NODE TableNode;
PEPROCESS Process;
NETWORK_USAGE_INFORMATION NetworkUsageInformation;

TableNode = (PNETWORK_USAGE_GENERIC_TABLE_NODE)CurrentEntry;

Process = TableNode->Process;
NetworkUsageInformation = TableNode->NetworkUsageInformation;

//
// 输出信息。
//
// DbgPrint 在输出 64 位整数时似乎不能识别 %lld（至少不能识别 %10lld），因此
// 此处使用 %I64d。
//

DbgWriteLine("Process Id: %4d, Sent: %10I64d bytes / s Received: %10I64d bytes / s.",
PsGetProcessId(Process),
NetworkUsageInformation.BytesSent / MILLISECOND_TO_SECOND(NETWORK_USAGE_DUMP_INTERVAL),
NetworkUsageInformation.BytesReceived / MILLISECOND_TO_SECOND(NETWORK_USAGE_DUMP_INTERVAL));
}

//
// 清空统计数据。
//

DeleteNetworkUsageGenericTableUnsafe();

//
// 释放自旋锁。
//

KeReleaseInStackQueuedSpinLock(&LockHandle);

return;
}

//
// 用于输出当前网络使用情况的工作线程。
//

VOID NetworkUsageDumper(
__in PVOID Context
)
{
PNETWORK_USAGE_DUMPER_CONTROL_BLOCK ControlBlock;

PAGED_CODE();

//
// 获得控制块指针。
//

ControlBlock = (PNETWORK_USAGE_DUMPER_CONTROL_BLOCK)Context;

while (!ControlBlock->Exit) {
LARGE_INTEGER DumpInterval;

//
// 每隔 NETWORK_USAGE_DUMP_INTERVAL 毫秒输出一次统计信息。
//
// 有关于 SleepingEvent 的使用，参考 NETWORK_USAGE_DUMPER_CONTROL_BLOCK
// 的定义。
//

DumpInterval.QuadPart = MILLISECOND_TO_KERNEL_DELAY(NETWORK_USAGE_DUMP_INTERVAL);

(VOID)KeWaitForSingleObject((PVOID)&ControlBlock->SleepingEvent, Executive, KernelMode,
FALSE, &DumpInterval);

//
// 输出统计信息。
//

DumpNetworkUsage();
}

return;
}

//
// 附加到指定设备。
//

NTSTATUS AttachDevice(
__out PDEVICE_ATTACH_STATE DeviceAttachState,
__in PDEVICE_OBJECT AttachingDevice,
__in PUNICODE_STRING DeviceName
)
{
NTSTATUS Status;

//
// 附加到指定设备。
//

Status = IoAttachDevice(AttachingDevice, DeviceName, &DeviceAttachState->AttachedDevice);

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

//
// 填充附加状态信息。
//

DeviceAttachState->AttachingDevice = AttachingDevice;

return Status;
}

//
// 创建一个设备对象并将其附加到指定的设备上。
//

NTSTATUS CreateAndAttachDevice(
__out PDEVICE_OBJECT *AttachingDevice,
__in PUNICODE_STRING DeviceName
)
{
PATTACHING_DEVICE_EXTENSION DeviceExtension;
PDEVICE_OBJECT LocalAttachingDevice;
NTSTATUS Status;

//
// 创建用于附加操作的设备对象。
//

Status = IoCreateDevice(FlowWatcherDriverObject,
sizeof(ATTACHING_DEVICE_EXTENSION), NULL, FILE_DEVICE_NETWORK, 0, FALSE,
&LocalAttachingDevice);

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

DeviceExtension = (PATTACHING_DEVICE_EXTENSION)LocalAttachingDevice->DeviceExtension;

//
// 将设备附加到指定的设备上。
//

Status = AttachDevice(&DeviceExtension->AttachState, LocalAttachingDevice, DeviceName);

if (!NT_SUCCESS(Status)) {
IoDeleteDevice(LocalAttachingDevice);

return Status;
}

//
// 将新创建的设备对象返回给调用方。
//

*AttachingDevice = LocalAttachingDevice;

return Status;
}

//
// 从之前附加到的设备上分离。
//

VOID DetachDevice(
__in PDEVICE_ATTACH_STATE DeviceAttachState
)
{
PAGED_CODE();

//
// 分离。
//

IoDetachDevice(DeviceAttachState->AttachedDevice);

return;
}

//
// 从之前附加到的设备上分离并释放为附加操作创建的设备对象。
//

VOID DeleteAndDetachDevice(
__in PDEVICE_ATTACH_STATE DeviceAttachState
)
{
PAGED_CODE();

//
// 分离设备。
//

DetachDevice(DeviceAttachState);

//
// 删除设备。
//

IoDeleteDevice(DeviceAttachState->AttachingDevice);

return;
}

//
// 附加到各个（TCP / UDP / RAW IP） TDI 驱动的设备对象上。
//

NTSTATUS AttachDevices(
VOID
)
{
ULONG DeviceIndex;
NTSTATUS Status;

//
// 附加到 TDI 的设备上。
//

for (DeviceIndex = 0; DeviceIndex < CountOf(DeviceAttachments); DeviceIndex++) { ULONG DeviceIndex2; // // 附加到该设备上。 // Status = CreateAndAttachDevice(&DeviceAttachments[DeviceIndex].AttachingDevice, &DeviceAttachments[DeviceIndex].TargetDeviceName); if (!NT_SUCCESS(Status)) { // // 撤销已经成功的附加操作。 // for (DeviceIndex2 = 0; DeviceIndex2 < DeviceIndex; DeviceIndex2++) { // // 分离设备并删除创建的设备。 // DeleteAndDetachDevice(&ATTACHING_DEVICE_EXTENSION( DeviceAttachments[DeviceIndex2].AttachingDevice)->AttachState);
}

break;
}
}

return Status;
}

//
// 从 TDI 的各个设备对象上分离。
//

VOID DetachDevices(
VOID
)
{
ULONG DeviceIndex;

PAGED_CODE();

//
// 撤销附加到 TDI 的设备的操作。
//

for (DeviceIndex = 0; DeviceIndex < CountOf(DeviceAttachments); DeviceIndex++) { // // 分离设备并删除创建的设备。 // DeleteAndDetachDevice(&ATTACHING_DEVICE_EXTENSION( DeviceAttachments[DeviceIndex].AttachingDevice)->AttachState);
}

return;
}

//
// 创建用于输出网络使用情况的线程。
//

NTSTATUS CreateNetworkUsageDumper(
VOID
)
{
HANDLE ThreadHandle;
NTSTATUS Status;

//
// 初始化线程控制块。
//

NetworkUsageDumperControlBlock.Exit = FALSE;

//
// 用于线程休眠的事件。
//

KeInitializeEvent(&NetworkUsageDumperControlBlock.SleepingEvent, NotificationEvent, FALSE);

//
// 创建线程。
//

Status = PsCreateSystemThread(&ThreadHandle, 0, NULL, 0, NULL, NetworkUsageDumper,
(PVOID)&NetworkUsageDumperControlBlock);

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

//
// 获得线程的 ETHREAD。
//

Status = ObReferenceObjectByHandle(ThreadHandle, 0, *PsThreadType, KernelMode,
(PVOID *)&NetworkUsageDumperControlBlock.Worker, NULL);

ZwClose(ThreadHandle);

if (!NT_SUCCESS(Status)) {
NetworkUsageDumperControlBlock.Exit = TRUE;
KeSetEvent(&NetworkUsageDumperControlBlock.SleepingEvent, 0, FALSE);

return Status;
}

return Status;
}

//
// 结束用于输出网络使用情况的线程。
//

VOID TerminateNetworkUsageDumper(
VOID
)
{
PAGED_CODE();

//
// 通知线程退出。
//

NetworkUsageDumperControlBlock.Exit = TRUE;

//
// 唤醒线程。
//

KeSetEvent(&NetworkUsageDumperControlBlock.SleepingEvent, 0, TRUE);

//
// 等待线程退出。
//

KeWaitForSingleObject((PVOID)NetworkUsageDumperControlBlock.Worker, Executive, KernelMode,
FALSE, NULL);

return;
}

//
// 用于统计进程的网络使用情况。
//
// 此处只能处理 IRP_MJ_INTERNAL_DEVICE_CONTROL 类型的 IRP，对于
// IRP_MJ_DEVICE_CONTROL，需要使用 TdiMapUserRequest 将其映射为
// IRP_MJ_INTERNAL_DEVICE_CONTROL 类型的 IRP。
//

VOID AnalyzeNetworkIrp(
__in PIRP Irp
)
{
PEPROCESS Process;
PIO_STACK_LOCATION IoStackLocation;
NETWORK_USAGE_INFORMATION NetworkUsageInformation;

IoStackLocation = IoGetCurrentIrpStackLocation(Irp);

ASSERT(IoStackLocation->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL);

Process = PsGetCurrentProcess();

if (QueryProcessNetworkUsage(Process, &NetworkUsageInformation) == FALSE) {

//
// 若查询不到当前进程的网络使用情况，则表示通用表中还没有这个进程
// 相关的表项，那么就将该进程的网络使用情况初始化成 0。
//

NetworkUsageInformation.BytesReceived = 0;
NetworkUsageInformation.BytesSent = 0;
}

if (IoStackLocation->MinorFunction == TDI_RECEIVE) {
PTDI_REQUEST_KERNEL_RECEIVE Parameters;

Parameters = (PTDI_REQUEST_KERNEL_RECEIVE)&IoStackLocation->Parameters;

//
// 统计接受的字节数。
//

NetworkUsageInformation.BytesReceived += Parameters->ReceiveLength;
} else if (IoStackLocation->MinorFunction == TDI_SEND) {
PTDI_REQUEST_KERNEL_SEND Parameters;

Parameters = (PTDI_REQUEST_KERNEL_SEND)&IoStackLocation->Parameters;

//
// 统计发送的字节数。
//

NetworkUsageInformation.BytesSent += Parameters->SendLength;
} else if (IoStackLocation->MinorFunction == TDI_RECEIVE_DATAGRAM) {
PTDI_REQUEST_KERNEL_RECEIVEDG Parameters;

Parameters = (PTDI_REQUEST_KERNEL_RECEIVEDG)&IoStackLocation->Parameters;

//
// 统计接受的字节数。
//

NetworkUsageInformation.BytesReceived += Parameters->ReceiveLength;
} else if (IoStackLocation->MinorFunction == TDI_SEND_DATAGRAM) {
PTDI_REQUEST_KERNEL_SENDDG Parameters;

Parameters = (PTDI_REQUEST_KERNEL_SENDDG)&IoStackLocation->Parameters;

//
// 统计发送的字节数。
//

NetworkUsageInformation.BytesSent += Parameters->SendLength;
}

//
// 更新进程的网络使用情况。
//

UpdateProcessNetworkUsage(Process, &NetworkUsageInformation);

return;
}

//
// 通用的传递 IRP 到下层驱动的派遣函数。
//

NTSTATUS GeneralPassthroughDispath(
__in PDEVICE_OBJECT DeviceObject,
__inout PIRP Irp
)
{
PDEVICE_OBJECT LowerDeviceObject;
PIO_STACK_LOCATION IoStackLocation;

ASSERT(ATTACHING_DEVICE_EXTENSION(DeviceObject));

LowerDeviceObject = ATTACHING_DEVICE_EXTENSION(DeviceObject)->AttachState.AttachedDevice;

IoStackLocation = IoGetCurrentIrpStackLocation(Irp);

//
// 对于 IRP_MJ_DEVICE_CONTROL 类型请求，将其转换为 IRP_MJ_INTERNAL_DEVICE_CONTROL
// 类型后再处理。
//

if (IoStackLocation->MajorFunction == IRP_MJ_DEVICE_CONTROL) {

//
// 若转换失败，则 IRP 依然是 IRP_MJ_DEVICE_CONTROL 类型，那么就不会被提交给
// AnalyzeNetworkIrp 做进一步处理。
//
// 由于稍后会再次判断 IRP 的类型，因此此处可以忽略 TdiMapUserRequest 的返回值。
//

(VOID)TdiMapUserRequest(DeviceObject, Irp, IoStackLocation);
}

//
// 若是 IRP_MJ_INTERNAL_DEVICE_CONTROL 类型请求，则记录进程的网络使用情况。
//

if (IoStackLocation->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL) {
AnalyzeNetworkIrp(Irp);
}

//
// 跳过当前 IRP 参数栈。
//

IoSkipCurrentIrpStackLocation(Irp);

//
// 下发 IRP。
//

return IoCallDriver(LowerDeviceObject, Irp);
}

//
// 驱动入口点。
//

NTSTATUS DriverEntry(
__inout PDRIVER_OBJECT DriverObject,
__in PUNICODE_STRING RegistryPath
)
{
ULONG MajorFunctionIndex;
NTSTATUS Status;

FlowWatcherDriverObject = DriverObject;

//
// 填充回调表。
//

DriverObject->DriverUnload = DriverUnload;

for (MajorFunctionIndex = 0; MajorFunctionIndex <= IRP_MJ_MAXIMUM_FUNCTION; MajorFunctionIndex++) { DriverObject->MajorFunction[MajorFunctionIndex] = GeneralPassthroughDispath;
}

//
// 初始化记录进程网络使用情况的通用表。
//

InitializeNetworkUsageGenericTable();

//
// 附加到 TDI 的设备对象上。
//

Status = AttachDevices();

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

//
// 创建输出网络使用信息的线程。
//

Status = CreateNetworkUsageDumper();

if (!NT_SUCCESS(Status)) {
DetachDevices();

return Status;
}

return Status;
}

//
// 驱动卸载回调。
//

VOID DriverUnload(
__in PDRIVER_OBJECT DriverObject
)
{
PAGED_CODE();

//
// 结束输出网络使用情况的线程。
//

TerminateNetworkUsageDumper();

//
// 从 TDI 的设备对象上分离。
//

DetachDevices();

//
// 释放用于统计进程网络使用情况的通用表。
//

DeleteNetworkUsageGenericTable();

//
// 类型定义。
//

typedef struct _IP_HEADER {
UCHAR HeaderVersion;
UCHAR TypeOfService;
USHORT PacketSize;
USHORT PacketId;

//
// 分片偏移也在 Flags 中。参考 IP 报头。
//

USHORT Flags;
UCHAR TimeToLive;
UCHAR Protocol;
USHORT Checksum;
ULONG Source;
ULONG Destination;
} IP_HEADER, *PIP_HEADER;

typedef PacketFilterExtensionPtr PACKET_FILTER_EXTENSION_PTR;

//
// 字符串形式的 IP 地址的最大长度。
//

#define STRING_ADDRESS_LENGTH 16

//
// 输出调试信息。
//

#define PROJECT_NAME "FilterHookDriver"

#define DbgWrite(Format, ...) DbgPrint("[" PROJECT_NAME "] " __FUNCTION__ ": " Format, __VA_ARGS__)
#define DbgWriteLine(Format, ...) DbgWrite(Format "\n", __VA_ARGS__)

//
// 修改指针位置。
//

#define PointerAdd(Pointer, Increment, Type) ((Type)((ULONG_PTR)(Pointer) + (Increment)))

//
// 计算静态数组中的元素个数。
//

#define CountOf(Array) ((sizeof(Array) / sizeof((Array)[0])))

//
// 过滤器驱动的设备对象名称。
//

#define IP_FILTER_DEVICE_NAME L"\\Device\\Ipfilterdriver"

//
// 函数原型。
//

DRIVER_INITIALIZE DriverEntry;

DRIVER_UNLOAD DriverUnload;

NTSTATUS OpenFilterDeviceObject(
__out PDEVICE_OBJECT *DeviceObject
);

NTSTATUS RegisterFilterCallback(
__in PACKET_FILTER_EXTENSION_PTR FilterCallbackProc
);

//
// 函数所属区段标记。
//

#ifdef ALLOC_PRAGMA
#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(PAGE, DriverUnload)
#pragma alloc_text(PAGE, OpenFilterDeviceObject)
#pragma alloc_text(PAGE, RegisterFilterCallback)
#endif

//
// 将 ULONG 型 IP 地址转换为 xxx.xxx.xxx.xxx 形式的 IP 地址。
//
// StringAddress 的长度至少需要为 STRING_ADDRESS_LENGTH。
//

VOID ConvertNumericAddressToString(
__out PUCHAR StringAddress,
__in ULONG NumericAddress
)
{
ULONG CurrentPosition;
ULONG Index;

CurrentPosition = 0;

//
// 将 IP 地址按字节分开，分别转换为字符串写入 StringAddress。
//

for (Index = 0; Index < sizeof(ULONG); Index++) { LPTSTR DestinationEnd; UCHAR Current; NTSTATUS Status; // // 得到这一位的数值。 // Current = *PointerAdd(&NumericAddress, Index, PUCHAR); // // 转换为字符串。 // ASSERT(STRING_ADDRESS_LENGTH > CurrentPosition);

Status = RtlStringCchPrintfExA(PointerAdd(StringAddress, CurrentPosition, PCHAR),
STRING_ADDRESS_LENGTH - CurrentPosition, &DestinationEnd, NULL, 0, "%d",
Current);

ASSERT(NT_SUCCESS(Status));

CurrentPosition = (PUCHAR)DestinationEnd - StringAddress;

//
// 如果不是最后一段，则增加分隔符（“.”）。
//

if (Index != sizeof(ULONG) - 1) {
StringAddress[CurrentPosition++] = '.';
}
}

//
// 末尾补 0。
//

StringAddress[CurrentPosition++] = 0;

ASSERT(CurrentPosition <= STRING_ADDRESS_LENGTH); return; } // // 过滤回调。 // // Packet、PacketSize 均不包含 IP 报头。 // // 本应为 PacketFilterExtensionPtr 类型，为了方便处理此处稍微有点变化。 // PacketFilterExtensionPtr 的定义参考 // http://msdn.microsoft.com/en-us/library/windows/hardware/ff562312(v=vs.85).aspx。 // PF_FORWARD_ACTION PacketFilterCallbackProc( __in PIP_HEADER Header, __in PUCHAR Packet, __in ULONG PacketSize, __in ULONG ReceiveInterfaceIndex, __in ULONG SendInterfaceIndex, __in PVOID ReceiveLinkNextHoop, __in PVOID SendLinkNextHoop ) { static const CHAR *OperationDescription[] = {"Sent", "Received"}; UCHAR Source[STRING_ADDRESS_LENGTH], Destination[STRING_ADDRESS_LENGTH]; ULONG OperationType; // // 将 IP 地址转换为字符串形式。 // ConvertNumericAddressToString(Source, Header->Source);
ConvertNumericAddressToString(Destination, Header->Destination);

//
// 获取操作类型。
//
// 如果 ReceiveInterfaceIndex 为 INVALID_PF_IF_INDEX，则表示当前过滤的是发送操作。
// 如果 SendInterfaceIndex 为 INVALID_PF_IF_INDEX，则表示当前过滤的是接受操作。
//
// ReceiveInterfaceIndex 和 SendInterfaceIndex 中有且只可能有一个是 INVALID_PF_IF_INDEX。
//

ASSERT(((ReceiveInterfaceIndex == INVALID_PF_IF_INDEX ^
SendInterfaceIndex == INVALID_PF_IF_INDEX)) == TRUE);

if (ReceiveInterfaceIndex == INVALID_PF_IF_INDEX) {
OperationType = 0;
} else {
OperationType = 1;
}

ASSERT(OperationType < CountOf(OperationDescription)); // // 输出信息。 // DbgWriteLine("%s a packet from %s to %s with length %d", OperationDescription[OperationType], Source, Destination, PacketSize); // // 放行。 // return PF_FORWARD; } // // 打开 IP 过滤器驱动的设备对象。 // // 这个函数在返回之前会引用 IP 过滤器驱动的设备对象。稍后应释放引用。 // NTSTATUS OpenFilterDeviceObject( __out PDEVICE_OBJECT *DeviceObject ) { UNICODE_STRING DeviceName = RTL_CONSTANT_STRING(IP_FILTER_DEVICE_NAME); PDEVICE_OBJECT LocalDeviceObject; PFILE_OBJECT FileObject; NTSTATUS Status; PAGED_CODE(); // // 获得设备对象的指针。 // Status = IoGetDeviceObjectPointer(&DeviceName, SYNCHRONIZE | GENERIC_READ | GENERIC_WRITE, &FileObject, &LocalDeviceObject); if (!NT_SUCCESS(Status)) { return Status; } // // IoGetDeviceObjectPointer 只会引用 FileObject （FileObject 间接引用了 DeviceObject。） // 此处引用 DeviceObject 并释放对 FileObject 的引用。 // ObReferenceObject((PVOID)LocalDeviceObject); ObDereferenceObject((PVOID)FileObject); // // 返回设备对象给调用方。 // *DeviceObject = LocalDeviceObject; return Status; } // // 注册过滤回调。 // // FilterCallbackProc 输入 NULL 则会从过滤器驱动中清除已经注册的过滤回调。 // NTSTATUS RegisterFilterCallback( __in PACKET_FILTER_EXTENSION_PTR FilterCallbackProc ) { PF_SET_EXTENSION_HOOK_INFO SetExtensionHookInfo; PDEVICE_OBJECT FilterDeviceObject; IO_STATUS_BLOCK IoStatusBlock; KEVENT Event; PIRP Irp; NTSTATUS Status; PAGED_CODE(); // // 打开过滤器驱动的设备对象。 // Status = OpenFilterDeviceObject(&FilterDeviceObject); if (!NT_SUCCESS(Status)) { return Status; } // // 填充注册信息。 // SetExtensionHookInfo.ExtensionPointer = FilterCallbackProc; // // 初始化用于同步操作的事件。 // KeInitializeEvent(&Event, NotificationEvent, FALSE); // // 构造 IRP。 // Irp = IoBuildDeviceIoControlRequest(IOCTL_PF_SET_EXTENSION_POINTER, FilterDeviceObject, (PVOID)&SetExtensionHookInfo, sizeof(PF_SET_EXTENSION_HOOK_INFO), NULL, 0, FALSE, &Event, &IoStatusBlock); if (!Irp) { ObDereferenceObject((PVOID)FilterDeviceObject); return STATUS_INSUFFICIENT_RESOURCES; } // // 注册过滤回调。 // Status = IoCallDriver(FilterDeviceObject, Irp); if (Status == STATUS_PENDING) { (VOID)KeWaitForSingleObject((PVOID)&Event, Executive, KernelMode, FALSE, NULL); Status = IoStatusBlock.Status; } // // 释放资源。 // // IRP 会由过滤器驱动在 IoCompleteRequest 时被系统释放。 // ObDereferenceObject((PVOID)FilterDeviceObject); if (!NT_SUCCESS(Status)) { return Status; } return Status; } // // 入口点。 // NTSTATUS DriverEntry( __in PDRIVER_OBJECT DriverObject, __in PUNICODE_STRING RegistryPath ) { NTSTATUS Status; UNREFERENCED_PARAMETER(RegistryPath); DbgWriteLine("Driver loaded. Built time: %s %s.", __DATE__, __TIME__); // // 填充回调函数。 // DriverObject->DriverUnload = DriverUnload;

//
// 注册过滤回调。
//

Status = RegisterFilterCallback((PACKET_FILTER_EXTENSION_PTR)PacketFilterCallbackProc);

if (!NT_SUCCESS(Status)) {
DbgWriteLine("We failed to register our filter hook. Unloading.");

return Status;
}

return Status;
}

//
// 卸载回调。
//

VOID DriverUnload(
__in PDRIVER_OBJECT DriverObject
)
{
NTSTATUS Status;

PAGED_CODE();

//
// 注销过滤回调。
//

Status = RegisterFilterCallback(NULL);

ASSERT(NT_SUCCESS(Status));

DbgWriteLine("Driver unloaded.");

//
// 用到的宏 / 例程等的声明。
//

#define IOCTL_NDIS_QUERY_GLOBAL_STATS 0x00170002

#define OID_802_3_PERMANENT_ADDRESS 0x01010101
#define OID_802_3_CURRENT_ADDRESS 0x01010102

#define STATUS_OBJECT_NAME_NOT_FOUND ((NTSTATUS)0xC0000034L)
#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)

typedef struct _UNICODE_STRING
{
USHORT Length;
USHORT MaximumLength;
PWSTR Buffer;
} UNICODE_STRING, *PUNICODE_STRING;

NTSYSAPI
VOID
NTAPI
RtlInitUnicodeString(
IN OUT PUNICODE_STRING DestinationString,
IN PCWSTR SourceString
);

typedef struct _OBJECT_ATTRIBUTES
{
ULONG Length;
HANDLE RootDirectory;
PUNICODE_STRING ObjectName;
ULONG Attributes;
PVOID SecurityDescriptor;
PVOID SecurityQualityOfService;
} OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;

#define OBJ_CASE_INSENSITIVE 0x00000040L

#define InitializeObjectAttributes(p,n,a,r,s) { \
(p)->Length = sizeof(OBJECT_ATTRIBUTES); \
(p)->RootDirectory = (r); \
(p)->Attributes = (a); \
(p)->ObjectName = (n); \
(p)->SecurityDescriptor = (s); \
(p)->SecurityQualityOfService = NULL; \
}

typedef struct _IO_STATUS_BLOCK
{
union
{
NTSTATUS Status;
PVOID Pointer;
};
ULONG_PTR Information;
} IO_STATUS_BLOCK, *PIO_STATUS_BLOCK;

#define FILE_OPEN 0x00000001

NTSYSAPI
NTSTATUS
NTAPI
ZwCreateFile(
OUT PHANDLE FileHandle,
IN ACCESS_MASK DesiredAccess,
IN POBJECT_ATTRIBUTES ObjectAttributes,
OUT PIO_STATUS_BLOCK IoStatusBlock,
IN PLARGE_INTEGER AllocationSize OPTIONAL,
IN ULONG FileAttributes,
IN ULONG ShareAccess,
IN ULONG CreateDisposition,
IN ULONG CreateOptions,
IN PVOID EaBuffer OPTIONAL,
IN ULONG EaLength
);

NTSYSAPI
NTSTATUS
NTAPI
ZwDeviceIoControlFile(
IN HANDLE DeviceHandle,
IN HANDLE Event OPTIONAL,
IN PVOID UserApcRoutine OPTIONAL,
IN PVOID UserApcContext OPTIONAL,
OUT PIO_STATUS_BLOCK IoStatusBlock,
IN ULONG IoControlCode,
IN PVOID InputBuffer,
IN ULONG InputBufferSize,
OUT PVOID OutputBuffer,
IN ULONG OutputBufferSize
);

NTSYSAPI
NTSTATUS
NTAPI
ZwClose(
IN HANDLE Handle
);

//
// MAC 地址长度。（6字节）
//

#define PHYSICAL_ADDRESS_LENGTH 6

//
// 打开设备。
//

NTSTATUS ObtainDeviceHandle(
__out PHANDLE DeviceHandle,
__in PUNICODE_STRING DeviceName
)
{
OBJECT_ATTRIBUTES ObjectAttributes;
IO_STATUS_BLOCK IoStatusBlock;
HANDLE LocalDeviceHandle;
NTSTATUS Status;

//
// 初始化 OBJECT_ATTRIBUTES。
//

InitializeObjectAttributes(&ObjectAttributes, DeviceName, OBJ_CASE_INSENSITIVE,
NULL, NULL);

//
// 尝试打开设备。
//

Status = ZwCreateFile(&LocalDeviceHandle, FILE_READ_ACCESS, &ObjectAttributes,
&IoStatusBlock, NULL, 0, FILE_SHARE_READ | FILE_SHARE_WRITE, FILE_OPEN, 0, NULL, 0);

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

//
// 将句柄返回给调用方。
//

*DeviceHandle = LocalDeviceHandle;

return Status;
}

//
// 获取网卡信息。
//

NTSTATUS QueryNdisDeviceInformation(
__in HANDLE DeviceHandle,
__in ULONG ObjectIdentifier,
__out_bcount_opt(OutputBufferSize) PVOID OutputBuffer,
__in SIZE_T OutputBufferSize
)
{
IO_STATUS_BLOCK IoStatusBlock;
NTSTATUS Status;

//
// 发起请求。
//

Status = ZwDeviceIoControlFile(DeviceHandle, NULL, NULL, NULL, &IoStatusBlock,
IOCTL_NDIS_QUERY_GLOBAL_STATS, (PVOID)&ObjectIdentifier, sizeof(ULONG),
OutputBuffer, OutputBufferSize);

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

return Status;
}

//
// 获取网卡 MAC 地址。
//
// 网卡 MAC 地址可以是出厂时的 MAC 地址，也可以是管理员指派的 MAC 地址。
//
// 输出参数 PhysicalAddress 至少需要 PHYSICAL_ADDRESS_LENGTH 个字节大小。
//（MAC 地址占 PHYSICAL_ADDRESS_LENGTH 字节。）
//

NTSTATUS QueryNdisDevicePhysicalAddress(
__out PUCHAR PhysicalAddress,
__in HANDLE DeviceHandle,
__in BOOLEAN PermanentPhysicalAddress
)
{
//
// 向设备查询网卡地址。
//

return QueryNdisDeviceInformation(DeviceHandle,
(PermanentPhysicalAddress ? OID_802_3_PERMANENT_ADDRESS : OID_802_3_CURRENT_ADDRESS),
(PVOID)PhysicalAddress, PHYSICAL_ADDRESS_LENGTH);
}

//
// 显示网卡信息。
//

VOID DumpNdisDevicePhysicalAddress(
__in PUNICODE_STRING DeviceName
)
{
HANDLE DeviceHandle;
BOOLEAN PermanentPhysicalAddressCandidates[] = {TRUE, FALSE};
LPCSTR PhysicalAddressDescriptions[] = {"permanent physical address", "current physical address"};
ULONG Index;
NTSTATUS Status;

//
// 打开网卡设备。
//

Status = ObtainDeviceHandle(&DeviceHandle, DeviceName);

if (!NT_SUCCESS(Status)) {
printf("We failed to open that device, status = 0x%08X\n", Status);

return;
}

//
// 对所有 PermanentPhysicalAddress 可能的取值分别查询其对应的物理地址。
//

//
// 用到的各种宏、类型、函数的声明。
//

#define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
#define STATUS_INSUFFICIENT_RESOURCES ((NTSTATUS)0xC000009AL)

typedef __success(return >= 0) LONG NTSTATUS;
typedef NTSTATUS *PNTSTATUS;

#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)

typedef struct _STRING
{
USHORT Length;
USHORT MaximumLength;
__field_bcount_part_opt(MaximumLength, Length) PCHAR Buffer;
} STRING, *PSTRING, ANSI_STRING, *PANSI_STRING, OEM_STRING, *POEM_STRING;

typedef struct _UNICODE_STRING
{
USHORT Length;
USHORT MaximumLength;
__field_bcount_part(MaximumLength, Length) PWCH Buffer;
} UNICODE_STRING, *PUNICODE_STRING;

#define RTL_CONSTANT_STRING(s) { sizeof(s) - sizeof((s)[0]), sizeof(s), s }

typedef struct _OBJECT_ATTRIBUTES
{
ULONG Length;
HANDLE RootDirectory;
PUNICODE_STRING ObjectName;
ULONG Attributes;
PVOID SecurityDescriptor;
PVOID SecurityQualityOfService;
} OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;

#define DOS_MAX_COMPONENT_LENGTH 255
#define DOS_MAX_PATH_LENGTH (DOS_MAX_COMPONENT_LENGTH + 5)

typedef struct _CURDIR
{
UNICODE_STRING DosPath;
HANDLE Handle;
} CURDIR, *PCURDIR;

#define RTL_USER_PROC_CURDIR_CLOSE 0x00000002
#define RTL_USER_PROC_CURDIR_INHERIT 0x00000003

typedef struct _RTL_DRIVE_LETTER_CURDIR
{
USHORT Flags;
USHORT Length;
ULONG TimeStamp;
STRING DosPath;
} RTL_DRIVE_LETTER_CURDIR, *PRTL_DRIVE_LETTER_CURDIR;

#define RTL_MAX_DRIVE_LETTERS 32
#define RTL_DRIVE_LETTER_VALID (USHORT)0x0001

typedef struct _RTL_USER_PROCESS_PARAMETERS
{
ULONG MaximumLength;
ULONG Length;

ULONG Flags;
ULONG DebugFlags;

HANDLE ConsoleHandle;
ULONG ConsoleFlags;
HANDLE StandardInput;
HANDLE StandardOutput;
HANDLE StandardError;

CURDIR CurrentDirectory;
UNICODE_STRING DllPath;
UNICODE_STRING ImagePathName;
UNICODE_STRING CommandLine;
PVOID Environment;

ULONG StartingX;
ULONG StartingY;
ULONG CountX;
ULONG CountY;
ULONG CountCharsX;
ULONG CountCharsY;
ULONG FillAttribute;

ULONG WindowFlags;
ULONG ShowWindowFlags;
UNICODE_STRING WindowTitle;
UNICODE_STRING DesktopInfo;
UNICODE_STRING ShellInfo;
UNICODE_STRING RuntimeData;
RTL_DRIVE_LETTER_CURDIR CurrentDirectories[RTL_MAX_DRIVE_LETTERS];

ULONG EnvironmentSize;
ULONG EnvironmentVersion;
} RTL_USER_PROCESS_PARAMETERS, *PRTL_USER_PROCESS_PARAMETERS;

#define RTL_USER_PROC_PARAMS_NORMALIZED 0x00000001

// begin_rev
#define PS_ATTRIBUTE_NUMBER_MASK 0x0000ffff
#define PS_ATTRIBUTE_THREAD 0x00010000 // can be used with threads
#define PS_ATTRIBUTE_INPUT 0x00020000 // input only
#define PS_ATTRIBUTE_UNKNOWN 0x00040000
// end_rev

// private
typedef enum _PS_ATTRIBUTE_NUM
{
PsAttributeParentProcess, // in HANDLE
PsAttributeDebugPort, // in HANDLE
PsAttributeToken, // in HANDLE
PsAttributeClientId, // out PCLIENT_ID
PsAttributeTebAddress, // out PTEB *
PsAttributeImageName, // in PWSTR
PsAttributeImageInfo, // out PSECTION_IMAGE_INFORMATION
PsAttributeMemoryReserve, // in PPS_MEMORY_RESERVE
PsAttributePriorityClass, // in UCHAR
PsAttributeErrorMode, // in ULONG
PsAttributeStdHandleInfo, // 10, in PPS_STD_HANDLE_INFO
PsAttributeHandleList, // in PHANDLE
PsAttributeGroupAffinity, // in PGROUP_AFFINITY
PsAttributePreferredNode, // in PUSHORT
PsAttributeIdealProcessor, // in PPROCESSOR_NUMBER
PsAttributeUmsThread, // ? in PUMS_CREATE_THREAD_ATTRIBUTES
PsAttributeMitigationOptions, // in UCHAR
PsAttributeMax
} PS_ATTRIBUTE_NUM;

// begin_rev

#define PsAttributeValue(Number, Thread, Input, Unknown) \
(((Number) & PS_ATTRIBUTE_NUMBER_MASK) | \
((Thread) ? PS_ATTRIBUTE_THREAD : 0) | \
((Input) ? PS_ATTRIBUTE_INPUT : 0) | \
((Unknown) ? PS_ATTRIBUTE_UNKNOWN : 0))

#define PS_ATTRIBUTE_PARENT_PROCESS \
PsAttributeValue(PsAttributeParentProcess, FALSE, TRUE, TRUE)
#define PS_ATTRIBUTE_DEBUG_PORT \
PsAttributeValue(PsAttributeDebugPort, FALSE, TRUE, TRUE)
#define PS_ATTRIBUTE_TOKEN \
PsAttributeValue(PsAttributeToken, FALSE, TRUE, TRUE)
#define PS_ATTRIBUTE_CLIENT_ID \
PsAttributeValue(PsAttributeClientId, TRUE, FALSE, FALSE)
#define PS_ATTRIBUTE_TEB_ADDRESS \
PsAttributeValue(PsAttributeTebAddress, TRUE, FALSE, FALSE)
#define PS_ATTRIBUTE_IMAGE_NAME \
PsAttributeValue(PsAttributeImageName, FALSE, TRUE, FALSE)
#define PS_ATTRIBUTE_IMAGE_INFO \
PsAttributeValue(PsAttributeImageInfo, FALSE, FALSE, FALSE)
#define PS_ATTRIBUTE_MEMORY_RESERVE \
PsAttributeValue(PsAttributeMemoryReserve, FALSE, TRUE, FALSE)
#define PS_ATTRIBUTE_PRIORITY_CLASS \
PsAttributeValue(PsAttributePriorityClass, FALSE, TRUE, FALSE)
#define PS_ATTRIBUTE_ERROR_MODE \
PsAttributeValue(PsAttributeErrorMode, FALSE, TRUE, FALSE)
#define PS_ATTRIBUTE_STD_HANDLE_INFO \
PsAttributeValue(PsAttributeStdHandleInfo, FALSE, TRUE, FALSE)
#define PS_ATTRIBUTE_HANDLE_LIST \
PsAttributeValue(PsAttributeHandleList, FALSE, TRUE, FALSE)
#define PS_ATTRIBUTE_GROUP_AFFINITY \
PsAttributeValue(PsAttributeGroupAffinity, TRUE, TRUE, FALSE)
#define PS_ATTRIBUTE_PREFERRED_NODE \
PsAttributeValue(PsAttributePreferredNode, FALSE, TRUE, FALSE)
#define PS_ATTRIBUTE_IDEAL_PROCESSOR \
PsAttributeValue(PsAttributeIdealProcessor, TRUE, TRUE, FALSE)
#define PS_ATTRIBUTE_MITIGATION_OPTIONS \
PsAttributeValue(PsAttributeMitigationOptions, FALSE, TRUE, TRUE)

// end_rev

// begin_private

typedef struct _PS_ATTRIBUTE
{
ULONG Attribute;
SIZE_T Size;
union
{
ULONG Value;
PVOID ValuePtr;
};
PSIZE_T ReturnLength;
} PS_ATTRIBUTE, *PPS_ATTRIBUTE;

typedef struct _PS_ATTRIBUTE_LIST
{
SIZE_T TotalLength;
PS_ATTRIBUTE Attributes[1];
} PS_ATTRIBUTE_LIST, *PPS_ATTRIBUTE_LIST;

typedef struct _PS_MEMORY_RESERVE
{
PVOID ReserveAddress;
SIZE_T ReserveSize;
} PS_MEMORY_RESERVE, *PPS_MEMORY_RESERVE;

typedef enum _PS_STD_HANDLE_STATE
{
PsNeverDuplicate,
PsRequestDuplicate, // duplicate standard handles specified by PseudoHandleMask, and only if StdHandleSubsystemType matches the image subsystem
PsAlwaysDuplicate, // always duplicate standard handles
PsMaxStdHandleStates
} PS_STD_HANDLE_STATE;

// begin_rev
#define PS_STD_INPUT_HANDLE 0x1
#define PS_STD_OUTPUT_HANDLE 0x2
#define PS_STD_ERROR_HANDLE 0x4
// end_rev

typedef struct _PS_STD_HANDLE_INFO
{
union
{
ULONG Flags;
struct
{
ULONG StdHandleState : 2; // PS_STD_HANDLE_STATE
ULONG PseudoHandleMask : 3; // PS_STD_*
};
};
ULONG StdHandleSubsystemType;
} PS_STD_HANDLE_INFO, *PPS_STD_HANDLE_INFO;

// windows-internals-book:"Chapter 5"
typedef enum _PS_CREATE_STATE
{
PsCreateInitialState,
PsCreateFailOnFileOpen,
PsCreateFailOnSectionCreate,
PsCreateFailExeFormat,
PsCreateFailMachineMismatch,
PsCreateFailExeName, // Debugger specified
PsCreateSuccess,
PsCreateMaximumStates
} PS_CREATE_STATE;

typedef enum _PS_IFEO_KEY_STATE
{
PsReadIFEOAllValues,
PsSkipIFEODebugger,
PsSkipAllIFEO,
PsMaxIFEOKeyStates
} PS_IFEO_KEY_STATE, *PPS_IFEO_KEY_STATE;

typedef struct _PS_CREATE_INFO
{
SIZE_T Size;
PS_CREATE_STATE State;
union
{
// PsCreateInitialState
struct
{
union
{
ULONG InitFlags;
struct
{
UCHAR WriteOutputOnExit : 1;
UCHAR DetectManifest : 1;
UCHAR SpareBits1 : 6;
UCHAR IFEOKeyState : 2; // PS_IFEO_KEY_STATE
UCHAR SpareBits2 : 6;
USHORT ProhibitedImageCharacteristics : 16;
};
};
ACCESS_MASK AdditionalFileAccess;
} InitState;

// PsCreateFailOnSectionCreate
struct
{
HANDLE FileHandle;
} FailSection;

// PsCreateFailExeName
struct
{
HANDLE IFEOKey;
} ExeName;

// PsCreateSuccess
struct
{
union
{
ULONG OutputFlags;
struct
{
UCHAR ProtectedProcess : 1;
UCHAR AddressSpaceOverride : 1;
UCHAR DevOverrideEnabled : 1; // from Image File Execution Options
UCHAR ManifestDetected : 1;
UCHAR SpareBits1 : 4;
UCHAR SpareBits2 : 8;
USHORT SpareBits3 : 16;
};
};
HANDLE FileHandle;
HANDLE SectionHandle;
ULONGLONG UserProcessParametersNative;
ULONG UserProcessParametersWow64;
ULONG CurrentParameterFlags;
ULONGLONG PebAddressNative;
ULONG PebAddressWow64;
ULONGLONG ManifestAddress;
ULONG ManifestSize;
} SuccessState;
};
} PS_CREATE_INFO, *PPS_CREATE_INFO;

extern "C"
NTSYSCALLAPI
NTSTATUS
NTAPI
ZwCreateUserProcess(
__out PHANDLE ProcessHandle,
__out PHANDLE ThreadHandle,
__in ACCESS_MASK ProcessDesiredAccess,
__in ACCESS_MASK ThreadDesiredAccess,
__in_opt POBJECT_ATTRIBUTES ProcessObjectAttributes,
__in_opt POBJECT_ATTRIBUTES ThreadObjectAttributes,
__in ULONG ProcessFlags,
__in ULONG ThreadFlags,
__in_opt PRTL_USER_PROCESS_PARAMETERS ProcessParameters,
__inout PPS_CREATE_INFO CreateInfo,
__in_opt PPS_ATTRIBUTE_LIST AttributeList
);;

#define IMAGE_PATH L"\\??\\C:\\Windows\\System32\\cmd.exe"

#define PointerAdd(Pointer, Increment, ResultType) ((ResultType)((ULONG_PTR)(Pointer) \
+ (ULONG_PTR)(Increment)))

#define RoundUp(Value, Alignment) (((Value) + (Alignment) - 1) / (Alignment) * (Alignment))

//
// 创建用户进程参数。
//

static NTSTATUS PsCreateProcessParameters(
OUT PRTL_USER_PROCESS_PARAMETERS *ProcessParameters,
IN PUNICODE_STRING ImagePath,
IN PUNICODE_STRING CurrentDirectory,
IN PWSTR Environment,
IN ULONG EnvironmentSize
)
{
static UNICODE_STRING EmptyUnicodeString = RTL_CONSTANT_STRING(L"");
static UNICODE_STRING DefaultDesktop = RTL_CONSTANT_STRING(L"Winsta0\\Default");

PRTL_USER_PROCESS_PARAMETERS LocalProcessParameters;
PUNICODE_STRING CommandLine, WindowTitle, DesktopInfo, ShellInfo, RuntimeData, DllPath;
ULONG ProcessParametersSize;
PWSTR CurrentString;

CommandLine = ImagePath;
DllPath = &EmptyUnicodeString;
WindowTitle = &EmptyUnicodeString;
DesktopInfo = &DefaultDesktop;
ShellInfo = &EmptyUnicodeString;
RuntimeData = &EmptyUnicodeString;

//
// 计算用户进程参数的大小。
//

ProcessParametersSize = sizeof(RTL_USER_PROCESS_PARAMETERS) +
RoundUp(ImagePath->Length + sizeof(UNICODE_NULL), sizeof(ULONG)) +
RoundUp(CommandLine->Length + sizeof(UNICODE_NULL), sizeof(ULONG)) +
RoundUp(DllPath->MaximumLength, sizeof(ULONG)) +
RoundUp(WindowTitle->MaximumLength, sizeof(ULONG)) +
RoundUp(DesktopInfo->MaximumLength, sizeof(ULONG)) +
RoundUp(ShellInfo->MaximumLength, sizeof(ULONG)) +
RoundUp(RuntimeData->MaximumLength, sizeof(ULONG)) +
RoundUp(DOS_MAX_PATH_LENGTH / sizeof(CHAR) * sizeof(WCHAR), sizeof(ULONG));

//
// 为用户进程参数分配内存。
//

LocalProcessParameters = (PRTL_USER_PROCESS_PARAMETERS)malloc(ProcessParametersSize);

if (!LocalProcessParameters) {
return STATUS_INSUFFICIENT_RESOURCES;
}

//
// 初始化用户进程参数。
//

RtlZeroMemory(LocalProcessParameters, ProcessParametersSize);

LocalProcessParameters->Length = LocalProcessParameters->MaximumLength =
ProcessParametersSize;
LocalProcessParameters->Flags = RTL_USER_PROC_PARAMS_NORMALIZED;
LocalProcessParameters->DebugFlags = 0;
LocalProcessParameters->Environment = (PVOID)Environment;
LocalProcessParameters->EnvironmentSize = EnvironmentSize;
LocalProcessParameters->CurrentDirectory.Handle = 0;
LocalProcessParameters->ConsoleFlags = 0;

//
// 复制字符串。
//

#define PspCopyUnicodeString(BufferPointer, Destination, Source, AllocationSize) \
(Destination)->Buffer = BufferPointer; \
(Destination)->Length = (Source)->Length; \
(Destination)->MaximumLength = AllocationSize; \
RtlCopyMemory(BufferPointer, (PVOID)(Source)->Buffer, Source->Length); \
BufferPointer = PointerAdd(BufferPointer, RoundUp( \
AllocationSize, sizeof(ULONG)), PWSTR);

CurrentString = (PWCHAR)(LocalProcessParameters + 1);

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->ImagePathName,
ImagePath, ImagePath->Length + sizeof(UNICODE_NULL));

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->CommandLine,
CommandLine, CommandLine->Length + sizeof(UNICODE_NULL));

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->DllPath,
DllPath, DllPath->MaximumLength);

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->WindowTitle,
WindowTitle, WindowTitle->MaximumLength);

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->DesktopInfo,
DesktopInfo, DesktopInfo->MaximumLength);

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->ShellInfo,
ShellInfo, ShellInfo->MaximumLength);

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->RuntimeData,
RuntimeData, RuntimeData->MaximumLength);

PspCopyUnicodeString(CurrentString, &LocalProcessParameters->CurrentDirectory.DosPath,
CurrentDirectory, DOS_MAX_PATH_LENGTH / sizeof(CHAR) * sizeof(WCHAR));

#undef PspCopyUnicodeString

//
// 返回结果给调用方。
//

*ProcessParameters = LocalProcessParameters;

return STATUS_SUCCESS;
}

//
// 释放用户进程参数。
//

static VOID PsFreeProcessParameters(
IN PRTL_USER_PROCESS_PARAMETERS ProcessParameters
)
{
free((PVOID)ProcessParameters);

return;
}

//
// 创建环境变量信息。
//

static NTSTATUS PsCreateEnvironment(
OUT PWSTR *Environment,
OUT ULONG *EnvironmentSize
)
{
static PWSTR LocalEnvironment = L"\0";

//
// FIX: 目前只是返回一个空的环境变量信息。应改为从注册表（
// Registry\Machine\System\CurrentControlSet\Control\Session Manager\Environment）
// 查询。
//
// 环境变量信息也可以从 PEB 中获取。
//

*Environment = LocalEnvironment;
*EnvironmentSize = sizeof(LocalEnvironment);

return STATUS_SUCCESS;
}

//
// 释放环境变量信息。
//

static VOID PsFreeEnvironment(
IN PWSTR Environment
)
{
//
// 因为目前环境变量信息是一个静态变量，因此不需释放。
//

UNREFERENCED_PARAMETER(Environment);

return;
}

//
// 创建进程。
//

NTSTATUS PsCreateProcess(
OUT HANDLE *ProcessHandle,
OUT HANDLE *ThreadHandle,
IN PUNICODE_STRING ImagePath
)
{
static UNICODE_STRING CurrentDirectory = RTL_CONSTANT_STRING(L"C:\\");

PRTL_USER_PROCESS_PARAMETERS ProcessParameters;
PWSTR Environment;
ULONG EnvironmentSize;
PS_CREATE_INFO CreateInfo = {0};
PS_ATTRIBUTE_LIST AttributeList = {0};
NTSTATUS Status;

//
// 创建环境变量。
//

Status = PsCreateEnvironment(&Environment, &EnvironmentSize);

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

//
// 创建进程参数。
//

Status = PsCreateProcessParameters(&ProcessParameters, ImagePath, &CurrentDirectory,
Environment, EnvironmentSize);

if (!NT_SUCCESS(Status)) {
PsFreeEnvironment(Environment);

return Status;
}

//
// 初始化 PS_CREATE_INFO 结构体。
//

CreateInfo.Size = sizeof(PS_CREATE_INFO);
CreateInfo.State = PsCreateInitialState;

//
// 将文件名通过 PS_ATTRIBUTE_LIST 传给 ZwCreateUserProcess。
//

AttributeList.TotalLength = sizeof(PS_ATTRIBUTE_LIST);
AttributeList.Attributes[0].Attribute = PS_ATTRIBUTE_IMAGE_NAME;
AttributeList.Attributes[0].Size = ImagePath->Length;
AttributeList.Attributes[0].ValuePtr = (PVOID)ImagePath->Buffer;

//
// 创建进程。
//

Status = ZwCreateUserProcess(ProcessHandle, ThreadHandle, MAXIMUM_ALLOWED, MAXIMUM_ALLOWED,
NULL, NULL, 0, 0, ProcessParameters, &CreateInfo, &AttributeList);

PsFreeProcessParameters(ProcessParameters);
PsFreeEnvironment(Environment);

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

return Status;
}

int main() {
static UNICODE_STRING strImagePath = RTL_CONSTANT_STRING(IMAGE_PATH);

HANDLE hProcess, hThread;
NTSTATUS Status;

Status = PsCreateProcess(&hThread, &hProcess, &strImagePath);

if (NT_SUCCESS(Status)) {
CloseHandle(hThread);
CloseHandle(hProcess);
}