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