监控系统所有进程的创建和销毁 (PsSetCreateProcessNotifyRoutine)

使用驱动方式，在原来的ProcObsrv.c基础上进行了完善，所有进程的创建和销毁都不会丢失，能完全捕获到。

具体代码如下：

//---------------------------------------------------------------------------
//
// ProcObsrv.c
//
// SUBSYSTEM:
// System monitor
// MODULE:
// Driver for monitoring NT process and DLLs mapping
// monitoring.
//
// DESCRIPTION:
// This code is based on the James Finnegan抯 sample
// (MSJ January 1999).
//
// Ivo Ivanov, January 2002
//
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
//
// Includes
//
//---------------------------------------------------------------------------
#include <ntddk.h>

//---------------------------------------------------------------------------
//
// Defines
//
//---------------------------------------------------------------------------

#define FILE_DEVICE_UNKNOWN 0x00000022
#define IOCTL_UNKNOWN_BASE FILE_DEVICE_UNKNOWN
#define IOCTL_PROCOBSRV_ACTIVATE_MONITORING /
CTL_CODE(IOCTL_UNKNOWN_BASE, 0x0800, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define IOCTL_PROCOBSRV_GET_PROCINFO /
CTL_CODE(IOCTL_UNKNOWN_BASE, 0x0801, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
//---------------------------------------------------------------------------
//
// Forward declaration
//
//---------------------------------------------------------------------------
void UnloadDriver(
PDRIVER_OBJECT DriverObject
);
NTSTATUS DispatchCreateClose(
IN PDEVICE_OBJECT DeviceObject,
IN PIRP Irp
);
NTSTATUS DispatchIoctl(
IN PDEVICE_OBJECT DeviceObject,
IN PIRP Irp
);
//
// Process function callback
//
VOID ProcessCallback(
IN HANDLE hParentId,
IN HANDLE hProcessId,
IN BOOLEAN bCreate
);
//
// Structure for holding info about activating/deactivating the driver
//
typedef struct _ActivateInfo
{
BOOLEAN bActivated;
} ACTIVATE_INFO, *PACTIVATE_INFO;

//
// Structure for process callback information
//
typedef struct _ProcessCallbackInfo
{
HANDLE hParentId;
HANDLE hProcessId;
BOOLEAN bCreate;
} PROCESS_CALLBACK_INFO, *PPROCESS_CALLBACK_INFO;

//
// Private storage for process retreiving
//
typedef struct _DEVICE_EXTENSION
{
PDEVICE_OBJECT DeviceObject;
//
// Shared section
//
HANDLE hProcessId;
//
// Process section data
//
PKEVENT ProcessEvent;
HANDLE hParentId;
BOOLEAN bCreate;
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;

typedef struct
{
LIST_ENTRY list_entry;

HANDLE hProcessId;
HANDLE hParentId;
BOOLEAN bCreate;
} CUSTOM_EXTENSION, *PCUSTOM_EXTENSION;

#define MEM_TAG 'Tony'
#define MAX_LIST_COUNT 1024

#define CONTAINING_RECORD(address, type, field) ((type *)( /
(PCHAR)(address) - /
(ULONG_PTR)(&((type *)0)->field)))

//
// Global variables
//
PDEVICE_OBJECT g_pDeviceObject;
ACTIVATE_INFO g_ActivateInfo;

// tony tang
KEVENT g_WaitEvent;
LIST_ENTRY g_list_head;
int g_list_count = 0;

//
// The main entry point of the driver module
//
NTSTATUS DriverEntry(
IN PDRIVER_OBJECT DriverObject,
IN PUNICODE_STRING RegistryPath
)
{
NTSTATUS ntStatus;
UNICODE_STRING uszDriverString;
UNICODE_STRING uszDeviceString;
UNICODE_STRING uszProcessEventString;
PDEVICE_OBJECT pDeviceObject;
PDEVICE_EXTENSION extension;
HANDLE hProcessHandle;

DbgPrint("==========================DriverEntry==========================");

// tony
KeInitializeEvent( &g_WaitEvent, SynchronizationEvent, TRUE );
InitializeListHead(&g_list_head);

//
// Point uszDriverString at the driver name
//
RtlInitUnicodeString(&uszDriverString, L"//Device//ProcObsrv");
//
// Create and initialize device object
//
ntStatus = IoCreateDevice(
DriverObject,
sizeof(DEVICE_EXTENSION),
&uszDriverString,
FILE_DEVICE_UNKNOWN,
0,
FALSE,
&pDeviceObject
);
if(ntStatus != STATUS_SUCCESS)
return ntStatus;
//
// Assign extension variable
//
extension = pDeviceObject->DeviceExtension;
//
// Point uszDeviceString at the device name
//
RtlInitUnicodeString(&uszDeviceString, L"//DosDevices//ProcObsrv");
//
// Create symbolic link to the user-visible name
//
ntStatus = IoCreateSymbolicLink(&uszDeviceString, &uszDriverString);

if(ntStatus != STATUS_SUCCESS)
{
//
// Delete device object if not successful
//
IoDeleteDevice(pDeviceObject);
return ntStatus;
}
//
// Assign global pointer to the device object for use by the callback functions
//
g_pDeviceObject = pDeviceObject;
//
// Setup initial state
//
g_ActivateInfo.bActivated = FALSE;
//
// Load structure to point to IRP handlers
//
DriverObject->DriverUnload = UnloadDriver;
DriverObject->MajorFunction[IRP_MJ_CREATE] = DispatchCreateClose;
DriverObject->MajorFunction[IRP_MJ_CLOSE] = DispatchCreateClose;
DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DispatchIoctl;
//
// Create event for user-mode processes to monitor
//
RtlInitUnicodeString(
&uszProcessEventString,
L"//BaseNamedObjects//ProcObsrvProcessEvent"
);
extension->ProcessEvent = IoCreateNotificationEvent(
&uszProcessEventString,
&hProcessHandle
);

// Clear it out
KeClearEvent(extension->ProcessEvent);

// Return success
return ntStatus;
}

//
// Create and close routine
//
NTSTATUS DispatchCreateClose(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
Irp->IoStatus.Status = STATUS_SUCCESS;
Irp->IoStatus.Information = 0;
IoCompleteRequest(Irp, IO_NO_INCREMENT);

return STATUS_SUCCESS;
}

NTSTATUS ProcessInfoAppendNode(HANDLE hProcessId, HANDLE hParentId, BOOLEAN bCreate)
{
DbgPrint(" ProcessInfoAppendNode.(hProcessId=%d, hParentId=%d, bCreate=%d)", hProcessId, hParentId, bCreate);

// 限制最大的链表节点数，这里默认为1024
if(g_list_count < MAX_LIST_COUNT)
{
PCUSTOM_EXTENSION extension = (PCUSTOM_EXTENSION)ExAllocatePoolWithTag(PagedPool, sizeof(CUSTOM_EXTENSION), MEM_TAG);
if(extension == NULL)
{
return STATUS_INTERNAL_ERROR;
}

extension->hProcessId = hProcessId;
extension->hParentId = hParentId;
extension->bCreate = bCreate;

// 插入到链表末尾。请注意这里没有使用任何锁。所以，这个函数不是多多线程安全的。在下面自旋锁的使用中讲解如何保证多线程安全性。
InsertTailList(&g_list_head, (PLIST_ENTRY)&extension->list_entry);

g_list_count++;
}
else
{
DbgPrint("!!!!! The list more than maximum number.(Default is 1024)");
}

return STATUS_SUCCESS;
}

//
// Process function callback
//
VOID ProcessCallback(
IN HANDLE hParentId,
IN HANDLE hProcessId,
IN BOOLEAN bCreate
)
{
PDEVICE_EXTENSION extension;

// Assign extension variable
extension = g_pDeviceObject->DeviceExtension;

// Assign current values into device extension.
// User-mode apps will pick it up using DeviceIoControl calls.
extension->hParentId = hParentId;
extension->hProcessId = hProcessId;
extension->bCreate = bCreate;

// 将进程信息插入到链表中
KeWaitForSingleObject(&g_WaitEvent, Executive, KernelMode, FALSE,(PLARGE_INTEGER)NULL);
ProcessInfoAppendNode(hProcessId, hParentId, bCreate);
KeSetEvent(&g_WaitEvent ,0 ,FALSE);

// Signal the event thus the user-mode apps listening will be aware
// that something interesting has happened.
KeSetEvent(extension->ProcessEvent, 0, FALSE);
KeClearEvent(extension->ProcessEvent);
}

//
// IOCTL handler for setting the callback
//
NTSTATUS ActivateMonitoringHanlder(
IN PIRP Irp
)
{
NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(Irp);
PACTIVATE_INFO pActivateInfo;

if (irpStack->Parameters.DeviceIoControl.InputBufferLength >= sizeof(ACTIVATE_INFO))
{
pActivateInfo = Irp->AssociatedIrp.SystemBuffer;
if (g_ActivateInfo.bActivated != pActivateInfo->bActivated)
{
if (pActivateInfo->bActivated)
{
// Set up callback routines
ntStatus = PsSetCreateProcessNotifyRoutine(ProcessCallback, FALSE);
if (ntStatus != STATUS_SUCCESS)
{
return ntStatus;
}

// Setup the global data structure
g_ActivateInfo.bActivated = pActivateInfo->bActivated;
} // if
else
{
// restore the call back routine, thus givinig chance to the user mode application to unload dynamically the driver
ntStatus = PsSetCreateProcessNotifyRoutine(ProcessCallback, TRUE);
if (ntStatus != STATUS_SUCCESS)
return ntStatus;
else
g_ActivateInfo.bActivated = FALSE;
}
ntStatus = STATUS_SUCCESS;
} // if
} // if

return ntStatus;
}

//
// The dispatch routine
//
NTSTATUS DispatchIoctl(
IN PDEVICE_OBJECT DeviceObject,
IN PIRP Irp
)
{
NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(Irp);
PDEVICE_EXTENSION extension = DeviceObject->DeviceExtension;
PPROCESS_CALLBACK_INFO pProcCallbackInfo;
PLIST_ENTRY pEntry = NULL;
PCUSTOM_EXTENSION pData = NULL;

// These IOCTL handlers are the set and get interfaces between the driver and the user mode app
switch(irpStack->Parameters.DeviceIoControl.IoControlCode)
{
case IOCTL_PROCOBSRV_ACTIVATE_MONITORING:
{
// 每次激活都清除一次列表，防止满了的情况
KeWaitForSingleObject(&g_WaitEvent, Executive, KernelMode, FALSE,(PLARGE_INTEGER)NULL);

DbgPrint(" DispatchIoctl(): IOCTL_PROCOBSRV_ACTIVATE_MONITORING.(Surplus number: %d)", g_list_count);

while(!IsListEmpty(&g_list_head))
{
pEntry = RemoveHeadList(&g_list_head);
pData = CONTAINING_RECORD(pEntry, CUSTOM_EXTENSION, list_entry);
ExFreePool(pData);
}

g_list_count = 0;

KeSetEvent(&g_WaitEvent ,0 ,FALSE);

ntStatus = ActivateMonitoringHanlder( Irp );
break;
}

case IOCTL_PROCOBSRV_GET_PROCINFO:
{
// 从队列里面取数据
KeWaitForSingleObject(&g_WaitEvent, Executive, KernelMode, FALSE,(PLARGE_INTEGER)NULL);

if(!IsListEmpty(&g_list_head))
{
pEntry = RemoveHeadList(&g_list_head);
if(pEntry)
{
pData = CONTAINING_RECORD(pEntry, CUSTOM_EXTENSION, list_entry);

extension->hProcessId = pData->hProcessId;
extension->hParentId = pData->hParentId;
extension->bCreate = pData->bCreate;

DbgPrint(" Remove Head List(Surplus number:%d). %d, %d, %d", g_list_count, pData->hProcessId, pData->hParentId, pData->bCreate);
g_list_count--;

ExFreePool(pData);
}
}

KeSetEvent(&g_WaitEvent ,0 ,FALSE);

if (irpStack->Parameters.DeviceIoControl.OutputBufferLength >= sizeof(PROCESS_CALLBACK_INFO))
{
pProcCallbackInfo = Irp->AssociatedIrp.SystemBuffer;
pProcCallbackInfo->hParentId = extension->hParentId;
pProcCallbackInfo->hProcessId = extension->hProcessId;
pProcCallbackInfo->bCreate = extension->bCreate;

ntStatus = STATUS_SUCCESS;
}
break;
}

default:
break;
}

Irp->IoStatus.Status = ntStatus;

// Set number of bytes to copy back to user-mode
if(ntStatus == STATUS_SUCCESS)
Irp->IoStatus.Information = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
else
Irp->IoStatus.Information = 0;

IoCompleteRequest(Irp, IO_NO_INCREMENT);

return ntStatus;
}

//
// Driver unload routine
//
void UnloadDriver(
IN PDRIVER_OBJECT DriverObject
)
{
UNICODE_STRING uszDeviceString;
//
// By default the I/O device is configured incorrectly or the
// configuration parameters to the driver are incorrect.
//
NTSTATUS ntStatus = STATUS_DEVICE_CONFIGURATION_ERROR;

PLIST_ENTRY pEntry = NULL;
PCUSTOM_EXTENSION pData = NULL;

if (g_ActivateInfo.bActivated)
// restore the call back routine, thus givinig chance to the user mode application to unload dynamically the driver
ntStatus = PsSetCreateProcessNotifyRoutine(ProcessCallback, TRUE);

IoDeleteDevice(DriverObject->DeviceObject);

RtlInitUnicodeString(&uszDeviceString, L"//DosDevices//ProcObsrv");
IoDeleteSymbolicLink(&uszDeviceString);

DbgPrint("==========================UnloadDriver(Surplus number: %d)==========================", g_list_count);

while(!IsListEmpty(&g_list_head))
{
pEntry = RemoveHeadList(&g_list_head);
pData = CONTAINING_RECORD(pEntry, CUSTOM_EXTENSION, list_entry);

DbgPrint(" Remove Head List. %d, %d, %d", pData->hProcessId, pData->hParentId, pData->bCreate);

ExFreePool(pData);
}

g_list_count = 0;
}

//----------------------------End of the file -------------------------------