// Tracer.cpp: implementation of the CTracer class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "TraceRoute.h"
#include "Tracer.h"
#include "TraceRouteDlg.h"

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

CTracer::CTracer()
{
	m_nSeq=1;
	icmpData=NULL;
	icmpRcvBuf=NULL;
    m_hSocket=INVALID_SOCKET;

	//初始化socket
    WSADATA wsaData;
    if(WSAStartup(MAKEWORD(2,2),&wsaData)!=0)
    {
	    AfxMessageBox("WSAStartup()出錯!");
    }
  
}

CTracer::~CTracer()
{
	//關(guān)閉Socket
	if (m_hSocket!=NULL)
		closesocket(m_hSocket);
	WSACleanup();
}

//填充ICMP報文首部

void CTracer::FillICMPData(char* icmpData,int size)
{
	memset(icmpData,0,size);
	ICMPHEADER* icmpHeader=NULL;
	icmpHeader=(ICMPHEADER*)icmpData;

	icmpHeader->i_type =ICMP_ECHO;
	icmpHeader->i_code =0;
	icmpHeader->i_id =(USHORT)GetCurrentProcessId();
	icmpHeader->i_seq =m_nSeq++; 
	//GetTickCount返回從0點到現(xiàn)在的毫秒數(shù)，作時間戳
	icmpHeader->timestamp=GetTickCount();
	char* datapart=icmpData+sizeof(ICMPHEADER);
    memset(datapart,'*',size-sizeof(ICMPHEADER));
	//填充校驗和
    icmpHeader->i_cksum =CheckSum(icmpData,size);
}

//由字符串轉(zhuǎn)化為地址

BOOL CTracer::FillAddress(char *addrDest)
{
	memset(&m_addrDest,0,sizeof(m_addrDest));
	m_addrDest.sin_family =AF_INET;
	if(inet_addr(addrDest)==INADDR_NONE)
	{
		//輸入的地址為計算機名字
		HOSTENT* hp=NULL;
		hp=gethostbyname(addrDest);
		if(hp)
		{
		    memcpy(&(m_addrDest.sin_addr),hp->h_addr,hp->h_length);
			m_addrDest.sin_family =hp->h_addrtype ;
		}
		else
		{
			AfxMessageBox("獲取地址失敗!");
			return FALSE;
		}
	}
	else
	{
		m_addrDest.sin_addr.s_addr=inet_addr(addrDest);
	}
	return TRUE;
}

void CTracer::Trace(char *destAddress)
{
	 int size=DEF_PACKET_SIZE+sizeof(ICMPHEADER);
	 //轉(zhuǎn)換地址
     if (!FillAddress(destAddress)) return ;

	 //分配必要的內(nèi)存空間
	 icmpData=(char*)HeapAlloc(GetProcessHeap(),HEAP_ZERO_MEMORY,MAX_PACKET);
	 icmpRcvBuf=(char*)HeapAlloc(GetProcessHeap(),HEAP_ZERO_MEMORY,MAX_PACKET);
	 if(!icmpData||!icmpRcvBuf)
	 {
		 AfxMessageBox("分配內(nèi)存空間失敗!");
		 TerminateProcess(GetCurrentProcess(),-1);
	 }
	 memset(icmpData,0,MAX_PACKET);
	 memset(icmpRcvBuf,0,MAX_PACKET);
	 


     //初始化套接字
	 m_hSocket=WSASocket(AF_INET,SOCK_RAW,IPPROTO_ICMP,NULL,0,WSA_FLAG_OVERLAPPED);
	 if(m_hSocket==INVALID_SOCKET)
	 {   
		 AfxMessageBox("套接字初始化失敗!");
		 TerminateProcess(GetCurrentProcess(),-1);
	 }
	
	 //設(shè)置超時選項
	 int nTimeOut=1000;
	 int result;
	 result=setsockopt(m_hSocket,SOL_SOCKET,SO_RCVTIMEO,(char*)&nTimeOut,sizeof(nTimeOut));
	 if(result==SOCKET_ERROR)
	 { 
		 AfxMessageBox("設(shè)置接收超時選項失敗!");
		 TerminateProcess(GetCurrentProcess(),-1);
	 }
	 result=setsockopt(m_hSocket,SOL_SOCKET,SO_SNDTIMEO,(char*)&nTimeOut,sizeof(nTimeOut));
	 if(result==SOCKET_ERROR)
	 {
		 AfxMessageBox("設(shè)置發(fā)送超時選項失敗!");
		 TerminateProcess(GetCurrentProcess(),-1);
	 }
	 //設(shè)置路由不查詢路由表選項
	 BOOL bDontRoute=TRUE;
	 result=setsockopt(m_hSocket,SOL_SOCKET,SO_DONTROUTE,(char*)&bDontRoute,sizeof(BOOL));
	 if(result==SOCKET_ERROR)
	 {
		 AfxMessageBox("設(shè)置不查詢路由表選項失敗!");
		 TerminateProcess(GetCurrentProcess(),-1);
	 }

	 for(int ttl=1;ttl<MAX_NOTES;ttl++)
	 {
		 //設(shè)定數(shù)據(jù)報的壽命
		 SetTTL(m_hSocket,ttl);
		 //發(fā)送數(shù)據(jù)報
		 if(SendData(icmpData,size)) continue; 
		 //接收數(shù)據(jù)報
		 if(RecvData(icmpRcvBuf,&result)) continue; 
		 //處理接收到的數(shù)據(jù)報
		 if (DecodeICMP(icmpRcvBuf,result,ttl)) break;
	 }
	 
	HeapFree(GetProcessHeap(), 0, icmpData);
    HeapFree(GetProcessHeap(), 0, icmpRcvBuf);

}

//設(shè)置數(shù)據(jù)報的壽命
BOOL CTracer::SetTTL(SOCKET hSocket, int ttl)
{
	int result;
	result=setsockopt(hSocket,IPPROTO_IP,IP_TTL,(LPSTR)&ttl,sizeof(ttl));
	if(result==SOCKET_ERROR)
	{
		AfxMessageBox("設(shè)置數(shù)據(jù)報壽命失敗!");
		TerminateProcess(GetCurrentProcess(),-1);
	}
	return TRUE;
}

//校驗和算法
USHORT CTracer::CheckSum(char* pBuffer,int size)
{   

	USHORT* buffer=(USHORT*)pBuffer;
	unsigned long cksum=0;
    while(size > 1) 
    {
        cksum += *buffer++;
        size -= sizeof(USHORT);
    }
    if(size )
        cksum += *(UCHAR*)buffer;
    cksum = (cksum >> 16) + (cksum & 0xffff);
    cksum += (cksum >> 16);
    return (USHORT)(~cksum);
}

//處理接收到的數(shù)據(jù)報
BOOL CTracer::DecodeICMP(char* pBuffer,int bytes,int ttl)
{
    IPHEADER       *ipHeader=NULL;
    ICMPHEADER     *icmpHeader=NULL;
    unsigned short  ipHeaderLen;
    HOSTENT *ph=NULL;
    in_addr inaddr=m_addrFrom.sin_addr;

    ipHeader=(IPHEADER*)pBuffer;
	ipHeaderLen=20;

    if (bytes<ipHeaderLen+ICMP_MIN) 
		AfxMessageBox("接收數(shù)據(jù)報長度不正確!");
	
	icmpHeader=(ICMPHEADER*)(pBuffer+20);
    
    switch (icmpHeader->i_type)
    {
		//目的站點的返回
        case ICMP_ECHOREPLY:     
            ph=gethostbyaddr((const char *)&inaddr,AF_INET, sizeof(in_addr));
            if (ph != NULL)
			{
				CString report;
				report.Format("%2d  %s (%s)",ttl,ph->h_name,inet_ntoa(inaddr));
				((CTraceRouteDlg*)m_pWnd)->InfoAdd(report);
			}
			return TRUE;
            break;
		//中途路由器的返回
        case ICMP_TIMEOUT:     
            {
				CString report;
				report.Format("%2d  %s", ttl, inet_ntoa(inaddr));
				((CTraceRouteDlg*)m_pWnd)->InfoAdd(report);
				return FALSE;
				break;
			}
		//錯誤 主機不可達
        case ICMP_DESTUNREACH:  
            {
				CString report;
				report.Format("%2d  %s   主機不可達",ttl,inet_ntoa(inaddr));
				((CTraceRouteDlg*)m_pWnd)->InfoAdd(report);
				return TRUE;
				break;
			}
		//收到一個不是回應(yīng)的報文
        default:
			{
				CString report;
				report.Format("非回應(yīng)報文");
				((CTraceRouteDlg*)m_pWnd)->InfoAdd(report);
				return TRUE;
			}
           
    }
    return FALSE;
}

//發(fā)送數(shù)據(jù)報
BOOL CTracer::SendData(char* icmpData,int size)
{
	//填充ICMP報頭
    FillICMPData(icmpData,size);

	//發(fā)送數(shù)據(jù)報
	int result;
	result=sendto(m_hSocket,icmpData,size,0,(SOCKADDR*)&m_addrDest,sizeof(m_addrDest));
		 
	if(result==SOCKET_ERROR)
	{
		if(WSAGetLastError()==WSAETIMEDOUT)
		{
			((CTraceRouteDlg*)m_pWnd)->InfoAdd ("發(fā)送超時");
			return TRUE;
		}
		AfxMessageBox("發(fā)送報文失敗!");
		TerminateProcess(GetCurrentProcess(),-1);
	}
	return FALSE;
}

BOOL CTracer::RecvData(char* icmpRcvBuf,int* presult)
{
	static int count=0;
	//總共6次出現(xiàn)接收超時，判斷存在連接問題。
	if(count>5)            
	{
		AfxMessageBox("連接存在問題!");
		TerminateProcess(GetCurrentProcess(),-1);
	}

	int fromlen=sizeof(SOCKADDR);
	*presult=SOCKET_ERROR;
	*presult=recvfrom(m_hSocket,icmpRcvBuf,MAX_PACKET,0,(SOCKADDR*)&m_addrFrom,&fromlen);
		 
	if(*presult==SOCKET_ERROR)
	{
		if(WSAGetLastError()==WSAETIMEDOUT)
		{
			((CTraceRouteDlg*)m_pWnd)->InfoAdd ("接收超時!");
			count++;
			return TRUE;
		}
		AfxMessageBox("接收數(shù)據(jù)報失敗!");
		TerminateProcess(GetCurrentProcess(),-1);
	}
	return FALSE;
}

void CTracer::SetWnd(CDialog *pWnd)
{
	//設(shè)置窗口指針。
	m_pWnd=pWnd;
}