/*
 * Copyright (c) 2005 Daniel Bryan
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or 
 * without modification, are permitted provided that the following 
 * conditions are met:
 * - Redistributions of source code must retain the above copyright 
 *   notice, this list of conditions and the following disclaimer.
 * - Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in 
 *   the documentation and/or other materials provided with the 
 *   distribution.
 * - Neither the name of the author nor the names of its contributors
 *   may be used to endorse or promote products derived from this 
 *   software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS 
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE 
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
 * POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/param.h>
#include <limits.h>
#include <paths.h>
#include <sys/time.h>
#include <stdarg.h>
#include <errno.h>

#include <sys/socket.h>
#include <net/if.h>
#include <net/if_var.h>
#include <net/if_types.h>

#include <fcntl.h>
#include <kvm.h>

#include <curses.h>

#define VER_MAJ 0
#define VER_MIN 7

#define DISP_BIT  0
#define DISP_BYTE 1
int disp_format = DISP_BYTE;

struct nlist netl[] = { {"_ifnet"}, {""} };
kvm_t *kvmd;
char *nlistf = NULL;
char *memf = NULL;
#if __FreeBSD_version >= 501113
char name[IFNAMSIZ];
#else
char name[32];
char tname[16];
#endif
unsigned long ifnetaddr = 0;

unsigned long long in_total = 0;
unsigned long long out_total = 0;
char *in_dev = NULL;

int winw,winh;
WINDOW *mainw;
WINDOW *inw;
WINDOW *outw;
unsigned int maxgh = 0;

struct timeval timea; /* update interval */
struct timeval now;
struct timeval last;

#define MAX_G 1024
unsigned int *logi[MAX_G];
unsigned int *logo[MAX_G];

/* defaults to 1mbit/1mbit connection */
/* these are set to be in bytes/s */
unsigned long in_max = 131072;
unsigned long out_max = 131072;

char *color_in = 0; /* in graph color */
char *color_out = 0; /* out graph color */

/* prints correct format, float */
void print_formatf(WINDOW *win, float x, int format) {
	if(format == DISP_BIT) {
		x *= 8;
		if(x < 1024) {
			wprintw(win,"%8.0f bit",x);
		} else if(x < 1048576) {
			wprintw(win,"%7.1f kbit",x / 1024);
		} else if(x < 1073741824) {
			wprintw(win,"%7.1f mbit",x / 1048576);
		} else {
			wprintw(win,"%7.1f gbit",x / 1073741824);
		}
	} else if(format == DISP_BYTE) {
		if(x < 1024) {
			wprintw(win,"%8.0f B",x);
		} else if(x < 1048576) {
			wprintw(win,"%7.1f kB",x / 1024);
		} else if(x < 1073741824) {
			wprintw(win,"%7.1f mB",x / 1048576);
		} else {
			wprintw(win,"%7.1f gB",x / 1073741824);
		}
	}
	return;
}

/* prints correct format, u_long */
void print_format(WINDOW *win, unsigned long x, int format) {
	
	if(format == DISP_BIT) {
		x *= 8;
		if(x < 1024) {
			wprintw(win,"%lu bit",x);
		} else if(x < 1048576) {
			wprintw(win,"%lu kbit",x / 1024);
		} else if(x < 1073741824) {
			wprintw(win,"%lu mbit",x / 1048576);
		} else {
			wprintw(win,"%lu gbit",x / 1073741824);
		}
	} else if(format == DISP_BYTE) {
		if(x < 1024) {
			wprintw(win,"%lu B",x);
		} else if(x < 1048576) {
			wprintw(win,"%lu kB",x / 1024);
		} else if(x < 1073741824) {
			wprintw(win,"%lu mB",x / 1048576);
		} else {
			wprintw(win,"%lu gB",x / 1073741824);
		}
	}
	return;
}

void free_logs() {
	int i;
	for(i=0;i<MAX_G;i++) {
		if(logi[i] != NULL) {
			free(logi[i]);
		}
		if(logo[i] != NULL) {
			free(logo[i]);
		}
        logo[i] = NULL;
        logi[i] = NULL;
	}
	return;
}

/* initial screen draw */
void screen_init() {
	int i = 0;
	getmaxyx(stdscr,winh,winw);
	if(winh < 20 || winw < 20) {
		endwin();
		kvm_close(kvmd);
		fprintf(stderr,"Screen size is too small, sorry\n");
		exit(1);
	}
	
	mainw = newwin(winh,winw,0,0);
	maxgh = (winh-3)/2;
	inw = newwin(maxgh,winw,maxgh+3,0);
	outw = newwin(maxgh,winw,2,0);

	overwrite(mainw,stdscr);
	overwrite(inw,stdscr);
	overwrite(inw,mainw);
	overwrite(outw,stdscr);
	overwrite(outw,mainw);
	nonl();
	cbreak();
	noecho();
	curs_set(0);
	nodelay(mainw,TRUE);
	refresh();
	wmove(mainw,0,0);
	wprintw(mainw,"cnd %d.%d [%s]",VER_MAJ,VER_MIN,name);
	wattron(mainw,A_REVERSE);
	wmove(mainw,1,0);
	for(i=0;i<winw;i++) {
		waddch(mainw,' ');
	}
	wmove(mainw,maxgh+2,0);
	for(i=0;i<winw;i++) {
		waddch(mainw,' ');
	}
	/* fill in extra line */
	wmove(mainw,winh-1,0);
	for(i=0;i<winw;i++) {
		waddch(mainw,' ');
	}	
	wmove(mainw,1,(winw/2)-9);
	wprintw(mainw," out[");
	wmove(mainw,maxgh+2,(winw/2)-9);
	wprintw(mainw,"  in[");

	wmove(mainw,1,0);
	print_format(mainw,out_max,disp_format);
	wmove(mainw,maxgh+2,0);
	print_format(mainw,in_max,disp_format);
	wattroff(mainw,A_REVERSE);
	
	wrefresh(mainw);
	return;
}

void screen_check() {
	int rwinh,rwinw;
	getmaxyx(stdscr, rwinh, rwinw);
	if(rwinh == winh && rwinw == winw) {
		return; /* all is good, return */
	} else {
		delwin(mainw);
		delwin(inw);
		delwin(outw);
		/* clear graph as it will be incorrect */
		free_logs();
		screen_init();
	}
	return;
}

/* read kernel memory, based off of netstat */
int kread(u_long addr,char *buf,int size) {
	if(kvmd != NULL) {
		if(kvm_nlist(kvmd,netl) < 0) {
			if(nlistf)
				fprintf(stderr,"error, kvm_nlist(%s): %s\n",nlistf,kvm_geterr(kvmd));
			else
				fprintf(stderr,"error, kvm_nlist: %s\n",kvm_geterr(kvmd));
			exit(1);
		}
		
		if(netl[0].n_type == 0) {
			if(nlistf)
				fprintf(stderr,"error, no namelist: %s\n",nlistf);
			else
				fprintf(stderr,"error, no namelist\n");
			exit(1);
		}
	} else {
		fprintf(stderr,"error, kvm not available\n");
		exit(1);
	}
	if(!buf)
		return 0;
	if(kvm_read(kvmd,addr,buf,size) != size) {
		fprintf(stderr,"error, %s\n",kvm_geterr(kvmd));
		exit(1);
	}
	return 0;
}

int main(int argc, char *argv[]) {

	unsigned long lasti = 0;
	unsigned long lasto = 0;
	unsigned long curi = 0;
	unsigned long curo = 0;
#define CURIS curi / timea.tv_sec
#define CUROS curo / timea.tv_sec
	/* all goods are defined in net/if.h and net/if_var.h */
	struct ifnet foonet;
	struct ifnethead ifnethead;
	int i = 0;
	int x = 0;
	int j = 0;
	char ch = '\0';
	timea.tv_usec = 0;
	timea.tv_sec = 2; /* default interval */
		
	while((ch = getopt(argc,argv,"i:m:n:vhs:c:bB")) != -1) {
		if(ch == 'v') {
			printf("cnd v%d.%d\n",VER_MAJ,VER_MIN);
			exit(0);
		} else if(ch == 'i') {
			in_dev = optarg;
			if(strlen(in_dev) > sizeof(name)) {
				fprintf(stderr,"error, device name is too long\n");
				exit(1);
			}
		} else if(ch == 'b') {
			disp_format = DISP_BIT;
		} else if(ch == 'B') {
			disp_format = DISP_BYTE;
		} else if(ch == 'm') {
			if(*optarg < '0' && *optarg > '9')
				break;
			if(atoi(optarg) < 1)
				break;		
			in_max = atoi(optarg);
			switch(*(optarg + strlen(optarg)-1)) {
			/* Byte -> Byte is pointless */
			case 'b': /* bit -> Byte */
				in_max /= 8;
				break;
			case 'k': /* kbit -> Byte */
				in_max *= 128;
				break;
			case 'K': /* KByte -> Byte */
				in_max *= 1024;
				break;
			case 'm': /* mbit -> Byte */
				in_max *= 131072;
				break;
			case 'M': /* MByte -> Byte */
				in_max *= 1048576;
				break;
			case 'g': /* gbit -> Byte */
				in_max *= 134217728;
				break;
			case 'G': /* GByte -> Byte */
				in_max *= 1073741824;
				break;		
			}
		} else if(ch == 'n') {
			if(*optarg < '0' && *optarg > '9')
				break;
	 		if(atoi(optarg) < 1)
				break;
			out_max = atoi(optarg);
			switch(*(optarg + strlen(optarg)-1)) {
			/* Byte -> Byte is pointless */
			case 'b': /* bit -> Byte */
				out_max /= 8;
				break;
			case 'k': /* kbit -> Byte */
				out_max *= 128;
				break;
			case 'K': /* KByte -> Byte */
				out_max *= 1024;
				break;
			case 'm': /* mbit -> Byte */
				out_max *= 131072;
				break;
			case 'M': /* MByte -> Byte */
				out_max *= 1048576;
				break;
			case 'g': /* gbit -> Byte */
				out_max *= 134217728;
				break;
			case 'G': /* GByte -> Byte */
				out_max *= 1073741824;
				break;		
			}
		} else if(ch == 'h') {
			printf("cnd %d.%d by daniel bryan\n",VER_MAJ,VER_MIN);
			printf("usage: cnd -hvbB -i INTERFACE -m SPEED -n SPEED -c [ COLOR | COLORIN,COLOROUT ]\n");
			printf("  i - select interface\n");
			printf("  b - display speed in bits\n");
			printf("  B - display speed in bytes, default\n");
			printf("  m SPEED - max input of connection, default is 1mbit\n");
			printf("  n SPEED - max output of connection, default is 1mbit\n");
			printf("  s SECONDS - update interval in seconds, default is 2\n");
			printf("  c COLOR - graph color, lowercase, ie. \"blue\"\n");
			printf("    COLORIN,COLOROUT - different in/out color, ie. \"red,blue\"\n");
			printf("  v - print version\n");
			printf("  h - print help (this)\n\n");
			
			printf("NOTE: SPEED can either be in bytes/s (uppsercase B/K/M/G), ie. \"4B\"\n");
			printf("      or bit (lowercase b/k/m/g), ie. \"4b\", defaults to bytes/s\n");
			exit(0);
		} else if(ch == 's') {
			if(*optarg >= '0' && *optarg <= '9')
				timea.tv_sec = atoi(optarg);
		} else if(ch == 'c') {
			color_in = optarg;
			if((color_out = strchr(optarg,',')) != NULL)
				color_out++;
		}
	}
	argc -= optind;
	argv += optind;
	
	for(i=0;i<sizeof(name);i++) {
		name[i] = '\0';
	}

	kvmd = kvm_openfiles(nlistf,memf,NULL,O_RDONLY,0);	
	setgid(getgid());
	kread(0,0,0);
	ifnetaddr = netl[0].n_value;
	
	if(kread(ifnetaddr,(char *)&ifnethead,sizeof ifnethead))
		return 1;
	
	ifnetaddr = (u_long)TAILQ_FIRST(&ifnethead);
	if(kread(ifnetaddr,(char *)&foonet,sizeof foonet))
		return 1;

#if __FreeBSD_version >= 501113
	strncpy(name,foonet.if_xname,sizeof(name));
#else
	if(kread((u_long)foonet.if_name, tname, 16))
		return 1;
	snprintf(name,32,"%s%d",tname,foonet.if_unit);
#endif

	while(in_dev != NULL && strncmp(in_dev,name,strlen(in_dev)) != 0) {
		ifnetaddr = (u_long)TAILQ_NEXT(&foonet,if_link);
		if(ifnetaddr < 1) {
			fprintf(stderr,"error, interface not found\n");
			exit(1);
		}
		if(kread(ifnetaddr,(char *)&foonet,sizeof foonet))
			return 1;

#if __FreeBSD_version >= 501113
		strncpy(name,foonet.if_xname,sizeof(name));
#else
		if(kread((u_long)foonet.if_name, tname, 16))
			return 1;
		snprintf(name,32,"%s%d",tname,foonet.if_unit);
#endif
			
	}

	/* start curses */
	initscr();
	/* start color setup */
	if(has_colors() != TRUE) {
		color_in = 0;
		color_out = 0;
	}
	if(color_in > 0) {
		start_color();
		use_default_colors();
		if(color_out == 0)
			color_out = color_in;
		
		if(strncmp(color_in,"black",5) == 0)
			init_pair(1, COLOR_BLACK,-1);
		else if(strncmp(color_in,"red",3) == 0)
			init_pair(1, COLOR_RED,-1);
		else if(strncmp(color_in,"green",5) == 0)
			init_pair(1, COLOR_GREEN,-1);
		else if(strncmp(color_in,"yellow",6) == 0)
			init_pair(1, COLOR_YELLOW,-1);
		else if(strncmp(color_in,"blue",4) == 0)
			init_pair(1, COLOR_BLUE,-1);
		else if(strncmp(color_in,"megenta",7) == 0)
			init_pair(1, COLOR_MAGENTA,-1);
		else if(strncmp(color_in,"cyan",4) == 0)
			init_pair(1, COLOR_CYAN,-1);
		else if(strncmp(color_in,"white",5) == 0)
			init_pair(1, COLOR_WHITE,-1);
	}
	if(color_out > 0) {
		if(strncmp(color_out,"black",5) == 0)
			init_pair(2, COLOR_BLACK,-1);
		else if(strncmp(color_out,"red",3) == 0)
			init_pair(2, COLOR_RED,-1);
		else if(strncmp(color_out,"green",5) == 0)
			init_pair(2, COLOR_GREEN,-1);
		else if(strncmp(color_out,"yellow",6) == 0)
			init_pair(2, COLOR_YELLOW,-1);
		else if(strncmp(color_out,"blue",4) == 0)
			init_pair(2, COLOR_BLUE,-1);
		else if(strncmp(color_out,"megenta",7) == 0)
			init_pair(2, COLOR_MAGENTA, -1);
		else if(strncmp(color_out,"cyan",4) == 0)
			init_pair(2, COLOR_CYAN,-1);
		else if(strncmp(color_out,"white",5) == 0)
			init_pair(2, COLOR_WHITE,-1);
	}
	/* screen init.. */
	screen_init();
	
	lasti = foonet.if_ibytes;
	lasto = foonet.if_obytes;
	gettimeofday(&last,NULL);
	
	for(i=0;i<MAX_G;i++) {
		logi[i] = NULL;
		logo[i] = NULL;
	}
	
	for(;;) {
		screen_check();
		ch = wgetch(mainw);
		/* quit when we get 'q' */
		if(ch == (int)'q' || ch == (int)' ') {
			
			endwin();
			kvm_close(kvmd);		

			exit(0);
		}
		/* clear screen */
		if(ch == (int)'c') {
            free_logs();
		}
		
		if(kread(ifnetaddr,(char *)&foonet,sizeof foonet))
			return 1;
		
		curi = foonet.if_ibytes - lasti;
		curo = foonet.if_obytes - lasto;
		in_total += curi;
		out_total += curo;
		
		wattron(mainw,A_REVERSE);
		wmove(mainw,1,(winw/2)-4);
		/* speed out */
		print_formatf(mainw,(float)CUROS,disp_format);
		if(disp_format == DISP_BYTE)
			wprintw(mainw,"/s");
		wprintw(mainw,"]");
		
		/* total out */
		wmove(mainw,1,(winw-17));
		wprintw(mainw,"total[");
		print_formatf(mainw,(float)out_total,DISP_BYTE);
		wprintw(mainw,"]");
		
		/* speed in */
		wmove(mainw,maxgh+2,(winw/2)-4);
		print_formatf(mainw,(float)CURIS,disp_format);
		if(disp_format == DISP_BYTE)
			wprintw(mainw,"/s");	
		wprintw(mainw,"]");
		
		/* total in */
		wmove(mainw,maxgh+2,(winw-17));
		wprintw(mainw,"total[");
		print_formatf(mainw,(float)in_total,DISP_BYTE);
		wprintw(mainw,"]");

		wattroff(mainw,A_REVERSE);
		
		gettimeofday(&now,NULL);
		
		wmove(mainw,winh,0);
		wrefresh(mainw);
		
		if(in_max/maxgh > 0)
			x = (curi/(in_max/maxgh)) / timea.tv_sec;
		else
			x = curi / timea.tv_sec;
		wmove(mainw,0,20);
		wrefresh(mainw);
		
		werase(inw);
		werase(outw);
		
		for(i=0;i<winw;i++) {
			if(logi[i] == NULL) {
				logi[i] = (unsigned int *)malloc(sizeof(unsigned int));
				if(logi[i] == NULL) {
					fprintf(stderr,"error,allocating memory\n");
					kvm_close(kvmd);
					endwin();
					exit(1);
				}
				*logi[i] = 0;
				break;
			}

		}
		for(i=0;i<winw;i++) {		
			if(logo[i] == NULL) {
				logo[i] = (unsigned int *)malloc(sizeof(unsigned int));
				if(logo[i] == NULL) {
					fprintf(stderr,"error,allocating memory\n");
					kvm_close(kvmd);
					endwin();
					exit(1);
				}
				*logo[i] = 0;
				break;
			}				
		}
		
		if(logi[1] != NULL && logo[1] != NULL) {
			for(i=winw-1;i>=1;i--) {
				if(logi[i] != NULL)
					*logi[i] = *logi[i-1];
			}
			for(i=winw-1;i>=1;i--) {
				if(logo[i] != NULL)
					*logo[i] = *logo[i-1];
			}			
		}
				
		*logi[0] = x;
		for(j=0;j<winw;j++) {
			if(logi[j] == NULL) {
				break;
			}
			if(logi[j] == 0)
				continue;
			if(*logi[j] > maxgh)
				*logi[j] = maxgh;
			for(i=0;i<*logi[j];i++) {
				wmove(inw,maxgh-i-1,j);
				wattron(inw,A_BOLD);
				wattron(inw,COLOR_PAIR(1));
				waddch(inw,'*');
				wattroff(inw,COLOR_PAIR(1));
				wattroff(inw,A_BOLD);
			}
		}
		wrefresh(inw);
		if(out_max/maxgh > 0)
			x = (curo/(out_max/maxgh)) / timea.tv_sec;
		else
			x = curo / timea.tv_sec;
		*logo[0] = x;
		for(j=0;j<winw;j++) {
			if(logo[j] == NULL) {
				break;
			}
			if(logo[j] == 0)
				continue;
			if(*logo[j] > maxgh)
				*logo[j] = maxgh;
			for(i=0;i<*logo[j];i++) {
				wmove(outw,maxgh-i-1,j);
				wattron(outw,A_BOLD);
				wattron(outw,COLOR_PAIR(2));
				waddch(outw,'*');
				wattroff(outw,COLOR_PAIR(2));
				wattroff(outw,A_BOLD);
			}
		}
		wrefresh(outw);
		
		gettimeofday(&last,NULL);
		select(0,NULL,NULL,NULL,&timea);
		lasti = foonet.if_ibytes;
		lasto = foonet.if_obytes;
		
	}
}