=head1 NAME
 
berleley_tokenbucketadd - This adds to the berkeley using a Algorithm::TokenBucket (see berkeley_ipblacklist), and could stop the connection or deny the access when.
 
=head1 DESCRIPTION
 
This module writes the BerkeleyDB.
When the bucket is empty it could log, deny, denysoft, denysoft_disconnect or deny_disconnect depending on the action parameter.
 
 
The client ip's can be whitelisted by the plugin L<berkeley_ipblacklist_ignore>
The plugin berkeley_ipblacklist must be configured with method tokenbucket
For an algorithm description L<http://en.wikipedia.org/wiki/Token_bucket>
For a database description read berkeley_ipblacklist
 
=head1 CONFIGURATION
 
=over 4
 
=item hook <deny,allrcptto>
 
When to substract a token from the bucket.
 
This is a comma separated string (without spaces)
 
This parameter is REQUIRED
 
=over 4
 
=item deny
 
Substract a token from the bucket each a command is denied (example: bad recipients, problems whith dns's, etc)
 
=back
 
=item allrcptto
 
Substract a token from the bucket each time a rcptto is received
 
=back
 
=item filename <file>
 
The name of the berkeley file. If it does'n exists, It creates the DB
 
This parameter is REQUIRED
 
=item type <btree|hash>
 
The type of Berkeley db. Hash or Btree. Default is btree
 
=item bucketrate <float>
 
C<rate of information> in items per second (see Algorithm::TokenBucket)
 
Default: 100 items/hour
 
=item bucketsize <integer>
 
C<burst size> in items (see Algorithm::TokenBucket)
 
Default: 10 items
 
=item action [string: deny, denysoft, deny_disconnect, denysoft_disconnect, log]
 
What to do when the bucket is empty -- the options are I<deny>,
I<denysoft> I<deny_disconnect>, I<denysoft_disconnect> or I<log>.
 
If I<log> is specified, the connection will be allowed to proceed as normal,
and only a warning will be logged.
 
The default is I<denysoft>.
 
Some actions may not be available for some hooks. 
See README.plugins in qpsmtpd sources.
 
=back
 
=head1 EXAMPLE
 
In the config/plugins file, add the next lines.
 
C<berkeley_ipblacklist filename /tmp/tokenbucket.bdb action denysoft method tokenbucket>
 
C<berkeley_tokenbucketadd filename /tmp/tokenbucket.bdb hook allrcptto bucketrate 0.00278 bucketsize 2 action denysoft_disconnect>
 
This is a bucket that every time ir receives a I<rcpt to>, it takes one token away. The bucket can contain only 2 items, and the filling rate is a little more than 10 per hour.
 
When the bucket is empty, the I<berkeley_ipblacklist> plugin, denies the attemped connections.
 
A denysoft_disconnect is issued when the buket is emptied.
 
 
=head1 NOTES
 
Remember to C<db_recover -h dbname> when the system is restarted
 
 
Remember to clean the databases using C<berkeley_operation --clean>
 
 
 
=head1 BUGS
 
We are using a read many, write one schema in BerkeleyDB, this allows for an IP to be counted once instead of twice (n).
 
 
=head1 AUTHOR
 
Written by Leonardo Helman <lhelman@pert(punto)com(punto)ar>.
Pert Consultores SRL
Argentina
 
=head1 COPYRIGHT AND LICENSE
 
Copyright (c) 2005 Leonardo Helman. Pert Consultores SRL Argentina
 
This plugin is licensed under the same terms as the qpsmtpd package itself.
Please see the LICENSE file included with qpsmtpd for details.
 
=head1 VERSION
 
$Id: berkeley_tokenbucketadd,v 1.16 2006/05/16 14:12:33 leoh Exp $
 
=cut
 
use warnings;
use strict;
 
use File::Basename;
use BerkeleyDB;
 
 
# VERY IMPORTANT TO KEEP THE SAME BETWEEN ALL THE berkeleys_...
sub BERKELEY_IPBLACKLIST_KEYLEN {15}  # 123.567.901.345
sub BERKELEY_IPBLACKLIST_DATALEN {90}  # TODO: MEJORAR
 
 
sub register {
	my ($self, $qp, @args) = @_;
 
	if (@args % 2) {
		$self->log(LOGERROR, "Unrecognized/mismatched arguments");
		return undef;
	}
	$self->{_args} = {
		'hook' => 0,
		'filename' => 0,
		'type' => 'btree',
		'bucketrate' => 1/36, # 100 items / 3600 seconds
		'bucketsize' => 10,
		@args,
		};
	if( $self->{_args}->{'filename'} ) {
		# untaint
		if( $self->{_args}->{'filename'} =~ /^([\w\:\=.\/_-]*)$/ ) {
			$self->{_args}->{'filename'}= $1;
		}
		else {
			$self->log(LOGCRIT, "The parameter 'filename' is invalid (" . $self->{_args}->{'filename'} . ")");
			return DECLINED;
		}
	}
	else {
		$self->log(LOGCRIT, "The parameter 'filename' is REQUIRED");
		return DECLINED;
	}
	if( $self->{_args}->{'hook'} ) {
		my @hooks=split(/,/,$self->{_args}->{'hook'} );
		$self->{_args}->{'hook'}={};
		for my $hook (@hooks) {
			$self->{_args}->{'hook'}->{$hook}=1;
		}
	}
	else {
		$self->log(LOGCRIT, "The parameter 'hook' is REQUIRED");
		return DECLINED;
	}
 
	$self->register_hook("rcpt", "hook_allrcptto") if( exists $self->{_args}->{'hook'}->{'allrcptto'} );
	# Esto lo necesito por que el hook deny solo tiene DONE, OK y DECLINE
	if( exists $self->{_args}->{'hook'}->{'deny'} ) {
		$self->register_hook("deny", "hook_deny") ;
		$self->register_hook("rcpt", "hook_rcptdeny") ;
		#$self->register_hook("data", "hook_datadeny") ;
	}
 
 
	1;
}
 
sub hook_allrcptto {
	my ($self, $transaction, $recipient) = @_;
 
	return DECLINED
		if ($self->qp->connection->notes('berkeley_ipblacklist_ignore'));
 
	if( $self->algorithmTokenBucket() ) {
		return $self->getReturnValueAndLog( $self->qp->connection->remote_ip, $self->{_args}->{'action'} );
	}
	return DECLINED;
}
 
sub hook_deny {
	my ($self, $transaction, $plugin, $level) = @_;
 
	# We're only interested in DENY or DENY_DISCONNECT
	unless ($level == DENY or $level == DENY_DISCONNECT) {
		return DECLINED;
	}
 
 
	return DECLINED
		if ($self->qp->connection->notes('berkeley_ipblacklist_ignore'));
 
	return DECLINED if $plugin eq $self->plugin_name;
 
	if( $self->algorithmTokenBucket() ) {
		$self->qp->connection->notes('berkeley_ipblacklist_hook_deny_true',1);
	}
	return DECLINED;
}
 
sub hook_rcptdeny {
	my ($self, $transaction, $recipient) = @_;
 
	return DECLINED
		if ($self->qp->connection->notes('berkeley_ipblacklist_ignore'));
	if( $self->qp->connection->notes('berkeley_ipblacklist_hook_deny_true')) {
		return $self->getReturnValueAndLog( $self->qp->connection->remote_ip, $self->{_args}->{'action'} );
	}
	return DECLINED;
}
 
sub hook_datadeny {
	my ($self, $transaction) = @_;
 
	return DECLINED
		if ($self->qp->connection->notes('berkeley_ipblacklist_ignore'));
	if( $self->qp->connection->notes('berkeley_ipblacklist_hook_deny_true')) {
		return $self->getReturnValueAndLog( $self->qp->connection->remote_ip, $self->{_args}->{'action'} );
	}
	return DECLINED;
}
 
sub getReturnValueAndLog {
	my( $self, $ip, $action )=@_;
 
	my @retval=(DECLINED);
	my $msg = "[$ip] is blacklisted";
	$self->log(LOGNOTICE, $msg);
	@retval=(DENY,$msg) if $action eq 'deny';
	@retval=(DENYSOFT,$msg) if $action eq 'denysoft';
	@retval=(DENY_DISCONNECT,$msg) if $action eq 'deny_disconnect';
	@retval=(DENYSOFT_DISCONNECT,$msg) if $action eq 'denysoft_disconnect';
	return @retval;
 
}
 
sub algorithmTokenBucket {
	my ($self) = @_;
 
	my @retval=();
	my $ip = $self->qp->connection->remote_ip;
 
	my $bucket;
	my $db= $self->openBDB();
	return DECLINED unless $db;
 
	my $value;
	my $key= sprintf( "%-*s", BERKELEY_IPBLACKLIST_KEYLEN, $ip );
	if( $db->db_get($key, $value) == 0 ) {
		$value =~ s/\s*$//;
		my ($timestamp, @state)=split(/\s*\|\s*/, $value);
		$bucket = new Algorithm::TokenBucket @state;
	}
	else {
		$bucket = new Algorithm::TokenBucket $self->{_args}->{'bucketrate'}, $self->{_args}->{'bucketsize'}, $self->{_args}->{'bucketsize'};
	}
	if( $bucket->conform(1) ) {
		$bucket->count(1);
		my $value= sprintf( "%-*s", BERKELEY_IPBLACKLIST_DATALEN, join( "|", time, $bucket->state ));
		$db->db_put( $key, $value );
	}
	else {
		@retval=(1);
	}
 
	$self->closeBDB();
 
	return @retval;
}
 
sub openBDB  {
	my ($self)=@_;
 
	my($name,$path) = fileparse($self->{_args}->{'filename'});
 
	# TODO: Poner esto en un modulo tipo pp= new db, pp->read, pp->close
 
	# Utilizo DB_INIT_CDB 
	# esto no es 100% preciso, si hay dos conexiones desde la misma IP,
	# cada uno va a hacer un db_get, modificar los datos, y escribirlos
	# pero el lock se realiza solo en la escritura.
	# Esto puede hacer que una ip pueda mandar mas cosas sin que el
	# sistema se de cuenta.
	# Pero no es una situacion catastrofica. 
	# A lo sumo esa IP podra multiplicar sus limites por el 
	# parametro cantidad de conexiones simultaneas desde una misma IP.
	# La forma correcta es mediante DB_INIT_LOCK
	my $env = new BerkeleyDB::Env  -Home => $path,
                                	-Cachesize => 204_800,
                                	-ErrFile => "$path/error.log",
                                	-Verbose=> 255,
                                	-Flags => DB_CREATE | DB_INIT_CDB | DB_INIT_MPOOL;
 
	my $db;
	if( $self->{_args}->{'type'} eq "btree" ) {
		$db= new BerkeleyDB::Btree(
			-Filename => $name,
			-Flags => BerkeleyDB::DB_CREATE(),
			-Env => $env );
	}
	else {  
		$db= new BerkeleyDB::Hash(
			-Filename => $name,
			-Flags => BerkeleyDB::DB_CREATE(),
			-Env => $env );
	}
	unless( $db ) {
		$self->log(LOGCRIT, "Can't open database " . $self->{_args}->{'filename'} . " failed: $!");
		return;
	}
	return $db;
}
 
sub closeBDB {
	my ($self, $db)=@_;
 
	undef $db;
}
 
 
 
1;
 
# vim:ft=perl: