#include <sys/types.h>
#include <sys/time.h>
#include <sys/param.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/io.h>
#include <fcntl.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <syslog.h>

/*
 * The Real Time Traffic Meter
 *
 * By Simon Kirby, 1997, 1998, 2002/03/19
 *
 * This used to use firewall accounting rules for traffic data, but iptables
 * doesn't use /proc files anymore and thus is easier to code for.  It would
 * be nice to do this, though.
 */

/****************************************************************************/

char acct_file[] = "/proc/net/dev";

typedef unsigned long long qword;
typedef unsigned long dword;
typedef unsigned short word;
typedef unsigned char byte;

/****************************************************************************/

#ifdef __builtin_expect
#define likely(x)       __builtin_expect((x),1)
#define unlikely(x)     __builtin_expect((x),0)
#else
#define likely(x)       x
#define unlikely(x)     x
#endif

#define ReadSize 32768
#define DEVICE_SIZE 16
#define MaxWords 12

struct words {
	char *arg[MaxWords];
	char *argc[MaxWords];
	char *string;
	char wordbuffer[8192+1+1];
	word wordcount;
};

struct units {
	double multiplier;
	char unit_name[8];
};

struct units byte_units[] = {
	{ 1,			"Bps", },
	{ 1024,		"kBps", },
	{ 1048576, 		"MBps", },
	{ 1073741824,	"GBps", },
	{ 0,			"", },
};

struct units bit_units[] = {
	{ 1,			"bps", },
	{ 1000,		"kbps", },
	{ 1000000, 		"Mbps", },
	{ 1000000000,	"Gbps", },
	{ 0,			"", },
};

struct units packet_units[] = {
	{ 1,			"pps", },
	{ 1000,		"kpps", },
	{ 1000000, 		"Mpps", },
	{ 1000000000,	"Gpps", },
	{ 0,			"", },
};

static struct words args;

static char readstring[ReadSize+1];
static char device[DEVICE_SIZE+1];
static enum direction {
	ALL,
	IN,
	OUT,
} direction;

static __inline__ void outport(u_short port,char value){
	__asm__("outb %0,%1" :: "a" ((char) value),"d" ((unsigned short) port));
}

static __inline__ short inport(unsigned short port){
	char _v;
	__asm__("inb %1,%0":"=a" (_v) : "d" ((unsigned short) port));
	return _v;
}

static __inline__ void cli(){__asm__("cli");}
static __inline__ void sti(){__asm__("sti");}

static char *strnzcpy(char *d,char *s,size_t n){
	strncpy(d,s,n);
	d[n] = '\0';
	return d;
}

static char *skipspace(char *s){
	while (*s == ' ')
		s++;
	return s;
}

static void eatwords(struct words *args,char *s){
	word wordlength,wordcount;
	register word i;
	register char *w;
	register char ls = '\0';

	for (i = 0;i < MaxWords;i++){
		args->arg[i] = &args->wordbuffer[8192];
		args->argc[i] = &args->wordbuffer[8192];
	}
	args->wordbuffer[8192] = '\0';
	args->wordcount = wordcount = 0;
	if (strlen(s) > 8192)
		s[8192] = '\0';
	w = args->wordbuffer;
	for (wordlength = 0;(*s != '\0') && (wordcount < MaxWords);s++){
		if ((*s == ' ') && (ls != ' ')){
			if (wordlength){
				*w++ = '\0';
				wordlength = 0;
			}
		} else {
			*w = *s;
			if (!wordlength++){
				args->arg[wordcount] = w;
				args->argc[wordcount++] = s;
			}
			w++;
			if (*s != ' ')
				args->wordcount = wordcount;
		}
		ls = *s;
	}
	*w = '\0';
	args->string = args->argc[0];
}

static qword atouq(char *s){
	qword r = 0;
	if (!s) return 0LL;
	if ((*s < '0') || (*s > '9')) return 0LL;
	r = *s++ - '0';
	while ((*s >= '0') && (*s <= '9')){
		r*= 10;
		r+= *s++ - '0';
	}
	return r;
}

static void setbg(char r,char g,char b){
	(void)iopl(3);
	outport(0x3c7,0);
	outport(0x3c8,0);
	outport(0x3c9,r);
	outport(0x3c9,g);
	outport(0x3c9,b);
	(void)iopl(0);
}

#define DESIRED_LENGTH 4

/*
 * Automatically select an appropriate unit and print number into strings
 *
 * Assume we have DESIRED_LENGTH characters to present something useful in,
 * excluding the unit.
 */
static void auto_unit(char *data,char *unit,struct units *up,double v){
	double multiplier;
	char *unitp,*s;
	int vlen,dlen;
	char lenstring[32+1];

	if (v < 0)
		v = -v;

	multiplier = up->multiplier;
	unitp = up->unit_name;
	for (;up->multiplier;up++){
		if ((up->multiplier / 1000 * 1000) == up->multiplier) {
			if (v < up->multiplier * 10)
				break;
		} else {
			if (v < up->multiplier)
				break;
		}
		multiplier = up->multiplier;
		unitp = up->unit_name;
	}
	strcpy(unit,unitp);
	v/= multiplier;
	sprintf(lenstring,"%.2f",v); 
	if ((s = strchr(lenstring,'.')) != NULL)
		*s = '\0';
	vlen = strlen(lenstring);
	dlen = DESIRED_LENGTH - 1 - vlen;
/* fprintf(stderr,"[%s]:%i\n",lenstring,dlen); */
	if (dlen > 0){
		sprintf(data,"%.*f",dlen,v);
	} else {
		strcpy(data,lenstring);	/* No decimals */
	}
}

static double interval;
static int goingberserk = 0;

static char field_data[16][16];
static char field_unit[16][16];

static struct winsize ws;

static double utilization,minimum_utilization,utilization_divisor;
static char *comment;


#define HISTORY_COUNT 3
#define DIRECTION_COUNT 2

static const dword history_time[HISTORY_COUNT] = { 1, 5, 15 };
static double history_exp[HISTORY_COUNT];
static qword hist_bytes[HISTORY_COUNT][DIRECTION_COUNT];
static qword hist_packets[HISTORY_COUNT][DIRECTION_COUNT];
static qword last_bytes[DIRECTION_COUNT];
static qword last_packets[DIRECTION_COUNT];

static void init_sample_history() {
	dword i;

	for (i = 0;i < HISTORY_COUNT;i++)
		history_exp[i] = 1. / exp(interval / history_time[i]);
	memset(&hist_bytes,0,sizeof(hist_bytes));
	memset(&hist_packets,0,sizeof(hist_packets));
	memset(&last_bytes,0,sizeof(last_bytes));
	memset(&last_packets,0,sizeof(last_packets));
}

qword import_counter(qword *last,qword now,qword rate) {
	qword change;

	if (now < *last) {
		/* Wrap master */
		if (((rate > 0x10000000) &&
		     (*last <= 0xffffffff)) ||
		    ((*last >= 0xe0000000) &&
		     (*last <= 0xffffffff)))
			change = now + 0x100000000ULL - *last;
		else
			change = 0;	 /* Interface likely reset */
	} else
		change = now - *last;
	*last = now;
	return change;
}

static void sample(int fd) {
	char *s,*c,*nexts;
	qword bytes,packets;
	static qword iter = 0;
	int i,l,r,screen_middle_row;
	double m;

	iter++;

	(void)lseek(fd,0,SEEK_SET);
	l = 0;
	do {
		r = read(fd,readstring + l,ReadSize - l);
		if (r == -1){
			perror("read()");
			exit(1);
		}
		l+= r;
	} while (r);
	readstring[l] = '\0';
	s = readstring;
	if (unlikely(!l)){
		fprintf(stderr,"Accouting file \"%s\" is empty.\n",acct_file);
		exit(1);
	}
	if (unlikely(!*s)){
		fprintf(stderr,"Huh?  Accounting file \"%s\" has NULL as first character.\n",acct_file);
		exit(1);
	}
	for (c = s = skipspace(s);s && *s;s = nexts){
		nexts = strchr(s,'\n');
		if (nexts){
			*nexts = '\0';
			nexts++;
		}
		s = skipspace(s);
		if (!*s)
			break;
		c = strchr(s,':');
		if (!c)
			continue;
		*c = '\0';
		if (strcmp(s,device) == 0)
			break;
		continue;
	}
	if (unlikely(!s || !*s)){
		fprintf(stderr,"Device \"%s\" not found.\n",device);
		exit(1);
	}

	eatwords(&args,c + 1);

	screen_middle_row = ws.ws_row >> 1;

	for (direction = 0;direction < 2;direction++) {
		s = skipspace(args.arg[direction << 3]);
		bytes = atouq(s);
		s = skipspace(args.arg[(direction << 3) + 1]);
		packets = atouq(s);

		if (unlikely(iter == 1)) {
			last_bytes[direction] = bytes;
			last_packets[direction] = packets;
			continue;
		}

		bytes = import_counter(&last_bytes[direction],bytes,hist_bytes[0][direction]);
		packets = import_counter(&last_packets[direction],packets,hist_packets[0][direction]);
	
		for (i = 0;i < HISTORY_COUNT;i++) {
			m = unlikely(iter == 2) ? 0 : history_exp[i];
			hist_bytes[i][direction]*= m;
			hist_packets[i][direction]*= m;
			hist_bytes[i][direction]+= (1 - m) * (bytes / interval);
			hist_packets[i][direction]+= (1 - m) * (packets / interval);
		}

		if (direction == 0) {
			printf("\033[1;%ur\033[%uf",screen_middle_row - 1,screen_middle_row - 1);
		} else
			printf("\033[%u;%ur\033[%uf",screen_middle_row + 1,ws.ws_row,ws.ws_row);


		utilization = (hist_bytes[0][direction] << 3) / utilization_divisor;
		if (iter > 10){
			if (utilization < minimum_utilization){
				setbg((iter & 1) * 63,((iter & 2) >> 1) * 63,((iter & 2) >> 1) * 63);
				goingberserk++;
				if ((iter & 7) == 0)
					printf("\007");
				if ((iter & 7) == 2)
					printf("\007");
				if ((iter & 1) == 1)
					(void)syslog(LOG_ALERT,"** NETWORK TROUBLE ** -> Device \"%s\": %.1f%% utilization (min: %.1f%%)",
					 device,utilization,minimum_utilization);
			} else
				if (goingberserk){
					setbg(0,0,0);
					goingberserk = 0;
				}
		}
		if (utilization >= 75)
			printf("\033[m\n\033[41m\033[K");
		else
			printf("\033[m\n\033[K");

		auto_unit(field_data[0],field_unit[0],packet_units,(double)packets / interval);
		auto_unit(field_data[1],field_unit[1],bit_units,(double)bytes / interval * 8);
		auto_unit(field_data[2],field_unit[2],byte_units,(double)bytes / interval);
		auto_unit(field_data[3],field_unit[3],packet_units,hist_packets[2][direction]);
		auto_unit(field_data[4],field_unit[4],bit_units,hist_bytes[2][direction] * 8);
		auto_unit(field_data[5],field_unit[5],byte_units,hist_bytes[2][direction]);

		printf("\033[0;1;36m%8Lu "
			"\033[0;1m%8Lu "
			" "
			"\033[1;36m%4s \033[34m%-4s "
			"\033[0;1m%4s \033[34m%-4s "
			"\033[1;33m%4s \033[34m%-4s "
			" "
			"\033[1;36m%4s \033[34m%-4s "
			"\033[0;1m%4s \033[34m%-4s "
			"\033[1;33m%4s \033[34m%-4s ",
			packets,bytes,
			field_data[0],field_unit[0],
			field_data[1],field_unit[1],
			field_data[2],field_unit[2],
			field_data[3],field_unit[3],
			field_data[4],field_unit[4],
			field_data[5],field_unit[5]
		);

		if (utilization < 1)
			printf("\033[0;37m");		/*   0 - 1 : Gray */
		else if (utilization < 10)
			printf("\033[1;32m");		/*  1 - 10 : Green */
		else if (utilization < 30)
			printf("\033[1;36m");		/* 10 - 30 : Cyan */
		else if (utilization < 50)
			printf("\033[1;33m");		/* 30 - 50 : Yellow */
		else if (utilization < 70)
			printf("\033[1;35m");		/* 50 - 70 : Red */
		else if (utilization < 80)
			printf("\033[1;5;33m");		/* 70 - 80 : Flashing red */
		else if (utilization < 90)
			printf("\033[1;5;33;41m");	/* 80 - 90 : Flashing yellow on red */
		else
			printf("\033[1;5;37;41m");	/*     90+ : Flashing white on red */

		printf("%6.1f%%\033[m\r",utilization);
		if (*comment)
			printf("\033[99A\033[1;44m[ %s ]\033[m\r\033[99B",comment);

	}
	r = utilization / 100. * ws.ws_row;
	printf("\033[%uf\033[0;7m\033[K Packets \033[0;7m  Bytes        /\\ IN /\\       \\/ OUT \\/\033[m",
		screen_middle_row);
	printf("\033[r\033[%uf",ws.ws_row);
	fflush(stdout);
}

void get_window_size() {
	memset(&ws,0,sizeof(ws));
	if ((ioctl(fileno(stdin), TIOCGWINSZ, &ws) != 0) || (ws.ws_row <= 0))
		ws.ws_row = 25,
		ws.ws_col = 80;
	printf("\nWINDOW SIZE: %ux%u\n",ws.ws_col,ws.ws_row);
}

void signal_winch(int si_number, siginfo_t *siginfo, void *bits) {
}

int main(int argc,char *argv[]) {
	char *s;
	int i,fd;
	qword integer_interval;
	struct itimerval it;
	struct sigaction sa;
	sigset_t sigset;
	siginfo_t siginfo;

	get_window_size();
	for (i = ws.ws_row / 2;i > 0;i--)
		printf("\n");

	if (argc < 1+1){
		fprintf(stderr,"Usage: %s <number of accounting rule to monitor> <interval> [<minimum utilization>] [<utilization divisor>] [<comment>]\n\n"
			"Example: %s eth0 0.2 1.5 100\n"
			"			`- Would monitor 1st line shown by ipchains -L acctin,\n"
			"				reporting every 0.2 seconds, and going berserk\n"
			"				if the utilization falls below 1.5%%.  Utilization\n"
			"				percentage would be based on a 100Mbit network.\n",
			argv[0],argv[0]);
		exit(1);
	}

	strnzcpy(device,argv[1],DEVICE_SIZE);

	if (argc >= 2+1) {
		interval = atof(argv[2]);
		if (!interval){
			fprintf(stderr,"Error: Invalid interval specified: \"%s\"\n",argv[2]);
			exit(1);
		}
	} else
		interval = .2;

	if (argc >= 3+1)
		minimum_utilization = atof(argv[3]);
	else
		minimum_utilization = 0;

	if (argc >= 4+1)
		utilization_divisor = atof(argv[4]) * 10000.;
	else
		utilization_divisor = 1000000;	/* 100Mbps default */

	if (argc >= 5+1){
		for (s = argv[5];*s;s++)
		  if (*s == '_')
			  *s = ' ';
		comment = argv[5];
	} else
		comment = "";

	setuid(0); seteuid(0); setgid(0);

	fprintf(stderr,"Starting dump of device \"%s\" at %g second intervals...\n",
		argv[1],interval);
	fd = open(acct_file,O_RDONLY);

	init_sample_history();

	memset(&sa,0,sizeof(sa));
	sa.sa_handler = SIG_IGN;
	sa.sa_sigaction = NULL;
	sigaction(SIGALRM,&sa,NULL);

	memset(&sa,0,sizeof(sa));
	sa.sa_handler = SIG_IGN;
	sa.sa_sigaction = signal_winch;
	sigaction(SIGWINCH,&sa,NULL);

	sigemptyset(&sigset);
	sigaddset(&sigset,SIGALRM);
	sigaddset(&sigset,SIGHUP);
	sigprocmask(SIG_BLOCK,&sigset,NULL);

	memset(&it,0,sizeof(it));
	integer_interval = interval * 1000000;
	it.it_interval.tv_sec = integer_interval / 1000000ULL;
	it.it_interval.tv_usec = integer_interval % 1000000ULL;
	it.it_value.tv_sec = it.it_interval.tv_sec;
	it.it_value.tv_usec = it.it_interval.tv_usec;
	setitimer(ITIMER_REAL,&it,NULL);

	sigemptyset(&sigset);
	sigaddset(&sigset,SIGALRM);
	sigaddset(&sigset,SIGWINCH);
	sigaddset(&sigset,SIGCHLD);
	sigaddset(&sigset,SIGHUP);

	for (;;) {
		sample(fd);
		for (;;) {
			sigwaitinfo(&sigset,&siginfo);
			if (siginfo.si_signo == SIGALRM)
				break;
			if (siginfo.si_signo == SIGWINCH) {
				get_window_size();
				printf("\033[2J");
			}
			if (siginfo.si_signo == SIGCHLD) {
				printf("CHLD\n");
				exit(123);
			}
			if (siginfo.si_signo == SIGHUP) {
				printf("HUP\n");
				exit(123);
			}
		}
	}

	exit(0);
}