#!/usr/bin/env python

#$Id: milter.py 44 2008-10-30 11:08:46Z valera $

import sys
import os
import StringIO
import popen2
import rfc822
from time import strftime
from subprocess import Popen

import Milter


class DspamMilter(Milter.Milter):

    def log(self, *msg):
        print "%s [%s]" % (strftime('%Y%b%d %H:%M:%S'), self.id),
        for i in msg:
            print i

    def __init__(self):
        self.tempname = None
        self.recipients = []
        self.fp_headers = StringIO.StringIO()
        self.fp_body = StringIO.StringIO()
        self.fp_dspam_in = None
        self.fp_dspam_out = None
        self.id = Milter.uniqueID()

    def header(self, field, value):
        self.fp_headers.write("%s: %s\n"%(field, value))
        return Milter.CONTINUE

    def envrcpt(self, to, *str):
        self.recipients.append(to)
        return Milter.CONTINUE

    def eoh(self):
        " at this point we have all headers"
        self.fp_headers.write('\n')
        return Milter.CONTINUE

    def body(self, chunk):
        self.fp_body.write(chunk)
        return Milter.CONTINUE

    def envfrom(self, f, *str):
        return Milter.CONTINUE

    def eom(self):
        " end of message "
        self.fp_headers.seek(0)
        self.fp_body.seek(0)
        self.pre_check_message="%s%s"%\
        (self.fp_headers.getvalue(), self.fp_body.getvalue())
        for user in self.recipients:
            args = \
            "dspamc --client --stdout --user %s --deliver=\"innocent,spam\""%\
            (user, )
            self.fp_dspam_out, self.fp_dspam_in = \
            popen2.popen2(args, bufsize=8192)
            self.fp_dspam_in.write(self.pre_check_message)
            self.fp_dspam_in.close()
        message = rfc822.Message(self.fp_dspam_out)
        for header, value in message.items():
            h_tuple = header.split('-')
            if h_tuple[0:2] == ['x', 'dspam']:
                c_header = "%s-%s-%s"%\
                (h_tuple[0].upper(),
                 h_tuple[1].upper(),
                 h_tuple[2].capitalize())
                self.addheader(c_header, value)
        return Milter.CONTINUE

if __name__ == '__main__':
    Milter.set_flags(Milter.CHGBODY + Milter.CHGHDRS + Milter.ADDHDRS)
    Milter.factory = DspamMilter
    Milter.runmilter("dspamfeeder", 'inet:20000@192.168.107.129', 240)