Add the capability to understand and filter on

userdata based on a launch-index (or leave userdata
alone if none is provided by the datasource). This
works by doing the following. 

1. Adjusting the userdata processor to attempt to
   inject a "Launch-Index" header into the messages
   headers (by either taking a header that already exists
   or by looking into the payload to see if it exists
   there). 
2. Adjust the get_userdata ds function to apply a filter
   on the returned userdata (defaulting to false) that
   will now use the datasources get_launch_index value
   to restrict the 'final' message used in consuming 
   user data (the same behavior if not existent). 
3. Further down the line processes that use the 'resultant'
   userdata now will only see the ones for there own launch
   index (ie cloud-config will be restricted automatically
   and so on) and are unaffected (although they can now
   ask the cloud object or the datasource for its launch index
   via the above new ds method.

5 files changed, 90 insertions, 21 deletions

diff --git a/cloudinit/cloud.py b/cloudinit/cloud.py
index 620b3c07..af69a541 100644
--- a/cloudinit/cloud.py
+++ b/cloudinit/cloud.py
@@ -70,12 +70,15 @@ 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()
 
+    def get_launch_index(self):
+        return self.datasource.get_launch_index()
+
     def get_public_ssh_keys(self):
         return self.datasource.get_public_ssh_keys()
 
diff --git a/cloudinit/sources/DataSourceEc2.py b/cloudinit/sources/DataSourceEc2.py
index 556dcafb..3e450e7e 100644
--- a/cloudinit/sources/DataSourceEc2.py
+++ b/cloudinit/sources/DataSourceEc2.py
@@ -77,6 +77,9 @@ class DataSourceEc2(sources.DataSource):
                         self.metadata_address)
             return False
 
+    def get_launch_index(self):
+        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..a1939d58 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
@@ -59,12 +61,53 @@ 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
 
+    def get_launch_index(self):
+        return None
+
+    def _filter_userdata(self, processed_ud):
+        idx = self.get_launch_index()
+        if idx is None:
+            return processed_ud
+        # First do a scan to see if any one with launch-index
+        # headers, if not just skip this....
+        launch_idxs = 0
+        for part in processed_ud.walk():
+            # multipart/* are just containers
+            if part.get_content_maintype() == 'multipart':
+                continue
+            launch_idx_h = part.get('Launch-Index', None)
+            if launch_idx_h is not None:
+                launch_idxs += 1
+        if not launch_idxs:
+            return processed_ud
+        # Reform a new message with those that either have
+        # no launch index or ones that have our launch index or ones
+        # that have some other garbage that we don't know what to do with
+        accumulating_msg = MIMEMultipart()
+        tot_attached = 0
+        for part in processed_ud.walk():
+            # multipart/* are just containers
+            if part.get_content_maintype() == 'multipart':
+                continue
+            try:
+                launch_idx_h = part.get('Launch-Index', None)
+                if launch_idx_h is None or int(launch_idx_h) == int(idx):
+                    accumulating_msg.attach(part)
+                    tot_attached += 1
+            except:
+                # If any int conversion fails (or other error), keep the part
+                accumulating_msg.attach(part)
+                tot_attached += 1
+        accumulating_msg[ud.ATTACHMENT_FIELD] = str(tot_attached)
+        return accumulating_msg
+
     @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..244e9223 100644
--- a/cloudinit/user_data.py
+++ b/cloudinit/user_data.py
@@ -58,10 +58,9 @@ class UserDataProcessor(object):
         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():
@@ -97,11 +96,38 @@ 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
+        try:
+            payload = util.load_yaml(msg.get_payload(decode=True))
+            if payload:
+                payload_idx = payload.get('launch-index')
+        except:
+            pass
+        # Header overrides contents...
+        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:
+                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>'
@@ -204,21 +230,15 @@ 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)
 
 
 # Coverts a raw string into a mime message