/*
 Copyright (C) JAEHYUK CHO
 All rights reserved.

 Author: JaeHyuk Cho <minzkn@minzkn.com>

 Tiny ping example source
*/

/* need for "struct addrinfo" */
#if 1L && (!defined(_POSIX_SOURCE))
# define _POSIX_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>

#define def_mzping_icmp_optimize (1)

#define def_mzping_icmp_min_packet_size (1 + 1 + 2 + 2 + 2) /* type(1) + code(1) + checksum(2) + identifier(2) + sequence_number(2) */

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)1)) {
#if def_mzping_icmp_optimize == (0)
        s_result += (unsigned int)(ntohs(*((const unsigned short int *)s_data)));
#else
        s_result += (unsigned int)(*((const unsigned short int *)s_data));
#endif
        s_data = ((const unsigned short int *)s_data) + ((size_t)1);
        s_size -= (size_t)2;
    }

    if(s_size > ((size_t)0)) {
        s_result += (unsigned int)(*((const unsigned char *)s_data));
    }

#if def_mzping_icmp_optimize == (0)
    while(s_result > 0xffffu) {
        s_result = (s_result >> 16) + (s_result & 0xffffu);
    }
#else
    s_result = (s_result >> 16) + (s_result & 0xffffu);
    s_result += s_result >> 16;
#endif

#if def_mzping_icmp_optimize == (0)
    return(htons((~s_result) & 0xffffu));
#else
    return((~s_result) & 0xffffu);
#endif
}

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) + (def_mzping_icmp_min_packet_size + 4) ];
    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_packet_size = (size_t)0;
    s_packet[s_packet_size] = 8u; /* ICMP_ECHO */
    s_packet_size += (size_t)1;
    s_packet[s_packet_size] = 0u; /* code */
    s_packet_size += (size_t)1;
    *((unsigned short int *)(&s_packet[s_packet_size])) = 0u; /* checksum */
    s_packet_size += (size_t)2;
    *((unsigned short int *)(&s_packet[s_packet_size])) = htons(s_myid); /* identifier */
    s_packet_size += (size_t)2;
    *((unsigned short int *)(&s_packet[s_packet_size])) = htons(s_sequence_number); /* sequence number */
    s_packet_size += (size_t)2;
#if def_mzping_icmp_optimize == (0)
    *((unsigned int *)(&s_packet[8])) = htonl(mzping_ts32_us()); /* optional: time stamp */
#else
    *((unsigned int *)(&s_packet[8])) = (unsigned int)mzping_ts32_us(); /* optional: time stamp */
#endif
    s_packet_size += (size_t)4;

    /* do checksum */
    *((unsigned short int *)(&s_packet[2])) = 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)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;
    }

    if(ntohs(*((unsigned short int *)(&s_packet[s_ip_header_size + 4]))) != s_myid) {
        /* (void)fprintf(stderr, "not my ping\n"); */
        goto l_need_echoreply;
    }

    if(s_packet[s_ip_header_size] == 8u /* ICMP_ECHO */) { /* maybe localhost loopback case */
        goto l_need_echoreply;
    }

    if(s_packet[s_ip_header_size] == 0u /* ICMP_ECHOREPLY */) {
        unsigned int s_trip_time;
        s_result = (int)ntohs(*((unsigned short int *)(&s_packet[s_ip_header_size + 1 + 1 + 2 + 2])));
#if def_mzping_icmp_optimize == (0)
        s_trip_time = mzping_ts32_us() - ntohl(*((unsigned int *)(&s_packet[s_ip_header_size + def_mzping_icmp_min_packet_size])));
#else
        s_trip_time = mzping_ts32_us() - (*((unsigned int *)(&s_packet[s_ip_header_size + def_mzping_icmp_min_packet_size])));
#endif
        (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],
            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[ (def_mzping_icmp_min_packet_size + 4) ];
    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_packet_size = (size_t)0;
    s_packet[s_packet_size] = 128u; /* ICMP6_ECHO_REQUEST */
    s_packet_size += (size_t)1;
    s_packet[s_packet_size] = 0u; /* code */
    s_packet_size += (size_t)1;
    *((unsigned short int *)(&s_packet[s_packet_size])) = 0u; /* checksum */
    s_packet_size += (size_t)2;
    *((unsigned short int *)(&s_packet[s_packet_size])) = htons(s_myid); /* identifier */
    s_packet_size += (size_t)2;
    *((unsigned short int *)(&s_packet[s_packet_size])) = htons(s_sequence_number); /* sequence number */
    s_packet_size += (size_t)2;
#if def_mzping_icmp_optimize == (0)
    *((unsigned int *)(&s_packet[8])) = htonl(mzping_ts32_us()); /* optional: time stamp */
#else
    *((unsigned int *)(&s_packet[8])) = (unsigned int)mzping_ts32_us(); /* optional: time stamp */
#endif
    s_packet_size += (size_t)4;

    /* do checksum */
    *((unsigned short int *)(&s_packet[2])) = 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)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;
    }

    if(ntohs(*((unsigned short int *)(&s_packet[4]))) != s_myid) {
        /* (void)fprintf(stderr, "not my ping\n"); */
        goto l_need_echoreply;
    }

    if(s_packet[0] == 128u /* ICMP6_ECHO_REQUEST */) {
        goto l_need_echoreply;
    }

    if(s_packet[0] == 129u /* ICMP6_ECHO_REPLY */) {
        unsigned int s_trip_time;
        s_result = (int)ntohs(*((unsigned short int *)(&s_packet[1 + 1 + 2 + 2])));
#if def_mzping_icmp_optimize == (0)
        s_trip_time = mzping_ts32_us() - ntohl(*((unsigned int *)(&s_packet[def_mzping_icmp_min_packet_size])));
#else
        s_trip_time = mzping_ts32_us() - (*((unsigned int *)(&s_packet[def_mzping_icmp_min_packet_size])));
#endif
        (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_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.2 - 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 */