/* Windows Ping Program.
 * Compiler: Visual C++.NET
 * by Kenji Aiko
 */
#define WIN32_LEAN_AND_MEAN

#include <winsock2.h>
#include <ws2tcpip.h>
#include <iostream>

using namespace std;

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

#define ICMP_ECHO_REPLY        0
#define ICMP_DEST_UNREACH      3
#define ICMP_TTL_EXPIRE        11
#define ICMP_ECHO_REQUEST      8
#define ICMP_MIN               8


struct IPHeader {
    BYTE    h_len:4;
    BYTE    version:4;
    BYTE    tos;
    USHORT  total_len;
    USHORT  ident;
    USHORT  flags;
    BYTE    ttl;
    BYTE    proto;
    USHORT  checksum;
    ULONG   source_ip;
    ULONG   dest_ip;
};


struct ICMPHeader {
    BYTE    type;
    BYTE    code;
    USHORT  checksum;
    USHORT  id;
    USHORT  seq;
    ULONG   timestamp;
};


#define DEFAULT_PACKET_SIZE     32
#define DEFAULT_TTL             30
#define MAX_PING_DATA_SIZE      1024
#define MAX_PING_PACKET_SIZE    (MAX_PING_DATA_SIZE + sizeof(IPHeader))


USHORT  ip_checksum(USHORT *data, int size);
int     host2ip(char *hostname, sockaddr_in &dest);
void    init_ping_packet(ICMPHeader *icmp_hdr, int packet_size, int seq_no);
int     send_ping(SOCKET sd, const sockaddr_in& dest, ICMPHeader *send_buf, int packet_size);
int     recv_ping(SOCKET sd, sockaddr_in &source, IPHeader *recv_buf, int packet_size);
int     decode_reply(IPHeader *reply, int bytes, sockaddr_in *from);
void    die(ICMPHeader *&send_buf, IPHeader *&recv_buf);



int main(int argc, char *argv[])
{
    ICMPHeader *send_buf = NULL;
    IPHeader   *recv_buf = NULL;
    
    if (argc < 2) {
        cerr << "usage: " << argv[0] << " <host> [data_size] [ttl]" << endl;
        cerr << "data_size can be up to " << MAX_PING_DATA_SIZE << " bytes. ";
        cerr << "Default is " << DEFAULT_PACKET_SIZE << "." << endl; 
        cerr << "ttl should be 255 or lower. ";
        cerr << "Default is " << DEFAULT_TTL << "." << endl;
        return 1;
    }

    int packet_size = DEFAULT_PACKET_SIZE;
    if (argc > 2) {
        packet_size = atoi(argv[2]);
    }

    int ttl = DEFAULT_TTL;
    if (argc > 3) { 
        ttl = atoi(argv[3]);
    }

    WSAData wsaData;
    if (WSAStartup(MAKEWORD(2, 1), &wsaData) != 0) {
        cerr << "Failed to find Winsock 2.1 or better." << endl;
        return 1;
    }

    SOCKET sd;
    if ((sd = WSASocket(AF_INET, SOCK_RAW, IPPROTO_ICMP, 0, 0, 0)) == INVALID_SOCKET) {
        cerr << "Failed to create raw socket: " << WSAGetLastError() << endl;
        die(send_buf, recv_buf);
    }

    if (setsockopt(sd, IPPROTO_IP, IP_TTL, (const char *)&ttl, sizeof(ttl)) == SOCKET_ERROR) {
        cerr << "TTL setsockopt failed: " << WSAGetLastError() << endl;
        die(send_buf, recv_buf);
    }

    sockaddr_in dest, source;
    memset(&dest, 0, sizeof(dest));
    if (host2ip(argv[1], dest) < 0) {
        cerr << "Failed to resolve " << argv[1] << endl;
    }

    packet_size = max(sizeof(ICMPHeader), min(MAX_PING_DATA_SIZE, (UINT)packet_size));

    if ((send_buf = (ICMPHeader *)new char[packet_size]) == 0) {
        cerr << "Failed to allocate output buffer." << endl;
        die(send_buf, recv_buf);
    }

    if ((recv_buf = (IPHeader *)new char[MAX_PING_PACKET_SIZE]) == 0) {
        cerr << "Failed to allocate output buffer." << endl;
        die(send_buf, recv_buf);
    }

    int seq_no = 0;
    init_ping_packet(send_buf, packet_size, seq_no);

    if (send_ping(sd, dest, send_buf, packet_size) < 0) {
        cerr << "send_ping function error." << endl;
        die(send_buf, recv_buf);
    }

    while (1) {
        if (recv_ping(sd, source, recv_buf, MAX_PING_PACKET_SIZE) < 0) {
            unsigned short header_len = recv_buf->h_len * 4;
            ICMPHeader *icmphdr = (ICMPHeader *)((char *)recv_buf + header_len);
            if (icmphdr->seq == seq_no) {
                break;
            }
            cerr << "bad sequence number!" << endl;
            continue;
        }

        if (decode_reply(recv_buf, packet_size, &source) != -2) {
            break;
        }
    }
    
    delete[]send_buf;
    delete[]recv_buf;
    WSACleanup();

    return 0;
}



USHORT ip_checksum(USHORT *data, int size)
{
    unsigned long sum = 0;

    while (size > 1) {
        sum += *(data++);
        size -= 2;
    }

    if (size > 0) {
        sum += (*data) & 0xff00;
    }

    sum = (sum & 0xffff) + (sum >> 16);
    return ((~(USHORT)((sum >> 16) + (sum & 0xffff))));
}



int host2ip(char *hostname, sockaddr_in &dest)
{
    UINT addr;
    if((addr = inet_addr(hostname)) == INADDR_NONE){

        struct hostent  *hp;
        if ((hp = gethostbyname(hostname)) == NULL) {
            return -1;
        }
        memcpy(&(dest.sin_addr), hp->h_addr, hp->h_length);
        dest.sin_family = hp->h_addrtype;
    
    }else{

        dest.sin_addr.s_addr = addr;
        dest.sin_family = AF_INET;
    }

    return 0;
}



void init_ping_packet(ICMPHeader *icmp_hdr, int packet_size, int seq_no)
{
    icmp_hdr->type        = ICMP_ECHO_REQUEST;
    icmp_hdr->code        = 0;
    icmp_hdr->checksum    = 0;
    icmp_hdr->id          = (USHORT)GetCurrentProcessId();
    icmp_hdr->seq         = seq_no;
    icmp_hdr->timestamp   = (ULONG)GetTickCount();

    const unsigned long int deadmeat = 0xDEADBEEF;
    char *datapart = (char *)icmp_hdr + sizeof(ICMPHeader);
    int bytes_left = packet_size - sizeof(ICMPHeader);
    
    while (bytes_left > 0) {
        memcpy(datapart, &deadmeat, min(int(sizeof(deadmeat)), bytes_left));
        bytes_left -= sizeof(deadmeat);
        datapart   += sizeof(deadmeat);
    }

    icmp_hdr->checksum = ip_checksum((USHORT *)icmp_hdr, packet_size);
}



int send_ping(SOCKET sd, const sockaddr_in& dest, ICMPHeader *send_buf, int packet_size)
{
    cout << "Sending " << packet_size << " bytes to " << inet_ntoa(dest.sin_addr) << "..." << flush;
    
    int bwrote = sendto(sd, (char *)send_buf, packet_size, 0, (sockaddr *)&dest, sizeof(dest));
    if (bwrote == SOCKET_ERROR) {
        cerr << "send failed: " << WSAGetLastError() << endl;
        return -1;
    }else if (bwrote < packet_size) {
        cout << "sent " << bwrote << " bytes..." << flush;
    }

    return 0;
}



int recv_ping(SOCKET sd, sockaddr_in &source, IPHeader *recv_buf, int packet_size)
{
    int fromlen = sizeof(source);
    int bread = recvfrom(sd, (char *)recv_buf, packet_size + sizeof(IPHeader), 0,
                (sockaddr *)&source, &fromlen);
   
    if (bread == SOCKET_ERROR) {
        cerr << "read failed: ";
        if (WSAGetLastError() == WSAEMSGSIZE) {
            cerr << "buffer too small" << endl;
        }else {
            cerr << "error #" << WSAGetLastError() << endl;
        }
        return -1;
    }

    return 0;
}



int decode_reply(IPHeader *reply, int bytes, sockaddr_in *from) 
{
    USHORT header_len = reply->h_len * 4;
    ICMPHeader *icmphdr = (ICMPHeader *)((char *)reply + header_len);

    if (bytes < header_len + ICMP_MIN) {
        cerr << "too few bytes from " << inet_ntoa(from->sin_addr) << endl;
        return -1;
    }
    
    if (icmphdr->type != ICMP_ECHO_REPLY && icmphdr->type != ICMP_TTL_EXPIRE) {
        cerr << "Unknown ICMP packet type " << (int)(icmphdr->type) << " received" << endl;
        return -1;
    }
   
    if (icmphdr->id != (USHORT)GetCurrentProcessId()) {
        return -2;
    }

    int nHops = int(256 - reply->ttl);
    if (nHops == 192) { 
        nHops = 1;
    }
    
    if (nHops == 128) {
        nHops = 0;
    }

    cout << endl << bytes << " bytes from ";
    cout << inet_ntoa(from->sin_addr) << ", icmp_seq " << icmphdr->seq << ", ";

    if (icmphdr->type == ICMP_TTL_EXPIRE) {
        cout << "TTL expired." << endl;
    }
    else {
        cout << nHops << " hop" << (nHops == 1 ? "" : "s");
        cout << ", time: " << (GetTickCount() - icmphdr->timestamp) << " ms." << endl;
    }

    return 0;
}



void die(ICMPHeader *&send_buf, IPHeader *&recv_buf)
{
    if (send_buf != NULL) {
        delete[]send_buf;
    }

    if (recv_buf != NULL) {
        delete[]recv_buf;
    }
    
    WSACleanup();
    exit(1);
}