#!/usr/local/bin/perl

# anycast.pl Copyright (c) 2007061501 Jeremy Kister
# http://jeremy.kister.net./code/anycast.pl
# Released under Perl's Artistic License.
# 
# Inject /32 routes into BGP for the purpose of anycast routing for our
# services.  When the service is broken, the route must go away.
# Currently supports UDP/DNS, TCP/DNS, and UDP/RADIUS.
#
# Recommended to run under djb's daemontools.  modify %our, %peer, %table.

use strict;
use Net::BGP::Process;
use Net::BGP::Peer;
use Net::BGP::Update;
use Net::BGP::ASPath;
my %modules = ('Net::DNS' => (eval{require Net::DNS}) ? 1 : 0,
               'Authen::Radius' => (eval{require Authen::Radius}) ? 1 : 0, );
 
my $DEBUG=1;

my %our   = (addr => '10.115.0.126', as => 65523);
my %peer  = (addr => '10.115.5.52', as => 65534);
my %table = ('10.115.0.9'      => { service  => 'DNS',
                                    qtype    => 'NS',
                                    host     => '.',
                                    recurse  => 1,
                                    protocol => 'TCP,UDP',
                                    expect   => '^[A-M]\.ROOT-SERVERS\.NET$', }, # a regex
             '10.115.0.40'     => { service  => 'DNS',
                                    qtype    => 'A',
                                    host     => 'ns1.example.net',
                                    recurse  => 0,
                                    protocol => 'UDP',
                                    expect   => '^10\.115\.0\.9$', }, # a regex
             '10.115.0.16'     => { service  => 'RADIUS',
                                    secret   => 'secret',
                                    username => 'username',
                                    password => 'password', },
            ); # 30 servers max


my @servers = keys %table;
foreach my $server (@servers){
	# taint the table - make sure all seems okay
	if($table{$server}{service} eq 'DNS'){
		die "Net::DNS module not found for service DNS\n" unless($modules{'Net::DNS'});
		unless($table{$server}{qtype} && $table{$server}{host} && $table{$server}{recurse} &&
		       $table{$server}{protocol} && $table{$server}{expect}){
			die "Configuration failed for server $server - check %table\n";
		}
		unless($table{$server}{protocol} =~ /^(?:TCP|UDP|TCP,UDP|UDP,TCP)$/){
			die "unknown protocol for server $server - check %table\n";
		}
	}elsif($table{$server}{service} eq 'RADIUS'){
		die "Authen::Radius module not found for service RADIUS\n" unless($modules{'Authen::Radius'});
		unless($table{$server}{secret} && $table{$server}{username} && $table{$server}{password}){
			die "Configuration failed for server $server - check %table\n";
		}
	}else{
		die "$server using unknown service: $table{$server}{service}\n";
	}
}

my $origin = ($our{as} == $peer{as}) ? 0 : 1; # 0 = igp, 1 = egp, 2 = incomplete
my $aspath = new Net::BGP::ASPath($our{as});
my %broken = map { $_ => 2 } @servers; # assume all broken, prove good.
my $timer = int(30/@servers); # check hosts equally during the 30 second period
my $loops_to_switch = (@servers * 2); # we check and confirm each server on start
my ($bgp,$peer,%retry,$run_loops);
my $id = 0;

connect_to_peer(); # infinite loop

sub connect_to_peer {
	warn "Connecting to Peer [$peer{addr}].\n";
	$bgp = new Net::BGP::Process();
	$peer = new Net::BGP::Peer( Start => 1,
	                            ThisID => $our{addr},
	                            ThisAS => $our{as},
	                            PeerID => $peer{addr},
	                            PeerAS => $peer{as},
	                          );
	$bgp->add_peer($peer);
	$peer->add_timer(\&run_on_timer, 1); # check each asap
	$run_loops = 1;
	$bgp->event_loop();
}


sub run_on_timer {
	if($run_loops == $loops_to_switch){
		# we've checked each host twice and have injected the routes.
		# stop checking asap, now check once per 30 seconds
		$peer->remove_timer(\&run_on_timer);
		$peer->add_timer(\&run_on_timer, $timer);
		# last time counter is incremented so we dont get here again
		$run_loops++; 
	}elsif($run_loops < $loops_to_switch){
		$run_loops++;
	}

	# make sure our session is established
	unless($peer->is_established){
		connect_to_peer(); # keep the infinite loop going
	}

	my $code;
	eval {
		local $SIG{ALRM} = sub { die "TIMEOUT.\n"; };
		alarm(1); # a full second is quite a lot of time
		
		if($table{$servers[$id]}{service} eq 'DNS'){
			# set the resolver to the host we're checking, set recurse appropriately
			my $res = Net::DNS::Resolver->new(nameservers => [$servers[$id]], recurse => $table{$servers[$id]}{recurse});
	
			foreach my $protocol (split /,/, $table{$servers[$id]}{protocol}){
				my $usevc = ($protocol eq 'TCP') ? 1 : 0;
				$res->usevc($usevc); # vc = 1 for tcp, 0 for udp
	
				my $query = $res->query($table{$servers[$id]}{host}, $table{$servers[$id]}{qtype});
				if($query){
					foreach my $rr (grep { $_->type eq $table{$servers[$id]}{qtype} } $query->answer){
						my $object = ($table{$servers[$id]}{qtype} eq 'NS') ? $rr->nsdname : $rr->address;
						unless($object =~ /$table{$servers[$id]}{expect}/i){
							warn "ERROR: [$object] does not match expected regex.\n" if($DEBUG);
							$code=0;
							last; # only one has to fail to know it's broken
						}
					}
				}else{
					warn "query failed: ", $res->errorstring, "\n" if($DEBUG);
					$code=0;
					last; # no use checking other protocols; we're calling it dead
				}
				last if(defined($code));
			}
		}elsif($table{$servers[$id]}{service} eq 'RADIUS'){
			my $radius = new Authen::Radius(Host => $servers[$id], Secret => $table{$servers[$id]}{secret});
			my $result = $radius->check_pwd($table{$servers[$id]}{username}, $table{$servers[$id]}{password});
		
			$code = ($result == 1) ? 1 : 0;
		}
		$code=1 unless(defined($code));

		alarm(0);
	};
	alarm(0);
	if($@){ $code=0; }; # if the eval fails but sets $code=1 -- needed??
	if($broken{$servers[$id]}){
		if($code){
			if($retry{$servers[$id]}){
				# server is working - inject into BGP ASAP
				delete $retry{$servers[$id]};
				delete $broken{$servers[$id]};
				warn "SENDING BGP UPDATE -> inject $servers[$id]/32\n";
	
				my $update = new Net::BGP::Update ( NLRI => [ "$servers[$id]/32" ],
				                                    AsPath => $aspath,
				                                    LocalPref => 100,
				                                    MED       => 0,
				                                    NextHop   => $our{addr},
				                                    Origin    => $origin,
				                                  );
				$peer->update($update);
			}else{
				warn "successful query to $servers[$id]; will retry to confirm.\n";
				$retry{$servers[$id]} = 1; # retry in 30 seconds to make sure its still ok
			}
		}
	}else{
		if($code){
			if($retry{$servers[$id]}){
				warn "successful query to $servers[$id] on retry.\n";
				delete $retry{$servers[$id]};
			}
		}else{
			warn "$servers[$id] down.\n" if($DEBUG);
			if($retry{$servers[$id]}){
				# server was dead before, is dead now - withdraw from BGP ASAP
				delete $retry{$servers[$id]};
				$broken{$servers[$id]} = 1;
				warn "SENDING BGP UPDATE -> withdraw $servers[$id]\n";

				my $update = new Net::BGP::Update ( Withdraw => [ "$servers[$id]/32" ], );
				$peer->update($update);
			}else{
				warn "failed query to $servers[$id]; will retry to confirm.\n";
				$retry{$servers[$id]} = 1; # we will retry once more, in 30 seconds
			}
		}
	}
	$id++; # next server in the loop to check
	$id = 0 if($id == @servers); # if we're out of the loop, start over (arrays 0-based)
}