#include "../include/traffic.h"

struct host		*hosts;				/* Pointer to array of hosts info */
u_int	   		h_num = 0;			/* Hosts number */
struct host		total;				/* Summarize */

struct ehost			*ehosts;	 /* Excluded hosts array */
struct ehost			*xhosts;	 /* Excluded hosts traffic array */
int						eh_num=0;	 /* Number of excluded hosts */
int						xh_num=0;	 /* Number of excluded hosts traffic */

const	char	*ver = "0.1rc4";	/* Version */

struct options opt = {
	NULL,					/* Device */
	NULL,					/* Log file */
	1,						/* Refresh timeout */
	0, 						/* Number of functions calls to estimate elapsed time */
	0,						/* Current page to display */
	{ 	COL_IP,				/* Diaplaying columns array */
		COL_DATA,
		COL_DATA_TOT,
		COL_PKT,
		COL_PKT_TOT,	
		0,0,0,0,0,0,0,0,0,0,0,0
	},
	MODE_STAT,						/* Displaying mode */
	COL_IP,							/* Sort column */
	0,								/* Disable reverse sorting */
	1,								/* Transfer unit */
	"B",							/* Unit name */
	1.0 / (1024 * 1024),			/* Total transfer unit */
	"MB",							/* Total unit name */
	0,								/* Total disabled */
	"",								/* Exclude file */
	0,								/* Daemon mode disabled by default */
	"/usr/local/etc/traffic.conf",	/* Default config file */
	0,								/* Discover hosts */
	0,								/* Only IP private classes */
	0,								/* Timeout */
	0,								/* RevDNS */
};

WINDOW			*win;				/* Curses window pointer */

/* Mutexes */
short	m_packet = 0;	/* Get packet function mutex */ 
short	m_display = 0;	/* Display function mutex */
short	m_file = 0;		/* File logging function mutex */

int
main(int argc, char *argv[])
{
	char           		*file = NULL;				/* File with hosts list */
	char				*class = NULL;				/* Network class */
	char           		errbuf[PCAP_ERRBUF_SIZE];	/* Error buffer */
	pcap_t         		*handle;					/* Device handler */
	FILE           		*fp;						/* File pointer */
	char           		buf[BUF_SIZE];				/* Buffor */
	struct itimerval	timer;						/* Timer characteristic */
	u_char				i = 1;						/* Counter */
	int					op;

	/* Init */
	opt.dev = pcap_lookupdev(errbuf);

	/* Init hosts and ehosts and xhosts arrays */
	if ((hosts = malloc(host_len)) == NULL) {
		perror("hosts array\n");
		exit(-1);
	}
	
	if ((ehosts = malloc(sizeof(struct ehost))) == NULL) {
		perror("ehosts array\n");
		exit(-1);
	}

	if ((xhosts = malloc(sizeof(struct ehost))) == NULL) {
		perror("xhosts array\n");
		exit(-1);
	}
	
	/* Parse arguments passed to program */
	while ((op = getopt(argc, argv, "f:i:u:U:l:r:c:s:tC:he:E:dRvDPT:Hx:X:")) != -1)
		switch (op) {
				/* Filename */
				case 'f':
					file = optarg;
					/* Read hosts list and create right array */
					/* Open file */
					if ((fp = fopen(file, "r")) == NULL) {
						perror(file);
						exit(1);
					}

					/* Proccess the file */
					while (fscanf(fp, "%s\n", buf) != EOF) {
						struct in_addr h;

						/* Convert ASCII address to BIN */
						if ((h.s_addr = inet_addr(buf)) == INADDR_ANY) {
							perror(buf);
							exit(1);
						}

						if (addhost(h)) {
							exit(-1);
						}
					}
					/* Close file */
					fclose(fp);
					break;

				/* Device */
				case 'i':
					opt.dev = optarg;
					break;

				/* Unit */
				case 'u':
					if (( opt.unit = get_unit(optarg) ) != 0 )
						strncpy(opt.unit_name, optarg, 5);
					else {
						printf("Invalid unit: %s\n", optarg);
						exit(1);
					}
					break;

				/* Total unit */
				case 'U':
					if (( opt.unit_t = get_unit(optarg) ) != 0 )
						strncpy(opt.unit_t_name, optarg, 5);
					else {
						printf("Invalid unit: %s\n", optarg);
						exit(1);
					}
					break;

				/* Log file */
				case 'l':
					opt.logfile = optarg;
					break;

				/* Refresh */
				case 'r':
					opt.delay = atoi(optarg);
					if (opt.delay == 0) {
						printf("Invalid refresh timeout !\n");
						exit(1);
					}
					break;

				/* Class */
				case 'c':
					{
					struct in_addr	ip;			/* Net addresss */
					char			tmp[16];	/* Temporary buffor */
					long			n;
					int				mask;		/* Netmask */
					int     		c1,c2,c3,c4;/* Parts of IP address */

					class = optarg;

					if (sscanf(class, "%d.%d.%d.%d/%d",&c1, &c2, &c3, &c4, &mask) != 5)
						invalid_class();

					sprintf(tmp, "%d.%d.%d.%d", c1, c2, c3, c4);

					/* Check validaity of net address */
					if (inet_addr(tmp) == INADDR_NONE)
						invalid_class();

					mask = pow(2, 32 - mask);
					ip.s_addr = inet_addr(tmp);

					/* Add hosts to array */
					for(n = 0; n < mask; n++) {
						if (addhost(ip))
							exit(-1);

						ip.s_addr = htonl(ntohl(ip.s_addr) + 1);
					}

					}
					break;

				/* Sort order */
				case 's':
					switch(optarg[0]) {
						case 'h':
							opt.sort = COL_IP;
							break;
						case 'c':
							opt.sort = COL_DATA;
							break;
						case 'm':
							opt.sort = COL_DATA_MAX;
							break;
						case 'a':
							opt.sort = COL_DATA_AVG;	
							break;
						case 's':
							opt.sort = COL_DATA_TOT;
							break;
						case 'p':
							opt.sort = COL_PKT;
							break;
						case 'O':
							opt.sort = COL_PKT_MAX;
							break;
						case 'o':
							opt.sort = COL_PKT_AVG;
							break;
						case 'P':
							opt.sort = COL_PKT_TOT;
							break;
						case 'M':
							opt.sort = COL_MAC;
							break;
						case 'H':
							opt.sort = COL_HOSTNAME;
							break;
						case 't':
							opt.sort = COL_TCP;
							break;
						case 'u':
							opt.sort = COL_UDP;
							break;
						case 'i':
							opt.sort = COL_ICMP;
							break;
						case 'T':
							opt.sort = COL_TCP_TOT;
							break;
						case 'U':
							opt.sort = COL_UDP_TOT;
							break;
						case 'I':
							opt.sort = COL_ICMP_TOT;
							break;
						default:
							printf("Invalid column : %s", optarg);
							exit(1);
							break;
					}
					break;

					/* Total enabled */
					case 't':
						opt.total = 1;
						break;
					
					/* Init columns */
					case 'C':
						{
							int i;

							bzero(opt.columns, sizeof(opt.columns));
							opt.columns[0] = COL_IP;
							for(i=0; i < strlen(optarg); i++)
								switch(optarg[i]) {
									case 'h':
										add_column(COL_IP) ;
										break;
									case 'c':
										add_column(COL_DATA) ;
										break;
									case 'm':
										add_column(COL_DATA_MAX) ;
										break;
									case 'a':
										add_column(COL_DATA_AVG) ;
										break;
									case 's':
										add_column(COL_DATA_TOT) ;
										break;
									case 'p':
										add_column(COL_PKT) ;
										break;
									case 'P':
										add_column(COL_PKT_TOT) ;
										break;
									case 'o':
										add_column(COL_PKT_AVG) ;
										break;
									case 'O':
										add_column(COL_PKT_MAX) ;
										break;
									case 'M':
										add_column(COL_MAC) ;
										break;
									case 'H':
										add_column(COL_HOSTNAME) ;
										break;
									case 't':
										add_column(COL_TCP);
										break;
									case 'u':
										add_column(COL_UDP);
										break;
									case 'i':
										add_column(COL_ICMP);
										break;
									case 'T':
										add_column(COL_TCP_TOT);
										break;
									case 'U':
										add_column(COL_UDP_TOT);
										break;
									case 'I':
										add_column(COL_ICMP_TOT);
										break;
									default:
										printf("Invalid column : %c\n", optarg[i]);
										exit(1);
									}
						}
						break;

					/* Exclude ip file list */
					case 'e':
						file = optarg;
						/* Read hosts list and create right array */
						/* Open file */
						if ((fp = fopen(file, "r")) == NULL) {
						perror(file);
						exit(1);
						}

						/* Proccess the file */
						{
							int line = 0;

							while (fscanf(fp, "%s\n", buf) != EOF) {
								eh_num++;
								line++;
								
								if (addehost(buf) != 0)
								{
									printf("Error in line %d : %s\n", line, buf);
									exit(-1);
								}
							}
						}

						/* Close file */
						fclose(fp);
						break;

					/* Exclude net or hosts */
					case 'E':
							eh_num++;
							
							if (addehost(optarg) != 0)
							{
								printf("Invalid class: %s\n", buf);
								exit(-1);
							}
							break;

					/* Exclude ip traffic file list */
					case 'x':
						file = optarg;
						/* Read hosts list and create right array */
						/* Open file */
						if ((fp = fopen(file, "r")) == NULL) {
						perror(file);
						exit(1);
						}

						/* Proccess the file */
						{
							int line = 0;

							while (fscanf(fp, "%s\n", buf) != EOF) {
								xh_num++;
								line++;
								
								if (addxhost(buf) != 0)
								{
									printf("Error in line %d : %s\n", line, buf);
									exit(-1);
								}
							}
						}

						/* Close file */
						fclose(fp);
						break;

					/* Exclude net or hosts traffic */
					case 'X':
							xh_num++;
							
							if (addxhost(optarg) != 0)
							{
								printf("Invalid class: %s\n", buf);
								exit(-1);
							}
							break;
						
					/* Enable daemon mode */
					case 'd':
						opt.daemon = 1;
						break;

					/* Reverse sorting */
					case 'R':
						opt.rsort = 1;
						break;

					/* Discover mode */
					case 'D':
						opt.discovery = 1;
						break;

					/* Private mode */
					case 'P':
						opt.private = 1;
						break;

					/* Exit after timeout */
					case 'T':
						opt.timeout = atoi(optarg);
						break;

					/* Enable RevDNS */
					case 'H':
						opt.revdns = 1;
						break;

					/* Print version */
					case 'v':
						printf("Version: %s\n", ver);
						exit(1);

					/* Help */
					case 'h':
						usage();
						break;
					default:
						usage();
						break;
		}	

	argc -= optind;
	argv += optind;

	/* Check */
	if ((h_num == 0) && (!opt.discovery)) {
		printf("Specify filename with IP list or class addreess !\n");
		exit(1);
	}
	
	if ((opt.daemon == 1) && (opt.logfile == NULL)) {
		printf("Daemon mode require log file !\n");
		exit(1);
	}

	bzero(&total, host_len);
	
	/* Exclude hosts from array */
	{
		int a,b;
	
		for(a=0; a < eh_num; a++)
			for(b=0; b < h_num; b++) 
				if (inclass(hosts[b].ip, ehosts[a]))
				{
					hosts[b] = hosts[h_num - 1];
					h_num--;
				}
	}

/*	sleep(1); */

	/* Daemonize if enabled */
	if (opt.daemon)
		daemonize();

	/* Open device and print error message if failed */
	if ((handle = pcap_open_live(opt.dev, BUFSIZ, 1, 0, errbuf)) == NULL) {
		perror(opt.dev);
		exit(2);
	}

	/* Set timer to show statistics every delay timeout */
	bzero(&timer, sizeof(timer));
	timer.it_interval.tv_sec = 1 * opt.delay;
	timer.it_value.tv_sec = 1 * opt.delay;
	if (setitimer(ITIMER_REAL, &timer, NULL) != 0) {
		perror("Timer");
		exit(5);
	}
	if (signal(SIGALRM, show_stats) == SIG_ERR) {
		perror("Signal");
		exit(6);
	}


	/* Initialize ncurses */
	if (!opt.daemon) {
		win = initscr();
		noecho(); 				/* Disable echo */
		nodelay(win, TRUE);		/* Non-blocking keyboard input */
		keypad(win, TRUE);		/* Enable functions keys */
		display();
	}

	/* Go into packet capturing */
	pcap_loop(handle, -1, got_packet, NULL);

	/* Deactivate curses */
	if (!opt.daemon) 
		endwin();

	/* Cleaning */
	pcap_close(handle);

	return (0);
}