#!/usr/bin/perl -w
#
# $Id: tinydns-readstats.txt,v 1.9 2005/02/01 14:40:22 nate Exp $
#
# Formatting functionality adapted from:
# tinydns log formatting utility
# based on Faried Nawaz's logfile formatter for dnscache
# by Kenji Rikitake <kenji.rikitake@acm.org> 29-JUL-2000
#
# The idea of opening up multilog for writing and sending
# log output to it was taken from tinydns-rrd by Ask Bjørn 
# Hansen, along with a code snippit or two.
#
# What's left was written by me, Nate Campi <nate@campin.net>
##################################################################
#
# Usage:
# 
# If you have a high traffic tinydns server leave out the --log
# option and pipe to multilog in order to minimize I/O. This is
# an example daemontools log run file (e.g. /service/tinydns/log/run):
#
# #!/bin/sh
# exec setuidgid dnslog tinydns-readstats.pl
#
# If you want pretty logs sent to multilog use the option --logpretty. 
# To run multilog you need a log run script something like this:
# 
# #!/bin/sh
# exec setuidgid dnslog tinydns-readstats.pl --logpretty -- multilog t ./main
#
# If you simply want to use this script to prettify your logs, use the
# --nostats option in your log/run script:
#
# #!/bin/sh
# exec setuidgid dnslog tinydns-readstats.pl --nostats --logpretty -- multilog t ./main
#
# You can use this with dnscache now as well,just add the --dnscache 
# option (/service/dnscache/log/run):
#
# #!/bin/sh
# exec setuidgid dnslog tinydns-readstats.pl --dnscache --log -- multilog t ./main
#
##################################################################
#
# To query these stats over SNMP use these lines in net-snmp snmpd.conf:
#
# exec VALUES /bin/echo A PTR ANY MX NS CNAME SOA SRV AAAA TOTAL
# exec bindstats /bin/cat /home/zoneaxfr/stats/stats_file
#
# See http://www.campin.net/DNS/graph.html for the rest of what you need to
# graph the stats.
#
##################################################################
#
# Hmm, each time I implemented this on a box, the stats file was already 
# in place and valid from a prototype version of these scripts. I totally
# spaced on whether it did the right thing when no file existed, the
# script should take some care to make sure things are in order.
#
# Make sure the file is there, readable, and has some valid values in
# it by running UNIX commands like this:
#
# $ mkdir -p /home/zoneaxfr/stats
# $ echo 0 0 0 0 0 0 0 0 0 0 > /home/zoneaxfr/stats/stats_file
# $ chown -R dnslog /home/zoneaxfr/stats
#
# Just be sure that if your logging account isn't named "dnslog" that
# you substitute the correct username in the chown command ("Gdnslog"
# perhaps).
#
##################################################################

use Getopt::Long;
use Fcntl qw(:DEFAULT :flock);
use strict;

my $stats_file = "/home/zoneaxfr/stats/stats_file";
my $stats_file_temp = "/home/zoneaxfr/stats/stats_file.temp";
my $stats_flush_interval = 60; # between 60 and 300 seconds is probably best
my $time = time();
my $stats_flush_time = ( $time + $stats_flush_interval );

# Scott Middlebrooks <scott DOT middlebrooks AT harrynorman DOT com> had
# a problem with zombies, contributed the following to reap them
$SIG{CHLD} = \&REAPER;
sub REAPER {
	my $waitedpid;
	while (($waitedpid = waitpid(-1, &WNOHANG)) > 0) {
	}
	$SIG{CHLD} = \&REAPER;
}

my ( 	$total, $srv, $any, $a, $ns,
	$cname, $soa, $aaaa, $mx, $ptr,
	$other, @line, 

	$oldtotal, $oldsrv, $oldany, $olda, $oldns,
	$oldcname, $oldsoa, $oldaaaa, $oldmx, $oldptr,
	
	$total_a, $total_any, $total_srv, $total_total,
	$total_ns, $total_soa, $total_cname, $total_aaaa,
	$total_mx, $total_ptr,
) = 0;

my ( 
	$DEBUG, $query_types, %opts, $pid, $i, @stats, $dnscache,
);


%opts = ('log'    => 0,
	'logpretty' => 0,
	'nostats' => 0,
	'dnscache' => $dnscache,
	'debug'  => $DEBUG,
);

GetOptions (\%opts,
            'log!',
            'nostats!',
            'dnscache!',
            'logpretty!',
            'debug!',
            )
  or exit 2;

$DEBUG = $opts{debug};
$dnscache = $opts{dnscache};

die "Can't use both --log and --logpretty at once\n" if $opts{log} and $opts{logpretty};

if ( $opts{log} || $opts{logpretty} ) { # pipe to multilog

	$| = 1;
	my $command = join " ", @ARGV; 
	open (MULTI, "|$command") or die "Could not open $command: $!";

	my $oldfh = select MULTI;
	$| = 1;
	select $oldfh;
}
 
while (<STDIN>) {

	$time = time();

	# increment the running total - unless it is a "starting tinydns" line,
	# I don't know of any other non-query lines, let me know if there are any
	if ($dnscache) {
		$total++ if /^query/ ; 
	} else { # then we're running tinydns
		$total++ unless /starting tinydns/ ; 
	}

	print "INPUT before transformation is $_\n" if $DEBUG;

	print MULTI "$_" if $opts{log}; # output for multilog's pleasure

	unless ($dnscache) {
	
		# convert addresses in hex to dotted decimal notation.
		s/\b([a-f0-9]{8})\b/join(".", unpack("C*", pack("H8", $1)))/eg;
	
		# clean up the rest
		s/^([\d.]+):(\w+):(\w+) ([\+\-\/]) \b([a-f0-9]+) \b([-.\w]+)/printQueryLine($1,$2,$3,$4,$5,$6)/e;
	
		print "INPUT after transformation is $_\n" if $DEBUG;
	
		print MULTI "$_" if $opts{logpretty}; # output in pretty format for multilog's pleasure
	
		@line = split(/\s+/);    # split it for easy parsing
	
		SWITCH: {
			if ( $line[2] eq "soa" )	{ $soa++; last SWITCH; }
			if ( $line[2] eq "ptr" )	{ $ptr++; last SWITCH; }
			if ( $line[2] eq "mx" )		{ $mx++; last SWITCH; }
			if ( $line[2] eq "a" )		{ $a++; last SWITCH; }
			if ( $line[2] eq "srv" )	{ $srv++; last SWITCH; }
			if ( $line[2] eq "ns" ) 	{ $ns++; last SWITCH; }
			if ( $line[2] eq "cname" )	{ $cname++; last SWITCH; }
			if ( $line[2] eq "any" )	{ $any++; last SWITCH; }
			if ( $line[2] eq "aaaa" )	{ $aaaa++; last SWITCH; }
			$other++;
		}
	}

	if ( !($opts{nostats}) && ($time >= $stats_flush_time) ) { #flush the stats with a child proc

		$stats_flush_time += $stats_flush_interval; # set the time to flush stats again

		$pid = fork();
		die "Cannot fork: $!" unless defined($pid);
		if ($pid == 0) { 
			# Child process
			updateStats();
			exit(0);   # Child process exits when it is done.
		}

		# clear out the stats now that we've flushed them to disk

		(	$total, $srv, $any, $a, $ns,
			$cname, $soa, $aaaa, $mx, $ptr,
			$other, @line, ) = 0;

		
	} # else 'tis the parent process, which goes back to processing logs

}


### subs

sub printQueryLine {
  my ($host, $port, $query_id, $flag, $query_type, $query) = @_;

  # pad hostname

  my $ret = "$host:";
  $ret .= hex($port);
  $ret .= ":" . hex($query_id);
  $ret .= " " . $flag;
  $ret .= " " . queryType(hex($query_type)) . " $query";
  
  return $ret;
}

sub queryType {
  my ($type) = shift;

  my $ret = "";
 
 # i only list the ones that are in dnscache's dns.h.
 SWITCH: {
    ($type == 1)        && do { $ret = "a";     last SWITCH; };
    ($type == 2)        && do { $ret = "ns";    last SWITCH; };
    ($type == 5)        && do { $ret = "cname"; last SWITCH; };
    ($type == 6)        && do { $ret = "soa";   last SWITCH; };
    ($type == 12)       && do { $ret = "ptr";   last SWITCH; };
    ($type == 13)       && do { $ret = "hinfo"; last SWITCH; };
    ($type == 15)       && do { $ret = "mx";    last SWITCH; };
    ($type == 16)       && do { $ret = "txt";   last SWITCH; };
    ($type == 17)       && do { $ret = "rp";    last SWITCH; };
    ($type == 24)       && do { $ret = "sig";   last SWITCH; };
    ($type == 25)       && do { $ret = "key";   last SWITCH; };
    ($type == 28)       && do { $ret = "aaaa";  last SWITCH; };
    ($type == 252)      && do { $ret = "axfr";  last SWITCH; };
    ($type == 255)      && do { $ret = "any";   last SWITCH; };
    do { $ret .= "$type "; last SWITCH; };
  }
  return $ret;
}

sub updateStats {

	sysopen(STATS_FILE,"$stats_file", O_RDWR|O_CREAT) ||
	    die "Sorry, I couldn't open $stats_file for writing: $!\n";

	flock(STATS_FILE, LOCK_EX)
		or die "Can't write-lock $stats_file: $!\n";
	
	sysopen(STATS_FILE_TEMP,"$stats_file_temp", O_RDWR|O_CREAT) ||
            die "Sorry, I couldn't open $stats_file_temp for writing: $!\n";

	flock(STATS_FILE_TEMP, LOCK_EX)
		or die "Can't write-lock $stats_file_temp: $!\n";

while (<STATS_FILE>) {
		chomp;
		@stats = split(/\s+/);    # split it for easy parsing
	
		$olda     = $stats[0];
		$oldptr   = $stats[1];
		$oldany   = $stats[2];
		$oldmx    = $stats[3];
		$oldns    = $stats[4];
		$oldcname = $stats[5];
		$oldsoa   = $stats[6];
		$oldsrv   = $stats[7];
		$oldaaaa  = $stats[8];
		$oldtotal = $stats[9];
	}
	
	print "oldA oldPTR oldANY oldMX oldNS oldCNAME oldSOA oldSRV oldAAAA oldTOTAL\n" if $DEBUG;
	print "$olda $oldptr $oldany $oldmx $oldns $oldcname $oldsoa $oldsrv $oldaaaa $oldtotal\n" if $DEBUG;
	
	print "A PTR ANY MX NS CNAME SOA SRV AAAA TOTAL\n" if $DEBUG;
	print "$a $ptr $any $mx $ns $cname $soa $srv $aaaa $total\n" if $DEBUG;
	
	$total_a     = ( $olda     + $a     );
	$total_ptr   = ( $oldptr   + $ptr   );
	$total_any   = ( $oldany   + $any   );
	$total_any   = ( $oldany   + $any   );
	$total_mx    = ( $oldmx    + $mx    );
	$total_ns    = ( $oldns    + $ns    );
	$total_cname = ( $oldcname + $cname );
	$total_soa   = ( $oldsoa   + $soa   );
	$total_srv   = ( $oldsrv   + $srv   );
	$total_aaaa  = ( $oldaaaa  + $aaaa  );
	$total_total = ( $oldtotal + $total );

	# be careful and truncate it
	seek(STATS_FILE_TEMP, 0, 0)		or die "can't rewind numfile : $!";
	truncate(STATS_FILE_TEMP, 0)		or die "can't truncate $stats_file: $!";
	
	print STATS_FILE_TEMP "$total_a $total_ptr $total_any $total_mx $total_ns $total_cname $total_soa $total_srv $total_aaaa $total_total\n";

	rename("$stats_file_temp","$stats_file") || die "Can't rename $stats_file_temp to $stats_file: $!";
	
	close(STATS_FILE);
	close(STATS_FILE_TEMP);
	
	print "$total_a $total_ptr $total_any $total_mx $total_ns $total_cname $total_soa $total_srv $total_aaaa $total_total\n" if $DEBUG;

}

__END__