diff --git a/Milter/greylist.py b/Milter/greylist.py index 58d5dbc..26b0dc8 100644 --- a/Milter/greylist.py +++ b/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() diff --git a/Milter/greysql.py b/Milter/greysql.py new file mode 100644 index 0000000..0c9fb13 --- /dev/null +++ b/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() diff --git a/testgrey.py b/testgrey.py index ee7fd0c..5309448 100644 --- a/testgrey.py +++ b/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')