'''
smtpd.py - SMTP daemon 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 *
import signal,time
import stats

__all__=['smtpd']

class smtps(smtp):
  '''
  SMTP server class
  '''
  reg_smtp_reply=re.compile("^[0-9]{3} ")
  reg_smtp_reply_ok=re.compile("^[23][0-9][0-9] ")
  reg_xforward=re.compile("^XFORWARD ",re.IGNORECASE)
  reg_mailfrom=re.compile("^MAIL +FROM:(.*)$",re.IGNORECASE)
  reg_rcptto=re.compile("^RCPT +TO:(.*)$",re.IGNORECASE)
  reg_rset=re.compile("^RSET",re.IGNORECASE)
  reg_data=re.compile("^DATA",re.IGNORECASE)
  reg_quit=re.compile("^QUIT",re.IGNORECASE)
  def __init__(self,scanners,conn,addr):
    self.SCANNERS=scanners
    self.f=conn.makefile('rw',self.bufsize)
    self.write=conn.sendall
    self.readline=self.f.readline
    self.RECV=RECV_SMTP
    self.t0=time.time()
    self.srv=smtpc()
    self.write(self.srv.readline())
    mail.__init__()
    mail.addr='NOOP CONNECT FROM: '+addr[0]+':'+str(addr[1])\
      +', '+mail.getnamebyaddr(addr[0])+'\r\n'
  def send_ok(self,r="Ok"):
    self.write("250 "+r+"\r\n")
  def quit_connect(self):
    debug.echo(5,"SMTPS: Closing connect to server")
    self.RECV=RECV_QUIT
    self.srv.close()
    return self.recv()
  def check_reply(self,line):
    '''Check for SMTP reply. Return True, if 2XX or 3XX returned.'''
    self.srv.conn.sendall(line)
    # send reply
    while 1:
      l=self.srv.readline()
      if not l:
        debug.echo(3,"SMTPS: Connection closed after:",line.rstrip())
        return False
      debug.echo(7,'SMTPC: '+l)
      self.write(l)
      if self.reg_smtp_reply.search(l):
        break
    if self.reg_smtp_reply_ok.search(l):
      return True
    debug.echo(4,"SMTPS: "+l.rstrip()+' --- '+line.rstrip())
    return False
  def recv(self):
    try:
      line=self.readline()
    except IOError,(ec,es):
      debug.echo(0, "smtps(): ",es)
      return 1
    if not line: return 1
    if self.RECV: # receiving data or header
      if self.RECV==RECV_QUIT: # waiting for quit?
        if self.reg_quit.search(line):
          self.write("221 Bye\r\n")
          self.f.close()
          return 1
        else:
          self.write("502 SMTP communication stopped by Sagator\r\n")
          return 0
      elif (self.RECV==RECV_HEADER) & ((line=='\r\n') | (line=='\n')):
        # receiving header
        debug.echo(5,"SMTPS: HEADER DONE")
        mail.df.write(line)
        self.RECV=RECV_BODY
      elif (line=='.\r\n') | (line=='.\n'):
        # end of DATA part
        mail.close()
        if len(mail.data)==len(mail.header):
          debug.echo(0,"SMTPS: Warning: No header?")
          mail.header=''
        debug.echo(3,"SMTPS: BODY DONE, size: ",len(mail.data)," B")
        self.RECV=RECV_SMTP
        v,virname=checkvir(self.SCANNERS)
        self.tw=time.time()-self.t0
        if v!=S_TEMPFAIL:
          debug.echo(2,"STATS: "+str(self.tw)+
            " seconds, ",len(mail.data)," bytes, status: ",virname)
        if (v==S_OK) | (v==S_FORCE_SEND):
          debug.echo(5,"SMTPS: Sending data")
          self.srv.conn.sendall(mail.header)
          self.srv.conn.sendall(mail.xheader)
          self.srv.conn.sendall(mail.data[mail.bodypos:])
          debug.echo(1,"SMTPS: OK: 250 Ok")
          self.srv.conn.sendall(line) # send .
          self.write(self.srv.readline())
          stats.update(len(mail.data),self.tw)
          mail.__init__() # reinit
        elif v==S_REJECT:
          debug.echo(1,"SMTPS: REJECT: ","550 Content rejected - VIRUS ",virname)
          self.write("550 Content rejected - VIRUS "+str(virname)+"\r\n")
          stats.update(len(mail.data),self.tw,1)
          return self.quit_connect()
        elif v==S_DROP:
          debug.echo(1,"SMTPS: DROP: 250 mail dropped, VIRUS ",virname)
          self.write("250 mail dropped, VIRUS "+str(virname)+"\r\n")
          stats.update(len(mail.data),self.tw,1)
          return self.quit_connect()
        else: # S_TEMPFAIL
          debug.echo(1,"SMTPS: TEMPFAIL: 451 ",virname)
          self.write("451 "+str(virname)+"\r\n")
          stats.update(0,0,0,1) # update fail statistics
          return self.quit_connect()
      else:
        if self.RECV==RECV_HEADER:
          mail.hf.write(line)
        else:
          mail.df.write(line)
    else:
      mail.comm=mail.comm+line
      rcptto=self.reg_rcptto.search(line)
      mailfrom=self.reg_mailfrom.search(line)
      if rcptto:
        if self.check_reply(line):
          debug.echo(2,'SMTPS: '+line.rstrip())
          try:
            mail.recip.append(rfc822.parseaddr(rcptto.group(1))[1])
          except:
            pass
        return 0
      elif mailfrom:
        if self.check_reply(line):
          debug.echo(2,'SMTPS: '+line.rstrip())
          mail.sender=mailfrom.group(1)
      elif self.reg_rset.search(line):
        if self.check_reply(line):
          debug.echo(2,'SMTPS: '+line.rstrip(),", Resetting mail class")
          mail.__init__() # reinit
      elif self.reg_data.search(line):
        if self.check_reply(line):
          debug.echo(5,'SMTPS: '+line.rstrip())
          self.RECV=RECV_HEADER
      elif self.reg_quit.search(line):
        debug.echo(5,'SMTPS: '+line.rstrip())
        self.check_reply(line) # 221 Bye
        return 1
      elif self.reg_xforward.search(line):
        debug.echo(2,'SMTPS: ',line.rstrip())
        self.check_reply(line)
      else:
        debug.echo(6,'SMTPS: '+line.rstrip())
        self.check_reply(line)
    return 0

class smtpd(service):
  '''
  smtpd() service for sagator.
  
  This service can be used to start sagator as separate filtering SMTP
  daemon. Is is useful for postfix and any other SMTP daemon, which
  can use these filters.
  
  Usage: smtpd(scanners, host, port [,prefork])
  
  Where: scanners is an array of scanners (see README.scanners for more info)
         host is a an ip address to bind
         port is a port to bind
         prefork is a number, which defines preforked process count
  
  Example: smtpd('127.0.0.1', 27)
  '''
  name='smtpd()'
  def __init__(self,scanners,host,port,prefork=5,listen=5):
      self.SCANNERS=scanners
      self.HOST=host
      self.PORT=port
      self.LISTEN=listen
      self.PREFORK=prefork
      self.EXITING=False
  def main(self):
      test_scanners(self.SCANNERS)
      # bind to socket
      self.s=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
      self.s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
      try:
        self.s.bind((self.HOST,self.PORT))
        self.s.listen(self.LISTEN)
      except socket.gaierror,(ec,es):
        debug.echo(0,"smtpd(): ERROR: BIND ",es)
        sys.exit(1)
      # and now start the daemon
      signal.signal(signal.SIGCHLD,self.sigchld)
      signal.signal(signal.SIGHUP,sighup)
      signal.signal(signal.SIGTERM,signal.SIG_DFL)
      pid=self.daemonize()
      if pid>0:
        self.s.close()
        signal.signal(signal.SIGCHLD,signal.SIG_DFL)
        return [pid]
      # chroot
      dochroot()
      debug.echo(1,"smtpd(): service started, waiting for connections ...")
      # prefork
      signal.signal(signal.SIGTERM,self.sigterm)
      signal.signal(signal.SIGHUP,self.sighup)
      self.PID=os.getpid()
      self.childs=[]
      for i in range(self.PREFORK):
        self.fork()
      debug.echo(3,"smtpd(): pids: ",str(self.childs))
      try:
        while True:
          time.sleep(60)
      except KeyboardInterrupt:
        self.sigterm(None,None) # kill childs
        raise
  def fork(self):
      try:
        pid=os.fork()
        if pid>0:
          self.childs.append(pid)
          return pid
      except OSError,(ec,es):
        debug.echo(1,"ERROR: Can't fork:",es)
        return -1
      signal.signal(signal.SIGCHLD,signal.SIG_DFL)
      signal.signal(signal.SIGHUP,sighup) # reopen log
      signal.signal(signal.SIGTERM,signal.SIG_DFL)
      time1,timec='',1
      while True:
        try:
          # accept
          conn,addr = self.s.accept()
          socket_settimeout(conn,120)
          # reinit scanners
          for scnr in self.SCANNERS:
            scnr.reinit()
          # generate ID
          time2=time.strftime("%Y%m%d-%H%M%S",time.localtime(time.time()))
          if time2!=time1:
            time1,timec=time2,1
          globals.id="%s-%04d-%05d-%s" % (time2,timec,os.getpid(),randomchars(6))
          debug.echo(1,"Connection from: "+addr[0]+" at "+\
            time.strftime("%c",time.localtime()))
          debug.echo(1,"Process id: %s" % globals.id)
          m=smtps(self.SCANNERS,conn,addr)
          while 1:
            if m.recv(): break
          conn.shutdown(2)
          debug.echo(1,"Closing connection.")
          conn.close()
        except SmtpcError,(ses,es):
          # error is written, accept new connection
          pass
        except socket.error,err:
          try:
            ec=err[0]
            es=err[1]
          except IndexError:
            ec=0
            es=str(err)
          if ec==107: # Transport endpoint is not connected
            debug.echo(1,"smtpd(): Connection closed by peer.")
          elif ec==104: # Connection reset by peer
            debug.echo(1,"smtpd(): Connection reset by peer.")
          elif ec==32: # 'Broken pipe'
            debug.echo(0,"smtpd(): ERROR: Broken pipe.")
          elif ec!=4:
            debug.echo(0,'smtpd(): ERROR: SocketError: ',es)
            debug.traceback(4,"smtpd(): ")
        except KeyboardInterrupt:
          self.sigterm(None,None)
        except:
          debug.echo(1,"ERROR: smtpd(): Unknow error!")
          debug.traceback(1,"smtpd(): ")
  def sighup(self,sn,stack):
      # reopen log for master process
      sighup(sn,stack)
      # reopen logs for all childs
      for pid in self.childs:
        try:
          os.kill(pid,signal.SIGHUP)
        except:
          pass
  def sigchld(self,sn,stack):
      # try to remove older zombies (sigchld lost?)
      try:
        while 1:
          pid,s=os.waitpid(0,os.WNOHANG)
          if pid==0:
            break
          if pid in self.childs:
            self.childs.remove(pid)
          debug.echo(1,"smtpd(): Child exited ... pid,status: ",(pid,s))
      except OSError: # no child processes?
        pass
      if not self.EXITING:
        # start another child
        debug.echo(1,"WARNING: smtpd(): An child exited? Starting new one ...")
        self.fork()
  def sigterm(self,sn,stack):
      self.EXITING=True
      if os.getpid()!=self.PID:
        debug.echo(1,"smtpd(child): Exiting - SIGTERM ...")
        os._exit(0)
      # kill all childs
      for pid in self.childs:
        debug.echo(5,'smtpd(master): killing children %d (TERM signal)' % pid)
        try:
          os.kill(pid,signal.SIGTERM)
        except OSError:
          pass
      time.sleep(0.1)
      if self.childs:
        time.sleep(1)
        for pid in self.childs:
          debug.echo(3,'smtpd(master): killing children %d (KILL signal)' % pid)
          try:
            os.kill(pid,signal.SIGKILL)
          except OSError:
            pass
      os._exit(0)
