#!/usr/bin/perl -w
#
# monthstats -- Retrieve and tally network usage statistics on a per-host
# 				basis over the current month.
#				Designed for use with the databases created 
#				and filled by the trafstats program; 
#				No guarantees of functionality are 
#				made when used with other programs.



# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

 
use strict;
use DBI;
use Pg;
use Socket;

my (%incoming,%outgoing);
my @allkeys;
my ($dbname,$dbuser);
my $connection;
my $timestamp_start;
my $timestamp_end;
my $result;
my @watchnets;
my $treshold;
my $net;
my $query;
my $month=undef;
my $netsfile;
my @monthnames=();
my @tmp=();
my $cmonth;
my $year;

#TODO: Read config file for username and database.
$dbname="trafstats";
$dbuser="accountreader";
$netsfile="/etc/trafstats/watchnets";

@monthnames= qw(January		February	March		April
		May		June		July		August
		September	October		November	December);

#@watchnets=("194.229.164.0/23","213.53.202.0/24");
$treshold=10485760; # that's 10 megabytes, for the non-binary inclined.

my $outfile=$ARGV[1];
my $line="+-----------------------------------+------------+------------+------------+\n";
# A sub
sub inmegs {
	my $value=shift;
	$value= int ($value / 1048576);         # Shadur. No. BAKA.
	return "$value Mb" unless $value == 0;
	return "";
}

sub testok {
        my ($host,$mask)=split /\//;

        return undef unless (inet_aton ($host));
        # $host is okay. Now for the mask part (if it exists)

        return 1 unless $mask; # No netmask; we're okay with that.

        if($mask =~ /^\d+$/) {
                return undef if $mask > 32;
        } else {
                return inet_aton($mask);
        }

        return 1;
}

sub readnetworks {
	my $netconffile=shift;
	my @nets=();
	open(WATCHNETS,"$netconffile") 
		or die "Can't open config file $netconffile for reading: $!\n";

	while(<WATCHNETS>) {
		chomp;
		next if /^\#/;	# Ignore comment lines
		next if /^\s*$/;	# ignore whitespace lines
		
		my $candidate=(split)[0];
		if (testok $candidate) {
			push(@nets,$candidate);
		}
	}
	return @nets;
}

# Main program start

#First off, let's see if the first arg (if there is one) is a month.
@tmp=localtime;

$cmonth=$tmp[4];
$year=$tmp[5]+1900;

if(@ARGV) {

	my $i;
	foreach $i (0..(scalar @monthnames)) {
		if ($ARGV[0] =~ m/$monthnames[$i]/i) {
			$month=$i;
			last;
		}
	}
	die qq("$ARGV[0]" is not a valid month.\n) unless defined $month;
	$year-- if($cmonth < $month); # Last year.
} else { 
	$month=$cmonth;
}

# Second off, get the list of networks to watch.
@watchnets=readnetworks($netsfile);

# Third, establish the connection
$connection=DBI->connect("dbi:Pg:dbname=$dbname;","$dbuser")
	or die "Connection to database $dbname failed: $DBI::errstr\n";

# Connection ok.

# Step 4: Get the start point.
$query=qq(SELECT date_trunc('month',timestamp '$monthnames[$month] 1, $year') AS timestamp);
my $sth=$connection->prepare($query) 
	or die "Query preparation failed: $connection->errstr\n";
$result=$sth->execute()
	or die "Query execution failed: $sth->errstr\n";

my @row=$sth->fetchrow_array;
$timestamp_start=$row[0];

$query=qq(SELECT date_trunc('month',timestamp '$monthnames[$month] 1, $year' + interval '1 month') AS timestamp);
$sth=$connection->prepare($query) 
	or die "Query preparation failed: $connection->errstr\n";
$result=$sth->execute()
	or die "Query execution failed: $sth->errstr\n";

@row=$sth->fetchrow_array;
$timestamp_end=$row[0];

# Start building the query for incoming.

$query ="SELECT dest_ip AS server, SUM(amount) AS incoming ";
$query.="FROM rawtraffic ";
$query.="WHERE timestamp > timestamp '$timestamp_start' ";
$query.="AND timestamp < timestamp '$timestamp_end' ";
# Do the subnets
$query.="AND (";
foreach $net (@watchnets) {
	$query .= "dest_ip << inet'$net' OR ";
}
# Have to do this. Should get optimized away anyways.
$query.="FALSE) ";

$query.="GROUP BY dest_ip ";
$query.="HAVING sum(amount) > $treshold ";

#print "Query: <$query>\n";

$sth=$connection->prepare($query)
	or die "Bad query <$query>: \n<$connection->errstr>\n";
$result=$sth->execute() 
	or die "Execution of <$query> failed:\n <$connection->errstr>\n";

my ($server,$incoming,$outgoing);
while(($server,$incoming)=$sth->fetchrow_array) {
	$incoming{$server}=$incoming;
	$incoming{"total"}+=$incoming;
}

@allkeys = keys(%incoming);
my $count=scalar @allkeys;
#print "$count entries in incoming.";

$query ="SELECT source_ip AS server, SUM(amount) AS outgoing ";
$query.="FROM rawtraffic ";
$query.="WHERE timestamp > timestamp '$timestamp_start' ";
$query.="AND timestamp < timestamp '$timestamp_end' ";
# Do the subnets
$query.="AND (";
foreach $net (@watchnets) {
	$query .= "source_ip << inet'$net' OR ";
}
$query.="FALSE) ";
$query.="GROUP BY source_ip ";
$query.="HAVING sum(amount) > $treshold ";

#print "Query: <$query>\n";

$sth=$connection->prepare($query)
	or die "Bad query <$query>: \n<$connection->errstr>\n";
$result=$sth->execute() 
	or die "Execution of <$query> failed:\n <$connection->errstr>\n";


while(($server,$outgoing)=$sth->fetchrow_array) {
	$outgoing{$server}=$outgoing;
	$outgoing{"total"}+=$outgoing;
}

push (@allkeys, keys %outgoing);
my %tmp=();
@tmp{@allkeys}=();
@allkeys=sort keys %tmp;

#print scalar(keys %outgoing) . " keys in outgoing\n";
#print scalar(@allkeys) . " keys total.\n";

my $now = scalar localtime;

print <<EOF;
+--------------------------------------------------------------------------+
|Trafstats traffic statistics                                              |
|Monthly stats generated $now                          |
EOF

$- -=3;
my $servname;

foreach $server (@allkeys) {
	next if( $server eq "total");
	$servname=gethostbyaddr(inet_aton($server),AF_INET);
	$servname=$server unless defined $servname;
	$incoming{$server}=0 unless defined $incoming{$server};
	$outgoing{$server}=0 unless defined $outgoing{$server};
	do { print $line ; $- --; } if ($- == 1);
	write;
}

print $line;$- --;
$servname=$server="total";
write;
print $line unless $-==0;

$sth->finish();
$connection->disconnect();

format STDOUT=
| @<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< | @>>>>>>>>> | @>>>>>>>>> | @>>>>>>>>> |
$servname, inmegs($incoming{$server}), inmegs($outgoing{$server}),inmegs ($incoming{$server} + $outgoing{$server})
.

format STDOUT_TOP=
+-----------------------------------+------------+------------+------------+
| Server                            |  Incoming  |  Outgoing  |  Total     |
+-----------------------------------+------------+------------+------------+
.

