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

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;


NETWORK_USAGE_GENERIC_TABLE NetworkUsageGenericTable;

思路如此,TDI的实现大家实现起来不难,我就不多说了,显得很唠叨烦人,嘿嘿

你可能感兴趣的:(驱动层 完成获取进程网络流量模块(总结))