#!/usr/local/bin/perl
#
# dnsstats.pl v1.11
# 20080729 00:01 (GMT-05:00) Jeremy Kister - http://jeremy.kister.net/
# Function: keep stats on tinydns and dnscache for use of cricket/cacti/etc
# Released under Perl's Artistic License
#

use strict;
use Getopt::Long;

chdir('/'); # be nice 
$SIG{ALRM} = \&_flush;

my %rrs = ('0001'=>'A','0002'=>'NS','0005'=>'CNAME','0006'=>'SOA','000c'=>'PTR',
           '000f'=>'MX','0010'=>'TXT','001c'=>'AAAA','0021'=>'SRV','0023'=>'NAPTR',
           '0026'=>'A6','00fb'=>'IXFR','00fc'=>'AXFR','00ff'=>'ANY',
           1=>'A',2=>'NS',5=>'CNAME',6=>'SOA',12=>'PTR',15=>'MX',16=>'TXT',28=>'AAAA',
           33=>'SRV',35=>'NAPTR',38=>'A6',251=>'IXFR',252=>'AXFR',255=>'ANY');
my @hexrrs = qw/0001 0002 0005 0006 000c 000f 0010 001c 0021 0023 0026 00fb 00fc 00ff/;
my @decrrs = qw/1 2 5 6 12 15 16 28 33 35 38 251 252 255/;

my %opts;
GetOptions(\%opts,
           'log!',
           'dnscache',
           'tinydns',
           'cricket',
           'cacti',
           'stats_file=s',
           'tmp_stats_file=s') || die "GetOptions Error: $!\n";

die "must specify either --tinydns or --dnscache\n" unless($opts{tinydns} || $opts{dnscache});
die "must specify either --cricket or --cacti\n" unless($opts{cricket} || $opts{cacti});
die "must specify --stats_file\n" unless($opts{stats_file});
$opts{tmp_stats_file} = $opts{stats_file} . '.tmp' unless(exists($opts{tmp_stats_file}));

foreach my $file ('stats_file','tmp_stats_file'){
	if(-f $opts{$file}){
		unless(-w $opts{$file}){
			die "could not write to $opts{$file}\n";
		}
	}else{
		if(open(FILE, ">$opts{$file}")){
			close FILE;
		}else{
			die "could not create $opts{$file}: $!\n";
		}
	}
}
chmod(0644,$opts{stats_file}) || die "couldnt chmod $opts{stats_file}: $!\n";
unlink($opts{tmp_stats_file}) || die "cannot unlink $opts{tmp_stats_file}: $!\n";

if($opts{log}){
	$|=1;
	my $command = join ' ', @ARGV;
	open(LOG, "| $command") || die "could not fork $command: $!\n";
	my $oldfh=select LOG;
	$|=1;
	select $oldfh;
}

my %total = ('total' => 0, 'other' => 0);

alarm(300);
if($opts{dnscache}){
	$total{motion} = 0;
	$total{querycount} = 1;
	foreach my $rr (@decrrs){
		$total{$rr} = 0;
	}

	while(<STDIN>){
		print LOG if($opts{log});

		my $f = substr($_,0,1);
		if($f eq 'q'){  # query
			my $rr = (split)[3];
			$total{total}++;
			if(exists($total{$rr})){
				$total{$rr}++;
			}else{
				$total{other}++;
			}
		}elsif($f eq 's'){ # stats, sent, servfail
			if(substr($_,1,1) eq 't'){ # stats
				($total{querycount},$total{motion}) = (split)[1,2];
			}
		}
	}
}else{
	$total{lame} = 0;
	$total{notimp} = 0;
	$total{formerr} = 0;
	foreach my $rr (@hexrrs){
		$total{$rr} = 0;
	}

	while(<STDIN>){
		print LOG if($opts{log});

		my @fields = split;
		next if($fields[0] eq 'starting');

		$total{total}++;
		if($fields[1] eq '-'){
			$total{lame}++;
		}elsif($fields[1] eq 'I'){
			$total{notimp}++;
		}elsif($fields[1] eq 'C'){
			$total{formerr}++;
		}elsif(exists($total{$fields[2]})){
			$total{$fields[2]}++;
		}else{
			$total{other}++;
		}
	}
}

sub _flush {
	# i dont care if this blocks for .001 seconds
	alarm(300);
	if(open(FILE, ">$opts{tmp_stats_file}")){
		if($opts{tinydns}){
			if($opts{cricket}){
				foreach my $rr (@hexrrs){
					print FILE "$total{$rr}\n";
				}
		 		print FILE "$total{lame}\n",
				           "$total{formerr}\n",
				           "$total{notimp}\n";
			}else{
				foreach my $rr (@hexrrs){
					print FILE "$rrs{$rr}:$total{$rr} ";
				}
		 		print FILE "lame:$total{lame} ",
				        "formerr:$total{formerr} ",
				        "notimp:$total{notimp} ";
			}
		}else{
			my $effectiveness = sprintf('%0.f', ($total{motion}/$total{querycount}));
			if($opts{cricket}){
				foreach my $rr (@decrrs){
					print FILE "$total{$rr}\n";
				}
				print FILE "${effectiveness}\n";
			}else{
				foreach my $rr (@decrrs){
					print FILE "$rrs{$rr}:$total{$rr} ";
				}
				print FILE "effectiveness:${effectiveness} ";
			}
		}
		if($opts{cricket}){
			print FILE "$total{other}\n",
			        "$total{total}\n";
		}else{
			print FILE "other:$total{other} ",
			        "total:$total{total}\n";
		}
		close FILE;
		if(rename($opts{tmp_stats_file},$opts{stats_file})){
			foreach my $key (keys %total){
				next if($key eq 'querycount' || $key eq 'motion'); # or a possible divide by zero error
				$total{$key} = 0;
			}
		}else{
			warn "could not rename $opts{tmp_stats_file} to $opts{stats_file}: $!\n";
		}
	}else{
		warn "cannot write to $opts{tmp_stats_file}: $!\n";
	}
}