#!/usr/bin/perl

# IP Accountant
# pedram amini <pedram@redhive.com, http://pedram.redhive.com>
#

use strict;

# configure these.
my $iptables    = "/usr/sbin/iptables";
my $rrdtool     = "/usr/local/sbin/rrdtool";
my $rrd_db_loc  = "/var/rrd/db2";
my $html_report = "/var/www/htdocs/rrd/stats.html";

my %LOKAL_UP  = ();
my %LOKAL_DOWN = ();
my %INTERNET_UP = ();
my %INTERNET_DOWN = ();

my ($key, $report_key, $report_count);

# we determine what "mode" we're in via the first argument.
#   rrd    - rrd updating mode
#   html   - html reporting mode
my $mode = shift;

# the default is to create a report in human readable form to stdout.
$mode = "default" if (!$mode);

if ($mode eq "default") {
    print "\n[ IP Accountant ]";
    print "\npedram amini";
    print "\npedram\@redhive.com, http://pedram.redhive.com\n";
}

my @lokal_up  = `$iptables -L LOKAL_UP -n -v -x -Z`;
my @lokal_down = `$iptables -L LOKAL_DOWN -n -v -x -Z`;
my @internet_up = `$iptables -L INTERNET_UP -n -v -x -Z`;
my @internet_down = `$iptables -L INTERNET_DOWN -n -v -x -Z`;

#my @iptables_output  = `$iptables -vxnL FORWARD`;
#my @iptables_input = `$iptables -vxnL INPUT`;
#my @iptables_output = `$iptables -vxnL OUTPUT`;


################################################################################
# load the /etc/hosts table
# 
my %hosts = ();
open (HOSTS, "/etc/hosts") || die "error opening /etc/hosts\n";

while (<HOSTS>) {
    chomp;      # remove newline.
    s/#.*//;    # remove comments.
    s/\s+/ /g;  # compress whitespace.
    s/^\s+//;   # remove leading whitespace.
    s/\s+$//;   # remove trailing whitespace.
    
    next if ( !$_ );
    
    # 0          1        2
    # ip_address hostname alias
    my @temp = split / /;
    $hosts{$temp[0]} = $temp[1];
}


################################################################################
# procesiranje LOKAL_UP-a
# 
# we start at 2 because the first 2 lines are part of the table header.
for (my $i = 2; $i <= $#lokal_up; $i++) {
    $_ = $lokal_up[$i];
    
    s/\s+/ /g;  # compress whitespace.
    s/^\s+//;   # remove leading whitespace.
    s/\s+$//;   # remove trailing whitespace.
    
    # 0    1     2      3    4   5  6   7      8
    # pkts bytes target prot opt in out source destination
    my @temp = split / /;
    if ( $hosts{$temp[7]} ){                    
    	    $LOKAL_UP{$hosts{$temp[7]}} = $temp[1];
    }
}


################################################################################
# procesiranje LOKAL_DOWN-a
# 
# we start at 2 because the first 2 lines are part of the table header.
for (my $i = 2; $i <= $#lokal_down; $i++) {
    $_ = $lokal_down[$i];
    
    s/\s+/ /g;  # compress whitespace.
    s/^\s+//;   # remove leading whitespace.
    s/\s+$//;   # remove trailing whitespace.
    
    # 0    1     2      3    4   5  6   7      8
    # pkts bytes target prot opt in out source destination
    my @temp = split / /;
    if ( $hosts{$temp[8]} ){                    
    	    $LOKAL_DOWN{$hosts{$temp[8]}} = $temp[1];
    }
}

################################################################################
# procesiranje INTERNET_UP-a
# 
# we start at 2 because the first 2 lines are part of the table header.
for (my $i = 2; $i <= $#internet_up; $i++) {
    $_ = $internet_up[$i];
    
    s/\s+/ /g;  # compress whitespace.
    s/^\s+//;   # remove leading whitespace.
    s/\s+$//;   # remove trailing whitespace.
    
    # 0    1     2      3    4   5  6   7      8
    # pkts bytes target prot opt in out source destination
    my @temp = split / /;
    if ( $hosts{$temp[7]} ){                    
    	    $INTERNET_UP{$hosts{$temp[7]}} = $temp[1];
    }
}

################################################################################
# procesiranje INTERNET_DOWN-a
# 
# we start at 2 because the first 2 lines are part of the table header.
for (my $i = 2; $i <= $#internet_down; $i++) {
    $_ = $internet_down[$i];
    
    s/\s+/ /g;  # compress whitespace.
    s/^\s+//;   # remove leading whitespace.
    s/\s+$//;   # remove trailing whitespace.
    
    # 0    1     2      3    4   5  6   7      8
    # pkts bytes target prot opt in out source destination
    my @temp = split / /;
    if ( $hosts{$temp[8]} ){                    
    	    $INTERNET_DOWN{$hosts{$temp[8]}} = $temp[1];
    }
}

################################################################################
# Pravljenje zbirnog stanja
################################################################################

foreach $key(sort keys %LOKAL_UP) {
	$LOKAL_UP{"_ukupno_.karaboorma.net"} += $LOKAL_UP{$key};
	$LOKAL_DOWN{"_ukupno_.karaboorma.net"} += $LOKAL_DOWN{$key};
	$INTERNET_UP{"_ukupno_.karaboorma.net"} += $INTERNET_UP{$key};
	$INTERNET_DOWN{"_ukupno_.karaboorma.net"} += $INTERNET_DOWN{$key};
}

################################################################################
# rrd update mode: update the rrd databases
#
if ($mode eq "rrd") {
    # the keys for INPUT and OUTPUT should be the same so we just step through
    # one of them. (DS order is rx:tx)
    foreach $key (sort keys %LOKAL_UP) {
        `$rrdtool update $rrd_db_loc/$key.rrd N:$LOKAL_UP{$key}:$LOKAL_DOWN{$key}:$INTERNET_UP{$key}:$INTERNET_DOWN{$key}`;
    }
}


################################################################################
# html reporting mode: generate an html report to the specified location
#
if ($mode eq "html") {
    open (HTML_REPORT, ">$html_report") || die "error opening $html_report\n";
    
#    my $total_up_lokal   = 0;
#    my $total_down_lokal  = 0;
#    my $total_lokal = 0;
#    my $total_up_net = 0;
#    my $total_down_net = 0;
#    my $total_net = 0;
    
    # tally the totals. the keys for INPUT and OUTPUT should be the same so we 
    # just step through one of them.
#    foreach $key (sort keys %LOKAL_UP) {
#       $total_up_lokal  += $LOKAL_UP{$key};
#    	$total_down_lokal  += $LOKAL_DOWN{$key};
#       $total_up_net  += $INTERNET_UP{$key};
#       $total_down_net  += $INTERNET_DOWN{$key};
#    }
    
#    $total_lokal = $total_up_lokal + $total_down_lokal;
#    $total_net = $total_up_net + $total_down_net;
    my $DATE = `date`;
#    
    print HTML_REPORT <<END_OF_HEADER;
    <html>
    <head><title>Statistike prometa</title></head>
    <body>
    	<center><h1>Statistika saobracaja</h1>
	<h2>Vreme generisanja: $DATE</h2>
	<h2>Obavezno pogledajte i <b><a href=pravilnik.html>pravilnik ponasanja na mrezi</a></b><h2>
	<h2>Imamo i neka <b><a href=obavestenja.html>obavestenja</a></b>.<h2>
	<hr>
END_OF_HEADER
#        <table border=1 cellpadding=1 cellspacing=0><tr><td bgcolor="#000000">
#        <table border=1 cellpadding=10 cellspacing=1 bgcolor="#FFFFFF">
#            <tr>
#                <td><font face="arial" size=2><b>domain</b></font></td>
#                <td align="center" colspan=2><font face="arial" size=2><b>Download</b></font></td>
#                <td align="center" colspan=2><font face="arial" size=2><b>Upload</b></font></td>
#                <td align="center" colspan=2><font face="arial" size=2><b>Ukupno</b></font></td>
#            </tr>
#            <tr>
#                <td><font face="arial" size=2><b>&nbsp;</b></font></td>
#                <td><font face="arial" size=2><b>ukupno bajta</b></font></td>
#                <td><font face="arial" size=2><b>% downloada</b></font></td>
#                <td><font face="arial" size=2><b>ukupno bajta</b></font></td>
#                <td><font face="arial" size=2><b>% uploada</b></font></td>
#                <td><font face="arial" size=2><b>ukupno bajta</b></font></td>
#                <td><font face="arial" size=2><b>ukupno %</b></font></td>
#            </tr>
#END_OF_HEADER

#    foreach $key (sort keys %INPUT) {
#        my $in_bytes    = convert_bytes($INPUT{$key});
#        my $in_percent  = sprintf("%.2f", $INPUT{$key} / $total_input * 200);
#        
#        my $out_bytes   = convert_bytes($OUTPUT{$key});
#        my $out_percent = sprintf("%.2f", $OUTPUT{$key} / $total_output * 200);
#        
#        my $tot_bytes   = convert_bytes($INPUT{$key} + $OUTPUT{$key});
#        my $tot_percent = sprintf("%.2f", ($INPUT{$key} + $OUTPUT{$key}) / 
#                                          $total_traffic * 200);
#        print HTML_REPORT <<END_OF_ROW;
#            <tr>
#                <td align="left"><font face="arial" size=2><a href="http://$key">$key</a></font></td>
#                <td align="right"><font face="arial" size=2>$in_bytes</font></td>
#                <td align="right"><font face="arial" size=2>$in_percent%</font></td>
#                <td align="right"><font face="arial" size=2>$out_bytes</font></td>
#                <td align="right"><font face="arial" size=2>$out_percent%</font></td>
#                <td align="right"><font face="arial" size=2>$tot_bytes</font></td>
#                <td align="right"><font face="arial" size=2>$tot_percent%</font></td>
#            </tr>
#END_OF_ROW
#    }

#    $total_input   = convert_bytes($total_input/2);
#    $total_output  = convert_bytes($total_output/2);
#    $total_traffic = convert_bytes($total_traffic/2);
    
#    print HTML_REPORT <<END_OF_FOOTER;
#            <tr>
#                <td align="left"><font face="arial" size=2><b>Ukupno</b></font></td>
#                <td align="right"><font face="arial" size=2><b>$total_input</b></font></td>
#                <td align="right"><font face="arial" size=2><b>&nbsp;</b></font></td>
#                <td align="right"><font face="arial" size=2><b>$total_output</b></font></td>
#                <td align="right"><font face="arial" size=2><b>&nbsp;</b></font></td>
#                <td align="right"><font face="arial" size=2><b>$total_traffic</b></font></td>
#                <td align="right"><font face="arial" size=2><b>&nbsp;</b></font></td> 
#            </tr>
#        </table>
#        </td></tr></table><hr>
#END_OF_FOOTER
	foreach $key (sort keys %LOKAL_UP) {
		print HTML_REPORT <<END_OF_HTML;
		<p><img src="$key.gif"><br><hr>
END_OF_HTML
	}
	print HTML_REPORT "</center></body></html>";
    
}


################################################################################
# byte conversion function, converts bytes to a human readable form
# 
sub convert_bytes () {
    my $bytes = shift;
    
    my $kbytes = $bytes / 1024;

    return sprintf("%.2f GB", $kbytes / 1048576) if ($kbytes > 1048576);
    return sprintf("%.2f MB", $kbytes / 1024)    if ($kbytes > 1024);
    return sprintf("%.2f KB", $kbytes);
}
