From 879e65bc3160a0776bb01e29a7190a9df522cc8d Mon Sep 17 00:00:00 2001 From: "Stuart D. Gathman" Date: Thu, 18 Jun 2020 16:40:54 -0400 Subject: [PATCH] bytes optimization --- Milter/__init__.py | 7 +++++-- Milter/test.py | 2 +- sample.py | 16 ++++++++-------- testsample.py | 15 ++++++++++----- 4 files changed, 24 insertions(+), 16 deletions(-) diff --git a/Milter/__init__.py b/Milter/__init__.py index f0d3347..5386a23 100755 --- a/Milter/__init__.py +++ b/Milter/__init__.py @@ -181,12 +181,12 @@ def noreply(func): wrapper.milter_protocol = nr_mask return wrapper -## Function decorator to set encoding error strategy. +## Function decorator to set decoding error strategy. # Current RFCs define UTF-8 as the standard encoding for SMTP # envelope and header fields. By default, Milter.Base decodes # envelope and header values with errors='surrogateescape'. # This decorator can change the error strategy to, e.g., 'ignore' or 'replace'. -def encodingerror(strategy): +def decode(strategy): def setstrategy(func): func.error_strategy = strategy return func @@ -391,6 +391,9 @@ class Base(object): def header_bytes(self,fld,val): try: e = getattr(self.header,'error_strategy','surrogateescape') + if e == 'bytes': + self.header_bytes = self.header + return self.header(fld,val) s = val.decode(encoding='utf-8',errors=e) except UnicodeDecodeError: s = val return self.header(fld,s) diff --git a/Milter/test.py b/Milter/test.py index 52e4520..02e01f4 100644 --- a/Milter/test.py +++ b/Milter/test.py @@ -183,7 +183,7 @@ class TestBase(object): if rc != Milter.CONTINUE: return rc # header for h,val in msg.items(): - rc = self.header(h,val) + rc = self.header_bytes(h,val.encode()) if rc != Milter.CONTINUE: return rc # eoh self._stage = Milter.M_EOH diff --git a/sample.py b/sample.py index 6476ba9..174ab40 100644 --- a/sample.py +++ b/sample.py @@ -67,7 +67,7 @@ class sampleMilter(Milter.Milter): self.log("rcpt to",to,str) return Milter.CONTINUE - @Milter.encodingerror('replace') + @Milter.decode('bytes') def header(self,name,val): lname = name.lower() if lname == 'subject': @@ -76,22 +76,22 @@ class sampleMilter(Milter.Milter): # (delete if you read chinese mail) #print('val=',val.encode(errors='surrogateescape')) print('val=',val) - if val.startswith('=?big5') or val.startswith('=?ISO-2022-JP'): + if val.startswith(b'=?big5') or val.startswith(b'=?ISO-2022-JP'): self.log('REJECT: %s: %s' % (name,val)) #self.setreply('550','','Go away spammer') return Milter.REJECT # check for common spam keywords - if val.find("$$$") >= 0 or val.find("XXX") >= 0 \ - or val.find("!!!") >= 0 or val.find("FREE") >= 0: + if val.find(b"$$$") >= 0 or val.find(b"XXX") >= 0 \ + or val.find(b"!!!") >= 0 or val.find(b"FREE") >= 0: self.log('REJECT: %s: %s' % (name,val)) #self.setreply('550','','Go away spammer') return Milter.REJECT # check for spam that pretends to be legal lval = val.lower() - if lval.startswith("adv:") or lval.startswith("adv.") \ - or lval.find('viagra') >= 0: + if lval.startswith(b"adv:") or lval.startswith(b"adv.") \ + or lval.find(b'viagra') >= 0: self.log('REJECT: %s: %s' % (name,val)) return Milter.REJECT @@ -103,7 +103,7 @@ class sampleMilter(Milter.Milter): # check for common bulk mailers if lname == 'x-mailer' and \ - val.lower() in ('direct email','calypso','mail bomber'): + val.lower() in (b'direct email',b'calypso',b'mail bomber'): self.log('REJECT: %s: %s' % (name,val)) #self.setreply('550','','Go away spammer') return Milter.REJECT @@ -112,7 +112,7 @@ class sampleMilter(Milter.Milter): if lname in ('subject','x-mailer'): self.log('%s: %s' % (name,val)) if self.fp: - self.fp.write(("%s: %s\n" % (name,val)).encode(errors='surrogateescape')) # add header to buffer + self.fp.write(b"%s: %s\n" % (name.encode(),val)) # add header to buffer return Milter.CONTINUE def eoh(self): diff --git a/testsample.py b/testsample.py index 175cea8..100d1a6 100644 --- a/testsample.py +++ b/testsample.py @@ -74,7 +74,8 @@ class BMSMilterTestCase(unittest.TestCase): self.assertTrue(rc == Milter.ACCEPT) self.assertTrue(ctx._bodyreplaced,"Message body not replaced") fp = ctx._body - open('test/'+fname+".tstout","wb").write(fp.getvalue()) + with open('test/'+fname+".tstout","wb") as f: + f.write(fp.getvalue()) #self.assertTrue(fp.getvalue() == open("test/virus1.out","r").read()) fp.seek(0) msg = mime.message_from_file(fp) @@ -97,7 +98,8 @@ class BMSMilterTestCase(unittest.TestCase): self.assertTrue(rc == Milter.ACCEPT) self.assertTrue(milter._bodyreplaced,"Message body not replaced") fp = milter._body - open('test/'+fname+".tstout","wb").write(fp.getvalue()) + with open('test/'+fname+".tstout","wb") as f: + f.write(fp.getvalue()) #self.assertTrue(fp.getvalue() == open("test/virus1.out","r").read()) fp.seek(0) msg = mime.message_from_file(fp) @@ -112,7 +114,8 @@ class BMSMilterTestCase(unittest.TestCase): self.assertTrue(rc == Milter.ACCEPT) self.assertFalse(milter._bodyreplaced,"Milter needlessly replaced body.") fp = milter._body - open('test/'+fname+".tstout","wb").write(fp.getvalue()) + with open('test/'+fname+".tstout","wb") as f: + f.write(fp.getvalue()) milter.close() def testDefang2(self): @@ -126,7 +129,8 @@ class BMSMilterTestCase(unittest.TestCase): self.assertTrue(rc == Milter.ACCEPT) self.assertTrue(milter._bodyreplaced,"Message body not replaced") fp = milter._body - open("test/virus3.tstout","wb").write(fp.getvalue()) + with open("test/virus3.tstout","wb") as f: + f.write(fp.getvalue()) #self.assertTrue(fp.getvalue() == open("test/virus3.out","r").read()) with self.zf.open("virus6") as fp: rc = milter.feedFile(fp) @@ -134,7 +138,8 @@ class BMSMilterTestCase(unittest.TestCase): self.assertTrue(milter._bodyreplaced,"Message body not replaced") self.assertTrue(milter._headerschanged,"Message headers not adjusted") fp = milter._body - open("test/virus6.tstout","wb").write(fp.getvalue()) + with open("test/virus6.tstout","wb") as f: + f.write(fp.getvalue()) milter.close() def suite(): return unittest.makeSuite(BMSMilterTestCase,'test')