#!/usr/bin/perl -w

# Simple script that read message from STDIN and test it on rspamd server
# using specified command.
#
# Usage: rspamc.pl [-c conf_file] [command] [-s statfile]
#
# By default rspamc.pl would read ./rspamd.conf and default command is SYMBOLS

use Socket qw(:DEFAULT :crlf);
use Getopt::Std;
use Term::Cap;
use Mail::Rspamd::Client;

my %cfg = (
    'conf_file' => '@CMAKE_INSTALL_PREFIX@/etc/rspamd.conf',
    'command'   => 'SYMBOLS',
    'hosts'      => ['localhost:11333', ],
	'require_input'	=> 0,
    'password'  =>  '',
    'control'   =>  0,
    'statfile'  =>  '',
	'deliver_to'=>  '',
	'weight'    =>  1,
	'imap_search' => 'ALL',
	'ip'		=> '127.0.0.1',
);

my $terminal;

$main::VERSION = '@RSPAMD_VERSION@';

sub HELP_MESSAGE {
    print <<EOD;
Usage: rspamc.pl [-h host] [-H hosts_list] [-P password] [-c conf_file] [-s statfile] [-d user\@domain] [command] [path]
-h         host to connect (in format host:port) or unix socket path 
-H         path to file that contains list of hosts
-P         define control password
-c         config file to parse
-s         statfile to use for learn commands

Additional options:
-d         define deliver-to header
-w         define weight for fuzzy operations
-S         define search string for IMAP operations
-i         emulate that message was send from specified IP

Notes:
imap format: imap:user:<username>:password:[<password>]:host:<hostname>:mbox:<mboxname>
Password may be omitted and then it would be asked in terminal
imaps requires IO::Socket::SSL

IMAP search strings samples:
ALL - All messages in the mailbox;
FROM <string> - Messages that contain the specified string in the envelope structure's FROM field;
HEADER <field-name> <string> - Messages that have a header with the specified field-name and that 
             contains the specified string in the text of the header (what comes after the colon);
NEW - Messages that have the \\Recent flag set but not the \\Seen flag. 
             This is functionally equivalent to "(RECENT UNSEEN)".
OLD - Messages that do not have the \\Recent flag set.
SEEN - Messages that have the \\Seen flag set.
SENTBEFORE <date> - Messages whose [RFC-2822] Date: header (disregarding time and timezone) 
             is earlier than the specified date.
TO <string> - Messages that contain the specified string in the envelope structure's TO field.
TEXT <string> - Messages that contain the specified string in the header or body of the message.
OR <search-key1> <search-key2> - Messages that match either search key (same for AND and NOT operations).

Version:   @RSPAMD_VERSION@
EOD
exit;
};

sub load_hosts_file {
	my $file = shift;

	open (HOSTS, "< $file") or die "cannot open file $file";
	$cfg{'hosts'} = [ ];
	while (<HOSTS>) {
		chomp;
		next if $_ =~ /^\s*#/;	
		if ($_ =~ /^\s*(([^:]+):(\d+))\s*$/) {
			push (@{ $cfg{'hosts'} }, $1);
		}
		elsif ($_ =~ /^\s*([^:]+)\s*$/) {
			if ($cfg{'control'}) {
				push (@{ $cfg{'hosts'} }, "$1:11334");
			}
			else {
				push (@{ $cfg{'hosts'} }, "$1:11333");
			}
		}
		elsif ($_ =~ /^\s*(\/\S*)\s*$/) {
			push (@{ $cfg{'hosts'} }, "$1");
		}
	}
	close HOSTS;
}

# Load rspamd config params
sub parse_config {
    my ($is_ctrl) = @_;
	
	if (! open CONF, "< $cfg{'conf_file'}") {
		print STDERR "Config file $cfg{'conf_file'} cannot be opened\n";
		return;
	}

    my $ctrl = 0, $skip = 0;
    while (<CONF>) {
        if ($_ =~ /^.*type.*=.*controller.*$/i) {
            $ctrl = 1;
        }
        if ($ctrl && $_ =~ /}/) {
            $ctrl = 0;
        }
        if ($_ =~ /^.*type.*=.*(?:lmtp|delivery|fuzzy).*$/i) {
            $skip = 1;
        }
        if ($skip && $_ =~ /}/) {
            $skip = 0;
        }
        if (!$skip && ((!$is_ctrl && !$ctrl) || ($ctrl && $is_ctrl))
                && $_ =~ /^\s*bind_socket\s*=\s*((([^:]+):(\d+))|(\/\S*))/i) {
            if ($3 && $4) {
				if ($3 eq '*') {
					$cfg{'hosts'} = [ "127.0.0.1:$4" ];
				}
				else {
					$cfg{'hosts'} = [ "$3:$4" ];
				}
            }
            else {
                $cfg{'hosts'} = [ "$5" ];
            }
        }
        if ($ctrl && $is_ctrl && $_ =~ /^\s*password\s*=\s*"(\S+)"/) {
            $cfg{'password'} = $1;
        }
    }

    close CONF;

}

sub print_control_result {
	my ($host, $res) = @_;

	$terminal->Tputs( 'md', 1, *STDOUT );
	print "Results for host $host:\n\n";
	$terminal->Tputs( 'me', 1, *STDOUT );
	if ($res->{error_code} == 0) {
		print "$res->{error}\n";
	}
	else {
		print "Error occured: $res->{error_code}:\n$res->{error}\n";
	}
}

sub print_rspamc_result {
	my ($host, $res) = @_;

	$terminal->Tputs( 'md', 1, *STDOUT );
	print "Results for host $host:\n\n";
	$terminal->Tputs( 'me', 1, *STDOUT );

	if (defined($res->{error})) {
		print "Error occured: $res->{error}\n\n";
	}
	else {
		while (my ($metric, $result) = each (%{ $res })) {
			$terminal->Tputs( 'md', 1, *STDOUT );
			print "$metric: ";
			$terminal->Tputs( 'me', 1, *STDOUT );
			print "$result->{isspam}, [ $result->{score} / $result->{threshold} ]\n";

			$terminal->Tputs( 'md', 1, *STDOUT );
			print "Symbols: ";
			$terminal->Tputs( 'me', 1, *STDOUT );
			print join("; ", @{ $result->{symbols} }) . "\n";
			print "Urls: " . join(", ", @{ $result->{urls} }) . "\n";
			foreach my $msg (@{ $result->{messages} }) {
				print "Message: $msg\n";
			}
			print "\n\n";
		}
	}
}

############################# Main part ###########################################
my %args;

HELP_MESSAGE() unless scalar @ARGV >= 1;

getopt('c:h:P:s:d:w:S:H:i:', \%args);

my $cmd = shift;
my @path = shift;

if (!defined ($cmd) || $cmd eq '') {
    $cmd = 'SYMBOLS';
}

if (defined ($args{c})) {
     if (-r $args{c}) {
            $cfg{'conf_file'} = $args{c};
     }
     else {
        die "config file $args{c} is not readable";
     }

}

if ($cmd =~ /(SYMBOLS|PROCESS|CHECK|URLS|EMAILS)/i) {
    $cfg{'command'} = $1;
    $cfg{'control'} = 0;
}
elsif ($cmd =~ /(STAT|LEARN|SHUTDOWN|RELOAD|UPTIME|COUNTERS|FUZZY_ADD|FUZZY_DEL|WEIGHTS)/i) {
    $cfg{'command'} = $1;
    $cfg{'control'} = 1;
}
else {
    die "unknown command $cmd";
}

if (-r $cfg{'conf_file'}) {
	# Try to parse config
    parse_config ($cfg{'control'});
}

if (defined ($args{S})) {
	$cfg{'imap_search'} = $args{S};
}

if (defined ($args{s})) {
    if ($args{s}) {
        $cfg{'statfile'} = $args{s};
    }
    else {
        main::HELP_MESSAGE();
    }
}
if (defined ($args{h})) {
    $cfg{'hosts'} = [ $args{h} ];
}
if (defined ($args{P})) {
    $cfg{'password'} = $args{P};
}
if (defined ($args{d})) {
	$cfg{'deliver_to'} = $args{d};
}
if (defined ($args{w})) {
	$cfg{'weight'} = $args{w};
}
if (defined ($args{i})) {
	$cfg{'ip'} = $args{i};
}

if ($cmd =~ /SYMBOLS|SCAN|PROCESS|CHECK|REPORT_IFSPAM|REPORT|URLS|EMAILS|LEARN|FUZZY_ADD|FUZZY_DEL|WEIGHTS/i) {
	$cfg{'require_input'} = 1;
}

if (defined ($args{H})) {
	load_hosts_file ($args{H});
}

my $rspamd = Mail::Rspamd::Client->new(\%cfg);

$terminal = Tgetent Term::Cap { TERM => undef, OSPEED => 9600 };

if (!defined ($path[0]) || ! $cfg{'require_input'}) {
	my $input;
	if ($cfg{'require_input'}) {
		while (defined (my $line = <>)) {
			$input .= $line;
		}
	}
	
	my $res = $rspamd->do_all_cmd ($input);
	while (my ($host, $result) =  each (%{ $res })) {
		if ($cfg{control}) {
			print_control_result ($host, $result);
		}
		else {
			print_rspamc_result ($host, $result);
		}
	}
}
else {
	my $res = $rspamd->process_path (@path);

	while (my ($item, $result) =  each (%{ $res })) {
		print "Results for item $item:\n";
		while (my ($host, $r) =  each (%{ $result })) {
			if ($cfg{control}) {
				print_control_result ($host, $r);
			}
			else {
				print_rspamc_result ($host, $r);
			}
		}
	}
}