#!/usr/bin/python2

'''
milter.py - milter service for sagator

(c) 2003-2004 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 signal
import sys,time

__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.mail=mail_class()
        return Milter.CONTINUE

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

    def envfrom(self,f,*s):
        self.log('MAIL FROM: '+f,2)
        self.t0=time.time()
        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:
          self.mail.recip.append(rfc822.parseaddr(to)[1])
        except:
          pass
        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,virname=checkvir(self.SCANNERS)
        self.tw=time.time()-self.t0
        debug.echo(2,"STATS: "+str(self.tw)+" seconds, ",
                     len(mail.data)," bytes, status: ",virname)
        if (v==S_OK) | (v==S_FORCE_SEND):
          self.log("ACCEPT: 250 Ok",1)
          stats.update(len(mail.data),self.tw)
          return Milter.ACCEPT
        elif v==S_REJECT:
          msg="550 Content rejected - VIRUS "+virname;
          self.log("REJECT: "+msg,1)
          stats.update(len(mail.data),self.tw,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)
          stats.update(len(mail.data),self.tw,1)
          return Milter.DISCARD
        self.log("TEMPFAIL: 451 "+str(virname),1)
        stats.update(0,0,0,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,es:
  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 service for sagator.
  
  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()'
  def __init__(self,scanners,name,conn):
      self.SCANNERS=scanners
      self.m_name=name
      self.m_conn=conn
      self.EXITING=False
      self.childs=[]
  def main(self):
      test_scanners(self.SCANNERS)
      # chroot
      dochroot()
      # and now start the daemon
      signal.signal(signal.SIGCHLD,self.sigchld)
      signal.signal(signal.SIGHUP,sighup)
      signal.signal(signal.SIGTERM,self.sigterm)
      pid=self.daemonize()
      if pid>0:
        return [pid]
      self.PID=os.getpid()
      self.fork()
      # wait for child exit
      while True:
        time.sleep(60)
      sigterm('Terminating') # kill self
      sys.exit(0)
      return [] # required by pychecker
  def fork(self):
      if self.EXITING:
        return -1
      if self.childs!=[]:
        return -1
      p=os.fork()
      if p==0:
        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
