From 6a803e2fba17ab74973038c1a93654c3127c223f Mon Sep 17 00:00:00 2001
From: Joshua Harlow <harlowja@yahoo-inc.com>
Date: Sun, 26 Aug 2012 15:04:06 -0700
Subject: 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.
---
 cloudinit/sources/__init__.py | 49 ++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 46 insertions(+), 3 deletions(-)

(limited to 'cloudinit/sources/__init__.py')

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
-- 
cgit v1.2.3