'''
milter.py - milter service for sagator

(c) 2003-2016,2019 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 *
import sys, time, signal, stats

__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 = parseaddr(tostr(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), tostr(virname)))
        if (v==S_OK) | (v==S_FORCE_SEND):
          self.log("ACCEPT: 250 Ok",1)
          # add headers
          for key, value in mail.xheaders(' \r'):
            self.addheader(key, value)
          self.stats.update(len(mail.data))
          return Milter.ACCEPT
        elif v==S_REJECT:
          msg="550 Content rejected - VIRUS "+tostr(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 "+tostr(virname), 1)
          self.stats.update(len(mail.data), 1)
          return Milter.DISCARD
        self.log("TEMPFAIL: 451 "+tostr(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, umask):
      sgMilter.SCANNERS=scanners
      Milter.factory = sgMilter
      Milter.set_flags(Milter.CHGBODY + Milter.CHGHDRS + Milter.ADDHDRS)
      if umask is not None:
        os.umask(umask)
      Milter.runmilter(name, conn, 256)

except ImportError:
  def milterInit(scanners, name, conn, umask):
      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, umask)
  
  Where: scanners is an array of scanners (see README.scanners for more info)
         name is an string, milter service name
         connection in an string, which defines, where should milter service
           listen
         umask is an integer, which defines which umask should be set before
           creating local socket

  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, umask=None):
      self.SCANNERS = scanners
      self.m_name = name
      self.m_conn = conn
      self.m_umask = umask
      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, self.m_umask)
        debug.echo(1, "milter(): Exiting ...")
        sys.exit(0)
      else:
        self.childs.append(p)
      return p