← Back to branch summary

~smoser/ubuntu/xenial/cloud-init/lp1506187

« back to all changes in this revision

Viewing changes to cloudinit/UserDataHandler.py

  • Committer: Bazaar Package Importer
  • Author(s): Scott Moser
  • Date: 2010-02-04 03:00:05 UTC
  • Revision ID: james.westby@ubuntu.com-20100204030005-r2y1568be8rbslxo
Tags: upstream-0.5.3
Import upstream version 0.5.3

Show diffs side-by-side

added added

removed removed

Lines of Context:
cloudinit/UserDataHandler.py
 
1
import email
 
2
 
 
3
from email.mime.multipart import MIMEMultipart
 
4
from email.mime.text import MIMEText
 
5
 
 
6
 
 
7
starts_with_mappings={
 
8
    '#include' : 'text/x-include-url',
 
9
    '#!' : 'text/x-shellscript',
 
10
    '#cloud-config' : 'text/cloud-config',
 
11
    '#upstart-job'  : 'text/upstart-job',
 
12
    '#part-handler' : 'text/part-handler'
 
13
}
 
14
 
 
15
# if 'str' is compressed return decompressed otherwise return it
 
16
def decomp_str(str):
 
17
    import StringIO
 
18
    import gzip
 
19
    try:
 
20
        uncomp = gzip.GzipFile(None,"rb",1,StringIO.StringIO(str)).read()
 
21
        return(uncomp)
 
22
    except:
 
23
        return(str)
 
24
 
 
25
def do_include(str,parts):
 
26
    import urllib
 
27
    # is just a list of urls, one per line
 
28
    for line in str.splitlines():
 
29
        if line == "#include": continue
 
30
        if line.startswith("#"): continue
 
31
        content = urllib.urlopen(line).read()
 
32
        process_includes(email.message_from_string(decomp_str(content)),parts)
 
33
 
 
34
def process_includes(msg,parts):
 
35
    # parts is a dictionary of arrays
 
36
    # parts['content']
 
37
    # parts['names']
 
38
    # parts['types']
 
39
    for t in ( 'content', 'names', 'types' ):
 
40
        if not parts.has_key(t):
 
41
            parts[t]=[ ]
 
42
    for part in msg.walk():
 
43
        # multipart/* are just containers
 
44
        if part.get_content_maintype() == 'multipart':
 
45
            continue
 
46
 
 
47
        payload = part.get_payload()
 
48
 
 
49
        ctype = None
 
50
        for str, gtype in starts_with_mappings.items():
 
51
            if payload.startswith(str):
 
52
                ctype = gtype
 
53
                break
 
54
 
 
55
        if ctype is None:
 
56
            ctype = part.get_content_type()
 
57
 
 
58
        if ctype == 'text/x-include-url':
 
59
            do_include(payload,parts)
 
60
            continue
 
61
 
 
62
        filename = part.get_filename()
 
63
        if not filename:
 
64
            filename = 'part-%03d' % len(parts['content'])
 
65
 
 
66
        parts['content'].append(payload)
 
67
        parts['types'].append(ctype)
 
68
        parts['names'].append(filename)
 
69
 
 
70
def parts2mime(parts):
 
71
    outer = MIMEMultipart()
 
72
 
 
73
    i = 0
 
74
    while i < len(parts['content']):
 
75
        if parts['types'][i] is None:
 
76
            # No guess could be made, or the file is encoded (compressed), so
 
77
            # use a generic bag-of-bits type.
 
78
            ctype = 'application/octet-stream'
 
79
        else: ctype = parts['types'][i]
 
80
        maintype, subtype = ctype.split('/', 1)
 
81
        if maintype == 'text':
 
82
            msg = MIMEText(parts['content'][i], _subtype=subtype)
 
83
        else:
 
84
            msg = MIMEBase(maintype, subtype)
 
85
            msg.set_payload(parts['content'][i])
 
86
            # Encode the payload using Base64
 
87
            encoders.encode_base64(msg)
 
88
        # Set the filename parameter
 
89
        msg.add_header('Content-Disposition', 'attachment', 
 
90
            filename=parts['names'][i])
 
91
        outer.attach(msg)
 
92
 
 
93
        i=i+1
 
94
    return(outer.as_string())
 
95
 
 
96
# this is heavily wasteful, reads through userdata string input
 
97
def preprocess_userdata(data):
 
98
    parts = { }
 
99
    process_includes(email.message_from_string(decomp_str(data)),parts)
 
100
    return(parts2mime(parts))
 
101
 
 
102
# callbacks is a dictionary with:
 
103
#  { 'content-type': handler(data,content_type,filename,payload) }
 
104
def walk_userdata(str, callbacks, data = None):
 
105
    partnum = 0
 
106
    for part in email.message_from_string(str).walk():
 
107
        # multipart/* are just containers
 
108
        if part.get_content_maintype() == 'multipart':
 
109
            continue
 
110
 
 
111
        ctype = part.get_content_type()
 
112
        if ctype is None:
 
113
            ctype = 'application/octet-stream'
 
114
 
 
115
        filename = part.get_filename()
 
116
        if not filename:
 
117
            filename = 'part-%03d' % partnum
 
118
 
 
119
        if callbacks.has_key(ctype):
 
120
            callbacks[ctype](data,ctype,filename,part.get_payload())
 
121
 
 
122
        partnum = partnum+1
 
123
 
 
124
if __name__ == "__main__":
 
125
    import sys
 
126
    data = decomp_str(file(sys.argv[1]).read())
 
127
    parts = { }
 
128
    process_includes(email.message_from_string(data),parts)
 
129
    print "#found %s parts" % len(parts['content'])
 
130
    print parts2mime(parts)