#!/usr/bin/perl
#
# Получает подсчитаный трафик от 'trafd' и сохраняет
# его в базе MySQL.
# Вызывается: trafsave.pl --eth=ETH [--oftime=time]
# ETH - имя интерфейса откуда будет считываться трафик
# time - добавить смещение времени (в минутах) при
#			записи данных в базу.
#

use warnings;
use strict;

use DBI;
use Socket;
use Sys::Syslog qw(:standard :macros);
use File::Basename;
use File::Copy;
use Getopt::Long;

# Параметры командной строки.
our $eth = '';
my $help = '';
our $daemon = '';
my $config = '/etc/trafstat/trafstat.conf';
my $oftime = 0;

# Получим параметры командной строки.
GetOptions("eth=s" => \$eth, "help" => \$help, "daemon" => \$daemon,
	"config=s" => \$config, "oftime=i" => \$oftime);
if ($help) {help(); exit;}
if (!$eth) {usage(); exit(1);}

# Используемые константы
my $log = "/var/spool/trafd/trafd.$eth";
my $dump = "/var/tmp/" . basename($0) . ".$$";
my $trafctl = "/usr/sbin/trafctl";
my $traflog = "/usr/sbin/traflog";
my $sleep = 0;
# размер хеша для буферизирования трафика
my $h_size = 10000;

# получим текущую дату
my ($sec,$min,$hour,$day,$mon,$year) = localtime(time + ($oftime * 60));
$year += 1900; $mon++;
# Округлим значение времени до 30 минут. Это делается для того,
# что бы данные в базе были кратны одному часу. Это позволит
# базе данных быстрее работать с таблицей хранящей трафик.
$min = 30;
$sec = 0;
my $date = "$year-$mon-$day $hour:$min:$sec";

# получим пользователя и пароль для доступа к базе
my ($db_name, $db_user, $db_pass) = ("", "", "");
open(CONFIG, $config) or error("Failed open config file: $config");
while (<CONFIG>) {
	next if /^#|^$|^\[/;
	SWITCH: {
		/^db_name\s*=\s*(\w+)/ && do {$db_name = $1; last SWITCH;};
		/^db_user\s*=\s*(\w+)/ && do {$db_user = $1; last SWITCH;};
		/^db_pass\s*=\s*(\w+)/ && do {$db_pass = $1; last SWITCH;};
		/^trafsave_delay\s*=\s*(\d+)/ && do {$sleep = $1; last SWITCH;};
	}
}
close(CONFIG);

# Выгрузим подсчитанный трафик из 'trafd', сделаем задержку и
# подготовимся для обработки выгруженного дампа.
system("$trafctl save $eth 2>/dev/null") and
	error("Failed execute: $trafctl save $eth");
sleep($sleep) if $sleep;
move($log, $dump);
open(TRAFFIC, "$traflog -na -o default -i $dump 2>/dev/null |") or
	error("Failed execute: $traflog -na -o default -i $dump");

# Подготовим подключение к базе для выгрузки в нее
# информации о трафике.
my ($dbh, $sth, $rv, $query);
$dbh = DBI->connect("DBI:mysql:database=$db_name", $db_user, $db_pass,
	{PrintError => 0}) or error("Connect to database failed: $DBI::errstr");

# Получим идентификаторы протоколов из базы
$query = 'select name, id from proto';
$sth = $dbh->prepare($query) or
	error("Failed prepare query '$query': " . $dbh->errstr);
$sth->execute() or
	error("Failed execute query '$query': " . $dbh->errstr);

my %proto;
while (my @val = $sth->fetchrow_array()) {
	$proto{$val[0]} = $val[1];
}

$sth->finish();

# подготовим запрос на вставку данных
$query = 'insert into traffic_cur ' .
	'(date, srcIp, dstIp, dstPort, pId, inBytes, outBytes) ' .
	'values(?, ?, ?, ?, ?, ?, ?) ' .
	'on duplicate key update inBytes = inBytes + ?, outBytes = outBytes + ?';
$sth = $dbh->prepare($query) or
	error("Failed prepare query '$query': " . $dbh->errstr);

# Сохраним все записи из дампа в базе.
LOOP_: {
	my $i = 0;
	my %traffic = ();

	while (<TRAFFIC>) {
		# подготовим строку для разбиения
		chomp;
		s/\bclient\b/-1/go;
		s/\bnone\b/0/go;

		# разобъем полученную строку
		my ($srcIp, $srcPort, $dstIp, $dstPort, $pId, $bytes, $psize) = split;
		my ($inBytes, $outBytes);
		my ($warnMsg, $protoName);

		$srcIp = ip2long($srcIp);
		$dstIp = ip2long($dstIp);
		$protoName = $pId;
		$pId = exists($proto{$pId}) ? $proto{$pId} : $proto{"unkn"};

		# проверим вносимые записи
		if (($srcPort == $dstPort) || (($srcPort != -1) && ($dstPort != -1))) {
			# не удается определить клиента в полученном пакете
			$warnMsg = "Can not determine the client: srcIp = " . long2ip($srcIp) . ", srcPort = $srcPort, " .
				"dstIp = " . long2ip($dstIp) . ", dstPort = $dstPort, pId = $protoName";
			warning($warnMsg);
			next;
		}

		# подготовим данные для внесения в базу
		if ($dstPort == -1) {
			($srcIp, $dstIp, $dstPort, $inBytes, $outBytes) = ($dstIp, $srcIp, $srcPort, $psize, 0);
		} elsif ($srcPort == -1) {
			($inBytes, $outBytes) = (0, $psize);
		} else {
			# не удается определить клиента в полученном TCP или UDP пакете
			$warnMsg = "Can not determine the client: srcIp = " . long2ip($srcIp) . ", srcPort = $srcPort, " .
				"dstIp = " . long2ip($dstIp) . ", dstPort = $dstPort, pId = $protoName";
			warning($warnMsg);
			next;
		}

		# добавим данные в буфер
		my $key = join '@', $srcIp, $dstIp, $dstPort, $pId;
		if (exists $traffic{$key}) {
			$traffic{$key}[0] += $inBytes;
			$traffic{$key}[1] += $outBytes;
		} else {
			$traffic{$key} = [$inBytes, $outBytes];
		}

		last if ++$i == $h_size;
	}

	# перенесем информацию о трафике в базу
	foreach my $key (sort keys %traffic) {
		my ($inBytes, $outBytes) = ($traffic{$key}[0], $traffic{$key}[1]);
		$sth->execute($date, split('@', $key), $inBytes, $outBytes, $inBytes, $outBytes) or
			error("Failed execute query '$query': " . $dbh->errstr);
	}

	# перезапустим цикл
	redo LOOP_ if $i == $h_size;
}

# закроем все открытые дискрипторы и соединения
close(TRAFFIC);
$sth->finish();
$dbh->disconnect();

# удалим полученный от 'trafd' дамп с трафиком
unlink($dump);
exit(0);

# преобразует строковое представление IP адреса в цифровое
sub ip2long {
	return unpack("L*", pack("N*", unpack("L*", inet_aton(shift))));
}

# преобразует цифровое представление IP адреса в строковое
sub long2ip {
	return inet_ntoa(pack("N*", shift));
}

# вывод сообщения об ошибке
sub warning {
	my $str = shift;
	my $msg = "WARNING: $str";

	if ($daemon) {
		openlog(basename($0), "ndelay,pid", LOG_LOCAL0);
		syslog(LOG_WARNING, "($eth) $msg");
		closelog();
	} else {
		warn("$msg\n");
	}
}

# вывод сообщения об ошибке и завершение работы
sub error {
	my $str = shift;
	my $msg = "ERROR: $str";

	if ($daemon) {
		openlog(basename($0), "ndelay,pid", LOG_LOCAL0);
		syslog(LOG_ERR, "($eth) $msg");
		closelog();
		exit(1);
	} else {
		die("$msg\n");
	}
}

# выводит пометку об использовании
sub usage {
	my $progName = basename($0);
	print("$progName: missing eth\n" .
		"Try '$progName --help' for more information.\n");
}

# выводит справку
sub help {
	my $progName = basename($0);
	print("Usage: $progName --eth=ETH [OPTION]...\n" .
		"Unloads the traffic to the database.\n" .
		"\n" .
		"  --eth=ETH      the interface on which there is a collection of\n" .
		"                 traffic\n" .
		"  --config=NAME  use a configuration file\n" .
		"  --daemon       run in daemon mode\n" .
		"  --oftime=INT   offset time of writing data to the database (min.)\n" .
		"  --help         display this help and exit\n" .
		"\n");
}
