/*
 Copyright (C) JAEHYUK CHO
 All rights reserved.
 
 Author: JaeHyuk Cho <minzkn@minzkn.com>

 Tiny ping example source
*/

#if !defined(_ISOC99_SOURCE)
# define _ISOC99_SOURCE (1L)
#endif

#if !defined(_GNU_SOURCE)
# define _GNU_SOURCE (1L)
#endif

#include <sys/types.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>

#if 1L
#include <netinet/ip_icmp.h> /* struct icmp */
#include <netinet/icmp6.h> /* struct icmp6_hdr */
#endif

#if !defined(ICMP_MINLEN)
# warning ICMP_MINLEN not defined ! (need include netinet/ip_icmp.h)
# define ICMP_MINLEN (1 + 1 + 2 + 2 + 2) /* 8 bytes : type(1) + code(1) + checksum(2) + identifier(2) + sequence_number(2) */
#endif

#if !defined(ICMP_ECHO)
# warning ICMP_ECHO not defined ! (need include netinet/ip_icmp.h)
# define ICMP_ECHO 8
#endif

#if !defined(ICMP_ECHOREPLY)
# warning ICMP_ECHOREPLY not defined ! (need include netinet/ip_icmp.h)
# define ICMP_ECHOREPLY 0
#endif

#if !defined(ICMP6_ECHO_REQUEST)
# warning ICMP6_ECHO_REQUEST not defined ! (need include netinet/icmp6.h)
# define ICMP6_ECHO_REQUEST 128
#endif

#if !defined(ICMP6_ECHO_REPLY)
# warning ICMP6_ECHO_REPLY not defined ! (need include netinet/icmp6.h)
# define ICMP6_ECHO_REPLY 129
#endif

struct mzping_icmp_header { /* 4 bytes */
    uint8_t m_type;
    uint8_t m_code;
    uint16_t m_checksum;
};

struct mzping_icmp_echo_header { /* 8 bytes */
    struct mzping_icmp_header m_common;
    uint16_t m_identifier;
    uint16_t m_sequence_number;
};

struct mzping_icmp_mydata {
    uint32_t m_timestamp32;
};

static unsigned int mzping_icmp_checksum(const void *s_data, size_t s_size);
static unsigned int mzping_ts32_us(void);
static int mzping_icmp_v4(int s_socket, const char *s_hostname, const char *s_address_string, struct addrinfo *s_addrinfo, int s_sequence_number, int s_timeout);
static int mzping_icmp_v6(int s_socket, const char *s_hostname, const char *s_address_string, struct addrinfo *s_addrinfo, int s_sequence_number, int s_timeout);

int mzping(const char *s_hostname, int s_count);

int main(int s_argc, char **s_argv);

static unsigned int mzping_icmp_checksum(const void *s_data, size_t s_size)
{
    register unsigned int s_result = 0u;

    while(s_size > ((size_t)1u)) {
        s_result += (unsigned int)(*((const uint16_t *)s_data));
        s_data = (const void *)(((const uint16_t *)s_data) + ((size_t)1));
        s_size -= (size_t)2u;
    }
   
    if(s_size > ((size_t)0u)) { /* odd */
        s_result += (unsigned int)(*((const uint8_t *)s_data));
    }

    s_result = (s_result >> 16) + (s_result & 0xffffu);
    s_result += s_result >> 16;

    return((~s_result) & 0xffffu);
}

static unsigned int mzping_ts32_us(void)
{
    struct timeval s_timeval;
    
    (void)gettimeofday((struct timeval *)(&s_timeval), (void *)0);
    
    return((s_timeval.tv_sec * 1000000u) + s_timeval.tv_usec);
}

static int mzping_icmp_v4(int s_socket, const char *s_hostname, const char *s_address_string, struct addrinfo *s_addrinfo, int s_sequence_number, int s_timeout)
{
    int s_result, s_check, s_myid;
    unsigned char s_packet[ (20 + 40) + (ICMP_MINLEN + 4) ];
    struct mzping_icmp_echo_header *s_icmp_echo_header;
    struct mzping_icmp_mydata *s_mydata;
    size_t s_packet_size;
    ssize_t s_send_bytes;
    size_t s_ip_header_size;
    fd_set s_fd_rx;
    struct timeval s_timeval;
    ssize_t s_recv_bytes;
    socklen_t s_socklen_in;
    struct sockaddr_in s_sockaddr_in;

    s_result = (-1);
    s_myid = (int)(getpid() & 0xffff);

    s_icmp_echo_header = (struct mzping_icmp_echo_header *)(&s_packet[0]);
    s_icmp_echo_header->m_common.m_type = (uint8_t)(ICMP_ECHO);
    s_icmp_echo_header->m_common.m_code = (uint8_t)0u;
    s_icmp_echo_header->m_common.m_checksum = (uint16_t)0u;
    s_icmp_echo_header->m_identifier = (uint16_t)htons(s_myid);
    s_icmp_echo_header->m_sequence_number = (uint16_t)htons(s_sequence_number);
    s_packet_size = (size_t)ICMP_MINLEN /* 8 bytes */;

    s_mydata = (struct mzping_icmp_mydata *)(&s_icmp_echo_header[1]);
    s_mydata->m_timestamp32 = (uint32_t)mzping_ts32_us(); /* keep host endian */
    s_packet_size += sizeof(struct mzping_icmp_mydata);

    /* do checksum */
    s_icmp_echo_header->m_common.m_checksum = mzping_icmp_checksum((const void *)(&s_packet[0]), s_packet_size); /* checksum */

    s_send_bytes = sendto(s_socket, (const void *)(&s_packet[0]), s_packet_size, MSG_NOSIGNAL,
        (struct sockaddr *)s_addrinfo->ai_addr, s_addrinfo->ai_addrlen);
    if(s_send_bytes == ((ssize_t)(-1))) {
        perror("sendto");
        return(-1);
    }
    if(s_send_bytes != ((ssize_t)s_packet_size)) {
        (void)fprintf(stderr, "send: can not send %ld/%lu\n", (long)s_send_bytes, (unsigned long)sizeof(s_packet));
        return(-1);
    }

l_need_echoreply:;
    FD_ZERO(&s_fd_rx);
    FD_SET(s_socket, &s_fd_rx);
    s_timeval.tv_sec = s_timeout / 1000;
    s_timeval.tv_usec = (s_timeout % 1000) * 1000;
    s_check = select(s_socket + 1, (fd_set *)(&s_fd_rx), (fd_set *)0, (fd_set *)0, (struct timeval *)(&s_timeval));
    if(s_check == (-1)) {
        perror("select");
        return(-1);
    }
    if(s_check == 0) {
        (void)fprintf(stderr, "select: timeout\n");
        return(-1);
    }
    if(FD_ISSET(s_socket, &s_fd_rx) == 0) {
        (void)fprintf(stderr, "select: is not set\n");
        return(-1);
    }

    s_socklen_in = (socklen_t)sizeof(s_sockaddr_in);
    s_recv_bytes = recvfrom(s_socket, (void *)(&s_packet[0]), sizeof(s_packet), MSG_NOSIGNAL, (struct sockaddr *)(&s_sockaddr_in), (socklen_t *)(&s_socklen_in));
    if(s_recv_bytes == ((ssize_t)(-1))) {
        perror("recvfrom");
        return(-1);
    }
    
    s_ip_header_size = ((size_t)((s_packet[0] >> 0) & 0x0fu)) << 2;
    if(s_recv_bytes < (s_ip_header_size + s_packet_size)) {
        /* (void)fprintf(stderr, "too small packet\n"); */
        goto l_need_echoreply;
    }

    s_icmp_echo_header = (struct mzping_icmp_echo_header *)(&s_packet[s_ip_header_size]);
    if(ntohs(s_icmp_echo_header->m_identifier) != s_myid) {
        /* (void)fprintf(stderr, "not my ping\n"); */
        goto l_need_echoreply;
    }
    
    if(s_icmp_echo_header->m_common.m_type == ((uint8_t)ICMP_ECHO)) { /* maybe localhost loopback case */
        goto l_need_echoreply;
    }
    
    if(s_icmp_echo_header->m_common.m_type != ((uint8_t)ICMP_ECHOREPLY)) {
        goto l_need_echoreply;
    }
    else {
        unsigned int s_trip_time;

	s_mydata = (struct mzping_icmp_mydata *)(&s_icmp_echo_header[1]);
        s_result = (int)ntohs(s_icmp_echo_header->m_sequence_number);
        s_trip_time = mzping_ts32_us() - s_mydata->m_timestamp32;
        (void)fprintf(stdout, "%ld bytes from %s (%s): icmp_seq=%d ttl=%u time=%u.%03u ms;\n",
            (unsigned long)(s_recv_bytes - (s_ip_header_size)),
            s_hostname,
            s_address_string,
            s_result,
            (unsigned int)s_packet[8] /* iphdr->ttl */,
            s_trip_time / 1000u,
            s_trip_time % 1000u);
    }

    return(s_result);
}

static int mzping_icmp_v6(int s_socket, const char *s_hostname, const char *s_address_string, struct addrinfo *s_addrinfo, int s_sequence_number, int s_timeout)
{
    int s_result, s_check, s_myid;
    unsigned char s_packet[ (ICMP_MINLEN + 4) ];
    struct mzping_icmp_echo_header *s_icmp_echo_header;
    struct mzping_icmp_mydata *s_mydata;
    size_t s_packet_size;
    ssize_t s_send_bytes;

    fd_set s_fd_rx;
    struct timeval s_timeval;
    ssize_t s_recv_bytes;
    socklen_t s_socklen_in;
    struct sockaddr_in s_sockaddr_in;

    s_result = (-1);
    s_myid = (int)(getpid() & 0xffff);

    s_icmp_echo_header = (struct mzping_icmp_echo_header *)(&s_packet[0]);
    s_icmp_echo_header->m_common.m_type = (uint8_t)(ICMP6_ECHO_REQUEST);
    s_icmp_echo_header->m_common.m_code = (uint8_t)0u;
    s_icmp_echo_header->m_common.m_checksum = (uint16_t)0u;
    s_icmp_echo_header->m_identifier = (uint16_t)htons(s_myid);
    s_icmp_echo_header->m_sequence_number = (uint16_t)htons(s_sequence_number);
    s_packet_size = (size_t)ICMP_MINLEN /* 8 bytes */;

    s_mydata = (struct mzping_icmp_mydata *)(&s_icmp_echo_header[1]);
    s_mydata->m_timestamp32 = (uint32_t)mzping_ts32_us(); /* keep host endian */
    s_packet_size += sizeof(struct mzping_icmp_mydata);

    /* do checksum */
    s_icmp_echo_header->m_common.m_checksum = mzping_icmp_checksum((const void *)(&s_packet[0]), s_packet_size); /* checksum */

    s_send_bytes = sendto(s_socket, (const void *)(&s_packet[0]), s_packet_size, MSG_NOSIGNAL,
        (struct sockaddr *)s_addrinfo->ai_addr, s_addrinfo->ai_addrlen);
    if(s_send_bytes == ((ssize_t)(-1))) {
        perror("sendto");
        return(-1);
    }
    if(s_send_bytes != ((ssize_t)s_packet_size)) {
        (void)fprintf(stderr, "send: can not send %ld/%lu\n", (long)s_send_bytes, (unsigned long)sizeof(s_packet));
        return(-1);
    }

l_need_echoreply:;
    FD_ZERO(&s_fd_rx);
    FD_SET(s_socket, &s_fd_rx);
    s_timeval.tv_sec = s_timeout / 1000;
    s_timeval.tv_usec = (s_timeout % 1000) * 1000;
    s_check = select(s_socket + 1, (fd_set *)(&s_fd_rx), (fd_set *)0, (fd_set *)0, (struct timeval *)(&s_timeval));
    if(s_check == (-1)) {
        perror("select");
        return(-1);
    }
    if(s_check == 0) {
        (void)fprintf(stderr, "select: timeout\n");
        return(-1);
    }
    if(FD_ISSET(s_socket, &s_fd_rx) == 0) {
        (void)fprintf(stderr, "select: is not set\n");
        return(-1);
    }
    
    s_socklen_in = (socklen_t)sizeof(s_sockaddr_in);
    s_recv_bytes = recvfrom(s_socket, (void *)(&s_packet[0]), sizeof(s_packet), MSG_NOSIGNAL, (struct sockaddr *)(&s_sockaddr_in), (socklen_t *)(&s_socklen_in));
    if(s_recv_bytes == ((ssize_t)(-1))) {
        perror("recvfrom");
        return(-1);
    }
    
    if(s_recv_bytes < s_packet_size) {
        /* (void)fprintf(stderr, "too small packet\n"); */
        goto l_need_echoreply;
    }
    
    s_icmp_echo_header = (struct mzping_icmp_echo_header *)(&s_packet[0]);
    if(ntohs(s_icmp_echo_header->m_identifier) != s_myid) {
        /* (void)fprintf(stderr, "not my ping\n"); */
        goto l_need_echoreply;
    }

    if(s_icmp_echo_header->m_common.m_type == ((uint8_t)ICMP6_ECHO_REQUEST)) {
        goto l_need_echoreply;
    }
    
    if(s_icmp_echo_header->m_common.m_type != ((uint8_t)ICMP6_ECHO_REPLY)) {
        goto l_need_echoreply;
    }
    else {
        unsigned int s_trip_time;

	s_mydata = (struct mzping_icmp_mydata *)(&s_icmp_echo_header[1]);
        s_result = (int)ntohs(s_icmp_echo_header->m_sequence_number);
        s_trip_time = mzping_ts32_us() - s_mydata->m_timestamp32;
        (void)fprintf(stdout, "%ld bytes from %s (%s): icmp_seq=%d ttl=%u time=%u.%03u ms;\n",
            (unsigned long)s_recv_bytes,
            s_hostname,
            s_address_string,
            s_result,
            (unsigned int)0u /* TODO: Hops limit here */,
            s_trip_time / 1000u,
            s_trip_time % 1000u);
    }

    return(s_result);
}

int mzping(const char *s_hostname, int s_count)
{
    int s_sequence_number = 0, s_check, s_socket;
    struct addrinfo *s_addrinfo_result;
    struct addrinfo s_addrinfo_hints;
    struct addrinfo *s_addrinfo;
    char s_address_string[ 64 ];
   
    /* resolv name */
    (void)memset((void *)(&s_addrinfo_hints), 0, sizeof(s_addrinfo_hints));
    s_addrinfo_hints.ai_socktype = SOCK_RAW;
    s_addrinfo_hints.ai_family = AF_UNSPEC;
    s_check = getaddrinfo(s_hostname, (const char *)0, (const struct addrinfo *)(&s_addrinfo_hints), (struct addrinfo **)(&s_addrinfo_result));
    if(s_check != 0) {
        (void)fprintf(stderr, "getaddrinfo error: %s\n", gai_strerror(s_check));
        return(-1);
    }

    do {
        s_sequence_number++;

        for(s_addrinfo = s_addrinfo_result;s_addrinfo != ((struct addrinfo *)0);s_addrinfo = s_addrinfo->ai_next) {
            if(s_addrinfo->ai_family == AF_INET) { /* ICMP */
                struct sockaddr_in *s_in;

                s_socket = socket(s_addrinfo->ai_family, SOCK_RAW, IPPROTO_ICMP);
                if(s_socket == (-1)) {
                    perror("socket");
                    continue;
                }

                s_in = (struct sockaddr_in *)s_addrinfo->ai_addr;
                inet_ntop(s_addrinfo->ai_family, (const void *)(&s_in->sin_addr), (char *)(&s_address_string[0]), (socklen_t)sizeof(s_address_string));
                
                s_check = mzping_icmp_v4(s_socket, s_hostname, (const char *)(&s_address_string[0]), s_addrinfo, s_sequence_number, 20000);
                close(s_socket);
            }
            else if(s_addrinfo->ai_family == AF_INET6) { /* ICMPv6 */
                struct sockaddr_in6 *s_in6;
                
                s_socket = socket(s_addrinfo->ai_family, SOCK_RAW, IPPROTO_ICMPV6);
                if(s_socket == (-1)) {
                    perror("socket");
                    continue;
                }
                
                s_in6 = (struct sockaddr_in6 *)s_addrinfo->ai_addr;
                inet_ntop(s_addrinfo->ai_family, (const void *)(&s_in6->sin6_addr), (char *)(&s_address_string[0]), (socklen_t)sizeof(s_address_string));
                
                s_check = mzping_icmp_v6(s_socket, s_hostname, (const char *)(&s_address_string[0]), s_addrinfo, s_sequence_number, 20000);
                close(s_socket);
            }
        }

        (void)sleep(1);
    }while((s_count == 0) || (s_sequence_number < s_count));

    freeaddrinfo((struct addrinfo *)s_addrinfo_result);

    return(1);

}

int main(int s_argc, char **s_argv)
{
    int s_count = 4;
    
    (void)fprintf(stdout, "mzping v0.0.4 - Code by JaeHyuk Cho <minzkn@minzkn.com>\n");

    if(s_argc <= 1) {
        (void)fprintf(stdout, "usage: %s <host> <count>\n", (char *)s_argv[0]);
        
        /* (void)mzping("localhost", s_count); */

        return(EXIT_SUCCESS);
    }

    if(s_argc >= 3) {
        if(sscanf(s_argv[2], "%i", &s_count) != 1) {
            perror("count");
            return(EXIT_FAILURE);
        }
    }

    (void)setuid(getuid());
    (void)mzping(s_argv[1], s_count);

    return(EXIT_SUCCESS);
}

/* vim: set expandtab: */
/* End of source */
