#!/usr/pkg/bin/perl
# $Id: band.pl 286 2012-03-26 14:05:15Z bouyer $

use strict;
use IO::Socket::UNIX;
use IO::Socket::INET;
use IO::Select;
use DBI;
use Socket qw(SOL_SOCKET SO_KEEPALIVE);
use Proc::Daemon;
use Getopt::Std;
use Sys::Syslog qw(:standard :macros);
use POSIX qw(setuid setgid);

our($opt_f, $opt_p, $opt_s, $opt_d, $opt_u, $opt_g);
getopts("fp:s:d:u:g:") or usage();

my $sockpath = "/var/run/band.sock";
my $dbpath = "/var/tmp/band.sqlite";

my $uid = 0;
my $gid = 0;

my $isdaemon = 0;

if (defined($opt_s)) {
	$sockpath = $opt_s;
}
if (defined($opt_d)) {
	$dbpath = $opt_d;
}

if (defined($opt_u)) {
	my ($n, $p, $u, $g, $q, $c, $gcos, $dir, $shell) =
	    getpwnam($opt_u) or mydie("bad user " . $opt_u);
	$uid = $u;
	$gid = $g;
}
if (defined($opt_g)) {
	my ($n, $p, $g, $m) =
	    getgrnam($opt_g) or mydie("bad group " . $opt_g);
	$gid = $g;
}

# Create the client listening socket
my $clientssock = new IO::Socket::INET (
	LocalPort => 1781,
	Proto => 'tcp',
	Listen => 16,
	Reuse => 1,
);
mydie("Could not create network socket: $@\n") unless $clientssock;

# Create the local control socket
unlink $sockpath;
my $controls = new IO::Socket::UNIX (
	Type => SOCK_STREAM,
	Local => $sockpath,
	Listen => 16,
	Reuse => 1,
);
mydie("Could not create control socket: $@\n") unless $controls;
chown $uid, $gid, $sockpath;
chmod 0660, $sockpath;

#now daemonize if requested
if ($opt_f ne "1") {
	my $daemon;
	if (defined($opt_p)) {
		$daemon = Proc::Daemon->new(
		    pid_file => $opt_p,
		    dont_close_fh => [ $clientssock, $controls ]
		);
	} else {
		$daemon = Proc::Daemon->new(
		    dont_close_fh => [ $clientssock, $controls ]
		);
	}
	my $pid = $daemon->init;
	if ($pid) {
		exit(0);
	}
	$isdaemon = 1;
}
openlog('band', 'pid', 'local5');
mylog(LOG_INFO, "band starting with pid " . $$);
unlink $dbpath;

#switch user
POSIX::setuid($uid);
POSIX::setgid($gid);
# create database
my $dbargs = {AutoCommit => 1,
                  PrintError => 0};
my $dbh = DBI->connect( "dbi:SQLite:$dbpath", "","",$dbargs) or
    mydie("Cannot connect: $DBI::errstr");

$dbh->do( "CREATE TABLE banned ( ip TEXT, service TEXT, time DATETIME, expire DATETIME, PRIMARY KEY (ip, service) )" )
    or mydie("can't create table: $DBI::errstr");
$dbh->func( 'now', 0, sub { return time }, 'create_function' );
$dbh->func( 'expire', 1, sub { my ($delay) = @_;  return time + $delay; }, 'create_function' );

my $read_set = new IO::Select();
$read_set->add($clientssock);
$read_set->add($controls);

my @clients;

while (1) {
	my ($rh_set) = IO::Select->select($read_set, undef, undef, 60);
	foreach my $rh (@$rh_set) {
		if ($rh == $clientssock) {
			# new client
			my $ns = $rh->accept();
			mylog (LOG_INFO, "new client " . $ns->peerhost() .
			    ":" . $ns->peerport());
			setsockopt($ns, SOL_SOCKET, SO_KEEPALIVE, 1)
			    or mylog(LOG_ERR, "can't set keepalive: $!");
			push @clients, $ns;
			$read_set->add($ns);
		} elsif ($rh == $controls) {
			#new control
			my $ns = $rh->accept();
			mylog(LOG_DEBUG, "new control");
			do_control($ns);
		} else {
			# new client message
			printf("read client ". $rh->peerhost());
			do_client($rh)
		}
	}
	do_expires();
}
exit 0;

#read a control message an process it
sub do_control {
	my ($ns) = @_;

	my $buf = <$ns>;
	chomp $buf;
	if ($buf =~ /^BAN (\d+\.\d+\.\d+\.\d+) (\w+) (\d+)$/) {
		mylog(LOG_DEBUG, "client message " .  $buf);
		$dbh->do("INSERT INTO banned VALUES (\"$1\", \"$2\", now(), expire($3))");
		if ($dbh->err()) {
			mylog(LOG_CRIT, "insert failed for $buf: $DBI::errstr");
		} else {
			sendtoclients($buf);
		}
	} elsif ($buf =~ /^UNBAN (\d+\.\d+\.\d+\.\d+) (\w+)$/) {
		mylog(LOG_DEBUG, "client message " .  $buf);
		$dbh->do("DELETE FROM banned WHERE ip = \"$1\" and service = \"$2\"");
		if ($dbh->err()) {
			mylog(LOG_CRIT, "delete failed for $buf: $DBI::errstr");
		}
		sendtoclients($buf);
	} else {
		mylog(LOG_ERR, "bad control message $buf");
	}
	close ($ns);
}

# process a message from a network client
sub do_client {
	my ($ns) = @_;

	my $buf = <$ns>;
	if ($buf) {
		chomp $buf;
		$buf =~ s/\r//; # for telnet
		mylog(LOG_DEBUG, "client wrote <$buf>");
		if ($buf =~ /^QUIT$/) {
			mylog(LOG_INFO, "client " . $ns->peerhost() . " quits");
			$read_set->remove($ns);
			close ($ns);
		} elsif ($buf =~ /^PING$/) {
			print $ns "PONG\n";
		} elsif ($buf =~ /^GETALL$/) {
			my $all = $dbh->selectall_arrayref(
			    "SELECT ip, service, expire FROM banned");
			foreach my $row (@$all) {
				my ($ip, $service, $expire) = @$row;
				my $timeout = $expire - time();
				if ($timeout > 0) {
					print $ns "BAN $ip $service $timeout\n";				}
			}
		} elsif ($buf =~ /^STATUS$/) {
			my $all = $dbh->selectall_arrayref(
			    "SELECT ip, service, time, expire FROM banned");
			foreach my $row (@$all) {
				my ($ip, $service, $time, $expire) = @$row;
				my $datetime = localtime($time);
				my $dateexpire = localtime($expire);
				print $ns "$datetime BANNED $ip $service until $dateexpire\n";
			}
		}
	} else {
		mylog(LOG_INFO, "client " . $ns->peerhost() . " closed");
		$read_set->remove($ns);
		close ($ns);
	}
}

#send a message to all network clients
sub sendtoclients {
	my ($cmd) = @_;

	foreach my $c (@clients) {
		print $c $cmd . "\n";
	}
}

#process expire dates in database
sub do_expires {
	my $all = $dbh->selectall_arrayref(
	    "SELECT ip, service FROM banned where expire <= now()");
	foreach my $row (@$all) {
		my ($ip, $service) = @$row;
		mylog(LOG_INFO, "expire " .  $ip . " " . $service);
		$dbh->do("DELETE FROM banned WHERE ip = \"$ip\" and service = \"$service\"");
		if ($dbh->err()) {
			mylog(LOG_CRIT, "delete failed for $ip $service: $DBI::errstr");
		}
		sendtoclients("UNBAN " . $ip . " " . $service);
	}
}

sub mylog {
  my ($level, $str) = @_;
  if ($isdaemon != "1") {
	  print STDERR "$str\n";
  }
  syslog $level, $str;
}

sub mydie {
  mylog(LOG_CRIT, @_);
  exit(1);
}

sub usage {
	print STDERR "usage: $ARGV[0] -f -p<pid file> -s<socket> -d<database> -u< user> -g<group>\n";
	exit 1;
}