geoip/geoip_updater

#!/usr/bin/env perl
use strict;
use warnings;
use IO::Socket::SSL;
use File::stat;
use Archive::Tar;
use DBI;
#Remember to install the XS-version, else this is dog slow:)
use My::geoip qw();
use My::localcheck qw(islocal);
#No idea why this is pulled in.. It's not used
#use Thread::Pool::Simple;
use File::Fetch;
use JSON qw( );
#load.pm comes from libload-perl (in Debian). Cannot find it in FreeBSD
use load qw(AutoLoader now);
use Data::Dumper;
use Time::HiRes qw(gettimeofday tv_interval);

my $debug = 0;
my $base_uri	= 'https://download.maxmind.com/app/geoip_download?';
my $config_file	= '/usr/local/etc/geoip_updater.conf';
my $config	= load_config($config_file);

unless ($config->{'account_id'} && $config->{'license_key'}) {
	print 'Remember to add your account id and license key before running this program'."\n";
	exit;
}

unless ($config->{'sql_user'} && $config->{'sql_pwd'} && $config->{'sql_host'}) {
	print 'Remember to add sql user, password and host before running this program'."\n";
	exit;
}

my %arg = map { $_ => 1 } @ARGV;
if($arg{'db'}) {
	&update_db;
	exit;
}

$debug = $arg{'debug'} || 0;

my ($sec,$min,$hour,$mday,$mon,$year,$wday,$yday,$isdst) = gmtime();
my $today = ($year + 1900).'-'.sprintf("%02d", $mon + 1).'-'.sprintf("%02d", $mday);
my %files = (
	'GeoLite2-Country.tar.gz'	=> {
		edition => 'GeoLite2-Country',
		file => 'GeoLite2-Country.mmdb',
		},
	'GeoLite2-ASN.tar.gz'		=> {
		edition => 'GeoLite2-ASN',
		file => 'GeoLite2-ASN.mmdb',
		},
	);
my $u_tmp	= '/usr/local/share/GeoIP/tmp/up/';
my $d_tmp	= "/usr/local/share/GeoIP/tmp/dl/$today/";
my $b_dir	= '/usr/local/share/GeoIP/tmp/backup/';
my $dir		= '/usr/local/share/GeoIP/';
my $restart	= 0;

die "Target directory ($dir) does not exist, this is an updater, not an installer" unless(-e "$dir");
mkdir $d_tmp unless(-e "$d_tmp");
mkdir $u_tmp unless(-e "$u_tmp");
mkdir $b_dir unless(-e "$b_dir");

foreach my $file(keys %files) {
	my $d_file;
	my $u_file;
	my $i_file;
	my $new_time;
	my $uri = $base_uri.'edition_id='.$files{$file}->{'edition'}.'&license_key='.$config->{'license_key'}.'&suffix=tar.gz';
	my $tounpack = $files{$file}->{'file'};
	my $old_time	= stat("$dir/$tounpack")->mtime;
	print "$dir/$tounpack has timestamp $old_time\n" if($debug);
	unless(-e "$d_tmp/$file") {
		print "$d_tmp/$file does not exist, downloading\n" if($debug);
		$d_file = &download($uri);
		exit() unless defined $d_file;
		system('mv',$d_file,$d_tmp.'/'.$file);
		$d_file = $d_tmp.'/'.$file;
	} else {
		print "$d_tmp/$file exists, no need to download\n" if($debug);
		$d_file = "$d_tmp/$file";
	}
	if(-e "$d_tmp/$file") {
		$u_file = &unpack($d_file,$tounpack);
		$new_time	= stat("$u_file")->mtime;
		print "$u_file has timestamp $new_time\n" if($debug);
	}
	if(-e "$u_file") {
		unless($old_time == $new_time) {
			if($i_file = &install($tounpack,$u_file)) {
				$restart = 1;
			}
		} else {
			print "$u_file and $dir/$tounpack has same timestamp, do nothing\n" if($debug);
		}
	}
}

if($restart) {
	system("/usr/sbin/service geoip reload");
	&update_db;
}

sub download {
	my $file = shift;
	unless($file) {
		print "No file to download\n";
		return 0;
	}
	my $ff		= File::Fetch->new(uri => $file);
	my $where	= $ff->fetch(to => $d_tmp);
	print "Downloaded $file to $where\n" if($debug);
	return $where if($where);
	return undef;
}

sub unpack {
	my $file = shift;
	my $e_file = shift;
	unless($file && $e_file) {
		print "No file to unpack, got file:$file and e_file:$e_file\n";
		return 0;
	}
	my $tar = Archive::Tar->new;
	$tar->read($file);
	my @file_list = $tar->list_files;
	foreach my $uf(@file_list) {
		if($uf =~ m/$e_file/g) {
			$tar->extract_file( $uf, "$u_tmp/$e_file" );
			print "Extracting $uf to $u_tmp/$e_file\n" if($debug);
			if(-e "$u_tmp/$e_file") {
				if(my $rmed = unlink "$file") {
					rmdir "$d_tmp";
					print "Deleted $d_tmp\n" if($debug);
				}
				return "$u_tmp/$e_file";
			} else {
				return 0;
			}
		}
	}
}

sub install {
	my $org_file	= shift;
	my $new_file	= shift;
	unless($org_file && $new_file) {
		print "No file to install or no new file specified\n";
		return 0;
	}
	if((-e "$dir/$org_file") && (-e "$new_file")) {
		print "$dir/$org_file and $new_file exists\n" if($debug);
		if(rename "$dir/$org_file","$b_dir/$org_file") {
			print "Renamed $dir/$org_file to $b_dir/$org_file\n" if($debug);
			if(rename "$new_file","$dir/$org_file") {
				print "Renamed $new_file to $dir/$org_file\n" if($debug);
				return 1;
			}
		}
	} else {
		print "Specified files does no exist or is not readable, tried with $dir/$org_file and $new_file\n";
		return 0;
	}
}

sub update_db {
	my $now		= time();
	print gmtime($now).' Started db update'."\n";
	my $dbh		= DBI->connect("DBI:mysql:database=syslog;host=$config->{'sql_host'}",$config->{'sql_user'},$config->{'sql_pwd'},{'AutoCommit' => 0, RaiseError => 1 });
	#This takes longer the more things there is in the database.. Something needs to be done:)
	my @dbwork	= &fetch_db($dbh);
	my $geoip_p	= My::geoip->new;
	my @parsed;
	my $upd		= 0;
	my $parse_start	= [gettimeofday];
	foreach my $job(@dbwork) {
		my $result = &worker_parser($job, $geoip_p);
		push(@parsed, $result) if($result->{'retval'} == 1);
	}
	my $parse_time	= tv_interval ($parse_start, [gettimeofday]);
	print gmtime(time()).' Done parsing. Took '.$parse_time.' seconds.'."\n";


	my $update_start = [gettimeofday];
	foreach my $job(@parsed) {
		my $result = &worker_updater($job, $dbh);
		$upd++ if($result->{'retval'} == 1);
	}
	my $update_time = tv_interval ($update_start, [gettimeofday]);
	print gmtime(time()).' Done updating. Took '.$update_time.' seconds.'."\n";

	$dbh->commit;
	$dbh->disconnect;
	print gmtime(time()).' Updated '.$upd.' entries in db'."\n";

}

sub fetch_db {
	my $dbh	= shift;
	my @tables = ('reject','reject_iptables');
	my @ret;
	foreach my $tb(@tables) {
		my $sth	= $dbh->prepare("SELECT DISTINCT ip,asn,iso FROM $tb");
		$sth->execute || die 'Execute failed while selecting entries: '.DBI::errstr;
		while(my $ref = $sth->fetchrow_hashref) {
			my $host	= $$ref{'ip'};
			my $asn		= $$ref{'asn'} || '';
			my $iso		= $$ref{'iso'} || '';
			unless(&islocal($host)) {
				my $toret = { 'host' => $host, 'asn' => $asn, 'iso' => $iso, 'table' => $tb };
				push(@ret,$toret);
			}
		}
	}
	return @ret;

}

sub worker_parser {
	my $todo	= shift;
	my $geoip_p	= shift;
	my $tb		= $todo->{'table'};
	my $host	= $todo->{'host'};
	my $asn		= $todo->{'asn'};
	my $iso		= $todo->{'iso'};
	my $g_r		= $geoip_p->parse($host);
	my $upd_asn	= $g_r->{'asn'};
	my $upd_iso	= $g_r->{'iso'};
	$upd_asn	= '' unless $upd_asn;
	unless($asn eq $upd_asn && $iso eq $upd_iso) {
		my $toret	= { 'host' => $host, 'upd_asn' => $upd_asn, 'upd_iso' => $upd_iso, 'asn' => $asn, 'iso' => $iso, 'did' => 'parse', 'table' => $tb, 'retval' => 1};
		return $toret;
	}
	return { 'did' => 'parse', 'retval' => 0, 'host' => $host };

}

sub worker_updater {
	my $todo	= shift;
	my $dbh		= shift;
	my $upd_asn	= $todo->{'upd_asn'};
	my $upd_iso	= $todo->{'upd_iso'};
	my $iso		= $todo->{'iso'};
	my $asn		= $todo->{'asn'};
	my $tb		= $todo->{'table'};
	my $host	= $todo->{'host'};
	my @upd_vals;
	my $query	= 'UPDATE '.$tb.' SET ';

	#Values to set. If NULL, it is not ok to pass the string NULL as a param..
	if(($upd_asn ne $asn) && ($upd_iso ne $iso)) {
		$query	.= 'asn = ';
		if($upd_asn eq '') {
			$query	.= 'NULL';
		} else {
			$query	.= '?';
			push(@upd_vals, $upd_asn);
		}
		$query	.= ', iso = ';
		if($upd_iso eq '') {
			$query	.= 'NULL';
		} else {
			$query	.= '?';
			push(@upd_vals, $upd_iso);
		}
	} elsif($upd_asn ne $asn) {
		$query	.= 'asn = ';
		if($upd_asn eq '') {
			$query	.= 'NULL';
		} else {
			$query	.= '?';
			push(@upd_vals, $upd_asn);
		}
	} elsif($upd_iso ne $iso) {
		$query	.= 'iso = ';
		if($upd_iso eq '') {
			$query	.= 'NULL';
		} else {
			$query	.= '?';
			push(@upd_vals, $upd_iso);
		}
	}

	#Always host
	$query	.= ' WHERE ip = ? ';
	push(@upd_vals, $host);

	#In where, if NULL, we need to use 'IS NULL'/'IS NOT NULL' instead of '= NULL'
	if($upd_asn eq $asn) {
		if($upd_asn eq '') {
			$query	.= ' AND asn IS NULL';
		} else {
			$query	.= ' AND asn = ?';
			push(@upd_vals, $upd_asn);
		}
	}
	if($upd_iso eq $iso) {
		if($upd_iso eq '') {
			$query	.= ' AND iso IS NULL';
		} else {
			$query	.= ' AND iso = ?';
			push(@upd_vals, $upd_iso);
		}
	}

	my $ins		= $dbh->prepare($query) || die "Prepare failed for updating $host, query was $query: ".DBI::errstr;
	$ins->execute(@upd_vals) || die "Execute failed for updating $host, query was $query: ".DBI::errstr;
	return { 'host' => $host, 'did' => 'update', 'retval' => 1}
}

sub worker {
	my $what	= shift;
	my $todo	= shift;
	my $toret;
	if($what eq 'parse') {
		$toret = &worker_parser($todo);
	} elsif($what eq 'update') {
		$toret = &worker_updater($todo);
	} else {
		my $todothing = keys %{$todo};
		return { 'did' => 'nothing', 'retval' => 0, 'was' => $todothing };
	}
	return $toret;
}

sub load_config {
	my $file	= shift;
	die 'Could not find config file '.$file unless(-e $file);
	my $json	= JSON->new;
	open(FH, '<', $file) or die 'Could not open file: '.$!;
	my $fc		= do {
		local $/;
		<FH>;
	};
	close(FH);
	my $toret	= $json->decode($fc);
	return $toret;
}