#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <signal.h>
#include <sys/types.h>
#include <time.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <netdb.h>
#include <net/if.h>

#define SERVER_ECHO 1
#define SERVER_DNS 2
#define SERVER_NTP 3

struct server {
    char *ifname;
    int protocol;
    unsigned char addr[4];
    unsigned short port;
    unsigned short gotone;
    unsigned short reach;
};

#define MAXSERVERS 16

struct server servers[MAXSERVERS];
int numservers;

int protocol = 43;
int installed = 0;

sig_atomic_t exiting = 0;

int gettime(struct timeval *tv);
void doze(void);
void init_signals(void);
void install(void), uninstall(void);

#define BUF_LEN 512

int
main(int argc, char **argv)
{
    int i, server_protocol, rc, port;
    char *ifname;
    int s;
    int one = 1;
    int interval = 120;
    struct timeval now;
    int starting = 1;

    server_protocol = SERVER_ECHO;
    port = 7;
    ifname = NULL;
    numservers = 0;

    for(i = 1; i < argc; i++) {
        struct hostent *host;
        if(argv[i][0] == '-') {
            if(strcmp(argv[i], "-p") == 0) {
                i++;
                protocol = atoi(argv[i]);
                if(protocol < 1 || protocol > 255)
                    goto syntax;
                continue;
            } else if(strcmp(argv[i], "-P") == 0) {
                i++;
                if(strcmp(argv[i], "echo") == 0) {
                    server_protocol = SERVER_ECHO;
                    port = 7;
                } else if(strcmp(argv[i], "dns") == 0) {
                    server_protocol = SERVER_DNS;
                    port = 53;
                } else if(strcmp(argv[i], "ntp") == 0) {
                    server_protocol = SERVER_NTP;
                    port = 123;
                } else
                    goto syntax;
                continue;
            } else if(strcmp(argv[i], "-i") == 0) {
                i++;
                ifname = argv[i];
                continue;
            } else if(strcmp(argv[i], "-t") == 0) {
                i++;
                interval = atoi(argv[i]);
                if(interval < 2)
                    goto syntax;
                continue;
            } else
                goto syntax;
        }

        if(numservers >= MAXSERVERS) {
            fprintf(stderr, "Too many servers.\n");
            exit(1);
        }

        host = gethostbyname(argv[i]);
        if(host == NULL || host->h_addr == NULL) {
            fprintf(stderr, "Couldn't find host %s.\n", argv[i]);
            exit(1);
        }

        if(ifname == NULL) {
            fprintf(stderr, "No interface defined.\n");
            goto syntax;
        }

        servers[numservers].protocol = server_protocol;
        servers[numservers].ifname = strdup(ifname);
        memcpy(servers[numservers].addr, host->h_addr, 4);
        servers[numservers].port = port;
        servers[numservers].gotone = 0;
        servers[numservers].reach = 0;
        numservers++;
    }
        
    if(numservers == 0) {
        fprintf(stderr, "No usable servers -- aborting.\n");
        exit(1);
    }

    s = socket(PF_INET, SOCK_DGRAM, 0);
    if(s < 0) {
        perror("socket");
        exit(1);
    }

    rc = setsockopt(s, IPPROTO_IP, IP_PKTINFO, &one, sizeof(one));
    if(rc < 0) {
        perror("setsockopt(IP_PKTINFO)");
        exit(1);
    }

    init_signals();

    doze();
    gettime(&now);

    while(!exiting) {
        long endtime;
        struct sockaddr_in sin;
        int reached;

        for(i = 0; i < numservers; i++) {
            unsigned char buf[BUF_LEN];
            int len;
            struct in_pktinfo pktinfo;
            struct msghdr msg;
            struct iovec iov[1];
            char cbuf[CMSG_SPACE(sizeof(pktinfo))];
            struct cmsghdr *cmsg;

            doze();

            if(servers[i].protocol == SERVER_ECHO) {
                /* Empty datagram */
                len = 0;
            } else if(servers[i].protocol == SERVER_DNS) {
                /* A query with qdcount = 0 */
                unsigned short id = random() & 0xFFFF;
                memset(buf, 0, 12);
                buf[0] = id >> 8;
                buf[1] = id & 0xFF;
                len = 12;
            } else if(servers[i].protocol == SERVER_NTP) {
                /* A version 3 client request with all fields set to 0 */
                memset(buf, 0, 48);
                buf[0] = 0x1b;
                len = 48;
            } else {
                abort();
            }

            memset(&pktinfo, 0, sizeof(pktinfo));
            pktinfo.ipi_ifindex = if_nametoindex(servers[i].ifname);
            if(pktinfo.ipi_ifindex == 0)
                continue;

            memset(&sin, 0, sizeof(sin));
            sin.sin_family = PF_INET;
            memcpy(&sin.sin_addr, servers[i].addr, 4);
            sin.sin_port = htons(servers[i].port);

            iov[0].iov_base = buf;
            iov[0].iov_len = len;

            memset(&msg, 0, sizeof(msg));
            msg.msg_name = &sin;
            msg.msg_namelen = sizeof(sin);
            msg.msg_iov = iov;
            msg.msg_iovlen = 1;
            msg.msg_control = cbuf;
            msg.msg_controllen = sizeof(cbuf);

            cmsg = CMSG_FIRSTHDR(&msg);
            cmsg->cmsg_level = IPPROTO_IP;
            cmsg->cmsg_type = IP_PKTINFO;
            cmsg->cmsg_len = CMSG_LEN(sizeof(struct in_pktinfo));
            memcpy(CMSG_DATA(cmsg), &pktinfo, sizeof(pktinfo));
            msg.msg_controllen = cmsg->cmsg_len;

            rc = sendmsg(s, &msg, 0);

            if(rc < len)
                perror("send");
        }

        /* Speed things up at the beginning */
        if(starting) {
            endtime = now.tv_sec + 4;
            starting = 0;
        } else {
            endtime = now.tv_sec + interval;
        }
        while(!exiting) {
            struct timeval tv;
            fd_set readfds;
            
            gettime(&now);
            if(now.tv_sec >= endtime)
                break;

            FD_ZERO(&readfds);
            FD_SET(s, &readfds);

            tv.tv_sec = endtime - now.tv_sec;
            tv.tv_usec = random() % 1000000;
            rc = select(s + 1, &readfds, NULL, NULL, &tv);
            if(rc < 0 && errno != EINTR && errno != EAGAIN) {
                perror("select");
                doze();
                continue;
            }

            if(exiting)
                break;
            
            if(rc > 0) {
                unsigned char buf[BUF_LEN];
                struct msghdr msg;
                struct iovec iov[1];
                struct cmsghdr *cmsg;
                char cbuf[512];
                struct in_pktinfo *pktinfo;

                iov[0].iov_base = buf;
                iov[0].iov_len = BUF_LEN;
                
                memset(&msg, 0, sizeof(msg));
                msg.msg_name = &sin;
                msg.msg_namelen = sizeof(sin);
                msg.msg_iov = iov;
                msg.msg_iovlen = 1;
                msg.msg_control = cbuf;
                msg.msg_controllen = 512;
                
                rc = recvmsg(s, &msg, 0);
                if(rc < 0) {
                    perror("recvmsg");
                    doze();
                    continue;
                }

                pktinfo = NULL;
                for(cmsg = CMSG_FIRSTHDR(&msg); cmsg != NULL;
                    cmsg = CMSG_NXTHDR(&msg, cmsg)) {
                    if(cmsg->cmsg_level == IPPROTO_IP &&
                       cmsg->cmsg_type == IP_PKTINFO) {
                        pktinfo = (struct in_pktinfo *)CMSG_DATA(cmsg);
                        break;
                    }
                }

                if(pktinfo == NULL) {
                    fprintf(stderr, "No pktinfo received.\n");
                    continue;
                }

                for(i = 0; i < numservers; i++) {
                    if(memcmp(&sin.sin_addr, servers[i].addr, 4) == 0 &&
                       ntohs(sin.sin_port) == servers[i].port &&
                       pktinfo->ipi_ifindex ==
                       if_nametoindex(servers[i].ifname)) {
                        servers[i].gotone = 1;
                        /* This is redundant, but it speeds things up. */
                        if((servers[i].reach & 0xC000))
                            install();
                    }
                }
            }
        }

        for(i = 0; i < numservers; i++) {
            servers[i].reach =
                (servers[i].reach >> 1) | (!!servers[i].gotone) << 15;
            servers[i].gotone = 0;
        }

        reached = 0;
        for(i = 0; i < numservers; i++) {
            unsigned short reach = servers[i].reach;
            int count =
                !!(reach & 0x8000) + !!(reach & 0x4000) +
                !!(reach & 0x2000);
            reached = reached || (count >= 2);
        }
        if(reached) {
            install();
        } else {
            uninstall();
        }
    }

    uninstall();

    return 0;

 syntax:
    fprintf(stderr,
            "Syntax: babel-pinger "
            "[-p protocol number] [-P server protocol] [-t timeout]\n"
            "                     "
            "-i interface host...\n");
    exit(1);
}

int
gettime(struct timeval *tv)
{
#if defined(_POSIX_TIMERS) && _POSIX_TIMERS > 0 && defined(CLOCK_MONOTONIC)
    static int have_posix_clocks = -1;

    if(have_posix_clocks < 0) {
        struct timespec ts;
        int rc;
        rc = clock_gettime(CLOCK_MONOTONIC, &ts);
        if(rc < 0) {
            have_posix_clocks = 0;
        } else {
            have_posix_clocks = 1;
        }
    }

    if(have_posix_clocks) {
        struct timespec ts;
        int rc;
        rc = clock_gettime(CLOCK_MONOTONIC, &ts);
        if(rc < 0)
            return rc;
        tv->tv_sec = ts.tv_sec;
        tv->tv_usec = ts.tv_nsec / 1000;
        return rc;
    }
#endif

    return gettimeofday(tv, NULL);
}

void
doze()
{
    if(exiting)
        return;

    usleep(random() % 1000000);
}

static void
sigexit(int signo)
{
    exiting = 1;
}

void
init_signals(void)
{
    struct sigaction sa;
    sigset_t ss;

    sigemptyset(&ss);
    sa.sa_handler = sigexit;
    sa.sa_mask = ss;
    sa.sa_flags = 0;
    sigaction(SIGTERM, &sa, NULL);

    sigemptyset(&ss);
    sa.sa_handler = sigexit;
    sa.sa_mask = ss;
    sa.sa_flags = 0;
    sigaction(SIGHUP, &sa, NULL);

    sigemptyset(&ss);
    sa.sa_handler = sigexit;
    sa.sa_mask = ss;
    sa.sa_flags = 0;
    sigaction(SIGINT, &sa, NULL);
}

void
install()
{
    char buf[100];
    if(!installed) {
        snprintf(buf, 100,
                 "ip route add 0.0.0.0/0 dev lo metric 65534 proto %d",
                 protocol);
        system(buf);
    }
    installed = 1;
}

void
uninstall()
{
    char buf[100];
    if(installed) {
        snprintf(buf, 100,
                 "ip route flush 0.0.0.0/0 dev lo proto %d", protocol);
        system(buf);
    }
    installed = 0;
}