#include "MyPingip.h"
#include <iostream>
#include <cstdlib>
#include <netdb.h>
#include <sys/select.h>

CMyPingip::CMyPingip(void) : m_sockfd(-1)
{
}

CMyPingip::~CMyPingip(void)
{
    if (m_sockfd != -1) {
        close(m_sockfd);
    }
}

long CMyPingip::GetTickCount()
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (tv.tv_sec * 1000) + (tv.tv_usec / 1000);
}

bool CMyPingip::PingFunction(const char* pstrHost)
{
    unsigned int nRetries = 5;
    int nLostConnectCnt = 0;
    bool bRet = true;

    // 创建原始套接字
    m_sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (m_sockfd < 0) {
        std::cerr << "Create Socket Error: " << strerror(errno) << std::endl;
        return false;
    }

    // 设置超时选项
    struct timeval timeout;
    timeout.tv_sec = 1;
    timeout.tv_usec = 0;
    setsockopt(m_sockfd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

    struct sockaddr_in saDest;
    struct hostent* host;

    // 解析主机名
    if (host = gethostbyname(pstrHost)) {
        memcpy(&saDest.sin_addr, host->h_addr, host->h_length);
    }
    else {
        std::cerr << "Hostname resolution failed: " << pstrHost << std::endl;
        close(m_sockfd);
        return false;
    }

    saDest.sin_family = AF_INET;
        saDest.sin_port = 0;

        // ping 循环
        for (unsigned int nLoop = 0; nLoop < nRetries; nLoop++) {
            // 发送ICMP回应请求
            if (SendEchoRequest(m_sockfd, &saDest) < 0) {
                nLostConnectCnt++;
                continue;
            }

            // 等待回应
            if (WaitForEchoReply(m_sockfd)) {
                struct sockaddr_in saSrc;
                u_char cTTL;

                // 获取回应
                    long dwTimeSent = RecvEchoReply(m_sockfd, &saSrc, &cTTL);
                    if (dwTimeSent > 0) {
                        char* strIp = inet_ntoa(saSrc.sin_addr);
                        if (strcmp(strIp, pstrHost) == 0) {
                            long dwElapsed = GetTickCount() - dwTimeSent;
                            std::cout << "Reply from: " << strIp
                                << ": bytes=" << REQ_DATASIZE
                                << " time=" << dwElapsed << " ms"
                                << " TTL=" << static_cast<int>(cTTL) << std::endl;
                        }
                        else {
                            std::cout << "Reply from: " << strIp
                                << ": Destination host unreachable" << std::endl;
                            nLostConnectCnt++;
                        }
                    }
                    else {
                        nLostConnectCnt++;
                    }
            }
            else {
                std::cout << "Request Timed Out" << std::endl;
                nLostConnectCnt++;
            }
        }

    close(m_sockfd);
    m_sockfd = -1;

    return (nLostConnectCnt < 3);
}

int CMyPingip::SendEchoRequest(int sock, struct sockaddr_in* dest_addr)
{
    static ECHOREQUEST echo;
    static int nId = 1;
    static int nSeq = 1;

    // 填充ICMP头
    echo.icmpHdr.Type = ICMP_ECHO;
    echo.icmpHdr.Code = 0;
    echo.icmpHdr.Checksum = 0;
    echo.icmpHdr.ID = htons(nId++);
    echo.icmpHdr.Seq = htons(nSeq++);

    // 填充数据
    for (int i = 0; i < REQ_DATASIZE; i++) {
        echo.cData[i] = ' ' + i;
    }

    // 保存发送时间
    echo.dwTime = GetTickCount();

    // 计算校验和
    echo.icmpHdr.Checksum = in_cksum((unsigned short*)&echo, sizeof(ECHOREQUEST));

    // 发送请求
    int nRet = sendto(sock, &echo, sizeof(ECHOREQUEST), 0,
        (struct sockaddr*)dest_addr, sizeof(struct sockaddr_in));

    if (nRet < 0) {
        std::cerr << "Send error: " << strerror(errno) << std::endl;
    }

    return nRet;
}

long CMyPingip::RecvEchoReply(int sock, struct sockaddr_in* from_addr, u_char* pTTL)
{
    ECHOREPLY echoReply;
    socklen_t from_len = sizeof(struct sockaddr_in);

    // 接收回应
    int nRet = recvfrom(sock, &echoReply, sizeof(ECHOREPLY), 0,
        (struct sockaddr*)from_addr, &from_len);

    if (nRet < 0) {
        if (errno != EAGAIN && errno != EWOULDBLOCK) {
            std::cerr << "Receive error: " << strerror(errno) << std::endl;
        }
        return -1;
    }

    // 获取TTL值
    *pTTL = echoReply.iphdr.TTL;

    // 返回发送时间
    return echoReply.echorequest.dwTime;
}

int CMyPingip::WaitForEchoReply(int sock)
{
    fd_set readfds;
    FD_ZERO(&readfds);
    FD_SET(sock, &readfds);

    struct timeval timeout;
    timeout.tv_sec = 1;
    timeout.tv_usec = 0;

    return select(sock + 1, &readfds, NULL, NULL, &timeout);
}

unsigned short CMyPingip::in_cksum(unsigned short* addr, int len)
{
    int nleft = len;
    unsigned short* w = addr;
    unsigned int sum = 0;
    unsigned short answer = 0;

    while (nleft > 1) {
        sum += *w++;
        nleft -= 2;
    }

    if (nleft == 1) {
        *(unsigned char*)(&answer) = *(unsigned char*)w;
        sum += answer;
    }

    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);
    answer = ~sum;

    return answer;
}