VC 实现 Netstat 功能

VC 实现 Netstat 功能



#ifndef _PORTPROC_H_

#define _PORTPROC_H_



#if _MSC_VER > 1000

#pragma once

#endif // _MSC_VER > 1000



#include <afx.h>

#include <vector>

#include <string>

#include <iostream>

#include <iomanip>

#include <iphlpapi.h>

#include <Winsock2.h>

#include <Tlhelp32.h>

using std::vector;

using std::string;

using std::setw;

#pragma comment ( lib, "Ws2_32.lib" )



class PortProc  

{

public:

enum {TCP_PROTOCOL,UDP_PROTOCOL};

private:

typedef struct NETSTATSTRUCT

{

   string strName;

   string strLocalIp;

   DWORD dwLocalPort;

   string strRemoteIp;

   DWORD dwRemotePort;

   DWORD dwState;  

   DWORD dwType;  

   DWORD dwPID;

   NETSTATSTRUCT(){strLocalIp=strName=strRemoteIp="";dwLocalPort=dwRemotePort=dwPID=dwState=0;dwType=TCP_PROTOCOL;}

}NETSTAT,* PNETSTAT;



struct MIB_TCPROWEX:public MIB_TCPROW

{

   DWORD dwProcessId;

};

struct _MIB_TCPTABLEEX

{

   DWORD        dwNumEntries;

   MIB_TCPROWEX   table[ANY_SIZE]; 

};

struct MIB_UDPROWEX:public MIB_UDPROW

{

   DWORD dwProcessId;

};

struct _MIB_UDPTABLEEX

{

   DWORD        dwNumEntries;

   MIB_UDPROWEX   table[ANY_SIZE]; 

};

typedef _MIB_TCPTABLEEX MIB_TCPTABLEEX, *PMIB_TCPTABLEEX;

typedef _MIB_UDPTABLEEX MIB_UDPTABLEEX, *PMIB_UDPTABLEEX;

typedef DWORD (WINAPI *AllocateAndGetTcpExTableFromStackFun)(PMIB_TCPTABLEEX *,BOOL,HANDLE,DWORD,DWORD);

typedef DWORD (WINAPI *AllocateAndGetUdpExTableFromStackFun)(PMIB_UDPTABLEEX *,BOOL,HANDLE,DWORD,DWORD);



vector<PNETSTAT>m_netstatTable;



//IpHlpApi.dll未公布函数,核型函数********

AllocateAndGetTcpExTableFromStackFun AllocateAndGetTcpExTableFromStack;

AllocateAndGetUdpExTableFromStackFun AllocateAndGetUdpExTableFromStack;

private:

string GetStatus(UINT dwState); //根据状态获取中文描述

//添加规则作过滤

bool ViaTcpRules(MIB_TCPROWEX& tcprow); 

bool ViaUdpRules(MIB_UDPROWEX& tcprow);

private:

bool ReceiveTcpPortProcTable();

bool ReceiveUdpPortProcTable();

void AssociatePortAndProc();

public:

PortProc();

virtual ~PortProc();

bool RecivePortProcTable(UINT nFlags);     //获取进程端口。

void Clear();   //清空vector,释放内存。

void Print();   //打印内容,做测试时使用。

};



//标记查询的协议,目前只针对TCP&UDP

#define TCP_TABLE_INFO 0x000001

#define UDP_TABLE_INFO 0x000002



#endif // _PORTPROC_H_









// PortProc.cpp: implementation of the PortProc class.

//

//



#include "stdafx.h"

#include "PortProc.h"



//

// Construction/Destruction

//



PortProc::PortProc()

{

AllocateAndGetTcpExTableFromStack=NULL;

AllocateAndGetUdpExTableFromStack=NULL;



HMODULE hM=LoadLibrary("iphlpapi.dll");

if(hM==NULL)

   throw new CException(TRUE);

AllocateAndGetTcpExTableFromStack=(AllocateAndGetTcpExTableFromStackFun)GetProcAddress(hM,"AllocateAndGetTcpExTableFromStack"); 

AllocateAndGetUdpExTableFromStack=(AllocateAndGetUdpExTableFromStackFun)GetProcAddress(hM,"AllocateAndGetUdpExTableFromStack");

}



PortProc::~PortProc()

{

Clear();

}



string PortProc::GetStatus(UINT dwState)

{

switch(dwState)

{   

case MIB_TCP_STATE_CLOSED:   return "CLOSED";     

case MIB_TCP_STATE_LISTEN:   return "LISTEN";      

case MIB_TCP_STATE_SYN_SENT: return "SYN_SENT";     

case MIB_TCP_STATE_SYN_RCVD: return "SYN_RCVD";     

case MIB_TCP_STATE_ESTAB:   return "ESTAB";      

case MIB_TCP_STATE_FIN_WAIT1: return "FIN_WAIT1";     

case MIB_TCP_STATE_FIN_WAIT2: return "FIN_WAIT2";     

case MIB_TCP_STATE_CLOSE_WAIT: return "CLOSE_WAIT";   

case MIB_TCP_STATE_CLOSING:   return "CLOSING";  

case MIB_TCP_STATE_LAST_ACK: return "LAST_ACK"; 

case MIB_TCP_STATE_TIME_WAIT: return "TIME_WAIT"; 

case MIB_TCP_STATE_DELETE_TCB: return "DELETE_TCB";  

default:       return "UNKNOWN";    

}

}

bool PortProc::ViaTcpRules(MIB_TCPROWEX& tcprow)

{

//目前仅过滤system进程

static bool bSystemPid = false;

if(tcprow.dwProcessId==4)

{

   if(!bSystemPid)

   {

    tcprow.dwLocalPort = htons(445);

    bSystemPid = true;

    return true; //放行

   }

   else

   {

    return false; //阻止

   }

}

return true;    //放行

}



bool PortProc::ViaUdpRules(MIB_UDPROWEX& tcprow)

{

return true;

}



bool PortProc::ReceiveTcpPortProcTable()

{

PMIB_TCPTABLEEX ptcpmib=NULL;

if(AllocateAndGetTcpExTableFromStack(&ptcpmib,TRUE,GetProcessHeap(),2,2)!=NO_ERROR)

   return false;

for (int i=0;i<ptcpmib->dwNumEntries;i++)

{

   if(!ViaTcpRules(ptcpmib->table[i]))

    continue;

   PNETSTAT pnetStat = new NETSTAT;

   pnetStat->dwState = ptcpmib->table[i].dwState;

   pnetStat->dwPID = ptcpmib->table[i].dwProcessId;

   in_addr locIP,RemIP;

   locIP.s_addr=ptcpmib->table[i].dwLocalAddr; 

   RemIP.s_addr=ptcpmib->table[i].dwRemoteAddr;

   pnetStat->strLocalIp=inet_ntoa(locIP);

   pnetStat->dwLocalPort=ntohs(ptcpmib->table[i].dwLocalPort);

   pnetStat->strRemoteIp=inet_ntoa(RemIP);

   pnetStat->dwRemotePort=pnetStat->dwState==MIB_TCP_STATE_LISTEN?ntohs(ptcpmib->table[i].dwRemotePort):0;

   m_netstatTable.push_back(pnetStat);

}

return true;

}

bool PortProc::ReceiveUdpPortProcTable()

{

PMIB_UDPTABLEEX pudpmib=NULL;

if(AllocateAndGetUdpExTableFromStack(&pudpmib,TRUE,GetProcessHeap(),2,2)!=NO_ERROR) 

   return false; 

for(int i=0;i< pudpmib->dwNumEntries;i++)

{

   PNETSTAT pnetStat = new NETSTAT;

   struct in_addr locIP;

   locIP.s_addr=pudpmib->table[i].dwLocalAddr;

   pnetStat->strLocalIp=inet_ntoa(locIP);

   pnetStat->dwLocalPort=ntohs(pudpmib->table[i].dwLocalPort);

   pnetStat->dwPID=pudpmib->table[i].dwProcessId;

   m_netstatTable.push_back(pnetStat);

}

return true;

}

void PortProc::AssociatePortAndProc()

{

HANDLE hProcessSnap=::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS,0);

if(hProcessSnap==INVALID_HANDLE_VALUE)return; 

for (int i=0;i<m_netstatTable.size();i++)

{

   PROCESSENTRY32 proEn;

   proEn.dwSize=sizeof(proEn);

   if(::Process32First(hProcessSnap,&proEn)) 

   { 

    do 

    {  

     if(proEn.th32ProcessID==m_netstatTable[i]->dwPID) 

     { 

      m_netstatTable[i]->strName=proEn.szExeFile;

      break;

     }   

    }while(::Process32Next(hProcessSnap,&proEn));

   }

}

::CloseHandle(hProcessSnap);

}

bool PortProc::RecivePortProcTable(UINT nFlags)

{

bool bRet=true;

if( nFlags & TCP_TABLE_INFO )

   bRet = ReceiveTcpPortProcTable();

if( bRet == false )

   goto clear;

if( nFlags & UDP_TABLE_INFO )

   bRet = ReceiveUdpPortProcTable();

if( bRet == false )

   goto clear;



AssociatePortAndProc();

return true;

clear:

Clear();

}



void PortProc::Clear()

{

for(int i=0;i<m_netstatTable.size();i++)

   delete m_netstatTable[i];

m_netstatTable.clear();

}



void PortProc::Print()

{

for(int i=0;i<m_netstatTable.size();i++)

{

   std::cout<<std::left<<setw(20)<<m_netstatTable[i]->strName<<setw(10)<<m_netstatTable[i]->dwPID<<setw(15)

    <<m_netstatTable[i]->strLocalIp<<setw(8)<<m_netstatTable[i]->dwLocalPort<<setw(15)<<m_netstatTable[i]->strRemoteIp

    <<setw(8)<<m_netstatTable[i]->dwRemotePort<<std::endl;

}

}







#include "stdafx.h"

#include "PortProc.h"



int main(int argc, char* argv[])

{

PortProc pp;

pp.RecivePortProcTable(TCP_TABLE_INFO/*|UDP_TABLE_INFO*/);

pp.Print();

return 0;

}