#!/usr/bin/ruby
# Ruby version of dspam-retrain perl script.
# Perl version: http://dspamwiki.expass.de/DspamRetrainScript
# Author: wsr@rushforthnetworks.com
# License: BSD
# Abstract: setup postfix to pipe mail into this script
#           so that we may then pass it off to dspam in
#           an appropriate way.
#
#           dspam-retrain.rb handles spam-user@domain.tld, and
#                                    spam@domain.tld formats
#
#           If dspam-retrain.rb does not find a signature,
#           it will train dspam using corpus or innoculation
#           mode.
#
# Requirements: open4 gem is installed
# -----------------------------------------------------------

#### Configuration BEGIN ####################################

#enable logging?

@enable_logging = true
@log_file = '/tmp/dspam_retrain.log'

#what mode to train dspam in if we do not find signature?
#corpus or innoculation

@alternative_mode = 'corpus'

#### Configuration END #####################################


require 'rubygems'
require 'open4'

if @enable_logging
        require 'logger'
        @logfile = File.new(@log_file, 'a+')
        @log = Logger.new(@logfile)
end

def logthis(message)
        if @enable_logging
                @log.info(message)
        end
end

# Get arguments
spam_class  = ARGV[0]
sender = ARGV[1]
recip  = ARGV[2]

logthis("dspam-retrain Started. Arguments: #{spam_class}, #{sender}, #{recip}")

#see if we were passed spam-user@ or just user@
if match = recip.to_s.match(/^(spam|ham)-(\w+)@/)
        user = recip.gsub(/#{match[1]}\-/, '')
elsif match = recip.to_s.match(/^(\w+)@/)
        user = sender
else
        logthis("\tCant't determine user")
        exit 75
end

signature = String.new
message   = String.new

#loop through email (passed via stdinput)
#search for signature

$stdin.each do |line|
        if line.match(/X-DSPAM-Signature/)
                signature = line.gsub(/X-DSPAM-Signature:/, '')
                #remove any potential whitespace
                signature.strip!
                #since we found signature, break loop
                break
        end
        message << line
end


if signature.length.to_i == 0
        #we did not find a signature, do normal training

        mode = 'train'

        logthis("\tEmail did not have signature passed in. Attempting #{@alternative_mode} train.")

        #open up dspam with appropriate options
        pid, dspam_in, dspam_out, dspam_err = Open4::popen4 "/usr/bin/dspam --source=#{@alternative_mode} --class=#{spam_class} --user #{user}"

        #attempt to feed message in
        begin
                dspam_in << message
        rescue
                #means dspam closed stdinput because our
                #options failed
                dspam_err.each_line do |o|
                        logthis("\t" + o.gsub(/\n/, ''))
                end
        end
        #close dspams stdinput
        dspam_in.close_write

        #see if dspam left any messages for us
        dspam_out.each_line do |o|
                if o.strip.length == 0 then next end
                logthis("\t" + o.gsub(/\n/, ''))
        end

else
        #we found signature, so we will only pass that.

        mode = 'retrain'

        logthis("\tRetraining Signature: #{signature} for User: #{user} as: #{spam_class}")

        #open up dspam with appropriate options
        pid, dspam_in, dspam_out, dspam_err = Open4::popen4 "/usr/bin/dspam --source=error --signature=#{signature} --class=#{spam_class} --user #{user}"

        #see if dspam left any messages for us
        dspam_out.each_line do |o|
                if o.strip.length == 0 then next end
                logthis("\t" + o.gsub(/\n/,''))
        end
        dspam_err.each_line do |o|
                if o.strip.length == 0 then next end
                logthis("\t" + o.gsub(/\n/,''))
        end
end

ignored, status = Process::waitpid2 pid

#see if we exited cleanly and log it
if status.exitstatus == 0
        logthis("\tMessage successfully #{mode}ed as #{spam_class}")
else
        logthis("\tMessage NOT #{mode}ed")
end