For the userdata 'post-filtering' add in a new folder that

can contain filters that serve this purpose only and add in
the initial launch-index filter and replace the code in 
the datasource class that previously did this.

4 files changed, 100 insertions, 46 deletions

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..b9e108c4
--- /dev/null
+++ b/cloudinit/filters/launch_index.py
@@ -0,0 +1,63 @@
+# 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):
+        if ud.is_skippable(message):
+            return False
+        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 apply(self, base_message):
+        if not base_message.is_multipart() or self.wanted_idx is None:
+            return base_message
+        prev_msgs = base_message.get_payload(decode=False)
+        to_attach = []
+        for sub_msg in base_message.walk():
+            if self._select(sub_msg):
+                to_attach.append(sub_msg)
+        if len(prev_msgs) != len(to_attach):
+            LOG.debug(("Discarding %s multipart messages "
+                       "which do not match launch index %s"),
+                      (len(prev_msgs) - len(to_attach)), self.wanted_idx)
+        filtered_msg = copy.deepcopy(base_message)
+        filtered_msg.set_payload(to_attach)
+        filtered_msg[ud.ATTACHMENT_FIELD] = str(len(to_attach))
+        return filtered_msg
diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index 74944e38..3f611d44 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -29,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'
@@ -77,48 +79,13 @@ class DataSource(object):
         return None
 
     def _filter_userdata(self, processed_ud):
-        if not processed_ud:
-            return processed_ud
-        idx = self.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():
-            if ud.is_skippable(part):
-                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
-        tot_processed = 0
-        for part in processed_ud.walk():
-            if ud.is_skippable(part):
-                continue
-            try:
-                tot_processed += 1
-                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
-                else:
-                    LOG.debug(("Discarding multipart message %s, "
-                               "launch-index provided destined for %s "
-                               "and not %s"),
-                               tot_processed, launch_idx_h, idx)
-            except (TypeError, ValueError):
-                # If any int conversion fails keep the message
-                accumulating_msg.attach(part)
-                tot_attached += 1
-        accumulating_msg[ud.ATTACHMENT_FIELD] = str(tot_attached)
-        return accumulating_msg
+        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):
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)