'''
Antivir Gateway Library for Sagator

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

This program is a email antivir gate for an SMTP server.

 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.

'''

import os,sys,string,re,time
import pwd,grp,socket,rfc822

from avlib import *
import stats

TEST_MSG='''\
From: Test User <test@tester.salstar.sk>\r
To: Test User2 <test2@tester2.salstar.sk>\r
Date: Fri, 16 Apr 2004 15:38:51 +0200 (CEST)\r
Subject: testing mail\r
\r
simple mail\r
'''

#########################################################################
### service class template

class service:
  '''Sagator service template'''
  name='ASrv()'
  def __init__(self,scanners,host,port,prefork=2,listen=5):
      self.SCANNERS=scanners
      self.BINDTO=(host,port)
      self.LISTEN=listen
      self.PREFORK=prefork
      self.EXITING=False
  def daemonize(self):
      # redirect stout and strerr
      debug.dup()
      # fork
      pid=os.fork()
      if (pid<=0) and (globals.daemon):
        globals.pidf.close()
      return pid
  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.BINDTO)
        self.s.listen(self.LISTEN)
      except socket.gaierror,(ec,es):
        debug.echo(0,self.name+": 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,self.name+": 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,self.name+": pids: ",str(self.childs))
      try:
        while True:
          time.sleep(60)
      except KeyboardInterrupt:
        self.sigterm(None,None) # kill childs
        raise
  def accept(self,connects=0):
      debug.echo(0,self.name+": This method must be overwrtten in subclass!")
      sys.exit(1)
  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)
      self.time1,self.timec='',1
      while True:
        try:
          self.accept()
        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,self.name+": Connection closed by peer.")
          elif ec==104: # Connection reset by peer
            debug.echo(1,self.name+": Connection reset by peer.")
          elif ec==32: # 'Broken pipe'
            debug.echo(0,self.name+": ERROR: Broken pipe.")
          elif ec!=4:
            debug.echo(0,self.name+': ERROR: SocketError: ',es)
            debug.traceback(4,self.name+": ")
        except KeyboardInterrupt:
          self.sigterm(None,None)
        except MemoryError:
          debug.echo(1,self.name+": ERROR: OUT OF MEMORY! EXITING PROCESS!")
          # it's safe to exit process now
          sys.exit(1)
        except:
          debug.echo(1,self.name+": ERROR: Unknown error!")
          debug.traceback(1,self.name+": ")
      return 0 # child process
  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,self.name+": Child exited ... pid,status: ",(pid,s))
      except OSError: # no child processes?
        pass
      if not self.EXITING:
        # start another child
        debug.echo(1,self.name+": WARNING: An child exited? Starting new one ...")
        self.fork()
  def sigterm(self,sn,stack):
      self.EXITING=True
      if os.getpid()!=self.PID:
        debug.echo(1,self.name[:-2]+"(child): Exiting - SIGTERM ...")
        os._exit(0)
      # kill all childs
      for pid in self.childs:
        debug.echo(5,self.name[:-2]+'(master): killing child %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,self.name[:-2]+'(master): killing child %d (KILL signal)' % pid)
          try:
            os.kill(pid,signal.SIGKILL)
          except OSError:
            pass
      os._exit(0)

#########################################################################
### Some helpful functions

# chroot, setgid and setuid
def dochroot():
    '''
    Do chroot() and setuid(),setgid()
    '''
    def func_err(fn,f):
        try:
          eval(f)
          return 0
        except OSError,(ec,es):
          debug.echo(0,"WARNING! Can't do %10s: %s" % (fn,es))
        except AttributeError:
          debug.echo(0,"WARNING! Can't do "+f+": Your python don't have this function!")
        return 1
    uid,gid=0,0
    if globals.USER:
      try:
        uid=pwd.getpwnam(globals.USER)[2]
      except KeyError,es:
        debug.echo(0,"ERROR, getpwnam: ",[[globals.USER]],es)
        sys.exit(100)
    if globals.GROUP:
      try:
        gid=grp.getgrnam(globals.GROUP)[2]
      except KeyError,es:
        debug.echo(0,"ERROR, getgrnam: ",[[globals.GROUP]],es)
        sys.exit(100)
    if globals.CHROOT:
      if func_err('chroot',"os.chroot('"+globals.CHROOT+"')"):
        safe.ROOT_PATH=globals.CHROOT
    if gid:
      func_err('setgroups',"os.setgroups(["+str(gid)+"])")
      func_err('setgid',"os.setgid("+str(gid)+")")
    if uid:
      func_err('setuid',"os.setuid("+str(uid)+")")
    if (os.getuid()==0) or (os.getgid()==0
        or (os.geteuid()==0) or (os.getegid()==0)):
      debug.echo(0,"ERROR! I can't setuid and i am running as root!")
      sys.exit(10)

def test_scanners(scanners):
    for scnr in scanners:
      scnr.reinit()
    if os.fork()==0:
      dochroot()
      globals.scan_only=True
      scanarr,scannames=[],'None'
      mail.data=TEST_MSG
      for scnr in scanners:
        if scnr.name:
          debug.echo(1, "Testing "+scnr.name+"...")
        level,detected,virlist,scan_report,err=do_scan(scnr)
        if scan_report:
          debug.echo(0, "Scanner ",scnr.name," test failed! Disable it manually!")
          debug.echo(0, "  ",scan_report)
        scanarr=scanarr+[scnr.name]
      mail.data=''
      sys.exit(0)
    os.wait()

def do_scan(scnr,fn=''):
    '''
    Scan email (stored in mail class) for viruses.
    Use scnr scanner.
    '''
    level,detected,virlist=0.0,'',[]
    try:
      if fn:
        scnr.filename=fn
      scnr.prescan()
      level,detected,virlist=scnr.scanstream(mail.data)
      scnr.postscan(level,detected,virlist)
      scnr.destroy()
      return level,detected,virlist,'',None
    except socket.error,(ec,es):
      debug.traceback(4,"do_scan: ")
      return level,detected,virlist,\
        scnr.name+": SocketError: "+es,'Socket error: '+es
    except (IOError,OSError):
      debug.traceback(4,"do_scan: ")
      return level,detected,virlist,\
        scnr.name+": Error: "+debug.traceback_value_str(),'Error: '+debug.traceback_value_str()
    except (ScannerError,ValueError),es:
      debug.traceback(4,"do_scan: ")
      return level,detected,virlist,\
        scnr.name+": "+str(es),'Error: '+str(es)
    except KeyboardInterrupt:
      raise
    except:
      debug.traceback(4,"do_scan: ")
      return level,detected,virlist,scnr.name+": unknown error",'Unknown error'

def do_param(p,scanners):
    n=0
    for scnr in scanners:
      n+=scnr.param(p)
    return n

def checkvir(scanners):
    '''
    Check email (stored in mail class) for viruses.
    '''
    try:
      level,detected,virlist=0.0,'',[]
      scanarr,scannames=[],'None'
      for scnr in scanners:
        l,detected,virlist,scan_reply,err=do_scan(scnr)
        level+=l
        if scan_reply:
          debug.echo(0, "CheckVir: ERROR: ",scan_reply)
          return S_TEMPFAIL,scan_reply
        else:
          scanarr=scanarr+[scnr.name]
          scannames=string.join(scanarr,' ')
        if is_infected(level,detected):
          debug.echo(3,"Status: ",[detected,virlist,level])
          break
      mail.addvirusinfo(detected,scannames)
      if not is_infected(level,detected):
        return S_OK,'CLEAN' # mail delivered OK
      else: # VIRUS FOUND
        rstat=globals.ACTION
        debug.echo(5,"action: ",rstat)
        return rstat,detected
    except socket.error,(ec,es):
      debug.echo(0,"checkvir: ERROR: socket.error: ",[[ec]],es)
      debug.traceback(4,"checkvir: ")
      return S_TEMPFAIL,es
    except IOError,(ec,es):
      debug.echo(0,"checkvir: ERROR: IO: ",es)
      debug.traceback(4,"checkvir: ")
      return S_TEMPFAIL,'IOError'
    except:
      debug.traceback(3,"checkvir: ")
      return S_TEMPFAIL,'AnyError, see sagator logs'

#########################################################################
### Signal handlers

def sigchld(sn,stack):
    ''' signal: Child process terminated '''
    # try to remove older zombies (sigchld lost?)
    try:
      while 1:
        pid,s=os.waitpid(0,os.WNOHANG)
        if pid==0:
          break
        debug.echo(1,"Child exited ... pid,status: ",(pid,s))
    except OSError: # no child processes?
      pass

def sighup(sn,stack):
    ''' signal: HUP signal '''
    debug.echo(3,"SIGHUP: Reopening log ...")
    debug.reopen()

def sigusr2(sn,stack):
    ''' signal: USR2 signal '''
    debug.echo(1,"SIGUSR2: Tracelog request received ...")
    debug.debug_level=debug.trace_level
    debug.stack(1,"SIGUSR2: ")

def sigterm(sn,stack=None):
    ''' signal: Terminate process '''
    if type(sn)==type(''):
      debug.echo(1,"SERVER: Exiting - %s ..." % sn)
    else:
      debug.echo(1,"SERVER: Exiting - SIGTERM ...")
    try:
      os.unlink(globals.pid_file)
    except OSError:
      pass
    try:
      if globals.stat_coll.pid>0:
        os.kill(globals.stat_coll.pid,signal.SIGTERM)
    except:
      pass
    sigchld(None,None) # unzombie all childs
    sys.exit(0)
