#!/usr/bin/env python

# Copyright 2003, 2004 Petru Paler (petru@paler.net)
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

import sys
import Milter
from optparse import OptionParser
from Milter import CONTINUE, REJECT, ACCEPT, TEMPFAIL
from rfc822 import AddressList
from siq import siq_query, SIQD_UNKNOWN, SIQD_TEMPFAIL, listEndsWith, lookupAddresses
import re
from IPy import IP

import DNS
DNS.ParseResolvConf()

VARA_ADDR = re.compile(r'([^@]+)_([-a-zA-Z.]+)_@([^@]+)')

def vara_parse(address):
    m = VARA_ADDR.match(address)
    if m:
        return m.group(1, 2, 3)
    else:
        return None

def vara_check(vara_dom, ip):
    reverse = DNS.revlookup(ip)
    ips = lookupAddresses(reverse)
    if ip not in ips:
        return False
    if listEndsWith(reverse.lower().split('.'), vara_dom.lower().split('.')):
        return True
    else:
        return False

do_reject = False
siq_server = "db.outboundindex.net"
threshold = 50

class SIQMilter(Milter.Milter):

    def __init__(self):
        self.remote_ip = None
        self.envdom = None
        self.helodomain = None
        if do_reject:
            self.reject = REJECT
            self.accept = CONTINUE
            self.tempfail = TEMPFAIL
        else:
            self.reject = CONTINUE
            self.accept = CONTINUE
            self.tempfail = CONTINUE

    def connect(self, hostname, unused, hostaddr):
    	self.log("Connect from", hostaddr)
        ip, port = hostaddr
        self.remote_ip = ip
        return CONTINUE

    def hello(self, hostname):
        self.helodomain = hostname
        return CONTINUE

    def envfrom(self, f, *str):
        self.log("mail from", f, str)
        addr = AddressList(f)[0][1]
        try:
            mbox, dom = addr.split("@")
        except ValueError:
            # bounce message, has "<>" sender
            dom = ""
        self.envdom = dom
        return CONTINUE

    def envrcpt(self, to, *str):
        self.log("rcpt to", to)
        address = IP(self.remote_ip)
        if address.version() != 4:
            self.log("%s is an IPv6 address, skipping checks" % self.remote_ip)
            return CONTINUE
        if address.iptype() != 'PUBLIC':
            self.log("%s is a %s address, skipping checks" % (self.remote_ip, address.iptype()))
            return CONTINUE
        d = vara_parse(to)
        if d:
            rcpt, fromdom, dom = d
            self.log("VARA: %s@%s <- %s" % (rcpt, dom, fromdom))
            if vara_check(fromdom, self.remote_ip):
                score = 100
                text = "VARA accept"
            else:
                score = 0
                text = "VARA reject"
            hl = ml = False
        else:
            r = siq_query(self.envdom, self.remote_ip, server=siq_server, helodom=self.helodomain)
            score, text, s1, s2, s3, hl, ml = r
        global threshold
        if score == SIQD_TEMPFAIL:
            decision = "TEMPFAIL"
            if do_reject:
                self.setreply("451", "4.7.1", text)
            mcode = self.tempfail
        elif score == SIQD_UNKNOWN:
            decision = "UNKNOWN"
            mcode = CONTINUE
        elif score > threshold and score <= 100:
            decision = "YES"
            mcode = self.accept
        elif score <= threshold and score >= 0:
            decision = "NO"
            if do_reject:
                self.setreply("550", "5.7.1", text)
            mcode = self.reject
        else:
            self.log("Unexpected response code: %r" % (score,))
            return CONTINUE
        report = "pass=%s ip=%s dn=%s hl=%d ml=%d siq=%s score=%d text='%s'" % (decision, self.remote_ip, self.envdom, hl, ml, siq_server, score, text)
        self.log(report)
        self.report = report
        return mcode

    def header(self, *args):
        return CONTINUE

    def eoh(self):
        return CONTINUE

    def eom(self):
        if hasattr(self, 'report'):
            self.addheader("X-milter-siq-Report", self.report)
        return CONTINUE

    def abort(self):
        return CONTINUE

    def close(self):
        return CONTINUE

    def log(self, *msg):
    	for m in msg:
	    print m,
	print
	sys.stdout.flush()

def main():
    global do_reject, siq_server, threshold
    parser = OptionParser("%prog [-r] [-h] [-s <SIQ server>] <socket name>")
    parser.add_option("-r", "--reject", action="store_true", dest="do_reject",
                      default=False, help="actually reject unauthorized mail")
    parser.add_option("-s", "--server", type="string", dest="server",
                      default=siq_server, help="SIQ server to use")
    parser.add_option("-t", "--threshold", type="int", dest="threshold",
                        default=threshold, help="Scoring threshold")
    (options, args) = parser.parse_args()
    if len(args) != 1:
        parser.error("incorrect number of arguments")
    do_reject = options.do_reject
    siq_server = options.server
    threshold = options.threshold

    Milter.factory = SIQMilter
    Milter.set_flags(Milter.ADDHDRS)
    msg = "starting ("
    if not do_reject:
        msg += "non-"
    msg += "reject mode)"
    print msg
    sys.stdout.flush()
    Milter.runmilter("siqfilter", args[0], 240)
    print "shutdown"

if __name__ == "__main__":
    main()