驱动层完成获取进程网络流量模块(总结)

浏览数：23 / 时间：2015年06月20日

xp平台下是使用TDI,win 7平台采用的是WFP的模式

WFP框架下思路

获取进程ID 获取进程的上传数据大小和下载数据大小,然后存储到链表中

链表中节点结构体如下:

typedef struct tagFlowInfoItem
{
    /** 链表结构
    */
    LIST_ENTRY m_listEntry;

    /** 进程ID
    */
    HANDLE m_processID;

    /** 已上传流量,BYTE为单位
    */
    UINT64 m_uploadFlow;

    /** 已下载流量,BYTE为单位
    */
    UINT64 m_downloadFlow;

    /** 上传速度,BYTE为单位
    */
    ULONG m_uploadSpeed;

    /** 下载速度,BYTE为单位
    */
    ULONG m_downloadSpeed;

    /** 限制的上传速度,BYTE为单位
    */
    ULONG m_limitUploadSpeed;

    /** 限制的下载速度,BYTE为单位
    */
    ULONG m_limitDownloadSpeed;
} FlowInfoItem;

wfp中是需要注册callout函数的,进程id是在FWPM_LAYER_ALE_FLOW_ESTABLISHED_V4层来获取,而UDP和TCP流数据信息(比如大小,内容)是在FWPM_LAYER_DATAGRAM_DATA_V4和FWPM_LAYER_STREAM_V4层来获取的

,问题来了,在流注册的callout函数中我们可以获取到数据的大小,但是问题来了,怎么样把进程ID和对应的进程数据相对应???

//////////////////////////////////////////////////////////////////////////////////////////////////////

不断的查找资料,研究wdk文档,wdk示例,各种搜索和qq群请教,(*^__^*) 嘻嘻……

//////////////////////////////////////////////////////////////////////////////////////////////////////

最终发现wdk的示例中有相关的代码可以参考一下,

目录如下:C:\WinDDK\7600.16385.1\src\network\trans\msnmntr

主要是用到了一个FwpsFlowAssociateContext0函数,具体的用法看msdn或者示例代码使用,很简单

一定要很仔细的看哦,这个函数有个很需要注意的地方,不然就蓝蓝更健康了哦,(关于DeleteFn函数的注册和写法)

这个FwpsFlowAssociateContext0函数的作用就是把进程id从FWPM_LAYER_ALE_FLOW_ESTABLISHED_V4层传递到了FWPM_LAYER_DATAGRAM_DATA_V4和FWPM_LAYER_STREAM_V4层,然后我们就可以使用链表结构把进程id和进程的数据关联起来了.

那么如何识别是上传数据还是下载数据呢?好问题

解析layerData中的streamData,streamData中有个标志位,

上传数据判断:

和系统定义的宏FWPS_STREAM_FLAG_SEND和FWPS_STREAM_FLAG_SEND_EXPEDITED与(&)操作一下,不就可以得到是否是上传了吗?

下载数据判断:

和上传数据的方法类似.

最后说一句,在操作链表的时候,记得加锁.

xp下比win7下要简单的多,而且TDI的资料也多.实现起来还是容易的多

VOID AnalyzeNetworkIrp(PIRP Irp)
{
   PEPROCESS Process;
   PIO_STACK_LOCATION IoStackLocation;
   NETWORK_USAGE_INFORMATION NetworkUsageInformation;
   //memset(&NetworkUsageInformation,0,sizeof(NETWORK_USAGE_INFORMATION));

   IoStackLocation = IoGetCurrentIrpStackLocation(Irp);

   ASSERT(IoStackLocation->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL);

   Process = PsGetCurrentProcess();


   if (QueryProcessNetworkUsage(Process, &NetworkUsageInformation) == FALSE)
   {
       NetworkUsageInformation.BytesReceived = 0;
       NetworkUsageInformation.BytesSent = 0;
   }

   if (IoStackLocation->MinorFunction == TDI_RECEIVE)       //   TDI_RECEIVE
   {
       DbgPrint("Enter TCP 接收!\n");
       PTDI_REQUEST_KERNEL_RECEIVE Parameters;
       Parameters = (PTDI_REQUEST_KERNEL_RECEIVE)&IoStackLocation->Parameters;
       NetworkUsageInformation.BytesReceived += Parameters->ReceiveLength;
   }
   else if (IoStackLocation->MinorFunction == TDI_SEND)   //TDI_SEND
   {
       DbgPrint("Enter TCP 发送!\n");
       PTDI_REQUEST_KERNEL_SEND Parameters;
       Parameters = (PTDI_REQUEST_KERNEL_SEND)&IoStackLocation->Parameters;
       NetworkUsageInformation.BytesSent += Parameters->SendLength;
   }
   else if (IoStackLocation->MinorFunction == TDI_RECEIVE_DATAGRAM)       //TDI_RECEIVE_DATAGRAM
   {
       DbgPrint("Enter UDP 接收!\n");
       PTDI_REQUEST_KERNEL_RECEIVEDG Parameters;
       Parameters = (PTDI_REQUEST_KERNEL_RECEIVEDG)&IoStackLocation->Parameters;
       NetworkUsageInformation.BytesReceived += Parameters->ReceiveLength;
   }
   else if (IoStackLocation->MinorFunction == TDI_SEND_DATAGRAM)           //TDI_SEND_DATAGRAM
   {
       DbgPrint("Enter UDP 发送!\n");
       PTDI_REQUEST_KERNEL_SENDDG Parameters;
       Parameters = (PTDI_REQUEST_KERNEL_SENDDG)&IoStackLocation->Parameters;
       NetworkUsageInformation.BytesSent += Parameters->SendLength;
   }
   else if (IoStackLocation->MinorFunction == TDI_CONNECT) //本机向外界发起连接请求
   {
       DbgPrint("Enter Connect!\n");
       PTDI_REQUEST_KERNEL_SENDDG Parameters;
       Parameters = (PTDI_REQUEST_KERNEL_SENDDG)&IoStackLocation->Parameters;
       NetworkUsageInformation.BytesSent += Parameters->SendLength;
   }
   else if (IoStackLocation->MinorFunction == TDI_ACCEPT)//外界向本机发起连接请求
   {
       DbgPrint("Enter Accept!\n");
       PTDI_REQUEST_KERNEL_RECEIVEDG Parameters;
       Parameters = (PTDI_REQUEST_KERNEL_RECEIVEDG)&IoStackLocation->Parameters;
       NetworkUsageInformation.BytesReceived += Parameters->ReceiveLength;
   }

   // 更新进程的网络使用情况
   UpdateProcessNetworkUsage(Process, &NetworkUsageInformation);

   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 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;
}

//网络使用情况结构体定义
typedef struct _NETWORK_USAGE_INFORMATION {

   ULONGLONG BytesSent;
   ULONGLONG BytesReceived;
} NETWORK_USAGE_INFORMATION, *PNETWORK_USAGE_INFORMATION;

typedef struct _NETWORK_USAGE_GENERIC_TABLE_NODE {
   PEPROCESS Process;
   //UINT32 uProcessID;
   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;
   KSPIN_LOCK TrafficLock;
} NETWORK_USAGE_GENERIC_TABLE, *PNETWORK_USAGE_GENERIC_TABLE;