/*-
 * Copyright (c) 2005 Fredrik Lindberg. <fli at shapeshifter dot se>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. 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.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``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 AUTHOR 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>
#if __FreeBSD_version >= 500000
# include <stdint.h>
#else
# include <inttypes.h>
#endif
#include <string.h>
#include <unistd.h>
#include <ctype.h>
#include <signal.h>
#include <math.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include "mylibpng/mylibpng.h"
#include "mylibpng/font.h"
#include "bw.h"

#define VERSION "0.2.1"

/* Struct to hold our configuration data */
typedef struct yabm_t {
	char 		*iface;		/* Interface */
	in_addr_t 	alias; 		/* IP-address */
	uint32_t	seconds;	/* Seconds */
	char		*filename;	/* Output filename */
	uint64_t	maxval;		/* Device speed */
	uint32_t	background;	/* Background color */
	uint32_t	ibarcolor;	/* Incoming bar color */
	uint32_t	obarcolor;	/* Outgoing bar color */
	uint16_t	width;		/* Image width */
	uint16_t	height;		/* Image height */
	uint8_t		bflag;		/* Bits or bytes */
	uint8_t		tflag;		/* Include text */
	uint8_t		separators;	/* Separators */
} yabm_t; 

#define DEF_SECONDS	60
#define DEF_BACKGROUND	0xc0c0c0
#define DEF_IBARCOLOR	0x69ff34	
#define DEF_OBARCOLOR	0x0079b7
#define DEF_FILENAME	"bw.png"
#define DEF_WIDTH	640	
#define DEF_HEIGHT	11
#define DEF_MAXVAL	100000000	/* 100Mbps */
#define DEF_PIDFILE	"/var/run/yabm.pid"
#define DEF_SEPARATORS	3	

#define SUFFIX_KILO	'k'
#define SUFFIX_MEGA	'm'
#define SUFFIX_GIGA	'g'

#define MARKERS		4

#define BFLAG_BITS	0
#define BFLAG_BYTES	1	

static void usage(char *);
static void main_loop(yabm_t *);

static char running = 1;
static void
reset_running()
{

        running = 0;
}

static void
do_setuid(uid_t euid, uid_t ruid)
{

	setreuid(ruid, euid);
}

static void
undo_setuid(uid_t euid, uid_t ruid)
{

	setreuid(ruid, euid);
}

static void
get_uids(uid_t *euid, uid_t *ruid)
{

	*euid = geteuid();
	*ruid = getuid();
}

int
main(int argc, char *argv[])
{
	unsigned char suffix;
	char 	*p;
	int	ch, fd;
	yabm_t	cfg;
	uid_t	euid, ruid;
	pid_t	pid;
	FILE 	*fp;
	char	*pidfile;

	get_uids(&euid, &ruid);
	undo_setuid(euid, ruid);

	if (argc < 2) {
		usage(argv[0]);
		exit(EXIT_FAILURE);
	}

	/* Setup default values */
	cfg.iface 	= NULL;
	cfg.alias 	= 0;
	cfg.seconds 	= DEF_SECONDS;
	cfg.filename 	= DEF_FILENAME; 
	cfg.maxval 	= 0;
	cfg.background 	= DEF_BACKGROUND; 
	cfg.ibarcolor 	= DEF_IBARCOLOR; 
	cfg.obarcolor 	= DEF_OBARCOLOR; 
	cfg.width 	= DEF_WIDTH;
	cfg.height 	= DEF_HEIGHT;
	cfg.bflag 	= BFLAG_BITS;
	cfg.tflag 	= 0;
	cfg.separators	= DEF_SEPARATORS;
	pidfile		= DEF_PIDFILE;

	/* Parse command line arguments */
	while ((ch = getopt(argc, argv, "i:a:s:o:m:b:f:h:w:n:u:ytp:k:")) != -1) {
		switch (ch) {
			case 'i':
				if (cfg.alias == 0)
					cfg.iface = argv[optind - 1];
				else {
					fprintf(stderr, "Interface conflicts with alias, can't do both\n\n");
					usage(argv[0]);
					exit(EXIT_FAILURE);
				}
				break;
			case 'a':
				if (cfg.iface == NULL)
					cfg.alias = inet_addr(argv[optind - 1]);
				else {
					fprintf(stderr, "Alias conflicts with interface, can't do both\n\n");
					usage(argv[0]);
					exit(EXIT_FAILURE);
				}	
				break;
			case 's':
				cfg.seconds = (uint32_t)strtol(argv[optind - 1], NULL, 10);
				if (cfg.seconds <= 0) cfg.seconds = DEF_SECONDS;
				break;
			case 'o':
				cfg.filename = argv[optind - 1];
				break;
			case 'm':
				p = argv[optind - 1];
				suffix = *(p + strlen(p) - 1); 
				
				if (!isdigit(suffix)) 
					*(p + strlen(p) - 1) = '\0';

				cfg.maxval = (uint32_t)strtol(argv[optind - 1], NULL, 10);

				if (suffix == SUFFIX_KILO || (suffix + 32) == SUFFIX_KILO) 
					cfg.maxval *= 1000;
				else if (suffix == SUFFIX_MEGA || (suffix + 32) == SUFFIX_MEGA) 
					cfg.maxval *= 1000000;
				else if (suffix == SUFFIX_GIGA || (suffix + 32) == SUFFIX_GIGA) 
					cfg.maxval *= 1000000000;

				break;
			case 'b':
				cfg.background = (uint32_t)strtol(argv[optind - 1], NULL, 16);
				break;
			case 'n':
				cfg.ibarcolor = (uint32_t)strtol(argv[optind - 1], NULL, 16);
				break;
			case 'u':
				cfg.obarcolor = (uint32_t)strtol(argv[optind - 1], NULL, 16);
				break;
			case 'h':
				cfg.height = (uint16_t)strtol(argv[optind - 1], NULL, 10);
				if (cfg.height == 0) cfg.height = DEF_HEIGHT;
				break;
			case 'w':
				cfg.width = (uint16_t)strtol(argv[optind - 1], NULL, 10);
				if (cfg.width == 0) cfg.width = DEF_WIDTH;
				break;
			case 'y':
				cfg.bflag = BFLAG_BYTES;	
				break;
			case 't':
				cfg.tflag++;
				break;
			case 'p':
				pidfile = argv[optind - 1];
				break;	
			case 'k':
				cfg.separators = (uint8_t)strtol(argv[optind - 1], NULL, 10);
				break;
			default:
				usage(argv[0]);
				exit(EXIT_FAILURE);
		}
	}

	/* Did we get an interface/alias ? */
	if (cfg.iface == NULL && cfg.alias == 0) {
		fprintf(stderr, "You must specify either an interface or an IP-address\n\n");
		usage(argv[0]);
		exit(EXIT_FAILURE);
	}

	/* Ok, we got some interface and alias data from the user. How accurate is it? */
	if (cfg.iface != NULL) {
		if (!bw_exists(cfg.iface, TYPE_IF)) {
			fprintf(stderr, "Unable to use interface %s, %s\n",
				cfg.iface, strerror(errno));
			exit(EXIT_FAILURE);
		}
	}
	else if (cfg.alias != 0) {
		uint8_t exists;
		struct in_addr inaddr;
		inaddr.s_addr = cfg.alias;

		do_setuid(euid, ruid);
		exists = bw_exists((in_addr_t *)cfg.alias, TYPE_ADR);
		undo_setuid(euid, ruid);

		if (!exists) {
			fprintf(stderr, "Unable to use IP-address %s, %s\n",
				inet_ntoa(inaddr), strerror(errno));
			exit(EXIT_FAILURE);
		}	
	}

	/* Do we have write-permissions to cfg.filename ? */
	if ((fd = open(cfg.filename, O_WRONLY)) == -1 && errno != ENOENT) {
		fprintf(stderr, "Can't access %s, %s\n", cfg.filename, strerror(errno));	
		exit(EXIT_FAILURE);
	}
	close(fd);
	
	/* No max speed was given, try to read it */
	if (cfg.maxval == 0) {
		do_setuid(euid, ruid);

		/* We need uid 0 to do this */
		if (geteuid() == 0) {
			if (cfg.iface != NULL)
				cfg.maxval = bw_maxspeed(cfg.iface, TYPE_IF);
			else
				cfg.maxval = bw_maxspeed(&cfg.alias, TYPE_ADR);
		}
		else {
			cfg.maxval = DEF_MAXVAL;
			fprintf(stderr, "No permissions to read max speed from device, using %llu\n", cfg.maxval);
		}
		undo_setuid(euid, ruid);
	}


#ifndef DEBUG
	if (fork() == 0) 
#endif /* DEBUG */
	{
		/* Try to write to pidfile, just ignore it if we failed */
		pid = getpid();

		do_setuid(euid, ruid);
		fp = fopen(pidfile, "w");
		undo_setuid(euid, ruid);

		if (fp != NULL) {
			fprintf(fp, "%d", pid);
			fclose(fp);
		} 

		/* Jump into main loop */
		main_loop(&cfg);

		/* Truncate file on exit */
		do_setuid(euid, ruid);
		fp = fopen(pidfile, "w");
		if (fp != NULL) fclose(fp);
		undo_setuid(euid, ruid);

	}

	return 0;
}

static void
usage(char *exec)
{

	printf("yabm %s (Yet Another Bandwidth Meter)\n", VERSION);
	printf("Usage %s [-i iface|-a alias] [options]\n", exec);
	printf(" -i interface \t Interface-based measurement\n");
	printf(" -a ipnumber \t IP-address based measurement (only AF_INET)\n");
	printf("Options\n");
	printf(" -s seconds \t Calculate average bandwidth over this many seconds (Default %d)\n", DEF_SECONDS);
	printf("            \t 60 secs = 1 min/avg, 1 sec = realtime, 300 secs = 5 min/avg\n"); 
	printf(" -o file \t Output filename (and path) (Default %s)\n", DEF_FILENAME);
	printf(" -m bits[kmg] \t Device max speed (eg. 512k, 100m, 1g, etc) (Default ask device)\n"); 
	printf(" -b hex \t Background color (Default %.6x)\n", DEF_BACKGROUND);
	printf(" -n hex \t Incoming bar color (Default %.6x)\n", DEF_IBARCOLOR);
	printf(" -u hex \t Outgoing bar color (Default %.6x)\n", DEF_OBARCOLOR);
	printf(" -w pixels \t Image width (Default %d)\n", DEF_WIDTH);
	printf(" -h pixels \t Image height (Default %d)\n", DEF_HEIGHT);
	printf(" -t \t\t Include text display of bandwidth utilization\n");
	printf(" -y \t\t Display bytes instead of bits\n");
	printf(" -k number \t How many markers to draw (Default %d)\n", DEF_SEPARATORS);
	printf(" -p file \t Pidfile (Default %s)\n", DEF_PIDFILE);

}

/*
 * Main loop, brings the bits and pieces together
 */
static void
main_loop(yabm_t *cfg)
{
	pngcolor8	bgcol, ocol, icol, ctmp; 
	pngimage	img;
	uid_t		euid, ruid;
	int		i;
	uint16_t	tmp; 
	uint64_t	bw[2] = {0}, maxval_bytes = cfg->maxval/8;
	uint32_t	i_barlen, o_barlen;
	double		bw_in, bw_out;

	/* 
	 * Some devices does apparently not work with kvm
	 * We downgrade them in silence to use bw_sysctl, this will only 
	 * work for interfaces.
	 * If working_kvm has changed from 0 bw_kvm has failed.
	 */
	char working_kvm = 0;

	get_uids(&euid, &ruid);

	/* Transform 32bit color values into pngcolor8 types */
	bgcol.r = (cfg->background & ~0x00ffff) >> 16;
	bgcol.g = (cfg->background & ~0xff00ff) >> 8;
	bgcol.b = (cfg->background & ~0xffff00);

	icol.r = (cfg->ibarcolor & ~0x00ffff) >> 16;
	icol.g = (cfg->ibarcolor & ~0xff00ff) >> 8;
	icol.b = (cfg->ibarcolor & ~0xffff00);
	icol.opacity = 0.7;

	ocol.r = (cfg->obarcolor & ~0x00ffff) >> 16;
	ocol.g = (cfg->obarcolor & ~0xff00ff) >> 8;
	ocol.b = (cfg->obarcolor & ~0xffff00);
	ocol.opacity = 0.5;

	/* Create our drawing area */
	img = mypng_new(cfg->width, cfg->height, 8);

        /* Setup signal handler to terminate loop */
        signal(SIGHUP,  reset_running);
        signal(SIGINT,  reset_running);
        signal(SIGKILL, reset_running);
        signal(SIGTERM, reset_running);
        signal(SIGQUIT, reset_running);

	while (running) {

		/* Fill the drawing area with black */
		ctmp.r = ctmp.g = ctmp.b = 0;
		mypng_fill8(img, &ctmp); 

		/* Draw our real background */
		mypng_box8(img, 1, 1, cfg->width - 1, cfg->height - 1, &bgcol);

		/* Draw some white "markers" */
		ctmp.r = ctmp.g = ctmp.b = 0xff;
		for (i = 1; i < (cfg->separators + 1); i++) {
			uint32_t j;
	
			tmp = (((cfg->width - 4) / (cfg->separators + 1)) * i) + i;
			for (j = 1; j < (cfg->height - 1); j++) {
				mypng_plot8(img, tmp, j, &ctmp);	
			}
		}

		/* Start the measure, try to increase our privileges */
		do_setuid(euid, ruid);

		/* We got uid 0 and kvm seems to work, go for it */
		if (geteuid() == 0 && working_kvm == 0) {
			if (cfg->iface != NULL)
				working_kvm = bw_kvm(cfg->iface, TYPE_IF, bw, cfg->seconds);
			else if (cfg->alias != 0) {
				working_kvm = bw_kvm((in_addr_t *)cfg->alias, TYPE_ADR, bw, cfg->seconds);
			}
		}
		else {
			if (cfg->iface != NULL) {
				bw_sysctl(cfg->iface, bw, cfg->seconds);
			}
			else if (working_kvm < 0) {
				fprintf(stderr, "kvm_read for your device failed, sorry there is nothing I can do\n");
				running = 0;
			}
			else {
				fprintf(stderr, "You need uid 0 to do ip-address based measurements\n");
				running = 0;
			}
		}
		/* Drop privileges */
		undo_setuid(euid, ruid);

		/* Calculate average bandwidth */ 
		bw_in  = bw[0] / (double)cfg->seconds;
		bw_out = bw[1] / (double)cfg->seconds;

		/* Calculate the length of the bars */
		i_barlen = (bw_in/maxval_bytes) * (cfg->width - 2);
		o_barlen = (bw_out/maxval_bytes) * (cfg->width - 2);

		/* Sanity checking to be on the safe side */ 
		if (i_barlen > (cfg->width - 2)) i_barlen = cfg->width - 2;
		if (o_barlen > (cfg->width - 2)) o_barlen = cfg->width - 2;

		/* Draw it */
		mypng_box8_opacity(img, 2, 2, i_barlen, cfg->height - 2, &icol);
		mypng_box8_opacity(img, 2, 2, o_barlen, cfg->height - 2, &ocol);

		/* Do we want text aswell? */
		if (cfg->tflag) {
			char *bw_out_str, *bw_in_str;

			ctmp.r = ctmp.g = ctmp.b = 0;
			ctmp.opacity = 0.4;

			/* The system gives us statistics in octets, recalculate to bits if needed */
			if (cfg->bflag == BFLAG_BITS) {
				bw_in *= 8;
				bw_out *= 8;

				bw_in_str = bw_in > 1000 ? (bw_in > 1000000 ? (bw_in > 1000000000 ? "Gbits" : "Mbits") : "Kbits") : "bits";
				bw_in = bw_in > 1000 ? (bw_in > 1000000 ? (bw_in > 1000000000 ? 
				   bw_in/1000000000 : bw_in/1000000) : bw_in/1000) : bw_in;

				bw_out_str = bw_out > 1000 ? (bw_out > 1000000 ? (bw_out > 1000000000 ? "Gbit" : "Mbits") : "Kbits") : "bits";
				bw_out = bw_out > 1000 ? (bw_out > 1000000 ? (bw_out > 1000000000 ? 
				   bw_out/1000000000 : bw_out/1000000) : bw_out/1000) : bw_out;
	
			}
			else {
				bw_in_str = bw_in > 1024 ? (bw_in > 1048576 ? (bw_in > 1073741824 ? "GiB" : "MiB") : "KiB") : "B";
				bw_in = bw_in > 1024 ? (bw_in > 1048576 ? (bw_in > 1073741824 ? 
				   bw_in/1073741824 : bw_in/1048576) : bw_in/1024) : bw_in;

				bw_out_str = bw_out > 1024 ? (bw_out > 1048576 ? (bw_out > 1073741824 ? "GiB" : "MiB") : "KiB") : "B";
				bw_out = bw_out > 1024 ? (bw_out > 1048576 ? (bw_out > 1073741824 ? 
				   bw_out/1073741824 : bw_out/1048576) : bw_out/1024) : bw_out;

			}
	
			mypng_printf(img, 3, (cfg->height - 8) / 2, &ctmp, "in:%0.2f %s/s out:%0.2f %s/s",
				bw_in, bw_in_str, bw_out, bw_out_str);
	
		}

		mypng_write(img,  cfg->filename);
	}

	mypng_destroy(img);
}