'''
Policy scanners for sagator

(c) 2005-2008 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 avlib import *
from interscan.match import match_any

__all__=['set_action',
         'greylist', 'not_listed', 'listed', 'elisted',
         'list_cleanup', 'auto_whitelist',
         'spf_check', 'dns_check', 'rbl_check',
         'add_listed']

class set_action(match_any):
  '''
  An interscanner to return a status to set an action to return.

  You can user this scanner to set default policy returned by policy scanner,
  or you can set policy according to scanner reply.

  Usage: set_action(action, *scanners)
  
  Where: action can be an action defined in postfix's SMTPD_POLICY_README,
           for example: dunno, defer_if_permit, ok, reject, ...

  New in version 0.8.0.
  '''
  is_policy_scanner=True
  name='set_action()'
  def __init__(self,action,*scanners):
      self.ACTION=action
      match_any.__init__(self,scanners)
  def scanbuffer(self,buffer,args={}):
      if not self.scanners:
        globals.ACTION=self.ACTION
        return 0.0,'',[]
      level,detected,ret=match_any.scanbuffer(self,buffer,args)
      if is_infected(level,detected):
        globals.ACTION=self.ACTION
      return level,detected,ret

class greylist(ascanner):
  '''
  A greylist policy scanner.

  This is an greylist implementation for sagator.
  You can learn more about greylisting at:
    http://projects.puremagic.com/greylisting/
  
  Usage: greylist(min=300,expire=48*3600)
  
  Where: min is minimum required seconds after an email is cleanly sent
           (default: 5 minutes)
         expire is record expiration time in seconds (default 48 hours).
           If this record will be passed after min time, expiration
           time will be disabled.
  
  Example: greylist()

  New in version 0.8.0.
  '''
  is_policy_scanner=True
  name='greylist()'
  def __init__(self,min=300,expire=5*3600):
      self.MIN_DELAY=min
      self.EXPIRE_TIME=expire
  def get_key(self,dbc):
      key=['','','']
      try:
        key=[mail.policy_request.get('client_address'),
             mail.policy_request.get('sender'),
             mail.policy_request.get('recipient')]
      except:
        debug.echo(4,'%s: wrong request: %s' % (self.name,mail.policy_request))
        raise
      return key
  def update_counter(self,dbc,cntr,flags,t0,key):
      dbc.execute_cycle(
        "UPDATE greylist SET %s_count=%s_count+1, last_update=%d WHERE flags LIKE %%s AND ip=%%s AND sender=%%s AND recipient=%%s" \
        % (cntr,cntr,t0),
        [flags]+key)
  def scanbuffer(self,buffer,args={}):
      key=self.get_key(args['dbc'])
      t0=time.time()
      q=args['dbc'].query(
          "SELECT timestamp,pass_count FROM greylist"
          " WHERE flags='GL' AND ip=%%s AND sender=%%s AND"
                " recipient=%%s AND ((expire>%d) OR (expire<1))" % (t0),
          key)
      debug.echo(8,"%s: %s %s" % (self.name,key,q))
      if q: # one row returned
        # found a time in database
        t=q[0][0]
        if t0-t > self.MIN_DELAY:
          # delay is acceptable
          args['dbc'].execute_cycle(
            "UPDATE greylist SET expire=-1, last_update=%d,"
                               " pass_count=pass_count+1"
            " WHERE flags='GL' AND ip=%%s AND sender=%%s"
                  " AND recipient=%%s" % (t0),
            key)
          # add header if this is first pass
          if q[0][1]==0:
            globals.PREPEND='X-Sagator-Greylist: delayed %d seconds at %s; %s' \
                            % (t0-t,socket.gethostname(),time.strftime("%c"))
          return 0.0,'',[]
        else:
          self.update_counter(args['dbc'],'block','GL',t0,key)
      else:
        # remove all possible expired records
        q=args['dbc'].query(
            "DELETE FROM greylist"
            " WHERE flags='GL' AND ip=%%s AND sender=%%s AND"
                  " recipient=%%s AND expire<=%d" % (t0),
            key)
        # add new record for nonexistent key
        t=t0
        args['dbc'].execute_cycle(
          "INSERT INTO greylist (flags,timestamp,expire,created,last_update,"
                                "ip,sender,recipient,block_count,pass_count)"
          " VALUES ('GL',%d,%d,%d,%d,%%s,%%s,%%s,1,0)" \
          % (t0,t0+self.EXPIRE_TIME,t0,t0), key)
      globals.ACTION='defer_if_permit Greylisted for %d seconds' \
        ' (see http://projects.puremagic.com/greylisting/)' \
        % (self.MIN_DELAY+t-t0)
      return 1.0,'greylisted',['%s greylisted' % key[2]]

class listed(greylist):
  '''
  SQL blacklist for a policy.
  
  You can use this scanner to chcek, if an record is present in SQL table.

  Usage: listed(flags='B',table='greylist')
  
  Where: flags is an character, which defines which flag will be searched
           in SQL table. By default 'B' for this scanner.
         table can define table alternative
  
  Returns: level=1.0, when client is in blacklist
           level=0.0, when client is not in blacklist
  
  Example: listed()

  New in version 0.8.0.
  '''
  name='listed()'
  def __init__(self,flags='B',repeats=0,table='greylist'):
      self.FLAGS=flags.upper()
      self.TABLE=table
      self.REPEATS=repeats
  def scanbuffer(self,buffer,args={}):
      key=self.get_key(args['dbc'])
      t0=time.time()
      args['dbc'].execute_cycle(
        "UPDATE %s SET block_count=block_count+1, last_update=%d \
           WHERE (expire<%d) AND (repeats>=%d) AND ( \
           (flags='%sA' AND %%s LIKE ip) OR \
           (flags='%sS' AND %%s LIKE sender) OR \
           (flags='%sR' AND %%s LIKE recipient))" \
        % (self.TABLE,t0,t0,self.REPEATS,self.FLAGS,self.FLAGS,self.FLAGS),
        key)
      if args['dbc'].rowcount>0:
        globals.ACTION='reject You are blacklisted!'
        return 1.0,'blacklisted',['blacklisted']
      return 0.0,'',[]

class elisted(greylist):
  '''
  SQL blacklist for a policy. Enhanced version.
  
  You can use this scanner to chcek, if an record is present in SQL table.

  Usage: listed(flags='B', table='greylist')
  
  Where: flags is an character, which defines which flag will be searched
           in SQL table. By default 'B' for this scanner.
         table can define table alternative
  
  Returns: level=1.0, when client is in blacklist
           level=0.0, when client is not in blacklist
  
  This is experimental and may change in future releases.

  Example: listed()

  New in version 1.1.1.
  '''
  name='elisted()'
  def __init__(self, flags='B', repeats=0, table='greylist'):
      self.FLAGS = flags.upper()
      self.TABLE = table
      self.REPEATS = repeats
  def scanbuffer(self, buffer, args={}):
      key = self.get_key(args['dbc'])
      t0 = time.time()
      args['dbc'].execute_cycle(
        "UPDATE %s SET block_count=block_count+1, last_update=%d \
           WHERE (expire<%d) AND (repeats>=%d) AND (flags='%s') \
             AND (%%s LIKE ip) \
             AND (%%s LIKE sender) \
             AND (%%s LIKE recipient)" \
        % (self.TABLE, t0, t0, self.REPEATS, self.FLAGS),
        key)
      if args['dbc'].rowcount>0:
        globals.ACTION = 'reject You are blacklisted!'
        return 1.0, 'blacklisted', ['blacklisted']
      return 0.0, '', []

class not_listed(greylist):
  '''
  SQL whitelist for a policy.
  
  Usage: not_listed(flags='W',table='greylist')
  
  Where: flags is an character, which defines which flag will be searched
           in SQL table. By default 'B' for this scanner.
         table can define table alternative
  
  Returns: level=1.0, when client is not in whitelist
           level=0.0, when client is in whitelist
  
  Example: not_listed()

  New in version 0.8.0.
  '''
  name='not_listed()'
  def __init__(self,flags='W',repeats=0,table='greylist'):
      self.FLAGS=flags.upper()
      self.TABLE=table
      self.REPEATS=repeats
  def scanbuffer(self,buffer,args={}):
      key=self.get_key(args['dbc'])
      t0=time.time()
      args['dbc'].execute_cycle(
        "UPDATE %s SET pass_count=pass_count+1, last_update=%d \
           WHERE (expire<%d) AND (repeats>=%d) AND ( \
           (flags='%sA' AND %%s LIKE ip) OR \
           (flags='%sS' AND %%s LIKE sender) OR \
           (flags='%sR' AND %%s LIKE recipient))" \
        % (self.TABLE,t0,t0,self.REPEATS,self.FLAGS,self.FLAGS,self.FLAGS),
        key)
      if args['dbc'].rowcount>0:
        debug.echo(4,"%s: %s whitelisted" % (self.name,key))
        return 0.0,'',[]
      return 1.0,'not_listed',['not listed']

class list_cleanup(greylist):
  '''
  Cleanup obsolete records from an SQL list.
  This cleanup is done not more than every 10 minutes to avoid server
  overload.
  
  Usage: list_cleanup(lifetime=36*24*3600,table='greylist')
  
  Where: lifetime is number of seconds after which records are removed (36 days)
         table can define table alternative
  
  Example: list_cleanup()

  New in version 0.8.0.
  '''
  name='list_cleanup()'
  def __init__(self, lifetime=36*24*3600, table='greylist'):
      self.LIFETIME = lifetime
      self.TABLE = table
      self.INTERVAL = 600 # every 10 minutes
      self.NEXT_ACTION = time.time() + self.INTERVAL
  def scanbuffer(self,buffer,args={}):
      t0=time.time()
      if t0>self.NEXT_ACTION:
        self.NEXT_ACTION = t0 + self.INTERVAL
        args['dbc'].execute_cycle(
          "DELETE FROM %s WHERE (last_update<%d) OR ((1<expire) AND (expire<%d))" \
          % (self.TABLE,t0-self.LIFETIME,t0))
      return 0.0,'',[]

class auto_whitelist(greylist):
  '''
  Whitelist IP addresses after sending some emails properly.
  This whitelisting is done not more than every 10 minutes to avoid server
  overload. This function is slow for large number of records, use it
  carefully.
  
  Usage: auto_whitelist(match_count=10,table='greylist')
  
  Where: match_count is number of record which will must match
         table can define table alternative
  
  Example: auto_whitelist(10)

  New in version 0.8.0.
  '''
  name='auto_whitelist()'
  def __init__(self,match_count=10,table='greylist'):
      self.MATCH_COUNT=match_count
      self.TABLE=table
      self.INTERVAL = 600 # every 10 minutes
      self.NEXT_ACTION = time.time()
  def scanbuffer(self,buffer,args={}):
      t0=time.time()
      if t0<self.NEXT_ACTION:
        return 0.0,'',[]
      self.NEXT_ACTION = t0 + self.INTERVAL
      for (ip,) in args['dbc'].query(
          "SELECT ip FROM %s WHERE flags='GL'"
                        " AND ip NOT IN (SELECT ip FROM %s WHERE flags='WA')"
          " GROUP BY ip HAVING sum(pass_count+repeats)>%d" \
          % (self.TABLE,self.TABLE,self.MATCH_COUNT)):
        try:
          debug.echo(4, "%s: Whitelisting %s" % (self.name, ip))
          args['dbc'].execute_cycle(
            "INSERT INTO %s (flags,timestamp,expire,ip,created,last_update) \
             VALUES ('WA',%d,-1,%%s,%d,%d)" \
            % (self.TABLE,t0,t0,t0), [ip])
        except Exception, e:
          # May be this record has been whitelisted from another process.
          # Ignore this error.
          debug.echo(8, "%s: Error auto-whitelisting %s: %s" \
                        % (self.name, ip, e))
      return 0.0,'',[]

# DNS and SPF tests

class dns_check(ascanner):
  '''
  IP to domain (and back) resolving checker.
  
  This scanner can be used to check senders IP address to resolve via
  DNS service. It can be uses asn real scanner or policy scanner.

  Usage: dns_check(reverse=False)
  
  Where: reverse can fe used to check also reverse records
  
  Example: dns_check()

  New in version 0.8.0.
  '''
  name='dns()'
  is_policy_scanner=True
  def __init__(self,reverse=False):
      self.REVERSE=reverse
  def getsender(self):
      # first try to use policy_request for policy scanners
      try:
        ip=mail.policy_request['client_address']
        sender=mail.policy_request['sender']
        helo=mail.policy_request['helo_name']
      except KeyError:
        # use standard method for content scanner
        ip=mail.getsender()['ADDR']
        sender=mail.sender
        try:
          helo=smtp.reg_helo_ehlo.search(mail.comm).group(1)
        except:
          helo='localhost'
      return ip,sender,helo
  def scanbuffer(self,buffer,args={}):
      ip,sender,helo=self.getsender()
      try:
        host=socket.gethostbyaddr(ip)
        debug.echo(4,"%s: %s" % (self.name,host))
        if self.REVERSE:
          rhost=socket.gethostbyname(host[0])
          debug.echo(4,"%s: %s" % (self.name,rhost))
      except socket.herror,(ec,es):
        if ec==1: # Unknown host
          return 1.0,"%s for %s" % (es,ip),[es,ip,sender,helo]
      return 0.0,'',[]

class spf_check(dns_check):
  '''
  SPF (Sender Permitted From) checker.
  
  This scanner can be used to scan for sender's SPF records in content
  scanner or policy scanner.
  For more information about SPF see: http://www.openspf.org/ .

  Usage: spf_check(hard_error=False)
  
  Where: hard_error can be set to True, if you want errors raise as exceptions,
           by default it is turned off.
  
  Example: spf_check()

  New in version 0.8.0.
  '''
  name='spf()'
  is_policy_scanner=True
  def __init__(self,hard_error=False):
      try:
        import spf
      except ImportError:
        print "You need pyspf (or python-spf) package to run spf_check() scanner."
        sys.exit(1)
      self.SPF=spf
      self.HARD_ERROR=hard_error
  def scanbuffer(self,buffer,args={}):
      ip,sender,helo=self.getsender()
      debug.echo(6,'spf_check(): %s[%s] %s' % (ip,helo,sender))
      result,status_code,explanation=self.SPF.check(ip,sender,helo)
      debug.echo(4,'spf_check(): %s[%s] %s: %s %d %s' \
                 % (ip,helo,sender,result,status_code,explanation))
      if result=='deny':
        return 1.0,'spf_hard_fail',['%s: %s' % (result,explanation)]
      elif result=='unknown':
        return 0.8,'spf_soft_fail',['%s: %s' % (result,explanation)]
      elif self.HARD_ERROR and (result=='error'):
        raise ScannerError,'%s: %s' % (result,explanation)
      return 0.0,'',[]

class rbl_check(dns_check):
  '''
  RBL (Real-time Blackhole list) checker.
  
  This scanner can be used to scan for sender's IP address contained
  in a blacklist.

  Usage: rbl_check(domains)
  
  Where: domains are blacklist domain names (multiple parameters allowed).
  
  Example: rbl_check('zen.spamhaus.org')

  New in version 1.1.0.
  '''
  name='rbl()'
  is_policy_scanner=True
  def __init__(self, *domains):
      self.DOMAINS = domains
  def scanbuffer(self, buffer, args={}):
      ip,sender,helo = self.getsender()
      debug.echo(6, '%s: %s' % (self.name, ip))
      reversed_ip = '.'.join(reversed(ip.split('.')))
      addrs = []
      for domain in self.DOMAINS:
        host = reversed_ip + '.' + domain
        debug.echo(6, '%s: Testing %s ...' % (self.name, host))
        try:
          addrs.append(
            "%s: %s" % (domain, socket.gethostbyname(host))
          )
        except socket.error:
          pass
      debug.echo(6, '%s: Result: %s' % (self.name, addrs))
      if addrs:
        return len(addrs), 'Found_in_RBL', addrs
      return 0.0, '', []

# This scanner is not a policy scanner, but it is a content scanner.
# It can set a whitelist or blacklist for policy scanners.

class add_listed(match_any):
  '''
  Add one record into SQL list for a policy.
  
  Usage: add_listed(dbc,flags,expire,table,scanners)
  
  Where: dbc is an database connection
         flags is an string, which defines what and where to add. It can be:
             BA - to blacklist sender's IP address
             BS - to blacklist sender's email address
             BR - to blacklist recipients
             WA - to whitelist sender's IP address
             WS - to whitelist sender's email address
             WR - to whitelist recipients
         expire is an integer, which defines validity of a new record
           in seconds. Set it to -1 for infinite time.
         table can define table alternative
  
  Example: add_listed(db.sqlite(),'BA',-1,'greylist',b2f(libclam()))

  New in version 0.8.0.
  '''
  name='add_listed()'
  def __init__(self,dbc,flags,expire=-1,table='greylist',*scanners):
      self.dbc=dbc
      self.FLAGS=flags.upper()
      self.EXPIRE_TIME=expire
      self.TABLE=table
      match_any.__init__(self,scanners)
  def scanbuffer(self,buffer,args={}):
      level,detected,ret=match_any.scanbuffer(self,buffer,args)
      if not is_infected(level,detected):
        return level,detected,ret
      if mail.getsender()['ADDR'].lower()=='unknown':
        debug.echo(4,"add_listed(): no sender address specified: %s" \
                   % (mail.getsender()['ADDR']))
        return level,detected,ret
      try:
        keys={
          'A': [[mail.getsender()['ADDR'],'','']],
          'S': [['',mail.sender,'']],
          'R': [['','',x] for x in mail.recip]
        }[self.FLAGS[1]]
      except KeyError:
        raise
      for key in keys:
        q=self.dbc.query(
          "SELECT count(*) FROM "+self.TABLE+" WHERE \
             flags=%s AND ip=%s AND sender=%s and recipient=%s",
          [self.FLAGS]+key)
        t0=time.time()
        if self.EXPIRE_TIME<0:
          t_expire=0
        else:
          t_expire=t0+self.EXPIRE_TIME
        if q[0][0]==0:
          self.dbc.execute_cycle(
            "INSERT INTO %s (flags,timestamp,expire,created,last_update,ip,sender,recipient) \
               VALUES ('%s',%d,%d,%d,%d,%%s,%%s,%%s)" \
            % (self.TABLE,self.FLAGS,t0,t_expire,t0,t0),
            key)
        else:
          self.dbc.execute_cycle(
            "UPDATE %s SET last_update=%d, expire=%d, repeats=repeats+1 WHERE \
               flags='%s' AND ip=%%s AND sender=%%s and recipient=%%s" \
            % (self.TABLE,t0,t_expire,self.FLAGS),
            key)
      return level,detected,ret
