diff --git a/email2pdf b/email2pdf index 7408b71..032260e 100755 --- a/email2pdf +++ b/email2pdf @@ -2,6 +2,7 @@ from datetime import datetime from subprocess import Popen, PIPE +from email.header import decode_header import argparse import email import os @@ -83,10 +84,29 @@ def main(): output = p.communicate(input=payload) else: payload = part.get_payload(decode=True) - with open(output_file_name, 'wb') as output_file: + + filename = part.get_filename() + if decode_header(filename)[0][1] is not None: + filename = str(decode_header(filename)[0][0]).decode(decode_header(filename)[0][1]) + + fullFilename = os.path.join(output_directory, filename) + + fullFilename = get_unique_version(fullFilename) + + with open(fullFilename, 'wb') as output_file: output_file.write(payload) +def get_unique_version(filename): + # From here: http://stackoverflow.com/q/183480/27641 + counter = 1 + file_name_parts = os.path.splitext(filename) + while os.path.isfile(filename): + filename = file_name_parts[0] + '_' + str(counter) + file_name_parts[1] + counter += 1 + return filename + + def find_depth_first(message, content_type): if message.is_multipart(): for part in message.get_payload(): diff --git a/email2pdf_unittest b/email2pdf_unittest index 5e764c8..b349e1a 100755 --- a/email2pdf_unittest +++ b/email2pdf_unittest @@ -3,8 +3,11 @@ from email.message import Message from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText +from email.mime.base import MIMEBase +from email import encoders from email.utils import formatdate from subprocess import Popen, PIPE, DEVNULL +from reportlab.pdfgen import canvas import os import sys @@ -64,6 +67,27 @@ def attachText(message, content): message.attach(MIMEText(content, 'plain')) +def attachPDF(message, string): + unused_f_handle, file_name = tempfile.mkstemp(prefix="email2pdf_unittest_file", suffix=".pdf") + + try: + cv = canvas.Canvas(file_name) + cv.drawString(0, 500, string) + cv.save() + + part = MIMEBase('application', "pdf") + openHandle = open(file_name, "rb") + part.set_payload(openHandle.read()) + encoders.encode_base64(part) + part.add_header('Content-Disposition', 'attachment; filename="%s"' % os.path.basename(file_name)) + message.attach(part) + openHandle.close() + + return file_name + finally: + os.unlink(file_name) + + class TestBasic(unittest.TestCase): def setUp(self): self.msg = Message() @@ -96,7 +120,7 @@ class TestBasic(unittest.TestCase): def test_plaincontent_fileexist(self): setPlainContent(self.msg, "Hello!") - f_handle, f_path = tempfile.mkstemp() + unused_f_handle, f_path = tempfile.mkstemp() try: self.assertNotEqual(invokeEmail2PDF(self.msg, outputFile=f_path, errorExpected=True), 0) finally: @@ -134,6 +158,14 @@ class TestMIME(unittest.TestCase): attachHTML(self.msg, "<p>Some basic textual content</p>") self.assertEqual(invokeEmail2PDF(self.msg), 0) + def test_pdf(self): + addHeaders(self.msg, "From", "To", "Subject") + attachText(self.msg, "Some basic textual content") + filename = attachPDF(self.msg, "Some PDF content") + self.assertEqual(invokeEmail2PDF(self.msg), 0) + self.assertTrue(os.path.exists(filename)) + os.unlink(filename) + def tearDown(self): time.sleep(DELAY)