Optional sqlite3 greylist implementation.

Milter/greylist.py

@@ -18,13 +18,19 @@ def quoteAddress(s):
 class Record(object):
   __slots__ = ( 'firstseen', 'lastseen', 'umis', 'cnt' )
 
-  def __init__(self):
-    now = time.time()
+  def __init__(self,timeinc=0):
+    now = time.time() + timeinc
     self.firstseen = now
     self.lastseen = now
     self.cnt = 0
     self.umis = None
 
+  def __str__(self):
+    return "Grey[%s:%s:%s:%d]" % (
+        time.ctime(self.firstseen),time.ctime(self.lastseen),
+        self.umis,self.cnt
+    )
+
 class Greylist(object):
 
   def __init__(self,dbname,grey_time=10,grey_expire=4,grey_retain=36):
@@ -34,6 +40,25 @@ class Greylist(object):
     self.greylist_retain = grey_retain * 24 * 3600   # days
     self.dbp = shelve.open(dbname,'c',protocol=2)
     self.lock = thread.allocate_lock()
+
+  def clean(self,timeinc=0):
+    "Delete records past the retention limit."
+    now = time.time() + timeinc
+    cnt = 0
+    dbp = self.dbp
+    for key, r in dbp.iteritems():
+      #print key,r,time.ctime(now)
+      if now > r.lastseen + self.greylist_retain:
+        self.lock.acquire()
+        try:
+          r = dbp[key]
+          now = time.time() + timeinc
+          if now > r.lastseen + self.greylist_retain:
+            del dbp[key]
+            cnt += 1
+        finally:
+          self.lock.release()
+    return cnt
   
   def check(self,ip,sender,recipient,timeinc=0):
     "Return number of allowed messages for greylist triple."
@@ -49,11 +74,11 @@ class Greylist(object):
         if now > r.lastseen + self.greylist_retain:
           # expired
           log.debug('Expired greylist: %s',key)
-          r = Record()
+          r = Record(timeinc)
         elif now < r.firstseen + self.greylist_time + 5:
           # still greylisted
           log.debug('Early greylist: %s',key)
-          #r = Record()
+          #r = Record(timeinc)
           r.lastseen = now
         elif r.cnt or now < r.firstseen + self.greylist_expire:
           # in greylist window or active
@@ -63,12 +88,15 @@ class Greylist(object):
         else:
           # passed greylist window
           log.debug('Late greylist: %s',key)
-          r = Record()
+          r = Record(timeinc)
         dbp[key] = r
       except:
-        r = Record()
+        r = Record(timeinc)
         dbp[key] = r
       dbp.sync()
     finally:
       self.lock.release()
     return r.cnt
+
+  def close(self):
+    self.dbp.close()

Milter/greysql.py

@@ -0,0 +1,76 @@
+import time
+import logging
+import urllib
+import sqlite3
+from datetime import datetime
+
+log = logging.getLogger('milter.greylist')
+
+class Greylist(object):
+
+  def __init__(self,dbname,grey_time=10,grey_expire=4,grey_retain=36):
+    self.ignoreLastByte = False
+    self.greylist_time = grey_time * 60         # minutes
+    self.greylist_expire = grey_expire * 3600   # hours
+    self.greylist_retain = grey_retain * 24 * 3600   # days
+    self.conn = sqlite3.connect(dbname)
+    self.conn.row_factory = sqlite3.Row
+    try:
+      self.conn.execute('''create table greylist(
+        ip text , sender text, recipient text,
+        firstseen timestamp, lastseen timestamp, cnt integer, umis text,
+        primary key (ip,sender,recipient))''')
+    except: pass
+
+  def clean(self,timeinc=0):
+    "Delete records past the retention limit."
+    now = time.time() + timeinc - self.greylist_retain
+    cur = self.conn.cursor()
+    cur.execute('delete from greylist where lastseen < ?',(now,))
+    cnt = cur.rowcount
+    self.conn.commit()
+    return cnt
+
+  def check(self,ip,sender,recipient,timeinc=0):
+    "Return number of allowed messages for greylist triple."
+    cur = self.conn.cursor()
+    cur.execute('''select firstseen,lastseen,cnt,umis from greylist where
+      ip=? and sender=? and recipient=?''',(ip,sender,recipient))
+    r = cur.fetchone()
+    now = time.time() + timeinc
+    cnt = 0
+    if not r:
+      cur.execute('''insert into 
+        greylist(ip,sender,recipient,firstseen,lastseen,cnt,umis)
+        values(?,?,?,?,?,?,?)''', (ip,sender,recipient,now,now,0,None))
+    elif now > r['lastseen'] + self.greylist_retain:
+      # expired
+      log.debug('Expired greylist: %s:%s:%s',ip,sender,recipient)
+      cur.execute('''update greylist set firstseen=?,lastseen=?,cnt=?,umis=?
+        where ip=? and sender=? and recipient=?''',
+        (now,now,0,None,ip,sender,recipient))
+    elif now < r['firstseen'] + self.greylist_time + 5:
+      # still greylisted
+      log.debug('Early greylist: %s:%s:%s',ip,sender,recipient)
+      #r = Record()
+      cur.execute('''update greylist set lastseen=?
+        where ip=? and sender=? and recipient=?''',
+        (now,ip,sender,recipient))
+    elif r['cnt'] or now < r['firstseen'] + self.greylist_expire:
+      # in greylist window or active
+      cnt = r['cnt'] + 1
+      cur.execute('''update greylist set lastseen=?,cnt=?
+        where ip=? and sender=? and recipient=?''',
+        (now,cnt,ip,sender,recipient))
+      log.debug('Active greylist(%d): %s:%s:%s',cnt,ip,sender,recipient)
+    else:
+      # passed greylist window
+      log.debug('Late greylist: %s:%s:%s',ip,sender,recipient)
+      cur.execute('''update greylist set firstseen=?,lastseen=?,cnt=?,umis=?
+        where ip=? and sender=? and recipient=?''',
+        (now,now,0,None,ip,sender,recipient))
+    self.conn.commit()
+    return cnt
+
+  def close(self):
+    self.conn.close()

testgrey.py

@@ -1,40 +1,51 @@
 import unittest
 import doctest
 import os
-import Milter.greylist
-from Milter.greylist import Greylist
+#from Milter.greylist import Greylist
+from Milter.greysql import Greylist
 
 class GreylistTestCase(unittest.TestCase):
 
   def setUp(self):
     self.fname = 'test.db'
+    os.remove(self.fname)
 
   def tearDown(self):
-    os.remove(self.fname)
+    #os.remove(self.fname)
+    pass
 
   def testGrey(self):
     grey = Greylist(self.fname)
     # first time
     rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com')
-    self.failUnless(rc == 0)
+    self.assertEqual(rc,0)
     # not in window yet
     rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=5*60)
-    self.failUnless(rc == 0)
+    self.assertEqual(rc,0)
     # within window
     rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=15*60)
-    self.failUnless(rc == 1)
+    self.assertEqual(rc,1)
     # new triple
     rc = grey.check('1.2.3.5','foo@bar.com','baz@spat.com',timeinc=15*60)
-    self.failUnless(rc == 0)
+    self.assertEqual(rc,0)
     # seen again
     rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=5*3600)
-    self.failUnless(rc == 2)
+    self.assertEqual(rc,2)
     # new one past expire
-    rc = grey.check('1.2.3.5','foo@bar.com','baz@spat.com',timeinc=5*3600)
-    self.failUnless(rc == 0)
+    rc = grey.check('1.2.3.5','foo@bar.com','baz@spat.com',timeinc=6*3600)
+    self.assertEqual(rc,0)
     # original past retain 
     rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=37*24*3600)
-    self.failUnless(rc == 0)
+    self.assertEqual(rc,0)
+    # new one for testing expire
+    rc = grey.check('1.2.3.5','flub@bar.com','baz@spat.com',timeinc=20*24*3600)
+    self.assertEqual(rc,0)
+    grey.close()
+    # test cleanup
+    grey = Greylist(self.fname)
+    rc = grey.clean(timeinc=37*24*3600)
+    self.assertEqual(rc,1)
+    grey.close()
 
 def suite(): 
   s = unittest.makeSuite(GreylistTestCase,'test')