diff --git a/Milter/test.py b/Milter/test.py index 5077f4c..f0dc58b 100644 --- a/Milter/test.py +++ b/Milter/test.py @@ -4,9 +4,9 @@ from __future__ import print_function import mime try: - from StringIO import StringIO + from io import BytesIO except: - from io import StringIO + from StringIO import StringIO as BytesIO import Milter Milter.NOREPLY = Milter.CONTINUE @@ -155,20 +155,20 @@ class TestBase(object): rc = self.eoh() if rc != Milter.CONTINUE: return rc header,body = msg.as_bytes().split(b'\n\n',1) - bfp = StringIO(body) + bfp = BytesIO(body) while 1: buf = bfp.read(8192) if len(buf) == 0: break rc = self.body(buf) if rc != Milter.CONTINUE: return rc self._msg = msg - self._body = StringIO() + self._body = BytesIO() rc = self.eom() if self._bodyreplaced: body = self._body.getvalue() - self._body = StringIO() + self._body = BytesIO() self._body.write(header) - self._body.write('\n\n') + self._body.write(b'\n\n') self._body.write(body) return rc @@ -177,7 +177,7 @@ class TestBase(object): # @param sender MAIL FROM # @param rcpts RCPT TO, multiple recipients may be supplied def feedMsg(self,fname,sender="spam@adv.com",*rcpts): - with open('test/'+fname,'r') as fp: + with open('test/'+fname,'rb') as fp: return self.feedFile(fp,sender,*rcpts) ## Call the connect and helo callbacks. diff --git a/Milter/utils.py b/Milter/utils.py index cd27b00..0637ee4 100644 --- a/Milter/utils.py +++ b/Milter/utils.py @@ -212,11 +212,11 @@ def parse_header(val): for s,enc in h: if enc: try: - u.append(unicode(s,enc,'replace')) + u.append(s.decode(enc,'replace')) except LookupError: - u.append(unicode(s)) + u.append(s.decode()) else: - u.append(unicode(s)) + u.append(s.decode()) u = u''.join(u) for enc in ('us-ascii','iso-8859-1','utf-8'): try: diff --git a/mime.py b/mime.py index 8106d2e..a770ba5 100644 --- a/mime.py +++ b/mime.py @@ -95,16 +95,20 @@ from __future__ import print_function try: - from StringIO import StringIO + from io import BytesIO, StringIO except: - from io import StringIO + from StringIO import StringIO + BytesIO = StringIO import socket import Milter import zipfile import email from email.message import Message -from email.generator import Generator +try: + from email.generator import BytesGenerator +except: + from email.generator import Generator as BytesGenerator from email.utils import quote if not getattr(Message,'as_bytes',None): @@ -113,7 +117,7 @@ if not getattr(Message,'as_bytes',None): ## Return a list of filenames in a zip file. # Embedded zip files are recursively expanded. def zipnames(txt): - fp = StringIO(txt) + fp = BytesIO(txt) zipf = zipfile.ZipFile(fp,'r') names = [] for nm in zipf.namelist(): @@ -124,7 +128,7 @@ def zipnames(txt): ## Fix multipart handling in email.Generator. # -class MimeGenerator(Generator): +class MimeGenerator(BytesGenerator): def _dispatch(self, msg): # Get the Content-Type: for the message, then try to dispatch to # self._handle__(). If there's no handler for the @@ -134,7 +138,7 @@ class MimeGenerator(Generator): if msg.is_multipart() and main.lower() != 'multipart': self._handle_multipart(msg) else: - Generator._dispatch(self,msg) + BytesGenerator._dispatch(self,msg) def unquote(s): """Remove quotes from a string.""" @@ -237,9 +241,9 @@ class MimeMessage(Message): g = MimeGenerator(file) g.flatten(self,unixfrom=unixfrom) - def as_string(self, unixfrom=False): + def as_bytes(self, unixfrom=False): "Return the entire formatted message as a string." - fp = StringIO() + fp = BytesIO() self.dump(fp,unixfrom=unixfrom) return fp.getvalue() @@ -300,7 +304,7 @@ class MimeMessage(Message): return None def message_from_file(fp): - msg = email.message_from_file(fp,MimeMessage) + msg = email.message_from_binary_file(fp,MimeMessage) for part in msg.walk(): part.modified = False assert not msg.ismodified() @@ -509,7 +513,7 @@ def check_html(msg,savname=None): htmlfilter.close() #except sgmllib.SGMLParseError: except: - #mimetools.copyliteral(msg.get_payload(),open('debug.out','w') + mimetools.copyliteral(msg.get_payload(),open('debug.out','wb')) htmlfilter.close() hostname = socket.gethostname() msg.set_payload( @@ -539,7 +543,7 @@ if __name__ == '__main__': return Milter.CONTINUE for fname in sys.argv[1:]: - fp = open(fname) + fp = open(fname,'rb') msg = message_from_file(fp) email.iterators._structure(msg) check_attachments(msg,_list_attach) diff --git a/sample.py b/sample.py index 7f92994..6c48296 100644 --- a/sample.py +++ b/sample.py @@ -8,9 +8,9 @@ from __future__ import print_function import sys import os try: - from StringIO import StringIO + from io import BytesIO except: - from io import StringIO + from StringIO import StringIO as BytesIO import mime import Milter import tempfile @@ -41,7 +41,7 @@ class sampleMilter(Milter.Milter): def envfrom(self,f,*str): "start of MAIL transaction" self.log("mail from",f,str) - self.fp = StringIO() + self.fp = BytesIO() self.tempname = None self.mailfrom = f self.bodysize = 0 @@ -97,12 +97,12 @@ 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)) # add header to buffer + self.fp.write(("%s: %s\n" % (name,val)).encode()) # add header to buffer return Milter.CONTINUE def eoh(self): if not self.fp: return Milter.TEMPFAIL # not seen by envfrom - self.fp.write("\n") + self.fp.write(b'\n') self.fp.seek(0) # copy headers to a temp file for scanning the body headers = self.fp.getvalue() @@ -144,7 +144,7 @@ class sampleMilter(Milter.Milter): msg.dump(out) out.seek(0) msg = mime.message_from_file(out) - fp = StringIO(msg.as_bytes().split(b'\n\n',1)[1]) + fp = BytesIO(msg.as_bytes().split(b'\n\n',1)[1]) while 1: buf = fp.read(8192) if len(buf) == 0: break diff --git a/testmime.py b/testmime.py index 03189bf..50110b3 100644 --- a/testmime.py +++ b/testmime.py @@ -62,18 +62,18 @@ class MimeTestCase(unittest.TestCase): self.failUnless(plist[0] == 'name="Jim&Girlz.jpg"') def testParse(self,fname='samp1'): - msg = mime.message_from_file(open('test/'+fname,"r")) + msg = mime.message_from_file(open('test/'+fname,"rb")) self.failUnless(msg.ismultipart()) parts = msg.get_payload() self.failUnless(len(parts) == 2) txt1 = parts[0].get_payload() self.failUnless(txt1.rstrip() == samp1_txt1,txt1) - msg = mime.message_from_file(open('test/missingboundary',"r")) + msg = mime.message_from_file(open('test/missingboundary',"rb")) # should get no exception as long as we don't try to parse # message attachments mime.defang(msg,scan_rfc822=False) - msg.dump(open('test/missingboundary.out','w')) - msg = mime.message_from_file(open('test/missingboundary',"r")) + msg.dump(open('test/missingboundary.out','wb')) + msg = mime.message_from_file(open('test/missingboundary',"rb")) try: mime.defang(msg) # python 2.4 doesn't get exceptions on missing boundaries, and @@ -85,12 +85,12 @@ class MimeTestCase(unittest.TestCase): def testDefang(self,vname='virus1',part=1, fname='LOVE-LETTER-FOR-YOU.TXT.vbs'): - msg = mime.message_from_file(open('test/'+vname,"r")) + msg = mime.message_from_file(open('test/'+vname,"rb")) mime.defang(msg,scan_zip=True) self.failUnless(msg.ismodified(),"virus not removed") oname = vname + '.out' - msg.dump(open('test/'+oname,"w")) - msg = mime.message_from_file(open('test/'+oname,"r")) + msg.dump(open('test/'+oname,"wb")) + msg = mime.message_from_file(open('test/'+oname,"rb")) txt2 = msg.get_payload() if type(txt2) == list: txt2 = txt2[part].get_payload() @@ -110,11 +110,11 @@ class MimeTestCase(unittest.TestCase): # virus6 has no parts - the virus is directly inline def testDefang6(self,vname="virus6",fname='FAX20.exe'): - msg = mime.message_from_file(open('test/'+vname,"r")) + msg = mime.message_from_file(open('test/'+vname,"rb")) mime.defang(msg) oname = vname + '.out' - msg.dump(open('test/'+oname,"w")) - msg = mime.message_from_file(open('test/'+oname,"r")) + msg.dump(open('test/'+oname,"wb")) + msg = mime.message_from_file(open('test/'+oname,"rb")) self.failIf(msg.ismultipart()) txt2 = msg.get_payload() self.failUnless(txt2 == mime.virus_msg % \ @@ -123,11 +123,11 @@ class MimeTestCase(unittest.TestCase): # honey virus has a sneaky ASP payload which is parsed correctly # by email package in python-2.2.2, but not by mime.MimeMessage or 2.2.1 def testDefang7(self,vname="honey",fname='story[1].scr'): - msg = mime.message_from_file(open('test/'+vname,"r")) + msg = mime.message_from_file(open('test/'+vname,"rb")) mime.defang(msg) oname = vname + '.out' - msg.dump(open('test/'+oname,"w")) - msg = mime.message_from_file(open('test/'+oname,"r")) + msg.dump(open('test/'+oname,"wb")) + msg = mime.message_from_file(open('test/'+oname,"rb")) parts = msg.get_payload() txt2 = parts[1].get_payload() txt3 = parts[2].get_payload() @@ -138,7 +138,7 @@ class MimeTestCase(unittest.TestCase): ('story[1].asp',hostname,None),txt3) def testParse2(self,fname="spam7"): - msg = mime.message_from_file(open('test/'+fname,"r")) + msg = mime.message_from_file(open('test/'+fname,"rb")) self.failUnless(msg.ismultipart()) parts = msg.get_payload() self.failUnless(len(parts) == 2) @@ -148,11 +148,11 @@ class MimeTestCase(unittest.TestCase): def testZip(self,vname="zip1",fname='zip.zip'): self.testDefang(vname,1,'zip.zip') # test scan_zip flag - msg = mime.message_from_file(open('test/'+vname,"r")) + msg = mime.message_from_file(open('test/'+vname,"rb")) mime.defang(msg,scan_zip=False) self.failIf(msg.ismodified()) # test ignoring empty zip (often found in DSNs) - msg = mime.message_from_file(open('test/zip2','r')) + msg = mime.message_from_file(open('test/zip2','rb')) mime.defang(msg,scan_zip=True) self.failIf(msg.ismodified()) # test corrupt zip (often an EXE named as a ZIP) @@ -177,10 +177,10 @@ class MimeTestCase(unittest.TestCase): def testCheckAttach(self,fname="test1"): # test1 contains a very long filename - msg = mime.message_from_file(open('test/'+fname,'r')) + msg = mime.message_from_file(open('test/'+fname,'rb')) mime.defang(msg,scan_zip=True) self.failIf(msg.ismodified()) - msg = mime.message_from_file(open('test/test2','r')) + msg = mime.message_from_file(open('test/test2','rb')) rc = mime.check_attachments(msg,self._chk_attach) self.assertEquals(self.filename,"7501'S FOR TWO GOLDEN SOURCES SHIPMENTS FOR TAX & DUTY PURPOSES ONLY.PDF") self.assertEquals(rc,Milter.CONTINUE) @@ -205,7 +205,7 @@ if __name__ == '__main__': unittest.main() else: for fname in sys.argv[1:]: - fp = open(fname,'r') + fp = open(fname,'rb') msg = mime.message_from_file(fp) mime.defang(msg,scan_zip=True) print(msg.as_string())