Ticket #1541: mail.diff

django/core/mail.py

@@ -4 +4 @@
 
 from django.conf import settings
 from email.MIMEText import MIMEText
+from email.mime.multipart import MIMEMultipart
+from email.mime.base 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.
@@ -154 +170 @@
     """
     A container for email information.
     """
-    def __init__(self, subject='', body='', from_email=None, to=None, bcc=None, connection=None):
+    def __init__(self, subject='', body='', from_email=None, to=None, bcc=None, attachments=None, connection=None):
         self.to = to or []
         self.bcc = bcc or []
         self.from_email = from_email or settings.DEFAULT_FROM_EMAIL
         self.subject = subject
         self.body = body
         self.connection = connection
+        self.attachments = attachments or []
 
     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 send_mail(subject, message, from_email, recipient_list, fail_silently=False, auth_user=None, auth_password=None):
+    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)
+
+    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
+
+def send_mail(subject, message, from_email, recipient_list, attachments=None, fail_silently=False, auth_user=None, auth_password=None):
     """
     Easy wrapper for sending a single message to a recipient list. All members
     of the recipient list will see the other recipients in the 'To' field.
@@ -202 +258 @@
     """
     connection = SMTPConnection(username=auth_user, password=auth_password,
                                  fail_silently=fail_silently)
-    return EmailMessage(subject, message, from_email, recipient_list, connection=connection).send()
+    return EmailMessage(subject, message, from_email, recipient_list, attachments=attachments, connection=connection).send()
 
 def send_mass_mail(datatuple, fail_silently=False, auth_user=None, auth_password=None):
     """
@@ -233 +289 @@
     EmailMessage(settings.EMAIL_SUBJECT_PREFIX + subject, message,
             settings.SERVER_EMAIL, [a[1] for a in
                 settings.MANAGERS]).send(fail_silently=fail_silently)
-