#!/usr/bin/perl -T

# $Id: process-quarantine.pl,v 1.4 2004/07/01 06:49:25 rjl Exp $

########################################################################
# MAIA MAILGUARD LICENSE v.1.0
#
# Copyright 2004 by Robert LeBlanc <rjl@renaissoft.com>
# All rights reserved.
#
# PREAMBLE
#
# This License is designed for users of Maia Mailguard
# ("the Software") who wish to support the Maia Mailguard project by
# leaving "Maia Mailguard" branding information in the HTML output
# of the pages generated by the Software, and providing links back
# to the Maia Mailguard home page.  Users who wish to remove this
# branding information should contact the copyright owner to obtain
# a Rebranding License.
#
# DEFINITION OF TERMS
#
# The "Software" refers to Maia Mailguard, including all of the
# associated PHP, Perl, and SQL scripts, documentation files, graphic
# icons and logo images.
#
# GRANT OF LICENSE
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. The end-user documentation included with the redistribution, if
#    any, must include the following acknowledgment:
#
#    "This product includes software developed by Robert LeBlanc
#    <rjl@renaissoft.com>."
#
#    Alternately, this acknowledgment may appear in the software itself,
#    if and wherever such third-party acknowledgments normally appear.
#
# 4. At least one of the following branding conventions must be used:
#
#    a. The Maia Mailguard logo appears in the page-top banner of
#       all HTML output pages in an unmodified form, and links
#       directly to the Maia Mailguard home page; or
#
#    b. The "Powered by Maia Mailguard" graphic appears in the HTML
#       output of all gateway pages that lead to this software,
#       linking directly to the Maia Mailguard home page; or
#
#    c. A separate Rebranding License is obtained from the copyright
#       owner, exempting the Licensee from 4(a) and 4(b), subject to
#       the additional conditions laid out in that license document.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
# TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
########################################################################

    use strict;
    use DBI;
    use Mail::SpamAssassin;
    my $sa_version = Mail::SpamAssassin::Version();
    if ($sa_version >= 3.0) {
        require Mail::SpamAssassin::Message;
        Mail::SpamAssassin::Message->import();
    } else {
        require Mail::SpamAssassin::NoMailAudit;
        Mail::SpamAssassin::NoMailAudit->import();
    }
    use Mail::SpamAssassin::PerMsgLearner;
    use Mail::SpamAssassin::Reporter;

    # CONFIGURE THIS: Train the Bayes engine? (1 = enable/yes,
    #                 0 = disable/no)
    my $learn = 1;

    # CONFIGURE THIS: Report to Razor2/DCC/Pyzor? (1 = enable/yes,
    #                 0 = disable/no)
    my $report = 1;

    # CONFIGURE THIS: Location of your database.cfg file
    my $cfg = "/var/amavisd/maia/scripts/database.cfg";

    # CONFIGURE THIS: Location of your encryption key file, or undef to disable
    # my $key_file = undef;
    my $key_file = "/var/amavisd/blowfish.key";

########################################################################
# End of user-configurable portion.  There should be no need to modify #
# anything below this point.                                           #
########################################################################

    # Retrieve the string value associated with a key in the database.cfg file.
    sub get_string_key($$)
    {
       my ($file, $key) = @_;

       if ($file =~ /\n[ \t]*$key[ \t]*=[ \t]*\"(.*)\"/) {
          return ($1);
       } else {
          die ("Maia: [get_string_key] Unable to find $key value in $file\n");
       }
    }


    # Read the encryption key from a file
    sub get_encryption_key($) {
        my ($key_file) = @_;
        my ($key);

        if (!defined $key_file) {
            return undef;
        }

        use IO::File;
        my $fh = new IO::File;

        $key_file = $1 if $key_file =~ /^(.+)$/si; # untaint

        # Key file exists, read key from file
        if ($fh->open("<" . $key_file)) {
            sysread($fh, $key, 56);
            $fh->close;
            return $key;
        } else {
            die("Can't open encryption key file " . $key_file);
        }
    }


    # Returns true (1) if the text is encrypted, false (0) otherwise.
    sub text_is_encrypted($) {
        my ($text) = @_;

        return ($text =~ /^RandomIV/);
    }


    # Decrypt a text string using a specified key.
    sub decrypt_text($$) {
        my ($key, $ciphertext) = @_;
        my ($cipher);

        use Crypt::CBC;

        $key = $1 if $key =~ /^([^\0]{56})$/si; # untaint
        if (text_is_encrypted($ciphertext)) {
    	    $cipher = Crypt::CBC->new( {'key' => $key,
    	                                'cipher' => 'Blowfish',
    	                                'regenerate_key' => 0,
    	                                'padding' => 'null',
    	                                'prepend_iv' => 1
    	                             } );

    	    return $cipher->decrypt($ciphertext);
        } else {
            return $ciphertext;
        }
    }


    # Delete one specific mail item and all references to it.
    sub delete_mail_item($$) {
    	my($dbh, $mail_id) = @_;
    	my($sth, $delete);

        $mail_id = $1 if $mail_id =~ /^([0-9]*)$/si; # untaint

        # Delete the mail item itself
        $delete = "DELETE FROM maia_mail WHERE id = ?";
        $sth = $dbh->prepare($delete)
                    or die (sprintf("Maia: [delete_mail_item] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute($mail_id)
            or die (sprintf("Maia: [delete_mail_item] Couldn't execute query: %s", $dbh->errstr));

        # Delete all recipient references to the mail item
        $delete = "DELETE FROM maia_mail_recipients WHERE mail_id = ?";
        $sth = $dbh->prepare($delete)
                    or die (sprintf("Maia: [delete_mail_item] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute($mail_id)
            or die (sprintf("Maia: [delete_mail_item] Couldn't execute query: %s", $dbh->errstr));

        # Delete any virus references to the mail item
        $delete = "DELETE FROM maia_viruses_detected WHERE mail_id = ?";
        $sth = $dbh->prepare($delete)
                    or die (sprintf("Maia: [delete_mail_item] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute($mail_id)
            or die (sprintf("Maia: [delete_mail_item] Couldn't execute query: %s", $dbh->errstr));

        # Delete any SpamAssassin rule references to the mail item
        $delete = "DELETE FROM maia_sa_rules_triggered WHERE mail_id = ?";
        $sth = $dbh->prepare($delete)
                     or die (sprintf("Maia: [delete_mail_item] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute($mail_id)
            or die (sprintf("Maia: [delete_mail_item] Couldn't execute query: %s", $dbh->errstr));

        # Delete any banned file attachment references to the mail item
        $delete = "DELETE FROM maia_banned_attachments_found WHERE mail_id = ?";
        $sth = $dbh->prepare($delete)
                    or die (sprintf("Maia: [delete_mail_item] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute($mail_id)
            or die (sprintf("Maia: [delete_mail_item] Couldn't execute query: %s", $dbh->errstr));

    }


    # Delete mail references for recipients who agree that
    # the mail item is [spam|ham].  Remove the mail item
    # itself, if no other recipients reference it.
    sub delete_mail($$$) {
    	my($dbh, $mail_id, $type) = @_;
    	my($sth, $sth2, $delete, $select);

        $mail_id = $1 if $mail_id =~ /^([0-9]*)$/si; # untaint

        # Delete mail references for all recipients who agree
        # that the mail item is [spam|ham].
        $delete = "DELETE FROM maia_mail_recipients " .
                  "WHERE mail_id = ? AND type = ?";
        $sth = $dbh->prepare($delete)
                   or die (sprintf("Maia: [delete_mail] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute($mail_id, $type)
            or die (sprintf("Maia: [delete_mail] Couldn't execute query: %s", $dbh->errstr));

        # See if any other mail references exist for this mail
        # item, i.e. any other recipients.
        $select = "SELECT recipient_id FROM maia_mail_recipients " .
                  "WHERE mail_id = ?";
        $sth = $dbh->prepare($select)
                   or die (sprintf("Maia: [process_spam] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute($mail_id)
            or die (sprintf("Maia: [process_spam] Couldn't execute query: %s", $dbh->errstr));

        # If no other recipients exist for this mail item,
        # delete the mail item itself.
        if (!$sth->fetchrow_array()) {
            delete_mail_item($dbh, $mail_id);
        }
        $sth->finish;
    }


    # Train the Bayes engine explicitly with the contents of
    # a single e-mail, forcing it to recognize the mail as
    # spam or ham.
    sub do_learn($$$) {
        my($sa, $msg, $isspam) = @_;
        my($ma, $learner, $learned);

        my @body = split (/^/m, $msg);
	my $dataref = \@body;
        if ($sa_version >= 3.0) {
            $ma = Mail::SpamAssassin::Message->new({message => $dataref, parse_now => 0});
        } else {
            $ma = Mail::SpamAssassin::NoMailAudit->new('data' => $dataref);
        }
        $ma->{noexit} = 1;
        $learner = $sa->learn($ma, undef, $isspam, 0);
	$learned = $learner->did_learn();
	$learner->finish();

	return $learned;
    }


    # Report spam to Vipul's Razor, the DCC, and/or Pyzor,
    # and optionally train the Bayes engine in the process
    # (recognizing the mail item as spam).
    sub do_report($$$) {
        my($sa, $msg, $learn) = @_;
        my($ma, $reported);

        my @body = split (/^/m, $msg);
	my $dataref = \@body;
        if ($sa_version >= 3.0) {
            $ma = Mail::SpamAssassin::Message->new({message => $dataref, parse_now => 0});
        } else {
            $ma = Mail::SpamAssassin::NoMailAudit->new ('data' => $dataref);
        }
        $ma->{noexit} = 1;
        $sa->{conf}->{bayes_learn_during_report} = $learn;
        $reported = !$sa->report_as_spam($ma);

        return $reported;
    }


    # Process all the mail items in the database that
    # recipients have confirmed as spam (type = 'C').
    sub process_spam($$$$) {
    	my($dbh, $learn, $report, $key) = @_;
    	my($sth, @row, $select);
    	my($mail_id, $mail_contents, $mail_subject);

    	# Assemble a list of all the confirmed spam (C) items
        $select = "SELECT DISTINCT maia_mail.id, " .
                         "maia_mail.contents, " .
                         "maia_mail.subject " .
                  "FROM maia_mail, maia_mail_recipients " .
                  "WHERE maia_mail.id = maia_mail_recipients.mail_id " .
                  "AND maia_mail_recipients.type = 'C'";
        $sth = $dbh->prepare($select)
                   or die (sprintf("Maia: [process_spam] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute()
            or die (sprintf("Maia: [process_spam] Couldn't execute query: %s", $dbh->errstr));

        my $count = 0;
        my $learnedcount = 0;
        my $reportedcount = 0;

        # Initialize SpamAssassin
        my $sa = new Mail::SpamAssassin();
        $sa->init (1);

        while (@row = $sth->fetchrow_array()) {
            $mail_id = $row[0];
            $mail_contents = $row[1];
            $mail_subject = $row[2];

            if (defined $key) {
                $mail_contents = decrypt_text($key, $mail_contents);
            }
            if ($learn && !$report) {
                if (do_learn($sa, $mail_contents, 1)) {
                    $learnedcount++;
                    print "Learned mail item " . $mail_id . " as SPAM\n";
                }
            } elsif ($report && !$learn) {
            	if (do_report($sa, $mail_contents, 0)) {
            	    $reportedcount++;
            	    print "Reported mail item " . $mail_id . "\n";
            	}
            } elsif ($report && $learn) {
            	if (do_report($sa, $mail_contents, 1)) {
            	    $reportedcount++;
            	    $learnedcount++;
            	    print "Learned mail item " . $mail_id . " as SPAM and reported it\n";
            	}
            }

            # Remove all recipient references to the mail item
            # (and the item itself, if no other recipients remain)
            delete_mail($dbh, $mail_id, "C");

            $count++;
        }
        $sth->finish;
        if ($learn) {

            # Resynchronize the Bayes engine after learning
            $sa->rebuild_learner_caches();
            $sa->finish_learner();
        }
        print $count . " spam items processed (" . $learnedcount . " learned, " .
              $reportedcount . " reported)\n";

    }


    # Process all the mail items in the database that
    # recipients have confirmed as ham (type = 'G').
    sub process_ham($$$) {
    	my($dbh, $learn, $key) = @_;
    	my($sth, @row, $select);
    	my($mail_id, $mail_contents, $mail_subject);

    	# Assemble a list of all the confirmed ham (G) items
        $select = "SELECT DISTINCT maia_mail.id, " .
                         "maia_mail.contents, " .
                         "maia_mail.subject " .
                  "FROM maia_mail, maia_mail_recipients " .
                  "WHERE maia_mail.id = maia_mail_recipients.mail_id " .
                  "AND maia_mail_recipients.type = 'G'";
        $sth = $dbh->prepare($select)
                   or die (sprintf("Maia: [process_ham] Couldn't prepare query: %s", $dbh->errstr));
        $sth->execute()
            or die (sprintf("Maia: [process_ham] Couldn't execute query: %s", $dbh->errstr));

        my $count = 0;
        my $learnedcount = 0;

        # Initialize SpamAssassin
        my $sa = new Mail::SpamAssassin();
        $sa->init (1);

        while (@row = $sth->fetchrow_array()) {
            $mail_id = $row[0];
            $mail_contents = $row[1];
            $mail_subject = $row[2];

            if ($learn) {
                if (defined $key) {
                    $mail_contents = decrypt_text($key, $mail_contents);
                }
                if (do_learn($sa, $mail_contents, 0)) {
                    $learnedcount++;
                    print "Learned mail item " . $mail_id . " as HAM\n";
                }
            }

            # Remove all recipient references to the mail item
            # (and the item itself, if no other recipients remain)
            delete_mail($dbh, $mail_id, "G");

            $count++;
        }
        $sth->finish;
        if ($learn) {

            # Resynchronize the Bayes engine after learning
            $sa->rebuild_learner_caches();
            $sa->finish_learner();
        }
        print $count . " ham items processed (" . $learnedcount . " learned)\n";

    }


    # Read the database configuration file into memory once
    open CONFIGFILE, "<" . $cfg
        or die ("Maia: [process-quarantine] Unable to open $cfg\n");
    my($config) = "";
    my $line;
    while ($line = <CONFIGFILE>) {
        $config .= $line;
    }
    close CONFIGFILE;

    # Connect to the database
    my $dsn = get_string_key($config, "dsn");
    my $username = get_string_key($config, "username");
    my $password = get_string_key($config, "password");
    my $dbh = DBI->connect($dsn, $username, $password)
                       or die("Can't connect to SQL database");

    my $key = get_encryption_key($key_file);
    process_spam($dbh, $learn, $report, $key);
    process_ham($dbh, $learn, $key);

    # Disconnect from the database
    $dbh->disconnect;

    # We're done.
    exit;