Unify filename, header replacement.

1 files changed, 32 insertions, 24 deletions

diff --git a/cloudinit/user_data.py b/cloudinit/user_data.py
index e17bcaee..454f3c06 100644
--- a/cloudinit/user_data.py
+++ b/cloudinit/user_data.py
@@ -70,6 +70,19 @@ ATTACHMENT_FIELD = 'Number-Attachments'
 EXAMINE_FOR_LAUNCH_INDEX = ["text/cloud-config"]
 
 
+def _replace_header(msg, key, value):
+    del msg[key]
+    msg[key] = value
+
+
+def _set_filename(msg, filename):
+    if not filename:
+        return
+    del msg['Content-Disposition']
+    msg.add_header('Content-Disposition',
+                   'attachment', filename=str(filename))
+
+
 class UserDataProcessor(object):
     def __init__(self, paths):
         self.paths = paths
@@ -83,14 +96,7 @@ class UserDataProcessor(object):
     def _process_msg(self, base_msg, append_msg):
 
         def find_ctype(payload):
-            ctype = handlers.type_from_starts_with(payload)
-            return ctype
-
-        def replace_header(part, key, value):
-            if key in part:
-                part.replace_header(key, value)
-            else:
-                part[key] = value
+            return handlers.type_from_starts_with(payload)
 
         for part in base_msg.walk():
             if is_skippable(part):
@@ -132,13 +138,17 @@ class UserDataProcessor(object):
                 maintype, subtype = ctype.split("/", 1)
                 n_part = MIMENonMultipart(maintype, subtype)
                 n_part.set_payload(payload)
-                if part.get_filename():
-                    n_part.add_header('Content-Disposition', 'attachment',
-                                      filename=part.get_filename())
+                # Copy various headers from the old part to the new one,
+                # but don't include all the headers since some are not useful
+                # after decoding and decompression.
+                _set_filename(n_part, part.get_filename())
+                for h in ('Launch-Index',):
+                    if h in part:
+                        _replace_header(n_part, h, str(part[h]))
                 part = n_part
 
             if ctype != ctype_orig:
-                replace_header(part, CONTENT_TYPE, ctype)
+                _replace_header(part, CONTENT_TYPE, ctype)
 
             if ctype in INCLUDE_TYPES:
                 self._do_include(payload, append_msg)
@@ -150,7 +160,7 @@ class UserDataProcessor(object):
 
             # TODO(harlowja): Should this be happening, shouldn't
             # the part header be modified and not the base?
-            replace_header(base_msg, CONTENT_TYPE, ctype)
+            _replace_header(base_msg, CONTENT_TYPE, ctype)
 
             self._attach_part(append_msg, part)
 
@@ -185,8 +195,7 @@ class UserDataProcessor(object):
 
     def _process_before_attach(self, msg, attached_id):
         if not msg.get_filename():
-            msg.add_header('Content-Disposition',
-                           'attachment', filename=PART_FN_TPL % (attached_id))
+            _set_filename(msg, PART_FN_TPL % (attached_id))
         self._attach_launch_index(msg)
 
     def _do_include(self, content, append_msg):
@@ -264,13 +273,15 @@ class UserDataProcessor(object):
                 msg.set_payload(content)
 
             if 'filename' in ent:
-                msg.add_header('Content-Disposition',
-                               'attachment', filename=ent['filename'])
+                _set_filename(msg, ent['filename'])
             if 'launch-index' in ent:
                 msg.add_header('Launch-Index', str(ent['launch-index']))
 
             for header in list(ent.keys()):
-                if header in ('content', 'filename', 'type', 'launch-index'):
+                if header.lower() in ('content', 'filename', 'type',
+                                      'launch-index', 'content-disposition',
+                                      ATTACHMENT_FIELD.lower(),
+                                      CONTENT_TYPE.lower()):
                     continue
                 msg.add_header(header, ent[header])
 
@@ -285,13 +296,13 @@ class UserDataProcessor(object):
             outer_msg[ATTACHMENT_FIELD] = '0'
 
         if new_count is not None:
-            outer_msg.replace_header(ATTACHMENT_FIELD, str(new_count))
+            _replace_header(outer_msg, ATTACHMENT_FIELD, str(new_count))
 
         fetched_count = 0
         try:
             fetched_count = int(outer_msg.get(ATTACHMENT_FIELD))
         except (ValueError, TypeError):
-            outer_msg.replace_header(ATTACHMENT_FIELD, str(fetched_count))
+            _replace_header(outer_msg, ATTACHMENT_FIELD, str(fetched_count))
         return fetched_count
 
     def _attach_part(self, outer_msg, part):
@@ -323,10 +334,7 @@ def convert_string(raw_data, headers=None):
     if "mime-version:" in data[0:4096].lower():
         msg = email.message_from_string(data)
         for (key, val) in headers.iteritems():
-            if key in msg:
-                msg.replace_header(key, val)
-            else:
-                msg[key] = val
+            _replace_header(msg, key, val)
     else:
         mtype = headers.get(CONTENT_TYPE, NOT_MULTIPART_TYPE)
         maintype, subtype = mtype.split("/", 1)