Fix a bug introduced in 4e58b51b

[snf-image-creator] / image_creator / os_type / __init__.py

diff --git a/image_creator/os_type/__init__.py b/image_creator/os_type/__init__.py
index f3fd173..6fe1336 100644 (file)
--- a/image_creator/os_type/__init__.py
+++ b/image_creator/os_type/__init__.py
@@ -31,8 +31,25 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
+from image_creator.util import FatalError
+
+import textwrap
 import re
-from clint.textui import indent, puts
+
+
+def os_cls(distro, osfamily):
+    module = None
+    classname = None
+    try:
+        module = __import__("image_creator.os_type.%s"
+            % distro, fromlist=['image_creator.os_type'])
+        classname = distro.capitalize()
+    except ImportError:
+        module = __import__("image_creator.os_type.%s"
+            % osfamily, fromlist=['image_creator.os_type'])
+        classname = osfamily.capitalize()
+
+    return getattr(module, classname)
 
 
 def add_prefix(target):
@@ -42,11 +59,94 @@ def add_prefix(target):
     return wrapper
 
 
+def sysprep(enabled=True):
+    def wrapper(func):
+        func.sysprep = True
+        func.enabled = enabled
+        return func
+    return wrapper
+
+
 class OSBase(object):
     """Basic operating system class"""
-    def __init__(self, rootdev, ghandler):
+
+    def __init__(self, rootdev, ghandler, output):
         self.root = rootdev
         self.g = ghandler
+        self.out = output
+
+        # Collect metadata about the OS
+        self.meta = {}
+        self.meta['ROOT_PARTITION'] = "%d" % self.g.part_to_partnum(self.root)
+        self.meta['OSFAMILY'] = self.g.inspect_get_type(self.root)
+        self.meta['OS'] = self.g.inspect_get_distro(self.root)
+        self.meta['DESCRIPTION'] = self.g.inspect_get_product_name(self.root)
+
+    def _is_sysprep(self, obj):
+        return getattr(obj, 'sysprep', False) and callable(obj)
+
+    def list_syspreps(self):
+
+        objs = [getattr(self, name) for name in dir(self) \
+            if not name.startswith('_')]
+
+        enabled = [x for x in objs if self._is_sysprep(x) and x.enabled]
+        disabled = [x for x in objs if self._is_sysprep(x) and not x.enabled]
+
+        return enabled, disabled
+
+    def _sysprep_change_status(self, name, status):
+
+        error_msg = "Syprep operation %s does not exist for %s" % \
+                (name, self.__class__.__name__)
+
+        method_name = name.replace('-', '_')
+        method = None
+        try:
+            method = getattr(self, method_name)
+        except AttributeError:
+            raise FatalError(error_msg)
+
+        if not self._is_sysprep(method):
+            raise FatalError(error_msg)
+
+        setattr(method.im_func, 'enabled', status)
+
+    def enable_sysprep(self, name):
+        """Enable a system preperation operation"""
+        self._sysprep_change_status(name, True)
+
+    def disable_sysprep(self, name):
+        """Disable a system preperation operation"""
+        self._sysprep_change_status(name, False)
+
+    def print_syspreps(self):
+        """Print enabled and disabled system preperation operations."""
+
+        enabled, disabled = self.list_syspreps()
+
+        wrapper = textwrap.TextWrapper()
+        wrapper.subsequent_indent = '\t'
+        wrapper.initial_indent = '\t'
+        wrapper.width = 72
+
+        self.out.output("Enabled system preperation operations:")
+        if len(enabled) == 0:
+            self.out.output("(none)")
+        else:
+            for sysprep in enabled:
+                name = sysprep.__name__.replace('_', '-')
+                descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
+                self.out.output('    %s:\n%s\n' % (name, descr))
+
+        self.out.output("Disabled system preperation operations:")
+        if len(disabled) == 0:
+            self.out.output("(none)")
+        else:
+            for sysprep in disabled:
+                name = sysprep.__name__.replace('_', '-')
+                descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
+                self.out.output('    %s:\n%s\n' % (name, descr))
 
     @add_prefix
     def ls(self, directory):
@@ -100,36 +200,18 @@ class OSBase(object):
             if has_ftype(f, ftype):
                 action(full_path)
 
-    def get_metadata(self):
-        """Returns some descriptive metadata about the OS."""
-        meta = {}
-        meta['ROOT_PARTITION'] = "%d" % self.g.part_to_partnum(self.root)
-        meta['OSFAMILY'] = self.g.inspect_get_type(self.root)
-        meta['OS'] = self.g.inspect_get_distro(self.root)
-        meta['description'] = self.g.inspect_get_product_name(self.root)
-
-        return meta
-
-    def data_cleanup(self):
-        """Cleanup sensitive data out of the OS image."""
-
-        puts('Cleaning up sensitive data out of the OS image:')
-        with indent(4):
-            for name in dir(self):
-                attr = getattr(self, name)
-                if name.startswith('data_cleanup_') and callable(attr):
-                    attr()
-        puts()
-
-    def sysprep(self):
+    def do_sysprep(self):
         """Prepere system for image creation."""
 
-        puts('Preparing system for image creation:')
-        with indent(4):
-            for name in dir(self):
-                attr = getattr(self, name)
-                if name.startswith('sysprep_') and callable(attr):
-                    attr()
-        puts()
+        self.out.output('Preparing system for image creation:')
+
+        tasks, _ = self.list_syspreps()
+        size = len(tasks)
+        cnt = 0
+        for task in tasks:
+            cnt += 1
+            self.out.output(('(%d/%d)' % (cnt, size)).ljust(7), False)
+            task()
+        self.out.output()
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :