/*
 * Copyright (c) 2006, Nate Nielsen
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *     * Redistributions of source code must retain the above
 *       copyright notice, this list of conditions and the
 *       following disclaimer.
 *     * Redistributions in binary form must reproduce the
 *       above copyright notice, this list of conditions and
 *       the following disclaimer in the documentation and/or
 *       other materials provided with the distribution.
 *     * The names of contributors to this software may not be
 *       used to endorse or promote products derived from this
 *       software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
 * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
 * THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 *
 *
 * CONTRIBUTORS
 *  Nate Nielsen <nielsen@memberwebs.com>
 */

#include "usuals.h"

#include <sys/types.h>
#include <sys/limits.h>
#include <sys/socket.h>
#include <sys/queue.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <sys/time.h>

#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <netinet/ip_var.h>
#include <arpa/inet.h>

#include <syslog.h>
#include <unistd.h>
#include <stdarg.h>
#include <ctype.h>
#include <fcntl.h>

#include <bsnmp/snmpmod.h>

#include "ping_tree.h"
#include "ping_oid.h"

/* our module handle */
static struct lmodule *module;

/* OIDs */
static const struct asn_oid oid_ping = OIDX_ping;

/* the Object Resource registration index */
static u_int reg_index = 0;

enum {
    RESPONSE_ERR = -1,
    RESPONSE_NONE = 0,
    RESPONSE_REPLY = 1
};

#define     A(t, bit)   (t)[(bit)>>3]	/* identify byte in array */
#define     B(bit)      (1 << ((bit) & 0x07))	/* identify bit in byte */
#define     SET(t, bit) (A((t), bit) |= B(bit))
#define     CLR(t, bit) (A((t), bit) &= (~B(bit)))
#define     TST(t, bit) (A((t), bit) & B(bit))
#define     MAX_DUP_CHK     8 * 64

struct pinger {
    uint32_t                index;
    TAILQ_ENTRY(pinger)     link;

    /* Configuration */
    u_char                  *host;
    struct sockaddr_in      addr;
    u_int                   interval;
    u_int                   history;

    /* Stats gathered */
    u_int                   at;
    int                     *latencies;

    /* For generating and filtering pings */
    void                    *timer;
    u_int                   transmitted;
    u_int                   received;
    char                    ck_table[MAX_DUP_CHK / 8];
};

TAILQ_HEAD(pinger_list, pinger);

/* list of regexes */
static struct pinger_list pingers = TAILQ_HEAD_INITIALIZER(pingers);
static u_int pinger_count = 0;

/* The ICMP socket */
static int icmp_sock = -1;
static void *icmp_sel = NULL;

/* The ICMP packet */
static u_char icmp_packet[IP_MAXPACKET];

/* -----------------------------------------------------------------------------
 * HELPERS
 */

static void
emsg(const char *format, ...)
{
    va_list va;
    va_start (va, format);
    vsyslog (LOG_ERR, format, va);
    va_end (va);
}

static uint64_t
get_current_ms ()
{
    struct timeval tv;

    if (gettimeofday (&tv, NULL) < 0) {
        emsg ("couldn't get current time: %s", strerror (errno));
        return get_ticks () * 10;
    }

    return ((uint64_t)(tv.tv_sec)) * 1000 +
           ((uint64_t)(tv.tv_usec)) / 1000;
}

/* -----------------------------------------------------------------------------
 * ICMP STUFF
 */

static void
process_icmp (u_char *packet, int len, struct sockaddr_in *from)
{
    struct pinger *ping;
    struct icmp *icp;
    struct ip *ip;
    uint64_t ptime;
    uint64_t ctime;
    const void *tp;
    int triptime;
    int hlen, seq;

    /* Check the IP header */
    ip = (struct ip*)packet;
    hlen = ip->ip_hl << 2;
    if (len < hlen + ICMP_MINLEN)
        return;

    /* Now the ICMP part */
    len -= hlen;
    icp = (struct icmp *)(packet + hlen);

    if (icp->icmp_type != ICMP_ECHOREPLY)
        return; /* An error or some such */

    for (ping = TAILQ_FIRST (&pingers); ping;
         ping = TAILQ_NEXT (ping, link)) {
        if (ping->index == icp->icmp_id)
            break;
    }

    if (!ping)
        return;     /* 'Twas not our ECHO */

    seq = ntohs (icp->icmp_seq);
    if (TST (ping->ck_table, seq % MAX_DUP_CHK))
        return;
    SET(ping->ck_table, seq % MAX_DUP_CHK);
    seq %= ping->history;

    ping->received++;
    triptime = -1;


    ctime = get_current_ms ();

#ifndef icmp_data
    tp = &icp->icmp_ip;
#else
    tp = icp->icmp_data;
#endif
    if (len - ICMP_MINLEN >= sizeof (ptime)) {
        /* Copy to avoid alignment problems: */
        memcpy(&ptime, tp, sizeof(ptime));
        triptime = ((int)(ctime - ptime));

        /* Some small rounding adjustments */
        if (triptime == 0)
            triptime = 1;
        if (triptime < 0)
            triptime = 0;
    }

    if (ping->latencies)
        ping->latencies[seq] = triptime;
}

static void
close_socket ()
{
    if (icmp_sel)
        fd_deselect (icmp_sel);
    icmp_sel = NULL;

    if (icmp_sock != -1)
        close (icmp_sock);
    icmp_sock = -1;
}

static void
io_icmp (int fd, void *unused)
{
    struct sockaddr_in from;
    socklen_t fromlen;
    int len;

    ASSERT (fd == icmp_sock);

    fromlen = sizeof (from);
    len = recvfrom (icmp_sock, icmp_packet, sizeof (icmp_packet), 0,
                    (struct sockaddr*)&from, &fromlen);
    if (len < 0) {
        if (errno == EINTR)
            return;
        emsg ("couldn't receive packet: %s", strerror (errno));
    }

    process_icmp (icmp_packet, len, &from);
}

static int
open_socket ()
{
    int hold;

    close_socket ();

    icmp_sock = socket (AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (icmp_sock == -1) {
        emsg ("couldn't open ICMP socket: %s", strerror (errno));
        return errno;
    }

    hold = IP_MAXPACKET + 128;
    setsockopt(icmp_sock, SOL_SOCKET, SO_RCVBUF, (char*)&hold, sizeof(hold));
    setsockopt(icmp_sock, SOL_SOCKET, SO_SNDBUF, (char*)&hold, sizeof(hold));

    icmp_sel = fd_select (icmp_sock, io_icmp, NULL, module);
    if (!icmp_sel) {
        emsg ("couldn't listen on ICMP socket: %s", strerror (errno));
        return errno;
    }

    return 0;
}

static u_short
in_cksum (u_short *addr, int len)
{
    int nleft, sum;
    u_short *w;
    union {
        u_short us;
        u_char uc[2];
    } last;
    u_short answer;

    nleft = len;
    sum = 0;
    w = addr;

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

    /* mop up an odd byte, if necessary */
    if (nleft == 1) {
        last.uc[0] = *(u_char *)w;
        last.uc[1] = 0;
        sum += last.us;
    }

    /* add back carry outs from top 16 bits to low 16 bits */
    sum = (sum >> 16) + (sum & 0xffff);     /* add hi 16 to low 16 */
    sum += (sum >> 16);                     /* add carry */
    answer = ~sum;                          /* truncate to 16 bits */
    return answer ;
}


static void
ping_fire (void *user_data)
{
    struct pinger *ping = (struct pinger*)user_data;
    uint64_t ptime;
    struct icmp *icp;
    int cc, i, seq;

    memset (icmp_packet, 0, sizeof (icmp_packet));

    ping->at = (ping->at + ((ping->transmitted > 1) ? 1 : 0)) % ping->history;
    seq = ping->transmitted % USHRT_MAX;
    ASSERT (ping->transmitted == 0 ||
            (ping->transmitted % ping->history) == ((ping->at + 1) % ping->history));

    if (ping->latencies)
        ping->latencies[seq % ping->history] = -1;

    icp = (struct icmp*)icmp_packet;
    icp->icmp_type = ICMP_ECHO;
    icp->icmp_code = 0;
    icp->icmp_cksum = 0;
    icp->icmp_seq = htons (seq);
    icp->icmp_id = ping->index; /* ID */

    /* Include the current time in the packet */
    ptime = get_current_ms ();
    memcpy (icmp_packet + ICMP_MINLEN, &ptime, sizeof (ptime));

    CLR (ping->ck_table, seq % MAX_DUP_CHK);

    /* compute ICMP checksum here */
    cc = ICMP_MINLEN + sizeof (ptime);
    icp->icmp_cksum = in_cksum ((u_short *)icp, cc);

    i = sendto(icmp_sock, icmp_packet, cc, 0, (struct sockaddr*)&(ping->addr),
               sizeof (ping->addr));

    if (i < 0) {
        switch (errno) {
        case EHOSTDOWN:
        case EHOSTUNREACH:
        case ENETDOWN:
        case ENETUNREACH:
            break;
        default:
            emsg ("couldn't send ICMP packet to: %s: %s", ping->host, strerror (errno));
            break;
        }
    } else if (i != cc)
        emsg ("partial write sending ICMP packet to: %", ping->host);

    ping->transmitted++;

    ping->timer = timer_start (ping->interval, ping_fire, ping, module);
    if (!ping->timer)
        emsg ("couldn't setup timer to ping host: %s", ping->host);
}

static void
ping_stop ()
{
    struct pinger *ping;

    for (ping = TAILQ_FIRST (&pingers); ping;
         ping = TAILQ_NEXT (ping, link)) {

        if (ping->timer)
            timer_stop (ping->timer);
        ping->timer = NULL;

        memset (ping->ck_table, 0, sizeof (ping->ck_table));
        ping->transmitted = 0;

        if (ping->latencies)
            memset (ping->latencies, 0, sizeof (u_int) * ping->history);
    }
}

static void
ping_start ()
{
    struct pinger *ping;

    ping_stop ();

    for (ping = TAILQ_FIRST (&pingers); ping;
         ping = TAILQ_NEXT (ping, link)) {

        /* Need to have an address configured */
        if (!ping->host || !ping->host[0])
            return;

        /* Allocate memory where needed */
        if (!ping->latencies)
            ping->latencies = (u_int*)calloc (ping->history, sizeof (u_int));
        if (!ping->latencies) {
            emsg ("out of memory");
            continue;
        }

        ping->at = 0;

        ping_fire (ping);
    }
}

/* -----------------------------------------------------------------------------
 * CALLBACKS/CONFIG
 */

static int
calc_average (struct pinger *ping)
{
    double sum = 0;
    double num = 0;
    int i;

    if (!ping->latencies)
        return -1;

    for (i = 0; i < ping->history; i++) {

        /* If not wrapped around yet, then don't wrap around */
        if (ping->transmitted <= ping->history && i >= ping->at)
            break;
        if (ping->latencies[i] < 0)
            continue;

        sum += ping->latencies[i];
        num ++;
    }

    return num ? (sum / num) : 0;
}

static int
calc_minimum (struct pinger *ping)
{
    int min = -1;
    int i;

    if (!ping->latencies)
        return -1;

    for (i = 0; i < ping->history; i++) {

        /* If not wrapped around yet, then don't wrap around */
        if (ping->transmitted <= ping->history && i >= ping->at)
            break;

        if (ping->latencies[i] < 0)
            continue;
        if (min > 0 && min < ping->latencies[i])
            continue;

        min = ping->latencies[i];
    }

    return min > 0 ? min : 0;
}

static int
calc_maximum (struct pinger *ping)
{
    int max = -1;
    int i;

    if (!ping->latencies)
        return -1;

    for (i = 0; i < ping->history; i++) {

        /* If not wrapped around yet, then don't wrap around */
        if (ping->transmitted <= ping->history && i >= ping->at)
            break;
        if (ping->latencies[i] < 0)
            continue;
        if (max > ping->latencies[i])
            continue;

        max = ping->latencies[i];
    }

    return max > 0 ? max : 0;
}

static int
calc_responses (struct pinger *ping)
{
    u_int num = 0;
    int i;

    if (!ping->latencies)
        return -1;

    for (i = 0; i < ping->history; i++) {

        /* If not wrapped around yet, then don't wrap around */
        if (ping->transmitted <= ping->history && i >= ping->at)
            break;
        if (ping->latencies[i] < 0)
            continue;

        num ++;
    }

    return num;
}

static int
calc_dropped (struct pinger *ping)
{
    u_int num = 0;
    int i;

    if (!ping->latencies)
        return -1;

    for (i = 0; i < ping->history; i++) {

        /* If not wrapped around yet, then don't wrap around */
        if (ping->transmitted <= ping->history && i >= ping->at)
            break;
        if (ping->latencies[i] >= 0)
            continue;

        num++;
    }

    return num;
}


static void
free_pinger (struct pinger *ping)
{
    if (ping->host)
        free (ping->host);
    ping->host = NULL;

    if (ping->latencies)
        free (ping->latencies);
    ping->latencies = NULL;

    TAILQ_REMOVE (&pingers, ping, link);
    free (ping);
}

static struct pinger*
alloc_pinger (u_int index)
{
    struct pinger *ping;

    ping = (struct pinger*)calloc (1, sizeof (struct pinger));
    if (!ping)
        return NULL;

    ping->host = strdup ("");
    if (!ping->host) {
        free (ping);
        return NULL;
    }

    ping->interval = 100;
    ping->history = 5;
    ping->index = index;
    ping->addr.sin_family = AF_INET;
    ping->addr.sin_len = sizeof (ping->addr);
    ping->addr.sin_port = 0;

    INSERT_OBJECT_INT (ping, &pingers);
    return ping;
}

static int
ping_config (struct pinger *ping, struct snmp_context *ctx,
             struct snmp_value *value, u_int sub, u_int iidx, enum snmp_op op)
{
    asn_subid_t which = value->var.subs[sub - 1];
    int index, r;

    /* Just return values, no creation */
    if (op == SNMP_OP_GET || op == SNMP_OP_GETNEXT) {

        if (!ping)
            return SNMP_ERR_NOSUCHNAME;

        switch (which) {
        case LEAF_pingIndex:
            value->v.integer = ping->index;
            return SNMP_ERR_NOERROR;
        case LEAF_pingHost:
            return string_get (value, ping->host, -1);
        case LEAF_pingHistory:
            value->v.integer = ping->history;
            return SNMP_ERR_NOERROR;
        case LEAF_pingInterval:
            value->v.integer = ping->interval;
            return SNMP_ERR_NOERROR;
        default:
            ASSERT (0);
            return SNMP_ERR_NOSUCHNAME;
        }
    }

    /* Remainder only at initialization */
    if (community != COMM_INITIALIZE)
        return ping ? SNMP_ERR_NOT_WRITEABLE : SNMP_ERR_NO_CREATION;

    /* No writing to pingIndex */
    if (which == LEAF_pingIndex)
        return SNMP_ERR_NOT_WRITEABLE;

    if (index_decode (&value->var, sub, iidx, &index))
        return SNMP_ERR_NO_CREATION;

    /* Create it if necessary */
    if (!ping) {
        ping = alloc_pinger (index);
        if (!ping) {
            emsg ("out of memory");
            return SNMP_ERR_GENERR;
        }
    }

    switch (which) {

    /* pingHistory */
    case LEAF_pingHistory:
        switch (op) {
        case SNMP_OP_SET:
            ctx->scratch->int1 = ping->history;
            ping->history = value->v.integer;
            return SNMP_ERR_NOERROR;
        case SNMP_OP_ROLLBACK:
            ping->history = ctx->scratch->int1;
            return SNMP_ERR_NOERROR;
        case SNMP_OP_COMMIT:
            return SNMP_ERR_NOERROR;
        default:
            ASSERT (0);
            return SNMP_ERR_GENERR;
        }
        break;

    /* pingInterval */
    case LEAF_pingInterval:
        switch (op) {
        case SNMP_OP_SET:
            ctx->scratch->int1 = ping->interval;
            ping->interval = value->v.integer;
            return SNMP_ERR_NOERROR;
        case SNMP_OP_ROLLBACK:
            ping->interval = ctx->scratch->int1;
            return SNMP_ERR_NOERROR;
        case SNMP_OP_COMMIT:
            return SNMP_ERR_NOERROR;
        default:
            ASSERT (0);
            return SNMP_ERR_GENERR;
        }
        break;

    /* pingHost */
    case LEAF_pingHost:
        switch (op) {
        case SNMP_OP_SET:
            if ((r = string_save (value, ctx, -1, &ping->host)) == SNMP_ERR_NOERROR) {
                if (!inet_aton (ping->host, &(ping->addr.sin_addr)))
                    r = SNMP_ERR_WRONG_VALUE;
            }
            if (r != SNMP_ERR_NOERROR)
                string_rollback (ctx, &ping->host);
            return r;
        case SNMP_OP_ROLLBACK:
            string_rollback (ctx, &ping->host);
            inet_aton (ping->host, &(ping->addr.sin_addr));
            return SNMP_ERR_NOERROR;
        case SNMP_OP_COMMIT:
            string_commit (ctx);
            return SNMP_ERR_NOERROR;
        default:
            ASSERT (0);
            return SNMP_ERR_GENERR;
        }
        break;

    /* Unknown OID */
    default:
        ASSERT (0);
        return SNMP_ERR_NOSUCHNAME;
    }
}

int
op_pingentry (struct snmp_context *ctx, struct snmp_value *value,
              u_int sub, u_int iidx, enum snmp_op op)
{
    asn_subid_t which = value->var.subs[sub - 1];
    struct pinger *ping = NULL;

    switch (op) {
    case SNMP_OP_GETNEXT:
        ping = NEXT_OBJECT_INT (&pingers, &value->var, sub);
        if (ping == NULL)
            return SNMP_ERR_NOSUCHNAME;
        value->var.len = sub + 1;
        value->var.subs[sub] = ping->index;
        break;

    case SNMP_OP_GET:
        ping = FIND_OBJECT_INT (&pingers, &value->var, sub);
        if (ping == NULL)
            return SNMP_ERR_NOSUCHNAME;
        break;

    default:
        ping = FIND_OBJECT_INT (&pingers, &value->var, sub);
        break;
    }

    /* Send configuration stuff off elsewhere */
    switch (which) {
    case LEAF_pingIndex:
    case LEAF_pingHost:
    case LEAF_pingHistory:
    case LEAF_pingInterval:
        return ping_config (ping, ctx, value, sub, iidx, op);
    }

    if (op != SNMP_OP_GET && op != SNMP_OP_GETNEXT)
        return SNMP_ERR_NOT_WRITEABLE;

    switch (which) {
    case LEAF_pingLatencyAvg:
        value->v.integer = calc_average (ping);
        return SNMP_ERR_NOERROR;
    case LEAF_pingLatencyMin:
        value->v.integer = calc_minimum (ping);
        return SNMP_ERR_NOERROR;
    case LEAF_pingLatencyMax:
        value->v.integer = calc_maximum (ping);
        return SNMP_ERR_NOERROR;
    case LEAF_pingResponses:
        value->v.integer = calc_responses (ping);
        return SNMP_ERR_NOERROR;
    case LEAF_pingDropped:
        value->v.integer = calc_dropped (ping);
        return SNMP_ERR_NOERROR;
    default:
        ASSERT (0);
        return SNMP_ERR_NOSUCHNAME;
    }
}

int
op_ping (struct snmp_context *ctx, struct snmp_value *value,
          u_int sub, u_int iidx, enum snmp_op op)
{
    asn_subid_t which = value->var.subs[sub - 1];

    switch (op) {
    case SNMP_OP_GET:
        break;

    case SNMP_OP_SET:
        return SNMP_ERR_NOT_WRITEABLE;

    case SNMP_OP_ROLLBACK:
    case SNMP_OP_COMMIT:
        return SNMP_ERR_NOERROR;

    default:
        ASSERT(0);
        break;
    }

    switch (which) {
    case LEAF_pingCount:
        value->v.integer = pinger_count;
        break;

    default:
        ASSERT(0);
        break;
    }

    return SNMP_ERR_NOERROR;
}


/* -----------------------------------------------------------------------------
 * MODULE
 */

/* the initialisation function */
static int
module_init (struct lmodule *mod, int argc, char *argv[])
{
    module = mod;

    if (argc != 0) {
        syslog(LOG_ERR, "bad number of arguments for %s", __func__);
        return EINVAL;
    }

    return open_socket ();
}

/* Module is started */
static void
module_start (void)
{
    reg_index = or_register (&oid_ping, "The MIB for ping data.", module);

    /* Start timers on them all */
    ping_start ();
}

/* Called, when the module is to be unloaded after it was successfully loaded */
static int
module_fini (void)
{
    struct pinger *ping;

    if (reg_index)
        or_unregister (reg_index);

    while ((ping = TAILQ_FIRST(&pingers)) != NULL)
        free_pinger (ping);

    return 0;
}

const struct snmp_module config = {
    .comment =      "This module implements SNMP listing of ping responses and latencies",
    .init =         module_init,
    .start =        module_start,
    .fini =         module_fini,
    .tree =         ping_ctree,
    .tree_size =    ping_CTREE_SIZE,
};