#!/usr/bin/perl 
use Socket;
use POSIX ;
use DBI;
use Switch;
use threads;
use threads::shared;

#CONFIGURATION SECTION
#my @allowedhosts = ('127.0.0.1', '10.0.0.1');
my @allowedhosts = ('127.0.0.1', '94.23.41.186');
my $LOGFILE="/var/log/ardeekpolicid.log";
chomp( my $vhost_dir = `pwd`);
my $port = 279;
my $listen_address = '0.0.0.0';
my $dsn = "DBI:mysql:keedra:localhost";
my $db_user = 'keedra';
my $db_passwd = 'k33drapwd';
my $db_table = 'emails';
my $db_quotacol = 'messagequota';
my $db_tallycol = 'messagetally';
my $db_timecol = 'timestamp';
my $db_wherecol = 'email';
my $deltaconf = 'monthly'; #daily, weekly, monthly
my $sql_getquota = "SELECT $db_quotacol, $db_tallycol, $db_timecol FROM $db_table WHERE $db_wherecol = ? AND $db_quotacol > 0";
my $sql_updatequota = "UPDATE $db_table SET $db_tallycol = ?, $db_timecol = ? WHERE $db_wherecol = ?";
#END OF CONFIGURATION SECTION

if($deltaconf eq 'monthly'){
	my $delta = 60*60*24*30;
}elsif($deltaconf eq 'weekly'){
	my $delta = 60*60*24*7;
}elsif($deltaconf eq 'daily'){
	my $delta = 60*60*24;
}else{
	my $delta = 60*60*24*30;
}
my %quotahash :shared;
my $lock:shared;
my $client;
my $proto = getprotobyname('tcp');
# create a socket, make it reusable
socket(SERVER, PF_INET, SOCK_STREAM, $proto) or die "socket: $!";
setsockopt(SERVER, SOL_SOCKET, SO_REUSEADDR, 1) or die "setsock: $!";
my $paddr = sockaddr_in($port, inet_aton($listen_address)); #Server sockaddr_in
bind(SERVER, $paddr) or die "bind: $!";# bind to a port, then listen
listen(SERVER, SOMAXCONN) or die "listen: $!";
my $client_addr;
my $client_ipnum;
my $cnt = 0;
&daemonize;
$SIG{TERM} = \&sigterm_handler;
$SIG{HUP} = \&print_cache;
while (1) {
	$client_addr = accept($client, SERVER) || next;
	$cnt++;
	if($cnt % 1500 == 0){
		lock($lock);
		&commit_cache;
		&flush_cache;
	}
	threads->new(\&start_thr)->detach();
}

sub start_thr {
	select($client);
	$|=1;

	my ($client_port, $client_ip) = unpack_sockaddr_in($client_addr);
	$client_ipnum = inet_ntoa($client_ip);
	if(grep $_ eq $client_ipnum, @allowedhosts){
		logger("Client $client_ipnum connected.");
		#my $client_host = gethostbyaddr($client_ip, AF_INET);
		#if (! defined ($client_host)) { $client_host=$client_ipnum;}
		my $message;
		my @buf;
		while(!eof($client)) {
			$message = <$client>;
			if($message =~ m/printshm/){
				print $client "Printing shm:\n";
				while(($k,$v) = each(%quotahash)){
					print $client "$k:\n";
					print $client "\tQuota: $quotahash{$k}{'quota'}\n";
					print $client "\tUsed: $quotahash{$k}{'tally'}\n";
				}
				next;
			}
			if($message =~ m/flush/){
				print $client "Flushing \r\n";
				&flush_cache;
				next;
			}
			if($message =~ m/=/){
				push(@buf, $message);
				next;
			}
			my $ret = &handle_req(@buf);
			if($ret =~ m/unknown/){
				shutdown($client,2);
				threads->exit(0);
			}else{
				print $client "action=$ret\n\n";
			}
			@buf = ();
		}
	}
	shutdown($client,2);
	threads->exit(0);
}

exit;

sub handle_req {
	my @buf = @_;
	my $protocol_state;
	my $sasl_username; 
	my $recipient_count;
	local $/ = "\n";
	foreach $aline(@buf){
		my @line = split("=", $aline);
		switch($line[0]){
			case "protocol_state" { 
				chomp($protocol_state = $line[1]);
			}
			case "sasl_username"{
				chomp($sasl_username = $line[1]);
			}
			case "recipient_count"{
				chomp($recipient_count = $line[1]);
			}
		}
	}
	if($protocol_state !~ m/DATA/ || $sasl_username eq "" ){
		return "ok";
	}
	#TODO: Maybe i should move to semaphore!!!
	lock($lock);
	if(!exists($quotahash{$sasl_username})){
		my $dbh = DBI->connect($dsn, $db_user, $db_passwd);
		my $sql_query = $dbh->prepare($sql_getquota);
		$sql_query->execute($sasl_username);
		if($sql_query->rows > 0){
			while(@row = $sql_query->fetchrow_array()){
				$quotahash{$sasl_username} = &share({});
				$quotahash{$sasl_username}{'quota'} = $row[0];
				$quotahash{$sasl_username}{'tally'} = $row[1];
				if($row[2]){
					$quotahash{$sasl_username}{'lastchange'} = $row[2];
				}else{
					$quotahash{$sasl_username}{'lastchange'} = time();
				}
			}
			$dbh->disconnect;
		}else{
			$dbh->disconnect;
			return "ok";
		}
	}
	if($quotahash{$sasl_username}{'lastchange'} + $delta > time()){
		$quotahash{$sasl_username}{'tally'} = 0;
		$quotahash{$sasl_username}{'lastchange'} += $delta;
	}
	if($quotahash{$sasl_username}{'tally'} + $recipient_count > $quotahash{$sasl_username}{'quota'}){
		return "471 Monthly message quota exceeded"; 
	}
	$quotahash{$sasl_username}{'tally'} += $recipient_count;
	return "ok";
}

sub sigterm_handler {
	shutdown(SERVER,2);
	lock($lock);
	logger("SIGTERM received.\nFlushing cache...\nExiting.");
	&commit_cache;
	exit(0);
}

sub commit_cache {
	my $dbh = DBI->connect($dsn, $db_user, $db_passwd);
    my $sql_query = $dbh->prepare($sql_updatequota);
    while(($k,$v) = each(%quotahash)){
        $sql_query->execute($quotahash{$k}{'tally'}, $quotahash{$k}{'lastchange'}, $k)
			or logger("Query error:".$sql_query->errstr);
    }
    $dbh->disconnect;
}

sub flush_cache {
	foreach $k(keys %quotahash){
		delete $quotahash{$k};
	}
}

sub print_cache {
	foreach $k(keys %quotahash){
        logger("$k: $quotahash{$k}{'quota'}, $quotahash{$k}{'tally'}");
    }
}

sub daemonize {
	my ($i,$pid);
	print "Keedra SMTP Policy Daemon. Logging to $LOGFILE\n";
	#Should i delete this??
	#$ENV{PATH}="/bin:/usr/bin";
	#chdir("/");
	close STDIN;
	if(!defined(my $pid=fork())){
		die "Impossible to fork\n";
	}elsif($pid >0){
		exit 0;
	}
	setsid();
	close STDOUT;
	open STDIN, "/dev/null";
	open LOG, ">>$LOGFILE" or die "Unable to open $LOGFILE: $!\n";
	select((select(LOG), $|=1)[0]);
	open STDERR, ">>$LOGFILE" or die "Unable to redirect STDERR to STDOUT: $!\n";
}

sub logger {
	my ($arg) = @_;
	my $time = localtime();
	chomp($time);
	print LOG  "$time $arg\n";
}