bytes optimization

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)

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

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):

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')