#!/usr/bin/perl -w
# Draw bandwidth pie graphs.
# Copyright 2002, 2003 Adam Sampson <ats@offog.org>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

use strict;
use GD;

# Set this to 1 to generate fake data (so you can see how the output will
# look).
my $simulate = 0;

my $statfile = 'net_stats';
my $ingraphfile = 'pie_in.png';
my $outgraphfile = 'pie_out.png';
my $ttffont = '/opt/lib/X11/fonts/TTF/luxisr.ttf';
my $width = 120;
my $height = 120;

open IN, "<$statfile" or die "can't open stats file";
my $last_time = int <IN>;
my %last_in = ();
my %last_out = ();
while (<IN>) {
	chomp;
	/^(\w+) (\d+) (\d+)$/ or die "bad line in stats file";
	$last_in{$1} = int $2;
	$last_out{$1} = int $3;
}
close IN;

my $time = time();
my %in = ();
my %out = ();
if ($simulate) {
	foreach ("cartman", "vindaloo", "naria") {
		my $period = $time - $last_time;
		unless (exists $last_in{$_}) {
			$last_in{$_} = $last_out{$_} = 0;
		}
		$in{$_} = $last_in{$_} + int rand(40000) * $period;
		$out{$_} = $last_out{$_} + int rand(10000) * $period;
	}
} else {
	my $current_chain = undef;
	my @lines = `iptables -vnxL`;
	foreach (@lines) {
		chomp;
		s/\s+/ /g;
		s/^ //;
		s/ $//;
		if (/^Chain (\w+) /) {
			$current_chain = $1;
		} elsif (/^$/) {
			# skip
		} elsif (/^pkts/) {
			# skip
		} else {
			my @bits = split / /;
			if ($bits[2] eq "RETURN"
				&& $current_chain =~ /^(.+)_(in|out)$/) {
				if ($2 eq "in") {
					$in{$1} = int $bits[1];
				} else {
					$out{$1} = int $bits[1];
				}
			}
		}
	}
}
my @machs = sort keys %in;

open OUT, ">$statfile" or die "can't open stats file for writing";
print OUT "$time\n";
foreach my $mach (@machs) {
	print OUT "$mach $in{$mach} $out{$mach}\n";
}
close OUT;

my $period = $time - $last_time;
my %delta_in = ();
my %delta_out = ();
foreach my $mach (@machs) {
	unless (defined $last_in{$mach}) {
		$last_in{$mach} = 0;
		$last_out{$mach} = 0;
	}
	$delta_in{$mach} = 1.0 * ($in{$mach} - $last_in{$mach});
	$delta_out{$mach} = 1.0 * ($out{$mach} - $last_out{$mach});
}

sub gd_center {
	my ($im, $size, $x, $y, $string, $color) = @_;
	my @b = GD::Image->stringFT($color, $ttffont, $size, 0, 0, 0, $string);
	my $w = $b[2] - $b[0];
	my $h = $b[3] - $b[5];
	$x -= $b[0];
	$y -= $b[1];
	$im->stringFT($color, $ttffont, $size, 0, $x - ($w / 2), $y + ($h / 2), $string);
}

sub graph {
	my ($machs, $delta, $period, $label, $file) = @_;

	my $im = new GD::Image($width, $height);
	my $white = $im->colorAllocate(255, 255, 255);
	my $black = $im->colorAllocate(0, 0, 0);
	$im->transparent($white);

	my %bother = ();
	my $total = 0;
	my %cols = ();
	my $i = 0;
	foreach my $mach (@$machs) {
		if ($delta->{$mach} > 1024) {
			$i = ($i + 1) % 6;
			my $j = $i + 1;
			$cols{$mach} = $im->colorAllocate(221 + 31 * ($j & 1),
			221 + 31 * (($j >> 1) & 1),
			221 + 31 * (($j >> 2) & 1));
			$total += $delta->{$mach};
			$bother{$mach} = 1;
		} else {
			$bother{$mach} = 0;
		}
	}

	my $cx = $width / 2;
	my $cy = $height / 2;
	my $xradius = ($width / 2) * 1.0;
	my $yradius = ($height / 2) * 1.0;

	my $factor = (2 * 3.1415927) / $total;
	my $pos = 0.0;
	foreach my $mach (@$machs) {
		next unless $bother{$mach};
		my $phi = $pos * $factor;
		my $phi_end = ($pos + $delta->{$mach}) * $factor;
		my $poly = new GD::Polygon;
		$poly->addPt($cx, $cy);
		my $frag = 1.0 / 200;
		for (my $i = $phi; $i < $phi_end; $i += $frag) {
			$poly->addPt($cx + cos($i) * $xradius,
				$cy + sin($i) * $yradius);
		}
		$im->filledPolygon($poly, $cols{$mach});
		$pos += $delta->{$mach};
	}
	$pos = 0.0;
	foreach my $mach (@$machs) {
		next unless $bother{$mach};
		my $phi = ($pos + $delta->{$mach} / 2) * $factor;
		my $x = $cx + cos($phi) * ($xradius * 0.6);
		my $y = $cy + sin($phi) * ($yradius * 0.6);
		&gd_center($im, 8, $x, $y - 5,
			"$mach", $black);
		&gd_center($im, 7, $x, $y + 5,
			sprintf("%.1fK", $delta->{$mach} / ($period * 1024)),
			$black);
		$pos += $delta->{$mach};
	}
	&gd_center($im, 10, $cx, $cy, $label, $black);

	open OUT, ">$file" or die "can't open image file $file";
	print OUT $im->png;
	close OUT;
}

&graph(\@machs, \%delta_in, $period, "in", $ingraphfile);
&graph(\@machs, \%delta_out, $period, "out", $outgraphfile);