/*
 * scamper
 *
 * $Id: scamper.c,v 1.143 2007/05/08 01:45:34 mjl Exp $
 *
 *        Matthew Luckie, WAND Group, Computer Science, University of Waikato
 *        mjl@wand.net.nz
 *
 * Copyright (C) 2003-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/uio.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/resource.h>

#include <netinet/in.h>

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

#include <assert.h>

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

#if defined(DMALLOC)
#include <dmalloc.h>
#endif

#include "scamper.h"
#include "scamper_debug.h"
#include "scamper_addr.h"
#include "scamper_list.h"
#include "scamper_tlv.h"
#include "scamper_trace.h"
#include "scamper_ping.h"
#include "utils.h"

#include "scamper_file.h"
#include "scamper_outfiles.h"
#include "scamper_task.h"
#include "scamper_target.h"
#include "scamper_queue.h"
#include "scamper_addresslist.h"
#include "scamper_getsrc.h"
#include "scamper_addr2mac.h"
#include "scamper_fds.h"
#include "scamper_icmp4.h"
#include "scamper_icmp6.h"
#include "scamper_udp4.h"
#include "scamper_udp6.h"
#include "scamper_tcp4.h"
#include "scamper_rtsock.h"
#include "scamper_dl.h"
#include "scamper_probe.h"
#include "scamper_privsep.h"
#include "scamper_control.h"
#include "scamper_do_trace.h"
#include "scamper_do_ping.h"

static uint32_t options = 0;
#define OPT_PPS         0x00000001 /* p: */
#define OPT_OUTFILE     0x00000002 /* o: */
#define OPT_OUTTYPE     0x00000004 /* O: */
#define OPT_VERSION     0x00000020 /* v: */
#define OPT_BGP         0x00000040 /* b: */
#define OPT_RTT         0x00000080 /* r: */
#define OPT_HOLDTIME    0x00000100 /* H: */
#define OPT_DAEMON      0x00000200 /* D: */
#define OPT_IPLIST      0x00000400 /* i: */
#define OPT_SPORT       0x00000800 /* s: */
#define OPT_DL          0x00001000 /* P: */
#define OPT_MONITORNAME 0x00002000 /* M: */
#define OPT_COMMAND     0x00004000 /* c: */
#define OPT_CYCLEID     0x00008000 /* C: */
#define OPT_LISTNAME    0x00010000 /* l: */
#define OPT_LISTID      0x00020000 /* L: */
#define OPT_HELP        0x00040000 /* ?: */

/*
 * parameters configurable by the command line:
 *
 * command:     default command to use with scamper
 * pps:         how many probe packets to send per second
 * sport:       source port to use by default.  (0x8000 & (getpid() & 0x7fff))
 * holdtime:    how long to hold tasks for late replies after completion
 * outfile:     where to send results by default
 * outtype:     format to use when writing results to outfile
 * daemon_port: port to use when operating as a daemon
 * monitorname: canonical name of monitor assigned by human
 * listname:    name of list assigned by human
 * listid:      id of list assigned by human
 * cycleid:     id of cycle assigned by human
 * arglist:     whatever is left over after getopt processing
 * arglist_len: number of arguments left over after getopt processing
 */
static char  *command      = NULL;
static int    pps          = SCAMPER_PPS_DEF;
static int    sport        = 0;
static int    holdtime     = SCAMPER_HOLDTIME_DEF;
static char  *outfile      = "-";
static char  *outtype      = "traceroute";
static int    daemon_port  = 0;
static char  *monitorname  = NULL;
static char  *listname     = NULL;
static int    listid       = -1;
static int    cycleid      = -1;
static char **arglist      = NULL;
static int    arglist_len  = 0;

/*
 * parameters calculated by scamper at run time:
 *
 * wait_between:   calculated wait between probes to reach pps, in microseconds
 * probe_window:   maximum extension of probing window before truncation
 * exit_when_done: exit scamper when current window of tasks is completed
 */
static int    wait_between   = 1000000 / SCAMPER_PPS_DEF;
static int    probe_window   = 250000;
static int    exit_when_done = 1;

/* central cache of addresses that scamper is dealing with */
scamper_addrcache_t *addrcache = NULL;

static void usage_str(char c, char *str)
{
  fprintf(stderr, "            -%c %s\n", c, str);
  return;
}

static void usage(uint32_t opt_mask)
{
  char buf[256];

  fprintf(stderr,
	  "usage: scamper [-?Pv] [-c command] [-p pps] [-M monitorname]\n"
	  "               [-s sport] [-H holdtime] [-o outfile] [-O outtype]\n"
	  "               [-i addr 1..N | listfile | -D port]\n"
	  "               [-l listname] [-L listid] [-C cycleid]\n");

  if(opt_mask == 0) return;

  fprintf(stderr, "\n");

  if(opt_mask & OPT_HELP)
    usage_str('?', "give an overview of the usage of scamper");

  if(opt_mask & OPT_COMMAND)
    {
      snprintf(buf, sizeof(buf),
	       "command string (default: %s)", SCAMPER_COMMAND_DEF);
      usage_str('c', buf);
    }

  if(opt_mask & OPT_CYCLEID)
    usage_str('C', "cycle id");

  if(opt_mask & OPT_DAEMON)
    usage_str('D', "start as a daemon listening for commands on a port");

  if(opt_mask & OPT_HOLDTIME)
    {
      snprintf(buf, sizeof(buf),
	       "time to hold trace for delayed responses (%d < holdtime %d)",
	       SCAMPER_HOLDTIME_MIN, SCAMPER_HOLDTIME_MAX);

      usage_str('H', buf);
    }

  if(opt_mask & OPT_IPLIST)
    usage_str('i', "IP addresses to trace provided on the command line");

  if(opt_mask & OPT_LISTID)
    usage_str('l', "name to assign to default list");

  if(opt_mask & OPT_LISTNAME)
    usage_str('L', "list id for default list");

  if(opt_mask & OPT_MONITORNAME)
    usage_str('M', "specify the canonical name of the monitor");

  if(opt_mask & OPT_OUTFILE)
    usage_str('o', "specify the file to write output to");

  if(opt_mask & OPT_OUTTYPE)
    usage_str('O', "specify the type of output [warts | traceroute]");

  if(opt_mask & OPT_PPS)
    {
      snprintf(buf, sizeof(buf),
	       "number of packets per second to send (%d <= pps <= %d)",
	       SCAMPER_PPS_MIN, SCAMPER_PPS_MAX);

      usage_str('p', buf);
    }

  if(opt_mask & OPT_DL)
    usage_str('P', "use a datalink to get tx timestamps for outgoing probes");

  if(opt_mask & OPT_SPORT)
    usage_str('s', "source port to use");

  if(opt_mask & OPT_VERSION)
    usage_str('v', "output the version of scamper this binary is");

  return;
}

static int set_opt(uint32_t opt,char *str,int (*setfunc)(int))
{
  long l;

  if(string_isnumber(str) == 0 || string_tolong(str, &l) == -1)
    {
      usage(opt);
      return -1;
    }

  return setfunc(l);
}

static int cycleid_set(const int cid)
{
  if(cid > 0 && cid <= 0x7fffffff)
    {
      cycleid = cid;
      return 0;
    }
  return -1;
}

int scamper_option_listid(void)
{
  return listid;
}

int scamper_option_cycleid(void)
{
  return cycleid;
}

const char *scamper_option_listname(void)
{
  return listname;
}

static int listid_set(const int lid)
{
  if(lid > 0 && lid <= 0x7fffffff)
    {
      listid = lid;
      return 0;
    }
  return -1;
}

static int check_options(int argc, char *argv[])
{
  char  ch;
  long  lo;
  char *opts = "c:C:D:H:il:L:M:o:O:p:Pv?";
  char *opt_cycleid = NULL, *opt_listid = NULL, *opt_listname = NULL;
  char *opt_daemon = NULL, *opt_holdtime = NULL, *opt_monitorname = NULL;
  char *opt_pps = NULL, *opt_sport = NULL, *opt_command = NULL;

  while((ch = getopt(argc, argv, opts)) != -1)
    {
      switch(ch)
	{
	case 'c':
	  options |= OPT_COMMAND;
	  opt_command = optarg;
	  break;

	case 'C':
	  options |= OPT_CYCLEID;
	  opt_cycleid = optarg;
	  break;

	case 'D':
	  options |= OPT_DAEMON;
	  opt_daemon = optarg;
	  break;

	case 'H':
	  options |= OPT_HOLDTIME;
	  opt_holdtime = optarg;
	  break;

	case 'i':
	  options |= OPT_IPLIST;
	  break;

	case 'l':
	  options |= OPT_LISTNAME;
	  opt_listname = optarg;
	  break;

	case 'L':
	  options |= OPT_LISTID;
	  opt_listid = optarg;
	  break;

	case 'M':
	  options |= OPT_MONITORNAME;
	  opt_monitorname = optarg;
	  break;

        case 'o':
          options |= OPT_OUTFILE; 
          outfile = optarg;
          break;

	case 'O':
	  options |= OPT_OUTTYPE;
	  outtype = optarg;
	  break;

	case 'p':
	  options |= OPT_PPS;
	  opt_pps = optarg;
	  break;

	case 'P':
	  options |= OPT_DL;
	  break;

	case 's':
	  options |= OPT_SPORT;
	  opt_sport = optarg;
	  break;

	case 'v':
	  options |= OPT_VERSION;
	  break;

	case '?':
	  options |= OPT_HELP;
	  usage(0xffffffff);
	  return -1;

	default:
	  return -1;
	}
    }

  if(options & OPT_VERSION)
    {
      fprintf(stderr, "scamper version %s\n", SCAMPER_VERSION);
      return -1;
    }

  if(options & OPT_PPS &&
     set_opt(OPT_PPS, opt_pps, scamper_pps_set) == -1)
    {
      usage(OPT_PPS);
      return -1;
    }

  if(options & OPT_HOLDTIME &&
     set_opt(OPT_HOLDTIME, opt_holdtime, scamper_holdtime_set) == -1)
    {
      usage(OPT_HOLDTIME);
      return -1;
    }

  if(options & OPT_LISTNAME &&
     (listname = strdup(opt_listname)) == NULL)
    {
      return -1;
    }

  if(options & OPT_LISTID &&
     set_opt(OPT_LISTID, opt_listid, listid_set) == -1)
    {
      usage(OPT_LISTID);
      return -1;
    }

  if(options & OPT_CYCLEID &&
     set_opt(OPT_CYCLEID, opt_cycleid, cycleid_set) == -1)
    {
      usage(OPT_CYCLEID);
      return -1;
    }

  if(options & OPT_MONITORNAME &&
     (monitorname = strdup(opt_monitorname)) == NULL)
    {
      return -1;
    }

  if(options & OPT_SPORT &&
     set_opt(OPT_SPORT, opt_sport, scamper_sport_set) == -1)
    {
      usage(OPT_SPORT);
      return -1;
    }

  if(options & OPT_IPLIST && options & OPT_DAEMON)
    {
      usage(OPT_IPLIST | OPT_DAEMON);
      return -1;
    }

  if(strcasecmp(outtype, "traceroute") != 0 &&
     strcasecmp(outtype, "warts") != 0)
    {
      usage(OPT_OUTTYPE);
      return -1;
    }

  /* set default command */
  if(scamper_command_set((options & OPT_COMMAND) ?
			 opt_command : SCAMPER_COMMAND_DEF) == -1)
    {
      usage(OPT_COMMAND);
      return -1;
    }

  /* these are the left-over arguments */
  arglist     = argv + optind;
  arglist_len = argc - optind;

  if(options & OPT_DAEMON)
    {
      /* if started as daemon, there should be no leftover arguments */
      if(arglist_len != 0)
	{
	  usage(OPT_DAEMON);
	  return -1;
	}

      /* port on which to run the daemon */
      if(string_isnumber(opt_daemon) == 0 ||
	 string_tolong(opt_daemon, &lo) == -1 ||
	 lo < 1 || lo > 65535)
	{
	  usage(OPT_DAEMON);
	  return -1;
	}

      daemon_port = lo;
    }
  else if(options & OPT_IPLIST)
    {
      /*
       * if a list of IP addresses is to be supplied, there has to be at
       * least one left over argument.
       */
      if(arglist_len < 1)
	{
	  usage(OPT_IPLIST);
	  return -1;
	}
    }
  else
    {
      /*
       * if a listfile is specified, then there may only be one left over
       * argument, which specifies the listfile.
       */
      if(arglist_len != 1)
	{
	  usage(0);
	  return -1;
	}
    }

  return 0;
}

const char *scamper_command_get(void)
{
  return command;
}

int scamper_command_set(const char *command_in)
{
  char *d;

  if(command_in == NULL || (d = strdup(command_in)) == NULL)
    {
      return -1;
    }

  if(command != NULL) free(command);

  command = d;
  return 0;
}

void scamper_exitwhendone(int on)
{
  if(on == 1 || on == 0)
    {
      exit_when_done = on;
    }
  return;
}

int scamper_sport_get()
{
  return sport;
}

int scamper_sport_set(const int sp)
{
  if(sp >= SCAMPER_SPORT_MIN && sp <= SCAMPER_SPORT_MAX)
    {
      sport = sp;
      return 0;
    }

  return -1;
}

int scamper_holdtime_get()
{
  return holdtime;
}

int scamper_holdtime_set(const int ht)
{
  if(ht >= SCAMPER_HOLDTIME_MIN && ht <= SCAMPER_HOLDTIME_MAX)
    {
      holdtime = ht;
      return 0;
    }

  return -1;
}

int scamper_pps_get()
{
  return pps;
}

int scamper_pps_set(const int p)
{
  if(p >= SCAMPER_PPS_MIN && p <= SCAMPER_PPS_MAX)
    {
      /*
       * reset the pps scamper is operating at.  re-calculate the inter-probe
       * delay, and the maximum size of the probe window.
       */

      pps = p;
      wait_between = 1000000 / pps;
      probe_window = (wait_between < 250000 ? 250000 : wait_between + 250000);

      return 0;
    }

  return -1;
}

const char *scamper_monitorname_get()
{
  return monitorname;
}

int scamper_monitorname_set(const char *mn)
{
  char *tmp;

  /*
   * before removing the old monitor name, get a copy of the monitor name
   * since that's what we'll be using to store afterward
   */
  if(mn != NULL)
    {
      if((tmp = strdup(mn)) == NULL)
	{
	  return -1;
	}
    }
  else
    {
      tmp = NULL;
    }

  if(monitorname != NULL)
    {
      free(monitorname);
    }

  monitorname = tmp;
  return 0;
}

int scamper_option_dl()
{
  if(options & OPT_DL) return 1;
  return 0;
}

static void scamper_hup(int sig)
{
  return;
}

static void scamper_chld(int sig)
{
  pid_t pid;
  int   status;

  for(;;)
    {
      if((pid = waitpid(-1, &status, WNOHANG)) == -1)
	{
	  break;
	}
    }

  return;
}

/*
 * scamper:
 * this bit of code contains most of the logic for driving the parallel
 * traceroute process.
 */
static int scamper(void)
{
  struct timeval        tv;
  struct timeval        lastprobe;
  struct timeval        nextprobe;
  struct timeval       *timeout;
  scamper_target_t     *target;
  scamper_task_t       *task;
  int64_t               diff;

  /*
   * this has to be done before priviledge separation, as if scamper is
   * running on a BPF system it has to open a BPF fd to establish
   * version compatibility
   */
  if(scamper_dl_init() == -1)
    {
      return -1;
    }

#ifndef WITHOUT_PRIVSEP
  /* revoke the root priviledges we started with */
  if(scamper_privsep_init() == -1)
    {
      return -1;
    }
#endif

  /* allocate the cache of addresses for scamper to keep track of */
  if((addrcache = scamper_addrcache_alloc()) == NULL)
    {
      return -1;
    }

  /* setup the file descriptor monitoring code */
  if(scamper_fds_init() == -1)
    {
      return -1;
    }

  /*
   * if we have been told to open a control socket and daemonise, then do
   * that now.
   */
  if(options & OPT_DAEMON)
    {
      if(scamper_control_init(daemon_port) == -1)
	{
	  return -1;
	}

      /*
       * scamper should wait for more tasks when it has finished with the
       * active window
       */
      exit_when_done = 0;
    }

  /* initialise the subsystem responsible for obtaining source addresses */
  if(scamper_getsrc_init() == -1)
    {
      return -1;
    }

  /* initialise the subsystem responsible for recording mac addresses */
  if(scamper_addr2mac_init() == -1)
    {
      return -1;
    }

  if(scamper_rtsock_init() == -1)
    {
      return -1;
    }

  /* initialise the structures necessary to keep track of addresses to probe */
  if(scamper_addresslist_init() == -1)
    {
      return -1;
    }

  /*
   * if we have an address list of some description on the command line,
   * read the addresses now
   */
  if(options & OPT_IPLIST)
    {
      if(scamper_source_do_array(NULL, command, arglist, arglist_len) == -1)
	{
	  return -1;
	}
    }
  else if((options & OPT_DAEMON) == 0)
    {
      if(scamper_source_do_file(NULL, arglist[0]) == -1)
	{
	  printerror(errno, strerror, __func__,
		     "could not parse %s", arglist[0]);
	  return -1;
	}
    }

  /*
   * initialise the data structures necessary to keep track of target
   * addresses currently being probed
   */
  if(scamper_targets_init() == -1)
    {
      return -1;
    }

  /* initialise the queues that hold the current tasks */
  if(scamper_queue_init() == -1)
    {
      return -1;
    }

  /*
   * initialise the data structures necessary to keep track of output files
   * currently being written to
   */
  if(scamper_outfiles_init(outfile, outtype) == -1)
    {
      return -1;
    }

  /* initialise scamper so it is ready to traceroute and ping */
  if(scamper_do_trace_init() == -1 || scamper_do_ping_init() == -1)
    {
      return -1;
    }

  gettimeofday_wrap(&lastprobe);

  for(;;)
    {
      if(scamper_queue_readycount() > 0 || scamper_addresslist_isready() == 1)
	{
	  /*
	   * if there is something ready to be probed right now, then set the
	   * timeout to go off when it is time to send the next probe
	   */
	  timeval_cpy(&nextprobe, &lastprobe);
	  timeval_add_usec(&nextprobe, wait_between);
	  timeout = &tv;
	}
      else if(scamper_queue_count() > 0)
	{
	  /*
	   * if there isn't anything ready to go right now, but we are
	   * waiting on a response from an earlier probe, then set the timer
	   * to go off when that probe expires.
	   */
	  scamper_queue_waittime(&nextprobe);
	  timeout = &tv;
	}
      else
	{
	  /*
	   * there is nothing to do, so block in select until a file
	   * descriptor supplies an address to probe.
	   */
	  timeout = NULL;
	}

      if(timeout != NULL)
	{
	  gettimeofday_wrap(&tv);
	  diff = timeval_diff_usec(&nextprobe, &tv);

	  memset(&tv, 0, sizeof(tv));
	  if(diff > (int64_t)0) timeval_add_usec(&tv, diff);
	  timeout = &tv;
	}
      else
	{
	  if(exit_when_done != 0 && scamper_addresslist_isempty() == 1)
	    {
	      break;
	    }
	  timeout = NULL;
	}

      /* listen until it is time to send the next probe */
      if(scamper_fds_poll(timeout) == -1)
	{
	  return -1;
	}

      /* take any 'done' traces and output them now */
      while((task = scamper_queue_getdone()) != NULL)
	{
	  /* write the trace out */
	  task->funcs->write(task);

	  target = scamper_target_find(task->dst);
	  scamper_target_free(target);

	  /* cleanup the task */
	  scamper_task_free(task);
	}

      /*
       * if there is something waiting to be probed, then find out if it is
       * time to probe yet
       */
      if(scamper_queue_readycount() > 0 || scamper_addresslist_isready() == 1)
	{
	  gettimeofday_wrap(&tv);

	  diff = timeval_diff_usec(&tv, &lastprobe);
	  if(diff < -((int64_t)probe_window) || diff > (int64_t)probe_window)
	    {
	      timeval_cpy(&lastprobe, &tv);
	      timeval_add_usec(&lastprobe, wait_between * -1);
	    }

	  /*
	   * when probing at > HZ, scamper might find that select blocks it
	   * from achieving the specified packets per second rate if it sends
	   * one probe per select.  Based on the time spent in the last call
	   * to select, send the necessary number of packets to fill that
	   * window where we sent no packets.
	   */
	  for(;;)
	    {
	      timeval_cpy(&nextprobe, &lastprobe);
	      timeval_add_usec(&nextprobe, wait_between);

	      /* if the next probe is not due to be sent, don't send one */
	      if(timeval_cmp(&nextprobe, &tv) > 0)
		{
		  break;
		}

	      /*
	       * look for an address that we can send a probe to.  if
	       * scamper doesn't have a task on the probe queue waiting
	       * to be probed, then get a fresh task. if there's absolutely
	       * nothing that scamper can probe, then break.
	       */
	      if((task = scamper_queue_select()) == NULL)
		{
		  if((task = scamper_addresslist_get()) == NULL)
		    {
		      break;
		    }

		  if((target = scamper_target_alloc(task->dst)) == NULL)
		    {
		      scamper_task_free(task);
		    }
		  target->task = task;
		}

	      task->funcs->probe(task);
	      timeval_cpy(&lastprobe, &nextprobe);
	    }
	}
    }
  
  return 0; 
}

/*
 * cleanup:
 *
 * be nice to the system and clean up all our mallocs
 */
static void cleanup(void)
{
#ifndef WITHOUT_PRIVSEP
  scamper_privsep_cleanup();
#endif
  scamper_getsrc_cleanup();
  scamper_rtsock_cleanup();

  scamper_icmp4_cleanup();
  scamper_icmp6_cleanup();
  scamper_udp4_cleanup();

  scamper_addr2mac_cleanup();

  scamper_do_trace_cleanup();
  scamper_do_ping_cleanup();

  scamper_addresslist_cleanup();

  scamper_dl_cleanup();

  if(options & OPT_DAEMON)
    {
      scamper_control_cleanup();
    }

  scamper_outfiles_cleanup();
  scamper_fds_cleanup();

  /* free the address cache, if one was used */
  if(addrcache != NULL)
    {
      scamper_addrcache_free(addrcache);
      addrcache = NULL;
    }

  if(monitorname != NULL)
    {
      free(monitorname);
      monitorname = NULL;
    }

  if(command != NULL)
    {
      free(command);
      command = NULL;
    }
  scamper_queue_cleanup();
  scamper_targets_cleanup();
  scamper_probe_cleanup();

  return;
}

int main(int argc, char *argv[])
{
  pid_t            pid;
  struct sigaction si_sa;

  /*
   * if we are using dmalloc, then we want to get it to register its
   * logdump function to occur after we have used cleanup to free up
   * scamper's core data structures.  this is a dirty hack.
   */
#if defined(DMALLOC)
  free(malloc(1));
#endif

  atexit(cleanup);

  /* set the port to bind to first! */
  pid = getpid();
  sport = (pid & 0x7fff) + 0x8000;

  if(check_options(argc, argv) == -1)
    {
      return -1;
    }

  sigemptyset(&si_sa.sa_mask);
  si_sa.sa_flags   = 0;
  si_sa.sa_handler = scamper_hup;
  if(sigaction(SIGHUP, &si_sa, 0) == -1)
    {
      printerror(errno, strerror, __func__,
		 "could not set sigaction for SIGHUP");
      return -1;
    }

  sigemptyset(&si_sa.sa_mask);
  si_sa.sa_flags   = 0;
  si_sa.sa_handler = scamper_chld;
  if(sigaction(SIGCHLD, &si_sa, 0) == -1)
    {
      printerror(errno, strerror, __func__,
		 "could not set sigaction for SIGCHLD");
      return -1;
    }

  return scamper();
}