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

void
got_packet(u_char *args, const struct pcap_pkthdr * header, const u_char * packet)
{

	eth_hdr 			*ethernet;	/* The ethernet header */
	ip_hdr 				*ip;		/* The IP header */

	u_int	          	i;								/* Counter */
	int					key;							/* Pressed key */
	static int			key_alt = 0;					/* Previous special key */

    u_char  			lpp = LINES - 6; 				/* Lines per page to display IP qtatistics */

	int					fs, fd;							/* found source and destination */

	ethernet = (eth_hdr *) (packet);			/* Pointer to ethernet header */
	ip = (ip_hdr *) (packet + sizeof(eth_hdr));	/* Pointer to IP header */

	/* Lock */
	m_packet = 1;

    /* Check exluded hosts traffic list */
    {   
        int a;
        
        for(a=0; a < xh_num; a++)
                if (inclass(ip->ip_src, xhosts[a]) || inclass(ip->ip_dst, xhosts[a])) {
					m_packet = 0;
                    return ;
				}
    }

	/* Init */	
	fs = fd = 0;

	/* Discover and add new host */
	if (opt.discovery) {
		/* Find host */
		for (i = 0; i < h_num; i++) {
			/* From host */
			if (hosts[i].ip.s_addr == ip->ip_src.s_addr)
			fs = 1;
			/* To host */
			if (hosts[i].ip.s_addr == ip->ip_dst.s_addr)
				fd = 1;
		}

		if (fs == 0) {
			if (opt.private) {
				uint32_t	t;
				
				t = ntohl(ip->ip_src.s_addr);
				/* 10.0.0.0/8 */
				if ((( t & 0xff000000 ) == 0x0a000000)
				/* 172.16.0.0/12 */
					|| (( t & 0xffe00000 ) == 0xab000000)
				/* 192.168.0.0/16 */
					|| (( t & 0xffff0000 ) == 0xc0a80000))
					addhost(ip->ip_src);
			}
			else
				addhost(ip->ip_src);
		}
			

		if (fd == 0) {
			if (opt.private) {
				uint32_t	t;
				
				t = ntohl(ip->ip_dst.s_addr);
				/* 10.0.0.0/8 */
				if ((( t & 0xff000000 ) == 0x0a000000)
				/* 172.16.0.0/12 */
					|| (( t & 0xffe00000 ) == 0xab000000)
				/* 192.168.0.0/16 */
					|| (( t & 0xffff0000 ) == 0xc0a80000))
					addhost(ip->ip_dst);
			}
			else
				addhost(ip->ip_dst);
		}
	}


	/* Find host */
	for (i = 0; i < h_num; i++) {
		/* From host */
		if (hosts[i].ip.s_addr == ip->ip_src.s_addr) {
			hosts[i].in += header->len;
			hosts[i].pkt_in++;
			
			/* Count data flow by protocol */
			switch (ip->ip_p) {
				case IPPROTO_TCP:
					hosts[i].tcp_in += header->len;
					break;
				case IPPROTO_UDP:
					hosts[i].udp_in += header->len;
					break;
				case IPPROTO_ICMP:
					hosts[i].icmp_in += header->len;
					break;
			}

			/* Save MAC */
			strncpy((char *)&hosts[i].mac, (char *)ether_ntoa((struct ether_addr *)&ethernet->ether_shost),18);
		}
		/* To host */
		else if (hosts[i].ip.s_addr == ip->ip_dst.s_addr) {
			hosts[i].out += header->len;
			hosts[i].pkt_out++;

			/* Count data flow by protocol */
			switch (ip->ip_p) {
				case IPPROTO_TCP:
					hosts[i].tcp_out += header->len;
					break;
				case IPPROTO_UDP:
					hosts[i].udp_out += header->len;
					break;
				case IPPROTO_ICMP:
					hosts[i].icmp_out += header->len;
					break;
			}

			/* Save MAC */
			strncpy((char *)&hosts[i].mac, (char *)ether_ntoa((const struct ether_addr *)&ethernet->ether_dhost),18);
		}
	}

	/* Remove unused hosts */
/*
	if (opt.discovery) {
		for (i = 0; i < h_num; i++) {
			if ((hosts[i].tot_pkt_in == 0) && (hosts[i].tot_pkt_out == 0) && (h_num > 1)) {
				h_num--;
				hosts[i] = hosts[h_num];
				hosts = realloc(hosts, host_len * h_num);
			}
		}	
	}
*/




	/* Keyboard control section */
	if (!opt.daemon) {
		key = getch();

		if (key_alt) {
			switch (key_alt) { 
				/* Add column */
				case 'a':
					switch(key) {
						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;
						}
					break;
	
				/* Delete column */
				case 'd':
					switch(key) {
						case 'h':
							del_column(COL_IP) ;
							break;
						case 'c':
							del_column(COL_DATA) ;
							break;
						case 'm':
							del_column(COL_DATA_MAX) ;
							break;
						case 'a':
							del_column(COL_DATA_AVG) ;
							break;
						case 's':
							del_column(COL_DATA_TOT) ;
							break;
						case 'p':
							del_column(COL_PKT) ;
							break;
						case 'P':
							del_column(COL_PKT_TOT) ;
							break;
						case 'o':
							del_column(COL_PKT_AVG) ;
							break;
						case 'O':
							del_column(COL_PKT_MAX) ;
							break;
						case 'M':
							del_column(COL_MAC) ;
							break;
						case 'H':
							del_column(COL_HOSTNAME) ;
							break;
						case 't':
							del_column(COL_TCP) ;
							break;
						case 'u':
							del_column(COL_UDP) ;
							break;
						case 'i':
							del_column(COL_ICMP) ;
							break;
						case 'T':
							del_column(COL_TCP_TOT) ;
							break;
						case 'U':
							del_column(COL_UDP_TOT) ;
							break;
						case 'I':
							del_column(COL_ICMP_TOT) ;
							break;
					}
					break;
	
				/* Sort column */
				case 's':
					switch(key) {
						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 'P':
							opt.sort = COL_PKT_TOT;
							break;
						case 'o':
							opt.sort = COL_PKT_AVG;
							break;
						case 'O':
							opt.sort = COL_PKT_MAX;
							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;
						}
					break;
	
				/* Change unit */
				case 'u':
					switch(key) {
						case 'b':
							strcpy(opt.unit_name, "Bit");
							opt.unit = get_unit("Bit");
							break;
						case 'B':
							strcpy(opt.unit_name, "B");
							opt.unit = get_unit("B");
							break;
						case 'k':
							strcpy(opt.unit_name, "kBit");
							opt.unit = get_unit("kBit");
							break;
						case 'K':
							strcpy(opt.unit_name, "KB");
							opt.unit = get_unit("KB");
							break;
						case 'm':
							strcpy(opt.unit_name, "mBit");
							opt.unit = get_unit("mBit");
							break;
						case 'M':
							strcpy(opt.unit_name, "MB");
							opt.unit = get_unit("MB");
							break;
						case 'g':
							strcpy(opt.unit_name, "gBit");
							opt.unit = get_unit("gBit");
							break;
						case 'G':
							strcpy(opt.unit_name, "GB");
							opt.unit = get_unit("GB");
							break;
						}
					break;
	
				/* Change total unit */
				case 'U':
					switch(key) {
						case 'b':
							strcpy(opt.unit_t_name, "Bit");
							opt.unit_t = get_unit("Bit");
							break;
						case 'B':
							strcpy(opt.unit_t_name, "B");
							opt.unit_t = get_unit("B");
							break;
						case 'k':
							strcpy(opt.unit_t_name, "kBit");
							opt.unit_t = get_unit("kBit");
							break;
						case 'K':
							strcpy(opt.unit_t_name, "KB");
							opt.unit_t = get_unit("KB");
							break;
						case 'm':
							strcpy(opt.unit_t_name, "mBit");
							opt.unit_t = get_unit("mBit");
							break;
						case 'M':
							strcpy(opt.unit_t_name, "MB");
							opt.unit_t = get_unit("MB");
							break;
						case 'g':
							strcpy(opt.unit_t_name, "gBit");
							opt.unit_t = get_unit("gBit");
							break;
						case 'G':
							strcpy(opt.unit_t_name, "GB");
							opt.unit_t = get_unit("GB");
							break;
					}
					break;
			}
			if (key != ERR)
				key_alt = 0;
		}
		else {
			switch (key) {
				/* Turn page left */
				case KEY_LEFT:
					if (opt.page > 0)
						opt.page--;
					break;
				/* Turn page right */
				case KEY_RIGHT:
					if (opt.page < h_num / lpp)
						opt.page++;
					break;
	
				/* Total line */
				case 't':
					/* Reverse */
					opt.total = (opt.total == 1)? 0 : 1;
					break;
	
				/* reverse sorting */
				case 'R':
					opt.rsort = (opt.rsort == 1)? 0 : 1;
					break;
	
				/* Reset data */
				case 'r':
					{
						struct in_addr tmp;
	
						for (i = 0; i < h_num; i++) {
							tmp = hosts[i].ip;
							bzero(&hosts[i], host_len);
							hosts[i].ip = tmp;
						}
					}
					break;
	
				/* Exit */
				case 'q':
					if (opt.mode == MODE_HELP) 
						opt.mode = MODE_STAT;
					else
						{
						endwin();
						exit(0);
							}
					break;
				/* Help */
				case '?':
					opt.mode = MODE_HELP;
					break;
				default:
					if (key != ERR) {
						key_alt = key ;
						m_display = 1;
					}
			}
		}
	}

	/* Display if enabled */
	if (m_display) {
		if (!opt.daemon)
			display();
		m_display = 0;
	}
	
	
	/* Log to file if enabled */
	if (m_file) {
		log_file();
		m_file = 0;
	}
	
	/* Unlock */
	m_packet = 0;
}