/*
 * scamper_do_ping.c
 *
 * $Id: scamper_do_ping.c,v 1.42 2007/05/14 04:50:23 mjl Exp $
 *
 * Copyright (C) 2005-2007 The University of Waikato
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 2.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 * 
 */

#include <sys/types.h>
#include <sys/time.h>
#include <sys/socket.h>

#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <netinet/ip6.h>
#include <netinet/icmp6.h>

#if defined(__APPLE__)
#include <stdint.h>
#endif

#include <stdlib.h>
#include <errno.h>
#include <unistd.h>
#include <string.h>

#include "scamper.h"
#include "scamper_addr.h"
#include "scamper_list.h"
#include "scamper_ping.h"
#include "scamper_getsrc.h"
#include "scamper_icmp_resp.h"
#include "scamper_dl.h"
#include "scamper_probe.h"
#include "scamper_task.h"
#include "scamper_queue.h"
#include "scamper_file.h"
#include "scamper_outfiles.h"
#include "scamper_addresslist.h"
#include "scamper_debug.h"
#include "scamper_do_ping.h"
#include "scamper_options.h"
#include "scamper_icmp4.h"
#include "scamper_icmp6.h"
#include "scamper_fds.h"
#include "utils.h"

#define SCAMPER_DO_PING_PROBECOUNT_MIN    1
#define SCAMPER_DO_PING_PROBECOUNT_DEF    4
#define SCAMPER_DO_PING_PROBECOUNT_MAX    65535

#define SCAMPER_DO_PING_PROBESIZE_V4_MIN  28
#define SCAMPER_DO_PING_PROBESIZE_V4_DEF  (28+56)
#define SCAMPER_DO_PING_PROBESIZE_V4_MAX  65535

#define SCAMPER_DO_PING_PROBESIZE_V6_MIN  48
#define SCAMPER_DO_PING_PROBESIZE_V6_DEF  (48+8)
#define SCAMPER_DO_PING_PROBESIZE_V6_MAX  65535

#define SCAMPER_DO_PING_PROBEWAIT_MIN     1
#define SCAMPER_DO_PING_PROBEWAIT_DEF     1
#define SCAMPER_DO_PING_PROBEWAIT_MAX     20

#define SCAMPER_DO_PING_PROBETTL_MIN      1
#define SCAMPER_DO_PING_PROBETTL_DEF      64
#define SCAMPER_DO_PING_PROBETTL_MAX      255

#define SCAMPER_DO_PING_PROBETOS_MIN      0
#define SCAMPER_DO_PING_PROBETOS_DEF      0
#define SCAMPER_DO_PING_PROBETOS_MAX      255

#define SCAMPER_DO_PING_REPLYCOUNT_MIN    0
#define SCAMPER_DO_PING_REPLYCOUNT_DEF    0
#define SCAMPER_DO_PING_REPLYCOUNT_MAX    65535

#define SCAMPER_DO_PING_PATTERN_MIN       1
#define SCAMPER_DO_PING_PATTERN_DEF       0
#define SCAMPER_DO_PING_PATTERN_MAX       32

/* the callback functions registered with the ping task */
static scamper_task_funcs_t ping_funcs;

/* ICMP ping probes are marked with the process' ID */
static pid_t pid;

/* packet buffer for generating the payload of an ICMP packet */
static uint8_t *pktbuf     = NULL;
static size_t   pktbuf_len = 0;

/* address cache used to avoid reallocating the same address multiple times */
extern scamper_addrcache_t *addrcache;

typedef struct ping_probe
{
  struct timeval tx;
  uint16_t       seq;
} ping_probe_t;

typedef struct ping_state
{
  scamper_fd_t  *fd;
  ping_probe_t **probes;
  uint16_t       replies;
  uint16_t       seq_min, seq_cur, seq_max;
} ping_state_t;

#define PING_OPT_PROBECOUNT    1
#define PING_OPT_PROBEWAIT     2
#define PING_OPT_PROBETTL      3
#define PING_OPT_REPLYCOUNT    4
#define PING_OPT_PATTERN       5
#define PING_OPT_PROBESIZE     6
#define PING_OPT_PROBETOS      7

static const scamper_option_in_t ping_opts_in[] = {
  {'c', NULL, PING_OPT_PROBECOUNT,   SCAMPER_OPTION_TYPE_NUM},
  {'i', NULL, PING_OPT_PROBEWAIT,    SCAMPER_OPTION_TYPE_NUM},
  {'m', NULL, PING_OPT_PROBETTL,     SCAMPER_OPTION_TYPE_NUM},
  {'o', NULL, PING_OPT_REPLYCOUNT,   SCAMPER_OPTION_TYPE_NUM},
  {'p', NULL, PING_OPT_PATTERN,      SCAMPER_OPTION_TYPE_STR},
  {'s', NULL, PING_OPT_PROBESIZE,    SCAMPER_OPTION_TYPE_NUM},
  {'z', NULL, PING_OPT_PROBETOS,     SCAMPER_OPTION_TYPE_NUM},
};

static const int ping_opts_cnt = SCAMPER_OPTION_COUNT(ping_opts_in);

/*
 * ping_abort
 *
 * some internal consistency check failed
 */
static void ping_abort(scamper_task_t *task)
{
  scamper_task_free(task);
  return;
}

static void ping_stop(scamper_task_t *task, uint8_t reason, uint8_t data)
{
  scamper_ping_t *ping = task->data;

  ping->stop_reason = reason;
  ping->stop_data   = data;

  scamper_queue_done(task->queue, scamper_holdtime_get()*1000);

  return;  
}

static void ping_handleerror(scamper_task_t *task, int error)
{
  ping_stop(task, SCAMPER_PING_STOP_ERROR, error);
  return;
}

/*
 * do_ping_probe
 *
 * it is time to send a probe for this task.  figure out the form of the
 * probe to send, and then send it.
 */
static int do_ping_probe(scamper_task_t *task)
{
  scamper_ping_t  *ping  = task->data;
  ping_state_t    *state = task->state;
  ping_probe_t    *pp = NULL;
  scamper_probe_t  probe;
  uint8_t         *buf;
  uint8_t          proto, type;
  uint8_t          icmp_len = (1 + 1 + 2 + 2 + 2);
  size_t           payload_len;
  int              i;

  if(ping->dst->type == SCAMPER_ADDR_TYPE_IPV4)
    {
      proto = IPPROTO_ICMP;
      type = ICMP_ECHO;
      payload_len = ping->probe_size - sizeof(struct ip) - icmp_len;
    }
  else if(ping->dst->type == SCAMPER_ADDR_TYPE_IPV6)
    {
      proto = IPPROTO_ICMPV6;
      type = ICMP6_ECHO_REQUEST;
      payload_len = ping->probe_size - sizeof(struct ip6_hdr) - icmp_len;
    }
  else
    {
      ping_abort(task);
      return -1;
    }

  /* make sure the global pktbuf is big enough for the probe we send */
  if(pktbuf_len < payload_len)
    {
      if((buf = realloc(pktbuf, payload_len)) == NULL)
	{
	  printerror(errno, strerror, __func__, "could not realloc");
	  ping_abort(task);
	  return -1;
	}
      pktbuf     = buf;
      pktbuf_len = payload_len;
    }

  probe.pr_dl        = NULL;
  probe.pr_dl_hdr    = NULL;
  probe.pr_dl_size   = 0;
  probe.pr_ip_src    = ping->src;
  probe.pr_ip_dst    = ping->dst;
  probe.pr_ip_ttl    = ping->probe_ttl;
  probe.pr_ip_proto  = proto;
  probe.pr_ip_id     = 0;
  probe.pr_ip_flow   = 0;
  probe.pr_icmp_type = type;
  probe.pr_icmp_code = 0;
  probe.pr_icmp_id   = pid & 0xffff;
  probe.pr_icmp_seq  = state->seq_cur;
  probe.pr_data      = pktbuf;
  probe.pr_len       = payload_len;
  probe.pr_ipoptc    = 0;
  probe.pr_ipopts    = NULL;
  probe.pr_fd        = scamper_fd_fd_get(state->fd);

  /* if the ping has to hold some pattern, then generate it now */  
  if(ping->pattern_bytes == NULL)
    {
      memset(pktbuf, 0, payload_len);
    }
  else
    {
      i = 0;
      while((size_t)(i + ping->pattern_len) < payload_len)
	{
	  memcpy(pktbuf+i, ping->pattern_bytes, ping->pattern_len);
	  i += ping->pattern_len;
	}
      memcpy(pktbuf+i, ping->pattern_bytes, payload_len - i);
    }

  /*
   * allocate a ping probe state record before we try and send the probe
   * as there is no point sending something into the wild that we can't
   * record
   */
  if((pp = malloc(sizeof(ping_probe_t))) == NULL)
    {
      ping_handleerror(task, errno);
      goto err;
    }

  if(scamper_probe(&probe) == -1)
    {
      ping_handleerror(task, probe.pr_errno);
      goto err;
    }

  /* fill out the details of the probe sent */
  pp->seq = state->seq_cur;
  timeval_cpy(&pp->tx, &probe.pr_tx);

  /* record the probe in the probes table */
  state->probes[state->seq_cur - state->seq_min] = pp;

  /* we've sent this sequence number now, so move to the next one */
  state->seq_cur++;

  /* increment the number of probes sent... */
  ping->ping_sent++;

  /* re-queue the ping task */
  scamper_queue_wait(task->queue, ping->probe_wait * 1000);

  return 0;

 err:
  if(pp != NULL) free(pp);
  return -1;
}

static int do_ping_handle_icmp(scamper_task_t *task, scamper_icmp_resp_t *ir)
{
  scamper_ping_t       *ping  = task->data;
  ping_state_t         *state = task->state;
  scamper_ping_reply_t *reply = NULL;
  ping_probe_t         *probe;
  uint16_t              seq;
  scamper_addr_t        addr;

  /* if this is an echo reply packet, then check the id and sequence */
  if(SCAMPER_ICMP_RESP_IS_ECHO_REPLY(ir))
    {
      /* if the response is not for us, then move on */
      if(ir->ir_icmp_id != (pid & 0xffff) ||
	 ir->ir_icmp_seq < state->seq_min ||
	 ir->ir_icmp_seq > state->seq_max)
	{
	  return 0;
	}

      seq = ir->ir_icmp_seq - state->seq_min;
    }
  /* if this is an ICMP response to an echo req, then check id and sequence */
  else if(SCAMPER_ICMP_RESP_INNER_IS_SET(ir) &&
	  SCAMPER_ICMP_RESP_INNER_IS_ICMP_ECHO_REQ(ir))
    {
      if(ir->ir_inner_icmp_id != (pid & 0xffff) ||
	 ir->ir_inner_icmp_seq < state->seq_min ||
	 ir->ir_inner_icmp_seq > state->seq_max)
	{
	  return 0;
	}

      seq = ir->ir_inner_icmp_seq - state->seq_min;
    }
  else
    {
      return 0;
    }

  /*
   * if the sequence number was in our range, but we have no record of the
   * probe, then just ignore the response
   */
  if((probe = state->probes[seq]) == NULL)
    {
      return 0;
    }

  /* allocate a reply structure for the response */
  if((reply = scamper_ping_reply_alloc()) == NULL)
    {
      goto err;
    }

  /* figure out where the response came from */
  if(scamper_icmp_resp_src(ir, &addr) != 0 ||
     (reply->addr = scamper_addrcache_get(addrcache,
					  addr.type, addr.addr)) == NULL)
    {
      goto err;
    }

  /* put together details of the reply */
  timeval_rtt(&reply->rtt, &probe->tx, &ir->ir_rx);
  reply->reply_size = ir->ir_ip_size;
  reply->icmp_type  = ir->ir_icmp_type;
  reply->icmp_code  = ir->ir_icmp_code;
  reply->probe_id   = seq;

  if(ir->ir_ip_ttl != -1)
    {
      reply->reply_ttl  = ir->ir_ip_ttl;
      reply->flags |= SCAMPER_PING_REPLY_FLAG_REPLY_TTL;
    }

  /*
   * if this is the first reply we have for this hop, then increment
   * the replies counter we keep state with
   */
  if(ping->ping_replies[seq] == NULL)
    {
      state->replies++;
    }

  /* put the reply into the ping table */
  scamper_ping_reply_append(ping, reply);

  /*
   * if only a certain number of replies are required, and we've reached
   * that amount, then stop probing
   */
  if(ping->reply_count != 0 && state->replies >= ping->reply_count)
    {
      ping_stop(task, SCAMPER_PING_STOP_COMPLETED, 0);
    }

  return 0;

 err:
  if(reply != NULL) scamper_ping_reply_free(reply);
  return -1;
}

static int do_ping_handle_dl(scamper_task_t *task, scamper_dl_rec_t *dl_rec)
{
  return -1;
}

/*
 * do_ping_handle_timeout
 *
 * the ping object expired on the pending queue
 * that means it is either time to send the next probe, or write the
 * task out
 */
static int do_ping_handle_timeout(scamper_task_t *task)
{
  ping_state_t *state = task->state;

  if(state->seq_cur == state->seq_max)
    {
      ping_stop(task, SCAMPER_PING_STOP_COMPLETED, 0);
    }

  return 0;
}

static int do_ping_write(scamper_task_t *task)
{
  scamper_outfile_t *outfile = scamper_source_getoutfile(task->source);
  scamper_file_t *sf = scamper_outfile_getfile(outfile);
  scamper_file_write_ping(sf, (scamper_ping_t *)task->data);
  return 0;
}

static void do_ping_free(scamper_task_t *task)
{
  scamper_ping_t *ping;
  ping_state_t *state;
  int i;

  /* free any ping data collected */
  if((ping = task->data) != NULL)
    {
      scamper_ping_free(ping);
    }

  if((state = task->state) != NULL)
    {
      /* close probe fd */
      if(state->fd != NULL)
	{
	  scamper_fd_free(state->fd);
	}

      if(state->probes != NULL)
	{
	  for(i=0; i<state->seq_max - state->seq_min; i++)
	    {
	      if(state->probes[i] != NULL)
		{
		  free(state->probes[i]);
		}
	    }
	  free(state->probes);
	}
      free(state);
    }

  return;
}

static uint8_t hex2byte(char a, char b)
{
  uint8_t out;

  if(a <= '9')      out = (((int)a - (int)'0') << 4);
  else if(a <= 'F') out = (((int)a - (int)'A') << 4);
  else              out = (((int)a - (int)'a') << 4);

  if(b <= '9')      out |= (((int)b - (int)'0') << 4);
  else if(b <= 'F') out |= (((int)b - (int)'A') << 4);
  else              out |= (((int)b - (int)'a') << 4);

  return out;
}

/*
 * scamper_do_ping_alloc
 *
 * given a string representing a ping task, parse the parameters and assemble
 * a ping.  return the ping structure so that it is all ready to go.
 *
 */
scamper_ping_t *scamper_do_ping_alloc(char *str)
{
  uint16_t  probe_count   = SCAMPER_DO_PING_PROBECOUNT_DEF;
  uint8_t   probe_wait    = SCAMPER_DO_PING_PROBEWAIT_DEF;
  uint8_t   probe_ttl     = SCAMPER_DO_PING_PROBETTL_DEF;
  uint8_t   probe_tos     = SCAMPER_DO_PING_PROBETOS_DEF;
  uint16_t  reply_count   = SCAMPER_DO_PING_REPLYCOUNT_DEF;
  uint16_t  probe_size    = 0; /* unset */
  uint16_t  pattern_len   = 0;
  uint8_t   pattern_bytes[SCAMPER_DO_PING_PATTERN_MAX/2];

  scamper_option_out_t *opts_out = NULL, *opt;
  scamper_ping_t *ping = NULL;
  char *addr;
  long tmp;
  int i;

  /* try and parse the string passed in */
  if(scamper_options_parse(str, ping_opts_in, ping_opts_cnt,
			   &opts_out, &addr) != 0)
    {
      goto err;
    }

  /* if there is no IP address after the options string, then stop now */
  if(addr == NULL)
    {
      goto err;
    }

  /* parse the options, do preliminary sanity checks */
  for(opt = opts_out; opt != NULL; opt = opt->next)
    {
      switch(opt->id)
	{
	/* number of probes to send */
	case PING_OPT_PROBECOUNT:
	  if(string_tolong(opt->str, &tmp) == -1  ||
	     tmp < SCAMPER_DO_PING_PROBECOUNT_MIN ||
	     tmp > SCAMPER_DO_PING_PROBECOUNT_MAX)
	    {
	      goto err;
	    }
	  probe_count = tmp;
	  break;

	/* how long to wait between sending probes */
	case PING_OPT_PROBEWAIT:
	  if(string_tolong(opt->str, &tmp) == -1 ||
	     tmp < SCAMPER_DO_PING_PROBEWAIT_MIN ||
	     tmp > SCAMPER_DO_PING_PROBEWAIT_MAX)
	    {
	      goto err;
	    }
	  probe_wait = tmp;
	  break;

	/* the ttl to probe with */
	case PING_OPT_PROBETTL:
	  if(string_tolong(opt->str, &tmp) == -1 ||
	     tmp < SCAMPER_DO_PING_PROBETTL_MIN  ||
	     tmp > SCAMPER_DO_PING_PROBETTL_MAX)
	    {
	      goto err;
	    }
	  probe_ttl = tmp;
	  break;

	/* how many unique replies are required before the ping completes */
	case PING_OPT_REPLYCOUNT:
	  if(string_tolong(opt->str, &tmp) == -1  ||
	     tmp < SCAMPER_DO_PING_REPLYCOUNT_MIN ||
	     tmp > SCAMPER_DO_PING_REPLYCOUNT_MAX)
	    {
	      goto err;
	    }
	  reply_count = tmp;
	  break;

	/* the pattern to fill each probe with */
	case PING_OPT_PATTERN:
	  /*
	   * sanity check that only hex characters are present, and that
	   * the pattern string is not too long.  then, compose the pattern
	   * bytes into the local array.
	   */
	  for(i=0; i<SCAMPER_DO_PING_PATTERN_MAX; i++)
	    {
	      if(opt->str[i] == '\0') break;
	      if(ishex(opt->str[i]) == 0) goto err;
	    }
	  if(i == SCAMPER_DO_PING_PATTERN_MAX) goto err;
	  if((i % 2) == 0)
	    {
	      pattern_len = i/2;
	      for(i=0; i<pattern_len; i++)
		{
		  pattern_bytes[i] = hex2byte(opt->str[i*2],opt->str[(i*2)+1]);
		}
	    }
	  else
	    {
	      pattern_len = (i/2) + 1;
	      pattern_bytes[0] = hex2byte('0', opt->str[0]);
	      for(i=1; i<pattern_len; i++)
		{
		  pattern_bytes[i] = hex2byte(opt->str[(i*2)-1],opt->str[i*2]);
		}
	    }
	  break;

	/* the size of each probe */
	case PING_OPT_PROBESIZE:
	  if(string_tolong(opt->str, &tmp) == -1 || tmp < 0 || tmp > 65535)
	    {
	      goto err;
	    }
	  probe_size = tmp;
	  break;

	/* the tos bits to include in each probe */
	case PING_OPT_PROBETOS:
	  if(string_tolong(opt->str, &tmp) == -1 ||
	     tmp < SCAMPER_DO_PING_PROBETOS_MIN  ||
	     tmp > SCAMPER_DO_PING_PROBETOS_MAX)
	    {
	      goto err;
	    }
	  probe_tos = tmp;
	  break;
	}
    }
  scamper_options_free(opts_out); opts_out = NULL;

  /* allocate the ping object and determine the address to probe */
  if((ping = scamper_ping_alloc()) == NULL)
    {
      goto err;
    }
  if((ping->dst = scamper_addrcache_resolve(addrcache,AF_UNSPEC,addr)) == NULL)
    {
      goto err;
    }

  /* ensure the probe size specified is suitable */
  if(ping->dst->type == SCAMPER_ADDR_TYPE_IPV4)
    {
      if(probe_size == 0) probe_size = SCAMPER_DO_PING_PROBESIZE_V4_MIN;
      else if(probe_size < SCAMPER_DO_PING_PROBESIZE_V4_MIN) goto err;
    }
  else if(ping->dst->type == SCAMPER_ADDR_TYPE_IPV6)
    {
      if(probe_size == 0) probe_size = SCAMPER_DO_PING_PROBESIZE_V6_MIN;
      else if(probe_size < SCAMPER_DO_PING_PROBESIZE_V6_MIN) goto err;
    }
  else goto err;

  /* copy in the pad bytes, if any */
  if(scamper_ping_setpattern(ping, pattern_bytes, pattern_len) != 0)
    {
      goto err;
    }

  ping->probe_count = probe_count;
  ping->probe_size  = probe_size;
  ping->probe_wait  = probe_wait;
  ping->probe_ttl   = probe_ttl;
  ping->probe_tos   = probe_tos;
  ping->reply_count = probe_count;

  return ping;

 err:
  if(ping != NULL) scamper_ping_free(ping);
  if(opts_out != NULL) scamper_options_free(opts_out);
  return NULL;
}

scamper_task_t *scamper_do_ping_alloctask(scamper_ping_t *ping,
					  scamper_list_t *list,
					  scamper_cycle_t *cycle)
{
  scamper_task_t *task;
  ping_state_t   *state;
  size_t          size;

  /* firstly, allocate the task structure */
  if((task = scamper_task_alloc(ping->dst, &ping_funcs)) == NULL)
    {
      goto err;
    }

  /* secondly, associate the ping structure with the task */
  task->data = ping;

  /* now, associate the list and cycle with the trace */
  ping->list  = scamper_list_use(list);
  ping->cycle = scamper_cycle_use(cycle);

  /* determine the source address used for sending probes */
  if((ping->src = scamper_getsrc(ping->dst)) == NULL)
    {
      goto err;
    }

  /* allocate the memory for ping replies */
  if(scamper_ping_replies_alloc(ping, ping->probe_count) == -1)
    {
      goto err;
    }

  /* allocate the necessary state to keep track of probes */
  if((task->state = malloc_zero(sizeof(ping_state_t))) == NULL)
    {
      goto err;
    }
  state = task->state;
  size = ping->probe_count * sizeof(ping_probe_t *);
  if((state->probes = malloc_zero(size)) == NULL)
    {
      goto err;
    }
  state->seq_max = state->seq_min + ping->probe_count;

  /* get the probe file descriptor */
  if(ping->dst->type == SCAMPER_ADDR_TYPE_IPV4)
    {
      state->fd = scamper_fd_icmp4();
    }
  else if(ping->dst->type == SCAMPER_ADDR_TYPE_IPV6)
    {
      state->fd = scamper_fd_icmp6();
    }
  else goto err;
  if(state->fd == NULL)
    {
      goto err;
    }

  /* timestamp the start time of the ping */
  gettimeofday_wrap(&ping->start);

  return task;

 err:
  if(task != NULL) scamper_task_free(task);
  return NULL;
}

void scamper_do_ping_cleanup()
{
  if(pktbuf != NULL)
    {
      free(pktbuf);
      pktbuf = NULL;
    }

  return;
}

int scamper_do_ping_init()
{
  ping_funcs.probe                  = do_ping_probe;
  ping_funcs.handle_icmp            = do_ping_handle_icmp;
  ping_funcs.handle_dl              = do_ping_handle_dl;
  ping_funcs.handle_timeout         = do_ping_handle_timeout;
  ping_funcs.write                  = do_ping_write;
  ping_funcs.task_free              = do_ping_free;

  pid = getpid();

  return 0;
}