'''
milter.py - milter service for sagator

(c) 2003-2007 Jan ONDREJ (SAL) <ondrejj(at)salstar.sk>

 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.

'''

from aglib import *
from StringIO import StringIO
import sys, time, signal, stats, rfc822

__all__=['milter']

try:
  import Milter

  class sgMilter(Milter.Milter):
    "Sagator's milter class."

    def log(self, msg, level=2):
        debug.echo(level, "milter(%s):" % self.CID,msg)

    def connect(self, hostname, unused, hostaddr):
        self.CID=randomchars(6,'0123456789abcdef')
        self.log("Connection from: "+hostname+" at "+\
                 time.strftime("%c",time.localtime()))
        self.comm1="XFORWARD ADDR=%s NAME=%s HELO= PROTO=SMTP\r\n" \
                   % (hostaddr, hostname)
        self.connect_from=(hostaddr, hostname)
        self.helo_name=''
        self.mail=mail_class()
        # reinitialize scanners
        for scanner in self.SCANNERS:
          scanner.reinit()
        return Milter.CONTINUE

    def hello(self,hostname):
        self.comm1+="HELO %s\r\n" % hostname
        self.helo_name=hostname
        return Milter.CONTINUE

    def envfrom(self,f,*s):
        self.log('MAIL FROM: '+f,2)
        self.stats=stats.statistics()
        self.mail.__init__() # reinit, may be a new email
        self.mail.comm=self.comm1+"MAIL FROM: %s\r\n" % f
        self.mail.sender=f
        return Milter.CONTINUE

    def envrcpt(self,to,*s):
        self.mail.comm+="RCPT TO: "+to+"\r\n"
        self.log('RCPT TO: '+to,2)
        try:
          recipient_addr=rfc822.parseaddr(to)[1]
          self.mail.recip.append(recipient_addr)
        except:
          recipient_addr=''
        mail.policy_request={
          'client_address': self.connect_from[0][0],
          'client_name':    self.connect_from[1],
          'helo_name':      self.helo_name,
          'sender':         mail.sender,
          'recipient':      recipient_addr
        }
        policy_reply=checkpolicy(globals.recipient_policy,True)
        self.stats.policy_update()
        if policy_reply[0]=='4':
          return Milter.TEMPFAIL
        elif policy_reply[0]=='5':
          return Milter.REJECT
        return Milter.CONTINUE

    def header(self,name,val):
        msg=name+": "+val+"\r\n"
        self.mail.df.write(msg)
        return Milter.CONTINUE

    def eoh(self):
        self.mail.df.write("\r\n")
        return Milter.CONTINUE

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

    def eom(self):
        self.mail.comm+="DATA\r\n"
        # reinit mail class and store mail data
        mail.__init__()
        mail.df=self.mail.df
        mail.comm=self.mail.comm
        mail.sender=self.mail.sender
        mail.recip=self.mail.recip
        mail.close()
        v,level,virname=checkvir(self.SCANNERS)
        debug.echo(2,"STATS: %s seconds, %s bytes, status: %s"
                     % (self.stats.end(), len(mail.data), virname))
        if (v==S_OK) | (v==S_FORCE_SEND):
          self.log("ACCEPT: 250 Ok",1)
          self.stats.update(len(mail.data))
          return Milter.ACCEPT
        elif v==S_REJECT:
          msg="550 Content rejected - VIRUS "+virname;
          self.log("REJECT: "+msg,1)
          self.stats.update(len(mail.data), 1)
          self.setreply(msg[0:3],None,msg[4:])
          return Milter.REJECT
        elif v==S_DROP:
          self.log("DROP: 250 mail dropped, VIRUS "+virname,1)
          self.stats.update(len(mail.data), 1)
          return Milter.DISCARD
        self.log("TEMPFAIL: 451 "+str(virname),1)
        self.stats.update(tempfail=1) # update fail statistics
        return Milter.TEMPFAIL

    def close(self):
        self.log("Connection closed.")
        return Milter.CONTINUE

    def abort(self):
        self.log("Connection aborted!")
        return Milter.CONTINUE

  def milterInit(scanners,name,conn):
      sgMilter.SCANNERS=scanners
      Milter.factory = sgMilter
      Milter.set_flags(Milter.CHGBODY + Milter.CHGHDRS + Milter.ADDHDRS)
      Milter.runmilter(name,conn,256)

except ImportError:
  def milterInit(scanners,name,conn):
      debug.echo(0,"ImportError: Can't import Milter module!")
      debug.echo(0,"Install sendmail milters in python from:")
      debug.echo(0,"  http://www.bmsi.com/python/milter.html")
      sys.exit(1)

class milter(service):
  '''
  Milter support service.
  
  This service can be used to start sagator as milter filter.
  
  Usage: milter(scanners, name, connection)
  
  Example: milter(SCANNERS, "sagator", "inet:3333@127.0.0.1")
  
  For more information about milter's parameters see milter documentation.
  
  You need python's milter module to run this service:
    http://www.bmsi.com/python/milter.html
  '''
  name = 'milter()'
  MIN_CHILDS = 1
  def __init__(self,scanners,name,conn):
      self.SCANNERS=scanners
      self.m_name=name
      self.m_conn=conn
      self.EXITING=False
      self.childs=[]
  def start(self):
      self.test_scanners(self.SCANNERS)
      pid=self.fork()
      return [pid]
  def fork(self):
      if self.EXITING:
        return -1
      if self.childs!=[]:
        return -1
      p=os.fork()
      if p==0:
        signal.signal(signal.SIGHUP,self.sighup)
        signal.signal(signal.SIGTERM,self.sigterm)
        dochroot()
        debug.echo(1, "milter(): service started, waiting for connections ...")
        milterInit(self.SCANNERS, self.m_name, self.m_conn)
        debug.echo(1, "milter(): Exiting ...")
        sys.exit(0)
      else:
        self.childs.append(p)
      return p
