#!/usr/bin/python

'''sgscan - sagator's mailbox/email scanner

(c) 2003-2010 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.
'''

__pychecker__ = 'unusednames=_ \
                 maxreturns=20 maxlines=300 \
                 no-local no-argsused no-stringiter no-abstract no-isliteral \
                 no-stdlib no-callinit no-badexcept'

SAVECLEAN = ''
SAVEINFECTED = ''
AV_ONLY = 0
VERBOSE = 0
QUIET = 0

import sys,os,fcntl,re,time,struct,shutil
from cStringIO import StringIO
from aglib import *

# load config
try:
  from etc import *
except ImportError,err_str:
  if str(err_str)[-6:]=="etc":
    print "ERROR! Config file not found! Exiting now."
    sys.exit(1)
  else:
    raise

def totime(t):
    '''Convert seconds to string time representation.'''
    return "%d:%02d:%02d" % (t/60/60, (t/60)%60, t%60)

def scan(data,fn,mailcount=-1):
    '''Scan data for viruses/spams.'''
    mail.data = data
    mail.xheader = ''
    debug.echo(2, "Object: ",fn)
    scanarr,scannames = [],'None'
    globals.reset()
    if fn[0:5]=="From ":
      fn2 = fn.split()
      mail.sender=fn2[1]
      fn2 = ' '.join([fn2[1]]+fn2[3:])
    else:
      fn2 = fn
    level = 0.0
    for scnr in SCANNERS:
      if (AV_ONLY>0) and (scnr.get('is_spamscan')>0): # if it contain a spamscan
        debug.echo(4,"Spam scanner skipped: ",scnr.name)
      else:
        l,detected,virlist,scan_reply,err=do_scan(scnr,os.path.basename(fn))
        if err:
          if mailcount>0:
            print "%2d: %s: %s" % (mailcount,fn2,err)
          else:
            print "%s: %s" % (fn,err)
          return False
        level += l
        if scan_reply:
          debug.echo(2, scan_reply)
        else:
          scanarr = scanarr+[scnr.name]
          scannames = ' '.join(scanarr)
        if is_infected(l,detected):
          if not QUIET:
            if mailcount>0:
              print "%2d: %s: %s [%s,%0.2f]" % (mailcount,fn2,detected,
                globals.found_by.name,level)
            else:
              print "%s: %s [%s,%0.2f]" % (fn,detected,
                globals.found_by.name,level)
          return True
    if VERBOSE:
      if mailcount>0:
        print "%2d: %s: CLEAN" % (mailcount,fn2)
      else:
        print fn+": CLEAN"
    return False

def usage():
    '''Show help.'''
    print "sgscan - sagator's file scanner"
    print ""
    print "Usage: sgscan [params...] files..."
    print ""
    print "Params: --help        this help"
    print "        --config=f    load \"f\" as alternate config (without .py extension)"
    print "        --logfile=l   filename for logging ('-' for stdout)"
    print "        --debug=l     set debug level to l"
    print "        --av-only     don't run antispam test, only antivir tests"
    print "        --clean       clean infected emails from mailboxes"
    print "        --progress    show progress"
    print "        --quiet       don't show INFECTED emails"
    print "        --verbose     show also CLEAN emails"
    for scnr in SCANNERS:
      for key,value in scnr.help().items():
        print "        --%11s %s" % (key,value)
    sys.exit(0)

class progress:
  hidden = False
  def start(self, max):
      self.max = max
      self.counter = 0
      self.time0 = time.time()
  def update(self, count=1):
      self.counter += count
      mps=self.counter/(time.time()-self.time0)
      sys.stderr.write(
        "Progress: %d/%d, infected/spams: %d [%5.1f%%], ETA: %s           \r" \
        % (self.counter, self.max, infected, 100.0*infected/self.counter,
           totime((self.max-self.counter)/mps))
      )
  def end(self):
      # overwrite progress indicator with spaces
      print " "*78, "\r",

class hidden_progress(progress):
  hidden = True
  def update(self, count=1):
      pass
  def end(self):
      pass

def scanfile(fn,progress_indicator):
    global infected, total
    while 1:
      try:
        try:
          os.chdir(cwd)
          if SAVECLEAN:
            f = open(fn,"r+")
          else:
            f = open(fn,"r")
        except IOError,(err_code,err_str):
          print "%s: IOError: %s [%d], skipping file" % (fn,err_str,err_code)
          break
        try:
          line1 = f.readline()
        except OverflowError:
          # too large line, not a mailbox, ignore
          return
        if line1[0:5]=="From ":
          print fn+": mailbox detected"
          if not progress_indicator.hidden:
            # count emails
            mailcount = 1
            while True:
              l = f.readline()
              if not l:
                break
              if l[:5]=='From ':
                mailcount += 1
            f.seek(0)
            f.readline()
            progress_indicator.start(mailcount)
          if SAVECLEAN:
            rv=fcntl.fcntl(f.fileno(), fcntl.F_SETFL, os.O_NDELAY)
            rv=fcntl.fcntl(f.fileno(), fcntl.F_SETLKW,
                           struct.pack('hhqqhh', fcntl.F_WRLCK, 0,0,0,0,0))
            fc = open(fn+SAVECLEAN, "w")
            os.chmod(fn+SAVECLEAN, 0600)
          if SAVEINFECTED:
            fi = open(fn+SAVEINFECTED, "w")
            os.chmod(fn+SAVEINFECTED, 0600)
          mailcount=0
          box_infected = 0
          time0 = time.time()
          while 1:
            progress_indicator.update()
            frm = line1
            lines = StringIO()
            while 1:
              line1 = f.readline()
              if line1=="": break
              if line1[0:5]=="From ": break
              lines.write(line1)
            if not scan(lines.getvalue(), frm.strip(), mailcount):
              if SAVECLEAN:
                fc.write(frm)
                fc.write(lines.getvalue())
            else:
              infected += 1
              box_infected += 1
              if SAVEINFECTED:
                fi.write(frm)
                fi.write(mail.xheader)
                fi.write(lines.getvalue())
            total += 1
            lines.close()
            if line1=="": break
          # remove .clean file if no viruses found
          if SAVECLEAN and (not SAVEINFECTED):
            os.chdir(cwd)
            if box_infected>0:
              # mailbox cleaned, viruses found
              fc.close()
              os.lseek(f.fileno(), 0, 0)
              os.ftruncate(f.fileno(), 0)
              shutil.copyfileobj(open(fn+SAVECLEAN,'r'), f, 1024*1024)
            try:
              os.unlink(fn+SAVECLEAN)
            except OSError:
              pass
        elif re.search('^NOOP CONNECT FROM: ',line1):
          print fn+": sagator's quarantine file detected"
          # skip to "data"
          while 1:
            line1=f.readline()
            if re.compile("^DATA",re.IGNORECASE).search(line1):
              break
          lines=StringIO()
          while 1:
            line1=f.readline()
            if (line1=='.\n') | (line1=='.\r\n'):
              break
            lines.write(line1)
          if scan(lines.getvalue(),os.path.basename(fn)):
            infected += 1
          total += 1
          lines.close()
          break # skip to next file
        else:
          # not mailbox, possible maildir file
          debug.echo(2, "Scanning file: "+fn)
          if scan(open(fn, 'r').read(50*1024*1024), fn):
            infected += 1
            if SAVECLEAN:
              os.unlink(fn)
          total += 1
        f.close()
        break
      except IOError, (err_code, err_str):
        if err_code!=35: # Resource deadlock avoided
          raise
        debug.echo(0, "Deadlock avoided, repeating last scan...")
      except KeyboardInterrupt:
        progress_indicator.end()
        debug.echo(0, "Interrupted! Exiting!")
        sys.exit(1)
    progress_indicator.end()

if __name__ == '__main__':
  debug.set_logfile('-')
  debug.set_level(0)
  globals.scan_only = True
  PROGRESS = hidden_progress()
  # parse command line arguments
  files=[]
  n=1
  try:
    opts=['help', 'config=', 'clean', 'separe', 'av-only', 'progress',
          'quiet', 'verbose', 'debug=']
    for scnr in SCANNERS:
      opts.extend(scnr.help().keys())
    opts,files=getopt.gnu_getopt(sys.argv[1:],'',opts)
  except getopt.GetoptError, (msg, opt):
    print "Error:",msg
    sys.exit(1)
    
  for key,value in opts:
    if key=="--help":
      usage()
    elif key=='--config':
      try:
        exec("from %s import *" % value)
      except ImportError,err_str:
        print "ImportError:",err_str
    elif key=="--clean":
      SAVECLEAN=".clean"
    elif key=="--separe":
      SAVECLEAN=".clean"
      SAVEINFECTED=".infected"
    elif key=="--av-only":
      AV_ONLY=1
    elif key=="--progress":
      PROGRESS=progress()
    elif key=="--quiet":
      QUIET=1
    elif key=="--verbose":
      VERBOSE=1
    elif key=="--debug":
      debug.set_level(int(value))
    else:
      for scnr in SCANNERS:
        err=scnr.param(key,value)
        if err:
          print "Unrecognized parameter (%s). %s" % ('='.join([key,value]),err)
          sys.exit(1)

  # reinit scanners
  for scnr in SCANNERS:
    scnr.reinit()
  safe.ROOT_PATH = CHROOT

  cwd = os.getcwd()
  total,infected = 0,0
  begintime = time.time()
  for fn in files:
    if os.path.isdir(fn):
      print "%s: directory detected" % fn
      PROGRESS.start(sum([len(c) for a,b,c in os.walk(fn)]))
      for fpath,fdirs,fnames in os.walk(fn):
        for fn in fnames:
          PROGRESS.update()
          scanfile(os.path.join(fpath,fn), hidden_progress())
      PROGRESS.end()
    else:
      scanfile(fn, PROGRESS)
  print "Total infected/spams emails: %d/%d [%5.1f%%] in: %1.2f seconds" \
        % (infected, total, 100.0*infected/total, time.time()-begintime)