Make the top level __init__ similar to the other top level __inits__ where just util functions and base classes are defined there.

5 files changed, 235 insertions, 30 deletions

diff --git a/cloudinit/handlers/__init__.py b/cloudinit/handlers/__init__.py
index 09163abb..0ff80a9e 100644
--- a/cloudinit/handlers/__init__.py
+++ b/cloudinit/handlers/__init__.py
@@ -20,3 +20,199 @@
 #    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 abc
+import os
+
+from cloudinit.settings import (PER_ALWAYS, PER_INSTANCE, FREQUENCIES)
+
+from cloudinit import importer
+from cloudinit import log as logging
+from cloudinit import url_helper
+from cloudinit import util
+
+LOG = logging.getLogger(__name__)
+
+# Used as the content type when a message is not multipart
+# and it doesn't contain its own content-type
+NOT_MULTIPART_TYPE = "text/x-not-multipart"
+
+# When none is assigned this gets used
+OCTET_TYPE = 'application/octet-stream'
+
+# Special content types that signal the start and end of processing
+CONTENT_END = "__end__"
+CONTENT_START = "__begin__"
+CONTENT_SIGNALS = [CONTENT_START, CONTENT_END]
+
+# Used when a part-handler type is encountered
+# to allow for registration of new types.
+PART_CONTENT_TYPES = ["text/part-handler"]
+PART_HANDLER_FN_TMPL = 'part-handler-%03d'
+
+# For parts without filenames
+PART_FN_TPL = 'part-%03d'
+
+# Different file beginnings to there content type
+INCLUSION_TYPES_MAP = {
+    '#include': 'text/x-include-url',
+    '#include-once': 'text/x-include-once-url',
+    '#!': 'text/x-shellscript',
+    '#cloud-config': 'text/cloud-config',
+    '#upstart-job': 'text/upstart-job',
+    '#part-handler': 'text/part-handler',
+    '#cloud-boothook': 'text/cloud-boothook',
+    '#cloud-config-archive': 'text/cloud-config-archive',
+}
+
+# Sorted longest first
+INCLUSION_SRCH = sorted(list(INCLUSION_TYPES_MAP.keys()),
+                        key=(lambda e: 0 - len(e)))
+
+
+class Handler(object):
+
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, frequency, version=2):
+        self.handler_version = version
+        self.frequency = frequency
+
+    def __repr__(self):
+        return "%s: [%s]" % (util.obj_name(self), self.list_types())
+
+    @abc.abstractmethod
+    def list_types(self):
+        raise NotImplementedError()
+
+    def handle_part(self, data, ctype, filename, payload, frequency):
+        return self._handle_part(data, ctype, filename, payload, frequency)
+
+    @abc.abstractmethod
+    def _handle_part(self, data, ctype, filename, payload, frequency):
+        raise NotImplementedError()
+
+
+def run_part(mod, data, ctype, filename, payload, frequency):
+    mod_freq = mod.frequency
+    if not (mod_freq == PER_ALWAYS or
+            (frequency == PER_INSTANCE and mod_freq == PER_INSTANCE)):
+        return
+    mod_ver = mod.handler_version
+    try:
+        if mod_ver == 1:
+            mod.handle_part(data, ctype, filename, payload)
+        else:
+            mod.handle_part(data, ctype, filename, payload, frequency)
+    except:
+        util.logexc(LOG, ("Failed calling mod %s (%s, %s, %s)"
+                         " with frequency %s"), 
+                    mod, ctype, filename,
+                    mod_ver, frequency)
+
+
+def call_begin(mod, data, frequency):
+    run_part(mod, data, CONTENT_START, None, None, frequency)
+
+
+def call_end(mod, data, frequency):
+    run_part(mod, data, CONTENT_END, None, None, frequency)
+
+
+def walker_handle_handler(pdata, _ctype, _filename, payload):
+    curcount = pdata['handlercount']
+    modname = PART_HANDLER_FN_TMPL % (curcount)
+    frequency = pdata['frequency']
+    modfname = os.path.join(pdata['handlerdir'], "%s" % (modname))
+    if not modfname.endswith(".py"):
+        modfname = "%s.py" % (modfname)
+    # TODO: Check if path exists??
+    util.write_file(modfname, payload, 0600)
+    handlers = pdata['handlers']
+    try:
+        mod = fixup_handler(importer.import_module(modname))
+        handlers.register(mod)
+        call_begin(mod, pdata['data'], frequency)
+        pdata['handlercount'] = curcount + 1
+    except:
+        util.logexc(LOG, "Failed at registered python file: %s", modfname)
+
+
+def _extract_first_or_bytes(blob, size):
+    # Extract the first line upto X bytes or X bytes from more than the
+    # first line if the first line does not contain enough bytes
+    first_line = blob.split("\n", 1)[0]
+    if len(first_line) >= size:
+        start = first_line[:size]
+    else:
+        start = blob[0:size]
+    return start
+
+
+def walker_callback(pdata, ctype, filename, payload):
+    if ctype in PART_CONTENT_TYPES:
+        walker_handle_handler(pdata, ctype, filename, payload)
+        return
+    handlers = pdata['handlers']
+    if ctype not in handlers:
+        # Extract the first line or 24 bytes for displaying in the log
+        start = _extract_first_or_bytes(payload, 24)
+        details = "'%s...'" % (start.encode("string-escape"))
+        if ctype == NOT_MULTIPART_TYPE:
+            LOG.warning("Unhandled non-multipart (%s) userdata: %s",
+                        ctype, details)
+        else:
+            LOG.warning("Unhandled unknown content-type (%s) userdata: %s",
+                        ctype, details)
+    else:
+        run_part(handlers[ctype], pdata['data'], ctype, filename,
+                 payload, pdata['frequency'])
+
+
+# Callback is a function that will be called with 
+# (data, content_type, filename, payload)
+def walk(msg, callback, data):
+    partnum = 0
+    for part in msg.walk():
+        # multipart/* are just containers
+        if part.get_content_maintype() == 'multipart':
+            continue
+
+        ctype = part.get_content_type()
+        if ctype is None:
+            ctype = OCTET_TYPE
+
+        filename = part.get_filename()
+        if not filename:
+            filename = PART_FN_TPL % (partnum)
+
+        callback(data, ctype, filename, part.get_payload(decode=True))
+        partnum = partnum + 1
+
+
+def fixup_handler(mod, def_freq=PER_INSTANCE):
+    if not hasattr(mod, "handler_version"):
+        setattr(mod, "handler_version", 1)
+    if not hasattr(mod, 'list_types'):
+        def empty_types():
+            return []
+        setattr(mod, 'list_types', empty_types)
+    if not hasattr(mod, 'frequency'):
+        setattr(mod, 'frequency', def_freq)
+    else:
+        freq = mod.frequency
+        if freq and freq not in FREQUENCIES:
+            LOG.warn("Handler %s has an unknown frequency %s", mod, freq)
+    if not hasattr(mod, 'handle_part'):
+        def empty_handler(_data, _ctype, _filename, _payload):
+            pass
+        setattr(mod, 'handle_part', empty_handler)
+    return mod
+
+
+def type_from_starts_with(payload, default=None):
+    for text in INCLUSION_SRCH:
+        if payload.startswith(text):
+            return INCLUSION_TYPES_MAP[text]
+    return default
+
+
diff --git a/cloudinit/handlers/boot_hook.py b/cloudinit/handlers/boot_hook.py
index b3aab366..10f60b8d 100644
--- a/cloudinit/handlers/boot_hook.py
+++ b/cloudinit/handlers/boot_hook.py
@@ -22,8 +22,8 @@
 
 import os
 
+from cloudinit import handlers
 from cloudinit import log as logging
-from cloudinit import user_data as ud
 from cloudinit import util
 
 from cloudinit.settings import (PER_ALWAYS)
@@ -31,31 +31,34 @@ from cloudinit.settings import (PER_ALWAYS)
 LOG = logging.getLogger(__name__)
 
 
-class BootHookPartHandler(ud.PartHandler):
+class BootHookPartHandler(handlers.Handler):
     def __init__(self, paths, instance_id, **_kwargs):
-        ud.PartHandler.__init__(self, PER_ALWAYS)
+        handlers.Handler.__init__(self, PER_ALWAYS)
         self.boothook_dir = paths.get_ipath("boothooks")
         self.instance_id = instance_id
 
     def list_types(self):
         return [
-            ud.type_from_starts_with("#cloud-boothook"),
+            handlers.type_from_starts_with("#cloud-boothook"),
         ]
 
-    def _handle_part(self, _data, ctype, filename, payload, _frequency):
-        if ctype in ud.CONTENT_SIGNALS:
-            return
-
+    def _write_part(self, payload, filename):
         filename = util.clean_filename(filename)
         payload = util.dos2unix(payload)
         prefix = "#cloud-boothook"
         start = 0
         if payload.startswith(prefix):
             start = len(prefix) + 1
-
         filepath = os.path.join(self.boothook_dir, filename)
         contents = payload[start:]
         util.write_file(filepath, contents, 0700)
+        return filepath
+
+    def _handle_part(self, _data, ctype, filename, payload, _frequency):
+        if ctype in handlers.CONTENT_SIGNALS:
+            return
+
+        filepath = self._write_part(payload, filename)
         try:
             env = os.environ.copy()
             if self.instance_id:
diff --git a/cloudinit/handlers/cloud_config.py b/cloudinit/handlers/cloud_config.py
index 12d1bd96..f6d95244 100644
--- a/cloudinit/handlers/cloud_config.py
+++ b/cloudinit/handlers/cloud_config.py
@@ -20,8 +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 cloudinit import handlers
 from cloudinit import log as logging
-from cloudinit import user_data as ud
 from cloudinit import util
 
 from cloudinit.settings import (PER_ALWAYS)
@@ -29,28 +29,34 @@ from cloudinit.settings import (PER_ALWAYS)
 LOG = logging.getLogger(__name__)
 
 
-class CloudConfigPartHandler(ud.PartHandler):
+class CloudConfigPartHandler(handlers.Handler):
     def __init__(self, paths, **_kwargs):
-        ud.PartHandler.__init__(self, PER_ALWAYS)
+        handlers.Handler.__init__(self, PER_ALWAYS)
         self.cloud_buf = []
         self.cloud_fn = paths.get_ipath("cloud_config")
 
     def list_types(self):
         return [
-            ud.type_from_starts_with("#cloud-config"),
+            handlers.type_from_starts_with("#cloud-config"),
         ]
 
+    def _write_cloud_config(self, buf):
+        if not self.cloud_fn:
+            return
+        lines = [str(b) for b in buf]
+        payload = "\n".join(lines)
+        util.write_file(self.cloud_fn, payload, 0600)
+
     def _handle_part(self, _data, ctype, filename, payload, _frequency):
-        if ctype == ud.CONTENT_START:
+        if ctype == handlers.CONTENT_START:
             self.cloud_buf = []
             return
-
-        if ctype == ud.CONTENT_END:
-            payload = "\n".join(self.cloud_buf)
-            util.write_file(self.cloud_fn, payload, 0600)
+        if ctype == handlers.CONTENT_END:
+            self._write_cloud_config(self.cloud_buf)
             self.cloud_buf = []
             return
 
         filename = util.clean_filename(filename)
-        entry = "\n".join(["#%s" % (filename), str(payload)])
-        self.cloud_buf.append(entry)
+        if not filename:
+            filename = '??'
+        self.cloud_buf.extend(["#%s" % (filename), str(payload)])
diff --git a/cloudinit/handlers/shell_script.py b/cloudinit/handlers/shell_script.py
index f6e2ef16..a9d8e544 100644
--- a/cloudinit/handlers/shell_script.py
+++ b/cloudinit/handlers/shell_script.py
@@ -22,8 +22,8 @@
 
 import os
 
+from cloudinit import handlers
 from cloudinit import log as logging
-from cloudinit import user_data as ud
 from cloudinit import util
 
 from cloudinit.settings import (PER_ALWAYS)
@@ -31,18 +31,18 @@ from cloudinit.settings import (PER_ALWAYS)
 LOG = logging.getLogger(__name__)
 
 
-class ShellScriptPartHandler(ud.PartHandler):
+class ShellScriptPartHandler(handlers.Handler):
     def __init__(self, paths, **_kwargs):
-        ud.PartHandler.__init__(self, PER_ALWAYS)
+        handlers.Handler.__init__(self, PER_ALWAYS)
         self.script_dir = paths.get_ipath_cur('scripts')
 
     def list_types(self):
         return [
-            ud.type_from_starts_with("#!"),
+            handlers.type_from_starts_with("#!"),
         ]
 
     def _handle_part(self, _data, ctype, filename, payload, _frequency):
-        if ctype in ud.CONTENT_SIGNALS:
+        if ctype in handlers.CONTENT_SIGNALS:
             # TODO: maybe delete existing things here
             return
 
diff --git a/cloudinit/handlers/upstart_job.py b/cloudinit/handlers/upstart_job.py
index 059a4851..411a5d68 100644
--- a/cloudinit/handlers/upstart_job.py
+++ b/cloudinit/handlers/upstart_job.py
@@ -23,8 +23,8 @@
 
 import os
 
+from cloudinit import handlers
 from cloudinit import log as logging
-from cloudinit import user_data as ud
 from cloudinit import util
 
 from cloudinit.settings import (PER_INSTANCE)
@@ -32,18 +32,18 @@ from cloudinit.settings import (PER_INSTANCE)
 LOG = logging.getLogger(__name__)
 
 
-class UpstartJobPartHandler(ud.PartHandler):
+class UpstartJobPartHandler(handlers.Handler):
     def __init__(self, paths, **_kwargs):
-        ud.PartHandler.__init__(self, PER_INSTANCE)
+        handlers.Handler.__init__(self, PER_INSTANCE)
         self.upstart_dir = paths.upstart_conf_d
 
     def list_types(self):
         return [
-            ud.type_from_starts_with("#upstart-job"),
+            handlers.type_from_starts_with("#upstart-job"),
         ]
 
     def _handle_part(self, _data, ctype, filename, payload, frequency):
-        if ctype in ud.CONTENT_SIGNALS:
+        if ctype in handlers.CONTENT_SIGNALS:
             return
 
         if not self.upstart_dir: