support launch index specific user-data

EC2 and openstack provide 'launch_index' in their metadata. This allows
the user to specify cloud-config or multipart mime data that includes the
'Launch-Index' header.

If launch index is available in the metadata service, then:
 * any part that contains a launch index other than the current launch-index
   of this instance will be ignored.
 * any part that does not contain a launch index will be considered as
   for this instance.

If there is no such header, or launch_index is not available in the metadata
service, then no such filtering will be done.

LP: #1023177

8 files changed, 195 insertions, 28 deletions

diff --git a/cloudinit/cloud.py b/cloudinit/cloud.py
index 620b3c07..95e0cfb2 100644
--- a/cloudinit/cloud.py
+++ b/cloudinit/cloud.py
@@ -70,12 +70,16 @@ class Cloud(object):
         return fn
 
     # The rest of thes are just useful proxies
-    def get_userdata(self):
-        return self.datasource.get_userdata()
+    def get_userdata(self, apply_filter=True):
+        return self.datasource.get_userdata(apply_filter)
 
     def get_instance_id(self):
         return self.datasource.get_instance_id()
 
+    @property
+    def launch_index(self):
+        return self.datasource.launch_index
+
     def get_public_ssh_keys(self):
         return self.datasource.get_public_ssh_keys()
 
diff --git a/cloudinit/filters/__init__.py b/cloudinit/filters/__init__.py
new file mode 100644
index 00000000..da124641
--- /dev/null
+++ b/cloudinit/filters/__init__.py
@@ -0,0 +1,21 @@
+# vi: ts=4 expandtab
+#
+#    Copyright (C) 2012 Canonical Ltd.
+#    Copyright (C) 2012 Hewlett-Packard Development Company, L.P.
+#    Copyright (C) 2012 Yahoo! Inc.
+#
+#    Author: Scott Moser <scott.moser@canonical.com>
+#    Author: Juerg Haefliger <juerg.haefliger@hp.com>
+#    Author: Joshua Harlow <harlowja@yahoo-inc.com>
+#
+#    This program is free software: you can redistribute it and/or modify
+#    it under the terms of the GNU General Public License version 3, as
+#    published by the Free Software Foundation.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU General Public License for more details.
+#
+#    You should have received a copy of the GNU General Public License
+#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
diff --git a/cloudinit/filters/launch_index.py b/cloudinit/filters/launch_index.py
new file mode 100644
index 00000000..4299fb46
--- /dev/null
+++ b/cloudinit/filters/launch_index.py
@@ -0,0 +1,75 @@
+# vi: ts=4 expandtab
+#
+#    Copyright (C) 2012 Canonical Ltd.
+#    Copyright (C) 2012 Hewlett-Packard Development Company, L.P.
+#    Copyright (C) 2012 Yahoo! Inc.
+#
+#    Author: Scott Moser <scott.moser@canonical.com>
+#    Author: Juerg Haefliger <juerg.haefliger@hp.com>
+#    Author: Joshua Harlow <harlowja@yahoo-inc.com>
+#
+#    This program is free software: you can redistribute it and/or modify
+#    it under the terms of the GNU General Public License version 3, as
+#    published by the Free Software Foundation.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU General Public License for more details.
+#
+#    You should have received a copy of the GNU General Public License
+#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import copy
+
+from cloudinit import log as logging
+from cloudinit import user_data as ud
+from cloudinit import util
+
+LOG = logging.getLogger(__name__)
+
+
+class Filter(object):
+    def __init__(self, wanted_idx, allow_none=True):
+        self.wanted_idx = wanted_idx
+        self.allow_none = allow_none
+
+    def _select(self, message):
+        msg_idx = message.get('Launch-Index', None)
+        if self.allow_none and msg_idx is None:
+            return True
+        msg_idx = util.safe_int(msg_idx)
+        if msg_idx != self.wanted_idx:
+            return False
+        return True
+
+    def _do_filter(self, message):
+        # Don't use walk() here since we want to do the reforming of the 
+        # messages ourselves and not flatten the message listings...
+        if not self._select(message):
+            return None
+        if message.is_multipart():
+            # Recreate it and its child messages
+            prev_msgs = message.get_payload(decode=False)
+            new_msgs = []
+            discarded = 0
+            for m in prev_msgs:
+                m = self._do_filter(m)
+                if m is not None:
+                    new_msgs.append(m)
+                else:
+                    discarded += 1
+            LOG.debug(("Discarding %s multipart messages "
+                       "which do not match launch index %s"),
+                       discarded, self.wanted_idx)
+            new_message = copy.copy(message)
+            new_message.set_payload(new_msgs)
+            new_message[ud.ATTACHMENT_FIELD] = str(len(new_msgs))
+            return new_message
+        else:
+            return copy.copy(message)
+
+    def apply(self, root_message):
+        if self.wanted_idx is None:
+            return root_message
+        return self._do_filter(root_message)
diff --git a/cloudinit/sources/DataSourceEc2.py b/cloudinit/sources/DataSourceEc2.py
index 7e845571..c7ad6d54 100644
--- a/cloudinit/sources/DataSourceEc2.py
+++ b/cloudinit/sources/DataSourceEc2.py
@@ -77,6 +77,12 @@ class DataSourceEc2(sources.DataSource):
                         self.metadata_address)
             return False
 
+    @property
+    def launch_index(self):
+        if not self.metadata:
+            return None
+        return self.metadata.get('ami-launch-index')
+
     def get_instance_id(self):
         return self.metadata['instance-id']
 
diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index 4719d254..3f611d44 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -20,6 +20,8 @@
 #    You should have received a copy of the GNU General Public License
 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+from email.mime.multipart import MIMEMultipart
+
 import abc
 
 from cloudinit import importer
@@ -27,6 +29,8 @@ from cloudinit import log as logging
 from cloudinit import user_data as ud
 from cloudinit import util
 
+from cloudinit.filters import launch_index
+
 DEP_FILESYSTEM = "FILESYSTEM"
 DEP_NETWORK = "NETWORK"
 DS_PREFIX = 'DataSource'
@@ -59,13 +63,31 @@ class DataSource(object):
         else:
             self.ud_proc = ud_proc
 
-    def get_userdata(self):
+    def get_userdata(self, apply_filter=False):
         if self.userdata is None:
-            raw_data = self.get_userdata_raw()
-            self.userdata = self.ud_proc.process(raw_data)
+            self.userdata = self.ud_proc.process(self.get_userdata_raw())
+        if apply_filter:
+            return self._filter_userdata(self.userdata)
         return self.userdata
 
     @property
+    def launch_index(self):
+        if not self.metadata:
+            return None
+        if 'launch-index' in self.metadata:
+            return self.metadata['launch-index']
+        return None
+
+    def _filter_userdata(self, processed_ud):
+        filters = [
+            launch_index.Filter(util.safe_int(self.launch_index)),
+        ]
+        new_ud = processed_ud
+        for f in filters:
+            new_ud = f.apply(new_ud)
+        return new_ud
+
+    @property
     def is_disconnected(self):
         return False
 
diff --git a/cloudinit/stages.py b/cloudinit/stages.py
index c9634a90..af902925 100644
--- a/cloudinit/stages.py
+++ b/cloudinit/stages.py
@@ -347,7 +347,7 @@ class Init(object):
             sys.path.insert(0, idir)
 
         # Ensure datasource fetched before activation (just incase)
-        user_data_msg = self.datasource.get_userdata()
+        user_data_msg = self.datasource.get_userdata(True)
 
         # This keeps track of all the active handlers
         c_handlers = helpers.ContentHandlers()
diff --git a/cloudinit/user_data.py b/cloudinit/user_data.py
index af98b488..5d550e1d 100644
--- a/cloudinit/user_data.py
+++ b/cloudinit/user_data.py
@@ -52,21 +52,23 @@ ARCHIVE_UNDEF_TYPE = "text/cloud-config"
 # Msg header used to track attachments
 ATTACHMENT_FIELD = 'Number-Attachments'
 
+# Only the following content types can have there launch index examined
+# in there payload, evey other content type can still provide a header
+EXAMINE_FOR_LAUNCH_INDEX = ["text/cloud-config", "text/cloud-config-archive"]
+
 
 class UserDataProcessor(object):
     def __init__(self, paths):
         self.paths = paths
 
     def process(self, blob):
-        base_msg = convert_string(blob)
-        process_msg = MIMEMultipart()
-        self._process_msg(base_msg, process_msg)
-        return process_msg
+        accumulating_msg = MIMEMultipart()
+        self._process_msg(convert_string(blob), accumulating_msg)
+        return accumulating_msg
 
     def _process_msg(self, base_msg, append_msg):
         for part in base_msg.walk():
-            # multipart/* are just containers
-            if part.get_content_maintype() == 'multipart':
+            if is_skippable(part):
                 continue
 
             ctype = None
@@ -97,11 +99,41 @@ class UserDataProcessor(object):
 
             self._attach_part(append_msg, part)
 
+    def _attach_launch_index(self, msg):
+        header_idx = msg.get('Launch-Index', None)
+        payload_idx = None
+        if msg.get_content_type() in EXAMINE_FOR_LAUNCH_INDEX:
+            try:
+                # See if it has a launch-index field
+                # that might affect the final header
+                payload = util.load_yaml(msg.get_payload(decode=True))
+                if payload:
+                    payload_idx = payload.get('launch-index')
+            except:
+                pass
+        # Header overrides contents, for now (?) or the other way around?
+        if header_idx is not None:
+            payload_idx = header_idx
+        # Nothing found in payload, use header (if anything there)
+        if payload_idx is None:
+            payload_idx = header_idx
+        if payload_idx is not None:
+            try:
+                msg.add_header('Launch-Index', str(int(payload_idx)))
+            except (ValueError, TypeError):
+                pass
+
     def _get_include_once_filename(self, entry):
         entry_fn = util.hash_blob(entry, 'md5', 64)
         return os.path.join(self.paths.get_ipath_cur('data'),
                             'urlcache', entry_fn)
 
+    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))
+        self._attach_launch_index(msg)
+
     def _do_include(self, content, append_msg):
         # Include a list of urls, one per line
         # also support '#include <url here>'
@@ -178,9 +210,11 @@ class UserDataProcessor(object):
             if 'filename' in ent:
                 msg.add_header('Content-Disposition',
                                'attachment', filename=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'):
+                if header in ('content', 'filename', 'type', 'launch-index'):
                     continue
                 msg.add_header(header, ent['header'])
 
@@ -204,21 +238,23 @@ class UserDataProcessor(object):
             outer_msg.replace_header(ATTACHMENT_FIELD, str(fetched_count))
         return fetched_count
 
-    def _part_filename(self, _unnamed_part, count):
-        return PART_FN_TPL % (count + 1)
-
     def _attach_part(self, outer_msg, part):
         """
-        Attach an part to an outer message. outermsg must be a MIMEMultipart.
-        Modifies a header in the message to keep track of number of attachments.
+        Attach a message to an outer message. outermsg must be a MIMEMultipart.
+        Modifies a header in the outer message to keep track of number of attachments.
         """
-        cur_c = self._multi_part_count(outer_msg)
-        if not part.get_filename():
-            fn = self._part_filename(part, cur_c)
-            part.add_header('Content-Disposition',
-                            'attachment', filename=fn)
+        part_count = self._multi_part_count(outer_msg)
+        self._process_before_attach(part, part_count + 1)
         outer_msg.attach(part)
-        self._multi_part_count(outer_msg, cur_c + 1)
+        self._multi_part_count(outer_msg, part_count + 1)
+
+
+def is_skippable(part):
+    # multipart/* are just containers
+    part_maintype = part.get_content_maintype() or ''
+    if part_maintype.lower() == 'multipart':
+        return True
+    return False
 
 
 # Coverts a raw string into a mime message
diff --git a/cloudinit/util.py b/cloudinit/util.py
index 6872cc31..33da73eb 100644
--- a/cloudinit/util.py
+++ b/cloudinit/util.py
@@ -1285,12 +1285,15 @@ def ensure_file(path, mode=0644):
     write_file(path, content='', omode="ab", mode=mode)
 
 
-def chmod(path, mode):
-    real_mode = None
+def safe_int(possible_int):
     try:
-        real_mode = int(mode)
+        return int(possible_int)
     except (ValueError, TypeError):
-        pass
+        return None
+
+
+def chmod(path, mode):
+    real_mode = safe_int(mode)
     if path and real_mode:
         with SeLinuxGuard(path):
             os.chmod(path, real_mode)