#!/usr/bin/perl -Tw
#
#  LMS version 1.11.7 Bastet
#
#  Copyright (C) 2001-2009 LMS Developers
#
#  Please, see the doc/AUTHORS for more information about authors!
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License Version 2 as
#  published by the Free Software Foundation.
#
#  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 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.
#
#  $Id: lms-traffic-logiptables,v 1.37 2009/01/13 07:45:28 alec Exp $

use strict;
use DBI;
use Config::IniFiles;
use Getopt::Long;
use vars qw($configfile $quiet $help $version);
use POSIX qw(strftime);

$ENV{'PATH'}='/sbin:/usr/sbin:/usr/local/sbin:/bin:/usr/bin:/usr/local/bin';

sub u32todotquad($)
{
	my $p = shift @_;
	return sprintf "%d.%d.%d.%d", ($p>>24)&0xff,($p>>16)&0xff, ($p>>8)&0xff,$p&0xff;
}		

my $_version = '1.11.7 Bastet';

my %options = (
	"--config-file|C=s"	=>	\$configfile,
	"--quiet|q"		=>	\$quiet,
	"--help|h"		=>	\$help,
	"--version|v"		=>	\$version
);

Getopt::Long::config("no_ignore_case");
GetOptions(%options);

if($help)
{
	print STDERR <<EOF;
lms-traffic-logiptables, version $_version
(C) 2001-2009 LMS Developers

-C, --config-file=/etc/lms/lms.ini	alternate config file (default: /etc/lms/lms.ini);
-h, --help			print this help and exit;
-v, --version			print version info and exit;
-q, --quiet			suppress any output, except errors;

EOF
	exit 0;
}

if($version)
{
	print STDERR <<EOF;
lms-traffic-logiptables, version $_version
(C) 2001-2009 LMS Developers

EOF
	exit 0;
}

if(!$configfile)
{
	$configfile = "/etc/lms/lms.ini";
}

if(!$quiet)
{
	print STDOUT "lms-traffic-logiptables, version $_version\n";
	print STDOUT "(C) 2001-2009 LMS Developers\n";
	print STDOUT "Using file $configfile as config.\n";
}

if(! -r $configfile)
{
	print STDERR "Fatal error: Unable to read configuration file $configfile, exiting.\n";
	exit 1;
}

my $ini = new Config::IniFiles -file => $configfile;
print @Config::IniFiles::errors;

my $outfile = $ini->val('traffic-logiptables','outfile') || '/etc/rc.d/rc.stat';
my $iptables_binary = $ini->val('traffic-logiptables','iptables_binary') || '/usr/sbin/iptables';
my $wan_interfaces = $ini->val('traffic-logiptables','wan_interfaces') || '';
my $local_ports = $ini->val('traffic-logiptables','local_ports') || '';
my $cuid = $ini->val('traffic-logiptables', 'script_owneruid') || '0';
my $cgid = $ini->val('traffic-logiptables', 'script_ownergid') || '0';
my $cperm = $ini->val('traffic-logiptables', 'script_permission') || '700';
my $networks_list = $ini->val('traffic-logiptables', 'networks') || '';
my $excluded_networks_list = $ini->val('traffic-logiptables', 'excluded_networks') || '';

my $dbtype = $ini->val('database', 'type') || 'mysql';
my $dbhost = $ini->val('database', 'host') || 'localhost';
my $dbuser = $ini->val('database', 'user') || 'root';
my $dbpasswd = $ini->val('database', 'password') || '';
my $dbname = $ini->val('database', 'database') || 'lms';

my $dbase;
my $utsfmt;

if(!$wan_interfaces)
{
	print STDERR "Fatal error: 'wan_interfaces' not defined, exiting.\n";
	exit 1;
}

if($dbtype =~ /mysql/)
{
	$dbase = DBI->connect("DBI:mysql:database=$dbname;host=$dbhost","$dbuser","$dbpasswd", { RaiseError => 1 });
	$utsfmt = "UNIX_TIMESTAMP()";
}
elsif($dbtype eq "postgres")
{
	$dbase = DBI->connect("DBI:Pg:dbname=$dbname;host=$dbhost","$dbuser","$dbpasswd", { RaiseError => 1 });
	$utsfmt = "EXTRACT(EPOCH FROM CURRENT_TIMESTAMP(0))";
}
else
{
	print STDERR "Fatal error: unsupported database type: $dbtype, exiting.\n";
	exit 1;
}

open(OUTFILE, ">$outfile") or die("Fatal error: Unable to write $outfile, exiting.\n");

my @local_ports_list = split(' ',$local_ports);
my @wan_interfaces_list = split(' ',$wan_interfaces);

print OUTFILE "#!/bin/bash\n";

foreach my $wan_interface (@wan_interfaces_list)
{
	print OUTFILE "$iptables_binary -t mangle -D FORWARD -i $wan_interface -j STAT >/dev/null 2>&1\n";
	print OUTFILE "$iptables_binary -t mangle -D FORWARD -o $wan_interface -j STAT >/dev/null 2>&1\n";
}

foreach my $port (@local_ports_list)
{
	print OUTFILE "$iptables_binary -t mangle -D INPUT -p tcp --dport $port -j STAT >/dev/null 2>&1\n";
	print OUTFILE "$iptables_binary -t mangle -D OUTPUT -p tcp --sport $port -j STAT >/dev/null 2>&1\n";
}

# 1) Lets create chain STAT
# 2) Flush it, it could exists
		
print OUTFILE "
$iptables_binary -t mangle -F STAT >/dev/null 2>&1
$iptables_binary -t mangle -X STAT >/dev/null 2>&1
$iptables_binary -t mangle -N STAT
";

foreach my $wan_interface (@wan_interfaces_list)
{
	print OUTFILE "$iptables_binary -t mangle -I FORWARD -i $wan_interface -j STAT\n";
	print OUTFILE "$iptables_binary -t mangle -I FORWARD -o $wan_interface -j STAT\n";
}

foreach my $port (@local_ports_list)
{
	print OUTFILE "$iptables_binary -t mangle -I INPUT -p tcp --dport $port -j STAT\n";
	print OUTFILE "$iptables_binary -t mangle -I OUTPUT -p tcp --sport $port -j STAT\n";
}

my %nodes;
my %down;
my %up;
my $networkswhere = '';
my @networks = split ' ', $networks_list;

if($excluded_networks_list)
{
    my @excluded_networks = split ' ', $excluded_networks_list;
    foreach my $excluded_network (@excluded_networks)
    {
	@networks = grep(!/^$excluded_network$/, @networks);
    }
}

foreach my $net (@networks)
{
        my $dbq = $dbase->prepare("SELECT INET_ATON(mask) AS mask, address FROM networks WHERE UPPER(name) = UPPER(?)");
	$dbq->execute($net);
	while (my $row = $dbq->fetchrow_hashref()) 
	{
		if (!$networkswhere)
		{
	    		$networkswhere = "WHERE $row->{'address'} = (ipaddr & $row->{'mask'})";
		}
		else
		{
	    		$networkswhere = "$networkswhere OR $row->{'address'} = (ipaddr & $row->{'mask'})";
		}
	}
}

my $dbq = $dbase->prepare("SELECT id, INET_NTOA(ipaddr) AS ip FROM nodes $networkswhere ORDER BY ipaddr");
$dbq->execute();
while (my $row = $dbq->fetchrow_hashref())
{
	print OUTFILE "$iptables_binary -t mangle -A STAT -s $row->{'ip'} -j RETURN\n";
	print OUTFILE "$iptables_binary -t mangle -A STAT -d $row->{'ip'} -j RETURN\n";
	$nodes{$row->{'ip'}} = $row->{'id'};
}
$dbq->finish();

close(OUTFILE);

chown $cuid, $cgid, $outfile or print "Warning! Unable to set owner of $outfile to $cuid.$cgid.\n";
chmod oct($cperm), $outfile or print "Warning! Unable to set permission $cperm to $outfile.\n";

my @info = `$iptables_binary -t mangle -L STAT -v -n -x`;

system("$outfile 1>/dev/null 2>&1");

foreach my $line (@info)
{
	my ($ipaddr, $downpkts, $uppkts, $downbytes, $upbytes);

	chomp $line;
	if($line =~ /^[ ]+([0-9]+)[ ]+([0-9]+) .*0.* 0\.0\.0\.0\/0[ ]+([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})/ )
	{
		$line =~ s/^[ ]+([0-9]+)[ ]+([0-9]+) .*0.* 0\.0\.0\.0\/0[ ]+([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})/$1 $2 $3/g;
		($downpkts, $downbytes, $ipaddr) = split ' ',$line;
	}
	elsif($line =~ /^[ ]+([0-9]+)[ ]+([0-9]+) .*0.* ([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})[ ]+0\.0\.0\.0\/0/ )
	{
		$line =~ s/^[ ]+([0-9]+)[ ]+([0-9]+) .*0.* ([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})[ ]+0\.0\.0\.0\/0/$1 $2 $3/g;
		($uppkts, $upbytes, $ipaddr) = split ' ',$line;
	}
	# iptables older than 1.3.7
	elsif($line =~ /^[ ]+([0-9]+)[ ]+([0-9]+) .*all.* 0\.0\.0\.0\/0[ ]+([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})/ )
	{
		$line =~ s/^[ ]+([0-9]+)[ ]+([0-9]+) .*all.* 0\.0\.0\.0\/0[ ]+([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})/$1 $2 $3/g;
		($downpkts, $downbytes, $ipaddr) = split ' ',$line;
	}
	elsif($line =~ /^[ ]+([0-9]+)[ ]+([0-9]+) .*all.* ([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})[ ]+0\.0\.0\.0\/0/ )
	{
		$line =~ s/^[ ]+([0-9]+)[ ]+([0-9]+) .*all.* ([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})[ ]+0\.0\.0\.0\/0/$1 $2 $3/g;
		($uppkts, $upbytes, $ipaddr) = split ' ',$line;
	}
	else
	{
		next;
	}

	if($nodes{$ipaddr})
	{
		$downbytes = $downbytes || 0;
		$upbytes = $upbytes || 0;

		$down{$ipaddr} = $down{$ipaddr} || $downbytes;
		$up{$ipaddr} = $up{$ipaddr} || $upbytes;
	}
}

foreach my $key (keys %down)
{
 	if($up{$key} || $down{$key} ) # don't need zeroes 
	{   
		$dbase->do("INSERT INTO stats (nodeid, dt, download, upload) 
			VALUES ( $nodes{$key} , $utsfmt, $down{$key}, $up{$key})");

		print "IP: $key\tSend: $up{$key}\t Recv: $down{$key}\n" if not $quiet;
	} 
	else
	{
		print "IP: $key\tSkipped - null data\n" if not $quiet;
	}
}

# finally 
$dbase->disconnect();
