bpo-32634: fix write in email/generator.py

Lib/email/generator.py

@@ -22,7 +22,6 @@
 fcre = re.compile(r'^From ', re.MULTILINE)
 
 
-
 class Generator:
     """Generates output from a Message object tree.
 
@@ -122,7 +121,7 @@ def clone(self, fp):
         """Clone this generator with the exact same options."""
         return self.__class__(fp,
                               self._mangle_from_,
-                              None, # Use policy setting, which we've adjusted
+                              None,  # Use policy setting, which we've adjusted
                               policy=self.policy)
 
     #
@@ -159,7 +158,7 @@ def _write_lines(self, lines):
         # XXX logic tells me this else should be needed, but the tests fail
         # with it and pass without it.  (NLCRE.split ends with a blank element
         # if and only if there was a trailing newline.)
-        #else:
+        # else:
         #    self.write(self._NL)
 
     def _write(self, msg):
@@ -186,8 +185,14 @@ def _write(self, msg):
         # If we munged the cte, copy the message again and re-fix the CTE.
         if munge_cte:
             msg = deepcopy(msg)
-            msg.replace_header('content-transfer-encoding', munge_cte[0])
-            msg.replace_header('content-type', munge_cte[1])
+            if msg.get('content-transfer-encoding', None) is not None:
+                msg.replace_header('content-transfer-encoding', munge_cte[0])
+            else:
+                msg.add_header('content-transfer-encoding', munge_cte[0])
+            if msg.get('content-type', None) is not None:
+                msg.replace_header('content-type', munge_cte[1])
+            else:
+                msg.add_header('content-type', munge_cte[1])
         # Write the headers.  First we see if the message object wants to
         # handle that itself.  If not, we'll do it generically.
         meth = getattr(msg, '_write_headers', None)
@@ -377,7 +382,8 @@ def _make_boundary(cls, text=None):
         b = boundary
         counter = 0
         while True:
-            cre = cls._compile_re('^--' + re.escape(b) + '(--)?$', re.MULTILINE)
+            cre = cls._compile_re('^--' + re.escape(b) + '(--)?$',
+                                  re.MULTILINE)
             if not cre.search(text):
                 break
             b = boundary + '.' + str(counter)
@@ -424,12 +430,12 @@ def _handle_text(self, msg):
         # just write it back out.
         if msg._payload is None:
             return
-        if _has_surrogates(msg._payload) and not self.policy.cte_type=='7bit':
+        if _has_surrogates(msg._payload) and self.policy.cte_type != '7bit':
             if self._mangle_from_:
                 msg._payload = fcre.sub(">From ", msg._payload)
             self._write_lines(msg._payload)
         else:
-            super(BytesGenerator,self)._handle_text(msg)
+            super(BytesGenerator, self)._handle_text(msg)
 
     # Default body handler
     _writeBody = _handle_text
@@ -439,17 +445,17 @@ def _compile_re(cls, s, flags):
         return re.compile(s.encode('ascii'), flags)
 
 
-
 _FMT = '[Non-text (%(type)s) part of message omitted, filename %(filename)s]'
 
+
 class DecodedGenerator(Generator):
     """Generates a text representation of a message.
 
     Like the Generator base class, except that non-text parts are substituted
     with a format string representing the part.
     """
-    def __init__(self, outfp, mangle_from_=None, maxheaderlen=None, fmt=None, *,
-                 policy=None):
+    def __init__(self, outfp, mangle_from_=None, maxheaderlen=None, fmt=None,
+                 *, policy=None):
         """Like Generator.__init__() except that an additional optional
         argument is allowed.
 
@@ -488,18 +494,17 @@ def _dispatch(self, msg):
                 pass
             else:
                 print(self._fmt % {
-                    'type'       : part.get_content_type(),
-                    'maintype'   : part.get_content_maintype(),
-                    'subtype'    : part.get_content_subtype(),
-                    'filename'   : part.get_filename('[no filename]'),
+                    'type': part.get_content_type(),
+                    'maintype': part.get_content_maintype(),
+                    'subtype': part.get_content_subtype(),
+                    'filename': part.get_filename('[no filename]'),
                     'description': part.get('Content-Description',
                                             '[no description]'),
-                    'encoding'   : part.get('Content-Transfer-Encoding',
-                                            '[no encoding]'),
+                    'encoding': part.get('Content-Transfer-Encoding',
+                                         '[no encoding]'),
                     }, file=self)
 
 
-
 # Helper used by Generator._make_boundary
 _width = len(repr(sys.maxsize-1))
 _fmt = '%%0%dd' % _width

Misc/NEWS.d/next/Library/2018-01-23-14-12-23.bpo-32634.KLH3kj.rst

@@ -0,0 +1 @@
+Fix chrashing email message parsing when incomplete headers given