← Back to branch summary

~smoser/ubuntu/vivid/cloud-init/2to3

« back to all changes in this revision

Viewing changes to cloudinit/user_data.py

  • Committer: Package Import Robot
  • Author(s): Scott Moser
  • Date: 2013-04-11 12:55:51 UTC
  • mfrom: (245.3.9 raring-proposed)
  • Revision ID: package-import@ubuntu.com-20130411125551-8k60jsoot7t21z4b
https://launchpad.net/bugs/1164147
* New upstream snapshot.
  * make apt-get invoke 'dist-upgrade' rather than 'upgrade' for
    package_upgrade. (LP: #1164147)
  * workaround 2.6 kernel issue that stopped blkid from showing /dev/sr0

Show diffs side-by-side

added added

removed removed

Lines of Context:
cloudinit/user_data.py
23
23
import os
24
24
 
25
25
import email
26
 
 
27
26
from email.mime.base import MIMEBase
28
27
from email.mime.multipart import MIMEMultipart
29
 
from email.mime.nonmultipart import MIMENonMultipart
30
28
from email.mime.text import MIMEText
31
29
 
32
30
from cloudinit import handlers
50
48
UNDEF_TYPE = "text/plain"
51
49
ARCHIVE_UNDEF_TYPE = "text/cloud-config"
52
50
 
53
 
# This seems to hit most of the gzip possible content types.
54
 
DECOMP_TYPES = [
55
 
    'application/gzip',
56
 
    'application/gzip-compressed',
57
 
    'application/gzipped',
58
 
    'application/x-compress',
59
 
    'application/x-compressed',
60
 
    'application/x-gunzip',
61
 
    'application/x-gzip',
62
 
    'application/x-gzip-compressed',
63
 
]
64
 
 
65
51
# Msg header used to track attachments
66
52
ATTACHMENT_FIELD = 'Number-Attachments'
67
53
 
70
56
EXAMINE_FOR_LAUNCH_INDEX = ["text/cloud-config"]
71
57
 
72
58
 
73
 
def _replace_header(msg, key, value):
74
 
    del msg[key]
75
 
    msg[key] = value
76
 
 
77
 
 
78
 
def _set_filename(msg, filename):
79
 
    del msg['Content-Disposition']
80
 
    msg.add_header('Content-Disposition',
81
 
                   'attachment', filename=str(filename))
82
 
 
83
 
 
84
59
class UserDataProcessor(object):
85
60
    def __init__(self, paths):
86
61
        self.paths = paths
92
67
        return accumulating_msg
93
68
 
94
69
    def _process_msg(self, base_msg, append_msg):
95
 
 
96
 
        def find_ctype(payload):
97
 
            return handlers.type_from_starts_with(payload)
98
 
 
99
70
        for part in base_msg.walk():
100
71
            if is_skippable(part):
101
72
                continue
103
74
            ctype = None
104
75
            ctype_orig = part.get_content_type()
105
76
            payload = part.get_payload(decode=True)
106
 
            was_compressed = False
107
 
 
108
 
            # When the message states it is of a gzipped content type ensure
109
 
            # that we attempt to decode said payload so that the decompressed
110
 
            # data can be examined (instead of the compressed data).
111
 
            if ctype_orig in DECOMP_TYPES:
112
 
                try:
113
 
                    payload = util.decomp_gzip(payload, quiet=False)
114
 
                    # At this point we don't know what the content-type is
115
 
                    # since we just decompressed it.
116
 
                    ctype_orig = None
117
 
                    was_compressed = True
118
 
                except util.DecompressionError as e:
119
 
                    LOG.warn("Failed decompressing payload from %s of length"
120
 
                             " %s due to: %s", ctype_orig, len(payload), e)
121
 
                    continue
122
 
 
123
 
            # Attempt to figure out the payloads content-type
 
77
 
124
78
            if not ctype_orig:
125
79
                ctype_orig = UNDEF_TYPE
 
80
 
126
81
            if ctype_orig in TYPE_NEEDED:
127
 
                ctype = find_ctype(payload)
 
82
                ctype = handlers.type_from_starts_with(payload)
 
83
 
128
84
            if ctype is None:
129
85
                ctype = ctype_orig
130
86
 
131
 
            # In the case where the data was compressed, we want to make sure
132
 
            # that we create a new message that contains the found content
133
 
            # type with the uncompressed content since later traversals of the
134
 
            # messages will expect a part not compressed.
135
 
            if was_compressed:
136
 
                maintype, subtype = ctype.split("/", 1)
137
 
                n_part = MIMENonMultipart(maintype, subtype)
138
 
                n_part.set_payload(payload)
139
 
                # Copy various headers from the old part to the new one,
140
 
                # but don't include all the headers since some are not useful
141
 
                # after decoding and decompression.
142
 
                if part.get_filename():
143
 
                    _set_filename(n_part, part.get_filename())
144
 
                for h in ('Launch-Index',):
145
 
                    if h in part:
146
 
                        _replace_header(n_part, h, str(part[h]))
147
 
                part = n_part
148
 
 
149
87
            if ctype != ctype_orig:
150
 
                _replace_header(part, CONTENT_TYPE, ctype)
 
88
                if CONTENT_TYPE in part:
 
89
                    part.replace_header(CONTENT_TYPE, ctype)
 
90
                else:
 
91
                    part[CONTENT_TYPE] = ctype
151
92
 
152
93
            if ctype in INCLUDE_TYPES:
153
94
                self._do_include(payload, append_msg)
157
98
                self._explode_archive(payload, append_msg)
158
99
                continue
159
100
 
160
 
            # TODO(harlowja): Should this be happening, shouldn't
 
101
            # Should this be happening, shouldn't
161
102
            # the part header be modified and not the base?
162
 
            _replace_header(base_msg, CONTENT_TYPE, ctype)
 
103
            if CONTENT_TYPE in base_msg:
 
104
                base_msg.replace_header(CONTENT_TYPE, ctype)
 
105
            else:
 
106
                base_msg[CONTENT_TYPE] = ctype
163
107
 
164
108
            self._attach_part(append_msg, part)
165
109
 
194
138
 
195
139
    def _process_before_attach(self, msg, attached_id):
196
140
        if not msg.get_filename():
197
 
            _set_filename(msg, PART_FN_TPL % (attached_id))
 
141
            msg.add_header('Content-Disposition',
 
142
                           'attachment', filename=PART_FN_TPL % (attached_id))
198
143
        self._attach_launch_index(msg)
199
144
 
200
145
    def _do_include(self, content, append_msg):
272
217
                msg.set_payload(content)
273
218
 
274
219
            if 'filename' in ent:
275
 
                _set_filename(msg, ent['filename'])
 
220
                msg.add_header('Content-Disposition',
 
221
                               'attachment', filename=ent['filename'])
276
222
            if 'launch-index' in ent:
277
223
                msg.add_header('Launch-Index', str(ent['launch-index']))
278
224
 
279
225
            for header in list(ent.keys()):
280
 
                if header.lower() in ('content', 'filename', 'type',
281
 
                                      'launch-index', 'content-disposition',
282
 
                                      ATTACHMENT_FIELD.lower(),
283
 
                                      CONTENT_TYPE.lower()):
 
226
                if header in ('content', 'filename', 'type', 'launch-index'):
284
227
                    continue
285
228
                msg.add_header(header, ent[header])
286
229
 
295
238
            outer_msg[ATTACHMENT_FIELD] = '0'
296
239
 
297
240
        if new_count is not None:
298
 
            _replace_header(outer_msg, ATTACHMENT_FIELD, str(new_count))
 
241
            outer_msg.replace_header(ATTACHMENT_FIELD, str(new_count))
299
242
 
300
243
        fetched_count = 0
301
244
        try:
302
245
            fetched_count = int(outer_msg.get(ATTACHMENT_FIELD))
303
246
        except (ValueError, TypeError):
304
 
            _replace_header(outer_msg, ATTACHMENT_FIELD, str(fetched_count))
 
247
            outer_msg.replace_header(ATTACHMENT_FIELD, str(fetched_count))
305
248
        return fetched_count
306
249
 
307
250
    def _attach_part(self, outer_msg, part):
333
276
    if "mime-version:" in data[0:4096].lower():
334
277
        msg = email.message_from_string(data)
335
278
        for (key, val) in headers.iteritems():
336
 
            _replace_header(msg, key, val)
 
279
            if key in msg:
 
280
                msg.replace_header(key, val)
 
281
            else:
 
282
                msg[key] = val
337
283
    else:
338
284
        mtype = headers.get(CONTENT_TYPE, NOT_MULTIPART_TYPE)
339
285
        maintype, subtype = mtype.split("/", 1)