Ticket #1541: mail.2.diff

django/common/core/mail.py

@@ -4 +4 @@
 
 from django.conf import settings
 from email.MIMEText import MIMEText
+from email.MIMEMultipart import MIMEMultipart
+from email.MIMEBase import MIMEBase
+from email import Encoders
+import mimetypes
 from email.Header import Header
 from email.Utils import formatdate
 from email import Charset
@@ -13 +17 @@
 import time
 import random
 
+# Default MIME type to use for attachments if it cannot be guessed.
+DEFAULT_ATTACHMENT_MIME_TYPE = 'application/octet-stream'
+
 # Don't BASE64-encode UTF-8 messages so that we avoid unwanted attention from
 # some spam filters.
 Charset.add_charset('utf-8', Charset.SHORTEST, Charset.QP, 'utf-8')
@@ -64 +71 @@
             val = Header(val, settings.DEFAULT_CHARSET)
         MIMEText.__setitem__(self, name, val)
 
+class SafeMIMEMultipart(MIMEMultipart):
+    def __setitem__(self, name, val):
+        "Forbids multi-line headers, to prevent header injection."
+        if '\n' in val or '\r' in val:
+            raise BadHeaderError, "Header values can't contain newlines (got %r for header %r)" % (val, name)
+        if name == 'Subject':
+            val = Header(val, settings.DEFAULT_CHARSET)
+        MIMEMultipart.__setitem__(self, name, val)
+
 class SMTPConnection(object):
     """
     A wrapper that manages the SMTP network connection.
@@ -161 +177 @@
         self.subject = subject
         self.body = body
         self.connection = connection
+        self.attachments = []
 
     def get_connection(self, fail_silently=False):
         if not self.connection:
@@ -169 +186 @@
 
     def message(self):
         msg = SafeMIMEText(self.body, 'plain', settings.DEFAULT_CHARSET)
+        if self.attachments:
+            body_msg = msg
+            msg = SafeMIMEMultipart()
+            msg.attach(body_msg)
+            for params in self.attachments:
+                msg.attach(self._compile_attachment(*params))
         msg['Subject'] = self.subject
         msg['From'] = self.from_email
         msg['To'] = ', '.join(self.to)
@@ -189 +212 @@
         """Send the email message."""
         return self.get_connection(fail_silently).send_messages([self])
 
+    def attach(self, filename, content, mimetype=None):
+        """Attaches a file with the given filename and content."""
+        self.attachments.append((filename, content, mimetype))
+
+    def attach_file(self, path, mimetype=None):
+        """Attaches a file from the filesystem."""
+        filename = os.path.basename(path)
+        f = open(path, 'rb')
+        content = f.read()
+        f.close()
+        self.attach(filename, content, mimetype)
+
+    def _compile_attachment(self, filename, content, mimetype):
+        """
+        Compiles the given attachment to a MIME object and returns the new
+        attachment.
+        """
+        if mimetype is None:
+            # Guess the mimetype based on the filename if possible.
+            mimetype, xx = mimetypes.guess_type(filename)
+            if mimetype is None:
+                mimetype = DEFAULT_ATTACHMENT_MIME_TYPE
+        basetype, subtype = mimetype.split('/', 1)
+        if basetype == 'text':
+            attachment = SafeMIMEText(content, subtype, settings.DEFAULT_CHARSET)
+        else:
+            # Encode non-text attachments with base64.
+            attachment = MIMEBase(basetype, subtype)
+            attachment.set_payload(content)
+            Encoders.encode_base64(attachment)
+        attachment.add_header('Content-Disposition', 'attachment', filename=filename)
+        return attachment
+
 
--- docs/email.txt       Mon Jun 25 09:36:30 2007
+++ docs/email.txt       Tue Jun 26 14:33:28 2007
@@ -198 +198 @@
 .. note::
     Not all features of the ``EmailMessage`` class are available through the
     ``send_mail()`` and related wrapper functions. If you wish to use advanced
-    features, such as BCC'ed recipients or multi-part e-mail, you'll need to
-    create ``EmailMessage`` instances directly.
+    features, such as BCC'ed recipients, file attachments, or multi-part
+    e-mail, you'll need to create ``EmailMessage`` instances directly.
 
 In general, ``EmailMessage`` is responsible for creating the e-mail message
 itself. ``SMTPConnection`` is responsible for the network connection side of
@@ -238 +238 @@
       SMTP server needs to be told the full list of recipients when the message
       is sent. If you add another way to specify recipients in your class, they
       need to be returned from this method as well.
+
+    * ``attach()`` creates a new file attachment and adds it to the message. It
+      takes three arguments: ``filename``, ``content`` and ``mimetype``.
+      ``filename`` is the name of the file attachment as it will appear in the
+      email. ``content`` is the data that will be contained inside the
+      attachment. ``mimetype`` is optional, and defines the attachment's MIME
+      content type. If you omit ``mimetype``, the MIME content type will be
+      guessed from the filename of the attachment.
+
+    * ``attach_file()`` creates a new attachment using a file on the
+      filesystem, and takes two arguments. The first is the path of the file to
+      attach. The second argument is the MIME content type to use for the
+      attachment, and is optional. If the MIME type is omitted then it will be
+      guessed from the filename.
 
 The ``SMTPConnection`` class is initialized with the host, port, username and
 password for the SMTP server. If you don't specify one or more of those