#!/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;

my %cfg = (
    'conf_file' => '@CMAKE_INSTALL_PREFIX@/etc/rspamd.conf',
    'command'   => 'SYMBOLS',
    'host'      => 'localhost',
    'port'      => '11333',
    'is_unix'   =>  0,
    'password'  =>  '',
    'control'   =>  0,
    'statfile'  =>  '',
);

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

sub HELP_MESSAGE {
    print <<EOD;
Usage: rspamc.pl [-h host] [-p port] [-P password] [-c conf_file] [-s statfile] [command]
-h         host to connect or unix socket path
-p         port to connect (not used with unix sockets)
-P         ask for control password
-c         config file to parse
-s         statfile to use for learn commands
Version:   @RSPAMD_VERSION@
EOD
};

# Load rspamd config params
sub parse_config {
    my ($is_ctrl) = @_;

    open CONF, "< $cfg{'conf_file'}" or die "config file $cfg{'conf_file'} cannot be opened";

    my $ctrl = 0, $skip = 0;
    while (<CONF>) {
        if ($_ =~ /^.*type.*=.*controller.*$/i) {
            $ctrl = 1;
        }
        if ($ctrl && $_ =~ /}/) {
            $ctrl = 0;
        }
        if ($_ =~ /^.*type.*=.*(?:lmtp|delivery).*$/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) {
                $cfg{'host'} = $3;
                $cfg{'port'} = $4;
                $cfg{'is_unix'} = 0;
            }
            else {
                $cfg{'host'} = $5;
                $cfg{'is_unix'} = 1;
            }
        }
        if ($ctrl && $is_ctrl && $_ =~ /^\s*password\s*=\s*"(\S+)"/) {
            $cfg{'password'} = $1;
        }
    }

    close CONF;

}

sub connect_socket {
    my $sock;

    if ($cfg{'is_unix'}) {
        socket ($sock, PF_UNIX, SOCK_STREAM, 0) or die "cannot create unix socket";
        my $sun = sockaddr_un($cfg{'host'});
        connect ($sock, $sun) or die "cannot connect to socket $cfg{'host'}";
    }
    else {
        my $proto = getprotobyname('tcp');
        my $sin;
        socket ($sock, PF_INET, SOCK_STREAM, $proto) or die "cannot create tcp socket";
        if ($cfg{'host'} eq '*') {
            $cfg{'host'} = '127.0.0.1';
        }
        if (inet_aton ($cfg{'host'})) {
            $sin = sockaddr_in ($cfg{'port'}, inet_aton($cfg{'host'}));
        }
        else {
            my $addr = gethostbyname($cfg{'host'});
            if (!$addr) {
                die "cannot resolve $cfg{'host'}";
            }
            $sin = sockaddr_in ($cfg{'port'}, $addr);
        }
        
        connect ($sock, $sin) or die "cannot connect to socket $cfg{'host'}:$cfg{'port'}";
    }
    return $sock;
}

# Currently just read stdin for user's message and pass it to rspamd
sub do_rspamc_command {
    my ($sock) = @_;

    my $input;
    while (defined (my $line = <>)) {
        $input .= $line;
    }

    print "Sending ". length ($input) ." bytes...\n";

    syswrite $sock, "$cfg{'command'} RSPAMC/1.0 $CRLF";
    syswrite $sock, "Content-Length: " . length ($input) . $CRLF . $CRLF;
    syswrite $sock, $input;
    syswrite $sock, $CRLF;
    while (defined (my $line = <$sock>)) {
        print $line;
    }
}

sub do_ctrl_auth {
    my ($sock) = @_;

    syswrite $sock, "password $cfg{'password'}" . $CRLF;
    if (defined (my $reply = <$sock>)) {
        my $end = <$sock>;
        if ($reply =~ /^password accepted/) {
            return 1;
        }
    }

    return 0;
}

sub do_control_command {
    my ($sock) = @_;

    # Read greeting first
    if (defined (my $greeting = <$sock>)) {
        if ($greeting !~ /^Rspamd version/) {
            die "not rspamd greeting line $greeting";
        }
    }
    if ($cfg{'command'} =~ /^learn$/i) {
        my $input;
        die "statfile is not specified to learn command" if !$cfg{'statfile'};

        while (defined (my $line = <>)) {
            $input .= $line;
        }
        
        if (do_ctrl_auth ($sock)) {
            my $len = length ($input);
            print "Sending $len bytes...\n";
            syswrite $sock, "learn $cfg{'statfile'} $len" . $CRLF;
            syswrite $sock, $input . $CRLF;
            if (defined (my $reply = <$sock>)) {
                if ($reply =~ /^learn ok/) {
                    print "Learn succeed\n";
                }
                else {
                    print "Learn failed\n";
                }
            }
        }
        else {
            print "Authentication failed\n";
        }
    }
    elsif ($cfg{'command'} =~ /(reload|shutdown)/i) {
        if (do_ctrl_auth ($sock)) {
            syswrite $sock, $cfg{'command'} . $CRLF;
            while (defined (my $line = <$sock>)) {
                last if $line =~ /^END/;
                print $line;
            }
        }
        else {
            print "Authentication failed\n";
        }
    }
    elsif ($cfg{'command'} =~ /(fuzzy_add|fuzzy_del)/i) {
        while (defined (my $line = <>)) {
            $input .= $line;
        }
        
        if (do_ctrl_auth ($sock)) {
            my $len = length ($input);
            print "Sending $len bytes...\n";
            syswrite $sock, $cfg{'command'} . " $len" . $CRLF;
            syswrite $sock, $input . $CRLF;
            if (defined (my $reply = <$sock>)) {
                if ($reply =~ /^OK/) {
                    print $cfg{'command'} . " succeed\n";
                }
                else {
                    print $cfg{'command'} . " failed\n";
                }
            }
        }
        else {
            print "Authentication failed\n";
        }
    
    }
    else {
        syswrite $sock, $cfg{'command'} . $CRLF;
        while (defined (my $line = <$sock>)) {
            last if $line =~ /^END/;
            print $line;
        }
    }
}

my %args;
getopt('c:h:p:P:s:', \%args);
my $cmd = shift;
my $do_parse_config = 1;

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

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

}
if (defined ($args{s})) {
    if ($args{s}) {
        $cfg{'statfile'} = $args{s};
    }
    else {
        main::HELP_MESSAGE();
        exit;
    }
}
if (defined ($args{h})) {
    $cfg{'host'} = $args{h};
    if ($args{h} =~ /^\/.*$/) {
        $cfg{'is_unix'} = 1;
    }
    $do_parse_config = 0;
}
if (defined ($args{p})) {
    $cfg{'port'} = $args{p};
}
if (defined ($args{P})) {
    $cfg{'password'} = $args{P};
}

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

if ($do_parse_config != 0) {
    parse_config ($cfg{'control'});
}

my $sock = connect_socket ();

if ($cfg{'control'}) {
    do_control_command ($sock);
}
else {
    do_rspamc_command ($sock);
}

close ($sock);