/*
	**
	** main.c
	**
	** Main module for traffic-vis program
	**
	** Copyright 1998 Damien Miller <dmiller@ilogic.com.au>
	**
	** This software is licensed under the terms of the GNU General 
	** Public License (GPL). Please see the file COPYING for details.
	** 
	** $Id: main.c,v 1.16 1998/12/02 13:34:52 dmiller Exp $
	**
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <getopt.h>
#include <unistd.h>
#include <signal.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <setjmp.h>

#include <pcap.h>

#include "packet-parse.h"
#include "packet-summary.h"
#include "summary-sort.h"

#include "textfe.h"
#include "htmlfe.h"
#include "psfe.h"

#define NET_READ_TIMEOUT	100

#define TEXTFE					1
#define HTMLFE					2
#define PSFE					3

/* Global variables */
static char			rcsid[] = "$Id: main.c,v 1.16 1998/12/02 13:34:52 dmiller Exp $";
static pcap_t		*pcap_handle;
static summary_t	*psum;
static int			exit_flag;
static jmp_buf		jbuf;

/* Prototypes */
int				main(int argc, char **argv);
void				cleanup(int signum);
void				configure_capture_file(pcap_t **p, const char *capture_file, int *link_type);
void				configure_live_interface(pcap_t **p, char *dev, int promisc_flag, bpf_u_int32 *localnet, bpf_u_int32 *netmask, int *link_type);
void				exit_main_loop(int signum);
void				pcap_packet_handler(u_char *user_data, const struct pcap_pkthdr* phdr, const u_char *pkt);
void				set_packet_filter(pcap_t *p, char *filter, bpf_u_int32 netmask);
void				usage(void);
void				version(void);

/* Commandline options */
static struct option long_options[] =
{
	{"no-promisc", 0, NULL, 'p'},
	{"no-dns", 0, NULL, 'd'},
	{"html", 0, NULL, 'W'},
	{"text", 0, NULL, 'T'},
	{"postscript", 0, NULL, 'C'},
	{"interface", 1, NULL, 'i'},
	{"filter", 1, NULL, 'f'},
	{"num-hosts", 1, NULL, 'n'},
	{"num-peers", 1, NULL, 'N'},
	{"read-file", 1, NULL, 'r'},
	{"host-sort", 1, NULL, 'H'},
	{"peer-sort", 1, NULL, 'P'},
	{"time", 1, NULL, 't'},
	{"reverse-hosts", 0, NULL, 'R'},
	{"reverse-peers", 0, NULL, 'E'},
	{"version", 0, NULL, 'V'},
	{"help", 0, NULL, 'h'},
	{NULL, 0, NULL, 0}
};

/* Called by pcap_dispatch for each packet processed */
void pcap_packet_handler(u_char *user_data, const struct pcap_pkthdr* phdr, const u_char *pkt)
{
	static pinfo_t pinfo;

	if (parse_packet(*(unsigned int *)user_data, pkt, phdr->caplen, &pinfo))
	{
		store_packet_info(psum, &pinfo);
	} else
	{
#if DEBUG
		printf("Packet parse error: %s\n", parse_error());
#endif
	}
}

/* Signal handler, shuts down pcap and exits cleanly */
void cleanup(int signum)
{
	if (pcap_handle != 0)
		pcap_close(pcap_handle);
		
	fprintf(stderr, "Caught signal %i - exiting.\n", signum);
	exit(1);
}

/* Signal handler for SIGINT, exits main loop */
void exit_main_loop(int signum)
{
	longjmp(jbuf, 1);
}

/* Sets up interface and returns information about network address, */
/* netmask and datalink type. Can be called with a NULL interface, */
/* in which case pcap will pick one. */
void configure_live_interface(pcap_t **p, char *dev, 
                              int promisc_flag, bpf_u_int32 *localnet, 
										bpf_u_int32 *netmask, int *link_type)
{
	char error_buf[1024];
	
	/* If no device selected, use pcap to select one */
	if (dev == NULL)
	{
		dev = pcap_lookupdev(error_buf);
		if (dev == NULL)
		{
			fprintf(stderr, "pcap_lookupdev: %s\n\n", error_buf);
			exit(-1);
		}
		fprintf(stderr, "No interface specified, using \'%s\'.\n", dev);
	}
		
	/* Get netmask of interface */
	if (pcap_lookupnet(dev, localnet, netmask, error_buf) < 0)
	{
		fprintf(stderr, "pcap_lookupnet: %s\n\n", error_buf);
		exit(-1);
	}
	
	/* Open interface */
	*p = pcap_open_live(dev, 256, promisc_flag, NET_READ_TIMEOUT, error_buf);
	if (pcap_handle == NULL)
	{
		fprintf(stderr, "pcap_open_live: %s\n\n", error_buf);
		exit(-1);
	}

	*link_type = pcap_datalink(*p);
}

/* Sets up reading from a previously recorded pcap capture file */
void configure_capture_file(pcap_t **p, const char *capture_file, 
                            int *link_type)
{
	char error_buf[1024];
	
	/* Open capture file */
	*p = pcap_open_offline(capture_file, error_buf);
	if (pcap_handle == NULL)
	{
		fprintf(stderr, "pcap_open_offline: %s\n\n", error_buf);
		exit(-1);
	}

	*link_type = pcap_datalink(*p);
}

/* Sets a filter on packet capture */
void set_packet_filter(pcap_t *p, char *filter, bpf_u_int32 netmask)
{
	struct bpf_program bpf_prog;
	
	if (filter == NULL)
		return;
	
	if (pcap_compile(p, &bpf_prog, filter, 1, netmask) == -1)
	{
		fprintf(stderr, "pcap_compile: %s\n\n", pcap_geterr(p));
		exit(-1);
	}
	
	if (pcap_setfilter(p, &bpf_prog) == -1)
	{
		fprintf(stderr, "pcap_setfilter: %s\n\n", pcap_geterr(p));
		exit(-1);
	}
}

void usage(void)
{
	fprintf(stderr, "Usage: traffic-vis [OPTION]\n");

	fprintf(stderr, "\n");

	fprintf(stderr, "Collect and print statistics on network traffic, including total traffic \n");
	fprintf(stderr, "(bytes and packets), traffic per host and which hosts communicated.\n");
	fprintf(stderr, "traffic-vis will run until stopped with a SIGINT or a SIGTERM.\n");

	fprintf(stderr, "\n");

	fprintf(stderr, "  -p, --no-promisc      Don\'t put the network interface into promiscious mode.\n");
	fprintf(stderr, "  -d, --no-dns          Do not attempt to resolve hostnames.\n");
	fprintf(stderr, "  -i, --interface xxx   Specify network interface to collect packets from.\n");
	fprintf(stderr, "  -f, --filter    xxx   Specify BPF filter to be applied before packet \n");
	fprintf(stderr, "                        processing. See tcpdump(8) for details.\n");
	fprintf(stderr, "  -n, --num-hosts xxx   Maximum number of hosts to display. (0 = unlimited)\n");
	fprintf(stderr, "  -N, --num-peers xxx   Maximum number of peers to display. (0 = unlimited)\n");
	fprintf(stderr, "  -r, --read-file xxx   Read a pcap capture file, such as those created by\n");
	fprintf(stderr, "                        tcpdump(8), instead of listening to live traffic.\n");
	fprintf(stderr, "  -H, --host-sort xxx   How to sort host list (see below).\n");
	fprintf(stderr, "  -P, --peer-sort xxx   How to sort peer lists (see below).\n");
	fprintf(stderr, "  -t, --time      xxx   Specify number of minutes to collect packets for.\n");
	fprintf(stderr, "  -R, --reverse-hosts   Reverse host list sort order.\n");
	fprintf(stderr, "  -E, --reverse-peers   Reverse peer lists sort order.\n");
	fprintf(stderr, "  -W, --html            Generate HTML format report.\n");
	fprintf(stderr, "  -T, --text            Generate ASCII format report (default).\n");
	fprintf(stderr, "  -C, --postscript      Generate Postscript format report (default).\n");
	fprintf(stderr, "  -V, --version         Print program version.\n");
	fprintf(stderr, "  -h, --help            Display this help text.\n");

	fprintf(stderr, "\n");

	fprintf(stderr, "Host sort keys:\n");
	fprintf(stderr, "  N, none         Do not sort\n");
	fprintf(stderr, "  b, bytes        Total bytes (sent + received)\n");
	fprintf(stderr, "  s, bsent        Bytes sent\n");
	fprintf(stderr, "  r, breceived    Bytes received\n");
	fprintf(stderr, "  p, packets      Total packets (sent + received)\n");
	fprintf(stderr, "  S, psent        Packets sent\n");
	fprintf(stderr, "  R, preceived    Packets received\n");
	fprintf(stderr, "  C, connects     Total TCP connections (sent + received)\n");
	fprintf(stderr, "  c, csent        TCP connections sent\n");
	fprintf(stderr, "  l, creceived    TCP connections received\n");
	fprintf(stderr, "  P, peers        Number of peers\n");
	fprintf(stderr, "  h, hostname     Hostname\n");
	fprintf(stderr, "  I, ipaddr       IP address\n");
	fprintf(stderr, "  F, firstseen    When host was first noticed\n");
	fprintf(stderr, "  L, lastseen     When host was last noticed\n");
	fprintf(stderr, "  X, random       Random order\n");

	fprintf(stderr, "\n");

	fprintf(stderr, "Peer sort keys:\n");
	fprintf(stderr, "  N, none         Do not sort\n");
	fprintf(stderr, "  b, bytes        Total bytes (sent + received)\n");
	fprintf(stderr, "  p, packets      Total packets (sent + received)\n");
	fprintf(stderr, "  C, connects     Total TCP connections (sent + received)\n");
	fprintf(stderr, "  F, firstseen    When host was first noticed\n");
	fprintf(stderr, "  L, lastseen     When host was last noticed\n");
	fprintf(stderr, "  X, random       Random order\n");

	fprintf(stderr, "\n");

	fprintf(stderr, "Please report bugs to dmiller@ilogic.com.au\n");
}

void version(void)
{
	fprintf(stderr, "traffic-vis 0.01\n");
	
	/* Avoid stupid compiler warning */
	(void)rcsid;
}

int main(int argc, char **argv)
{
	bpf_u_int32 	localnet;
	bpf_u_int32		netmask;
	unsigned int	link_type;
	int 				c;
	int				dispatch_result;
	extern char 	*optarg;
	
	/* Commandline options */
	static char		*dev = NULL;
	static char 	*filter = NULL;
	static char 	*capture_file = NULL;
	static int 		promisc_flag = 1;
	static int		max_hosts = 0;
	static int		max_peers = 0;
	static int		host_sort_mode = SORT_HOST_DEFAULT;
	static int		peer_sort_mode = SORT_PEER_DEFAULT;
	static int		host_reverse = 0;
	static int		peer_reverse = 0;
	static int		display_mode = TEXTFE;
	static int		resolve_dns = 1;

	/* Fetch commandline options */
	while (1)
	{
		c = getopt_long (argc, argv, "pdi:f:n:N:r:H:P:t:REWCTVh?", long_options, NULL);

		if (c == -1)
			break;
		
		switch (c)
		{
			case 'p':
				promisc_flag = 0;
				break;
			
			case 'd':
				resolve_dns = 0;
				break;
			
			case 'i':
				dev = strdup(optarg);
				break;

			case 'f':
				filter = strdup(optarg);
				break;
				
			case 'n':
				if (sscanf(optarg, "%u", &max_hosts) == 0)
				{
					fprintf(stderr, "Invalid commandline options.\n\n");
					usage();
					exit(-1);
				}
				break;
			
			case 'N':
				if (sscanf(optarg, "%u", &max_peers) == 0)
				{
					fprintf(stderr, "Invalid commandline options.\n\n");
					usage();
					exit(-1);
				}
				break;
			
			case 'r':
				capture_file = strdup(optarg);
				break;
			
			case 'H':
				host_sort_mode = parse_host_sort_mode(optarg);
				if (host_sort_mode == -1)
				{
					fprintf(stderr, "Invalid host sort mode.\n\n");
					usage();
					exit(-1);
				}
				break;
			
			case 'P':
				peer_sort_mode = parse_peer_sort_mode(optarg);
				if (peer_sort_mode == -1)
				{
					fprintf(stderr, "Invalid peer sort mode.\n\n");
					usage();
					exit(-1);
				}
				break;

			case 't':
				/* Schedule exit signal */
				alarm(60 * atoi(optarg));
				break;
			
			case 'T':
				display_mode = TEXTFE;
				break;
			
			case 'W':
				display_mode = HTMLFE;
				break;
			
			case 'C':
				display_mode = PSFE;
				break;
			
			case 'R':
				host_reverse = 1;
				break;
			
			case 'E':
				peer_reverse = 1;
				break;
			
			case 'V':
				version();
				exit(0);
				
			case 'h':
			case '?':
				usage();
				exit(0);
				
			default:
				fprintf(stderr, "Invalid commandline options.\n\n");
				usage();
				exit(-1);
		}
	}
					
	/* Seed the RNG in case we need it later */
	srandom(time(NULL) ^ getpid() ^ ~getppid());

	/* Initialise packet summary data structure */
	packet_summary_init(&psum);

	/* Exits main loop on user signal */
	signal(SIGINT, exit_main_loop);
	signal(SIGTERM, exit_main_loop);
	signal(SIGHUP, exit_main_loop);
	signal(SIGALRM, exit_main_loop);

	/* Shut down cleanly is something has gone wrong */
	signal(SIGABRT, cleanup);
/*	signal(SIGSEGV, cleanup); */ /* FIXME */
	signal(SIGQUIT, cleanup);
	signal(SIGILL, cleanup);
	
	/* Set up packet capture */
	if (capture_file != NULL)
	{
		configure_capture_file(&pcap_handle, capture_file, &link_type);
		localnet = 0;
		netmask = 0;
	} else
	{
		configure_live_interface(&pcap_handle, dev, promisc_flag, &localnet, 
	    								 &netmask, &link_type);
	}
	set_packet_filter(pcap_handle, filter, netmask);
	
	/* Test for user interrupt */
	if (setjmp(jbuf) == 0)
	{
		/* Main loop */
		exit_flag = 0;
		while(!exit_flag)
		{
			dispatch_result = pcap_dispatch(pcap_handle, 0, pcap_packet_handler, 
			                                (u_char*)&link_type);
			if (dispatch_result == -1)
			{
				/* An error has occurred */
				fprintf(stderr, "pcap_dispatch: %s\n\n", pcap_geterr(pcap_handle));
				exit(-1);
				pcap_close(pcap_handle);
			} else if ((dispatch_result == 0) && (capture_file != NULL))

			{
				/* An EOF has occurred */
				exit_flag = 1;
			}
		}
	}
	
	/* Close packet capture interface */
	pcap_close(pcap_handle);

	/* Free options variables */	
	if (dev != NULL)
		free(dev);
		
	if (filter != NULL)
		free(filter);
		
	if (capture_file != NULL)
		free(capture_file);

	/* Reset signals */
	signal(SIGINT, SIG_DFL);
	signal(SIGTERM, SIG_DFL);
	signal(SIGSEGV, SIG_DFL);
	signal(SIGABRT, SIG_DFL);
	signal(SIGQUIT, SIG_DFL);
	signal(SIGILL, SIG_DFL);

	/* Resolve host names */
	if (resolve_dns)
	{
		fprintf(stderr, "Resolving hostnames (this can take a while)...");
		resolve_host_names(psum);
		fprintf(stderr, " done.\n");
	}
	
	/* Display results */
	switch (display_mode)
	{
		case TEXTFE:
		text_display_summary(psum, max_hosts, host_sort_mode, host_reverse, 
									max_peers, peer_sort_mode, peer_reverse);
		break;

		case HTMLFE:
		html_display_summary(psum, max_hosts, host_sort_mode, host_reverse, 
									max_peers, peer_sort_mode, peer_reverse);
		break;

		case PSFE:
		ps_display_summary(psum, max_hosts, host_sort_mode, host_reverse, 
									max_peers, peer_sort_mode, peer_reverse);
		break;
	}	
	
	/* Free packet summary data structure */
	packet_summary_free(psum);

	return(0);
}
