#!/usr/bin/python3

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

(c) 2003-2018 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 __future__ import absolute_import
from __future__ import print_function

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

import sys, os, fcntl, re, time, struct, shutil
from avlib import BytesIO
from aglib import *

# load config
try:
  from etc import *
except ImportError as 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 = b''
    debug.echo(2, "Object: ",fn)
    scanarr,scannames = [],'None'
    globals.reset()
    if fn[0:5]==b"From ":
      fn2 = fn.split()
      mail.sender=fn2[1]
      fn2 = b' '.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" % (tostr(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, tostr(fn2),
                tostr(detected), globals.found_by.name, level))
            else:
              print("%s: %s [%s,%0.2f]" % (tostr(fn),
                tostr(detected), globals.found_by.name, level))
          return True
    if VERBOSE:
      if mailcount>0:
        print("%2d: %s: CLEAN" % (mailcount, tostr(fn2)))
      else:
        print("%s: CLEAN" % tostr(fn))
    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", end=' ')

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,"rb+")
          else:
            f = open(fn,"rb")
        except IOError as err:
          (err_code,err_str) = err.args
          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]==b"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]==b'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, "wb")
            os.chmod(fn+SAVECLEAN, 0o600)
          if SAVEINFECTED:
            fi = open(fn+SAVEINFECTED, "wb")
            os.chmod(fn+SAVEINFECTED, 0o600)
          mailcount=0
          box_infected = 0
          time0 = time.time()
          while 1:
            progress_indicator.update()
            frm = line1
            lines = BytesIO()
            while 1:
              line1 = f.readline()
              if line1==b"": break
              if line1[0:5]==b"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==b"": 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,'rb'), f, 1024*1024)
            try:
              os.unlink(fn+SAVECLEAN)
            except OSError:
              pass
        elif re.search(b'^NOOP CONNECT FROM: ',line1):
          print(fn+": sagator's quarantine file detected")
          # skip to "data"
          while 1:
            line1=f.readline()
            if re.compile(b"^DATA",re.IGNORECASE).search(line1):
              break
          lines=BytesIO()
          while 1:
            line1=f.readline()
            if (line1==b'.\n') | (line1==b'.\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, 'rb').read(50*1024*1024), fn):
            infected += 1
            if SAVECLEAN:
              os.unlink(fn)
          total += 1
        f.close()
        break
      except IOError as err:
        (err_code, err_str) = err.args
        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(list(scnr.help().keys()))
    opts,files=getopt.gnu_getopt(sys.argv[1:],'',opts)
  except getopt.GetoptError as err:
    (msg, opt) = err.args
    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 as 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)
  if total>0:
    print("Total infected/spams emails: %d/%d [%5.1f%%] in: %1.2f seconds" \
          % (infected, total, 100.0*infected/total, time.time()-begintime))