From f165adc0d7dd0816def505588dd184ed8c2d805f Mon Sep 17 00:00:00 2001
From: Nikos Skalkotos <skalkoto@grnet.gr>
Date: Tue, 10 Apr 2012 16:38:13 +0300
Subject: [PATCH] Add options for enabling/disabling sysprep tasks

Also merge data_cleanup tasks with syspreps and did a major cleanup
---
 image_creator/__init__.py          |    3 +-
 image_creator/disk.py              |    3 +-
 image_creator/main.py              |   47 +++++-----
 image_creator/os_type/__init__.py  |  171 +++++++++++++++---------------------
 image_creator/os_type/freebsd.py   |    2 +-
 image_creator/os_type/hurd.py      |    2 +-
 image_creator/os_type/linux.py     |   15 ++--
 image_creator/os_type/netbsd.py    |    2 +-
 image_creator/os_type/slackware.py |    5 +-
 image_creator/os_type/ubuntu.py    |    2 +-
 image_creator/os_type/unix.py      |   25 +++---
 image_creator/os_type/windows.py   |    2 +-
 image_creator/util.py              |    7 +-
 13 files changed, 133 insertions(+), 153 deletions(-)

diff --git a/image_creator/__init__.py b/image_creator/__init__.py
index 117c010..afdda63 100644
--- a/image_creator/__init__.py
+++ b/image_creator/__init__.py
@@ -33,6 +33,7 @@
 
 __version__ = '0.1'
 
+
 import image_creator.os_type
 
 
@@ -51,7 +52,5 @@ def get_os_class(distro, osfamily):
     return getattr(module, classname)
 
 
-class FatalError(Exception):
-    pass
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
diff --git a/image_creator/disk.py b/image_creator/disk.py
index 6b69aad..be728a2 100644
--- a/image_creator/disk.py
+++ b/image_creator/disk.py
@@ -32,8 +32,7 @@
 # or implied, of GRNET S.A.
 
 from image_creator.util import get_command
-from image_creator.util import warn, progress, success, output
-from image_creator import FatalError
+from image_creator.util import warn, progress, success, output, FatalError
 
 import stat
 import os
diff --git a/image_creator/main.py b/image_creator/main.py
index 336d795..f2aba56 100755
--- a/image_creator/main.py
+++ b/image_creator/main.py
@@ -35,9 +35,8 @@
 
 from image_creator import get_os_class
 from image_creator import __version__ as version
-from image_creator import FatalError
 from image_creator.disk import Disk
-from image_creator.util import get_command, error, success, output
+from image_creator.util import get_command, error, success, output, FatalError
 from image_creator import util
 import sys
 import os
@@ -65,9 +64,6 @@ def parse_options(input_args):
     parser.add_option("-f", "--force", dest="force", default=False,
         action="store_true", help="overwrite output files if they exist")
 
-    parser.add_option("--no-cleanup", dest="cleanup", default=True,
-        help="don't cleanup sensitive data", action="store_false")
-
     parser.add_option("--no-sysprep", dest="sysprep", default=True,
         help="don't perform system preperation", action="store_false")
 
@@ -78,16 +74,20 @@ def parse_options(input_args):
         default=None, action="callback", callback=check_writable_dir,
         help="dump image to FILE", metavar="FILE")
 
-    parser.add_option("--print-sysprep", dest="print_sysprep", default=False,
-        help="Print the enabled and disable sysprep actions for this image",
-        action="store_true")
+    parser.add_option("--enable-sysprep", dest="enabled_syspreps", default=[],
+        help="Run SYSPREP operation on the input media",
+        action="append", metavar="SYSPREP")
 
-    parser.add_option("--print-data-cleanup", dest="print_data_cleanup",
-        default=False, help="Print the enabled and disable data cleanup "
-        "operations actions for this source", action="store_true")
+    parser.add_option("--disable-sysprep", dest="disabled_syspreps",
+        help="Prevent SYSPREP operation from running on the input media",
+        default=[], action="append", metavar="SYSPREP")
+
+    parser.add_option("--print-sysprep", dest="print_sysprep", default=False,
+        help="Print the enabled and disabled sysprep operations for this "
+        "input media", action="store_true")
 
     parser.add_option("-s", "--silent", dest="silent", default=False,
-        help="silent mode, only output error", action="store_true")
+        help="silent mode, only output errors", action="store_true")
 
     parser.add_option("-u", "--upload", dest="upload", default=False,
         help="upload the image to pithos", action="store_true")
@@ -116,9 +116,9 @@ def image_creator():
         util.silent = True
 
     if options.outfile is None and not options.upload \
-        and not options.print_sysprep and not options.print_data_cleanup:
-        FatalError("At least one of the following: `-o', `-u', "
-        "`--print-sysprep' `--print-data-cleanup' must be set")
+                                            and not options.print_sysprep:
+        raise FatalError("At least one of `-o', `-u' or" \
+                            "`--print-sysprep' must be set")
 
     output('snf-image-creator %s\n' % version)
 
@@ -144,22 +144,21 @@ def image_creator():
 
         output()
 
-        if options.print_sysprep:
-            image_os.print_sysprep()
-            output()
+        for sysprep in options.disabled_syspreps:
+            image_os.disable_sysprep(sysprep)
+
+        for sysprep in options.enabled_syspreps:
+            image_os.enable_sysprep(sysprep)
 
-        if options.print_data_cleanup:
-            image_os.print_data_cleanup()
+        if options.print_sysprep:
+            image_os.print_syspreps()
             output()
 
         if options.outfile is None and not options.upload:
             return 0
 
         if options.sysprep:
-            image_os.sysprep()
-
-        if options.cleanup:
-            image_os.data_cleanup()
+            image_os.do_sysprep()
 
         dev.umount()
 
diff --git a/image_creator/os_type/__init__.py b/image_creator/os_type/__init__.py
index e9a8c3a..1383807 100644
--- a/image_creator/os_type/__init__.py
+++ b/image_creator/os_type/__init__.py
@@ -31,8 +31,9 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.util import output
+from image_creator.util import output, FatalError
 
+import textwrap
 import re
 
 
@@ -43,33 +44,85 @@ def add_prefix(target):
     return wrapper
 
 
-def exclude_task(func):
-    func.excluded = True
-    return func
+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):
         self.root = rootdev
         self.g = ghandler
 
-        self.sysprep_regexp = re.compile('^sysprep_')
-        self.data_cleanup_regexp = re.compile('^data_cleanup_')
+    def _is_sysprep(self, obj):
+        return getattr(obj, 'sysprep', False) and callable(obj)
 
-    def _print_task(self, task):
-        name = task.__name__
+    def list_syspreps(self):
 
-        if self.sysprep_regexp.match(name):
-            name = self.sysprep_regexp.sub("", name)
-        elif self.data_cleanup_regexp.match(name):
-            name = self.data_cleanup_regexp.sub("", name)
-        else:
-            raise FatalError("%s is not a task" % name)
+        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."""
 
-        name = name.replace('_', '-')
+        enabled, disabled = self.list_syspreps()
 
-        output("  %s:\n    %s" % (name, task.__doc__))
+        wrapper = textwrap.TextWrapper()
+        wrapper.subsequent_indent = '\t'
+        wrapper.initial_indent = '\t'
+
+        output("Enabled system preperation operations:")
+        if len(enabled) == 0:
+            output("(none)")
+        else:
+            for sysprep in enabled:
+                name = sysprep.__name__.replace('_', '-')
+                descr = wrapper.fill(sysprep.__doc__)
+                output('    %s:\n%s\n' % (name, descr))
+
+        output("Disabled system preperation operations:")
+        if len(disabled) == 0:
+            output("(none)")
+        else:
+            for sysprep in disabled:
+                name = sysprep.__name__.replace('_', '-')
+                descr = wrapper.fill(sysprep.__doc__)
+                output('    %s:\n%s\n' % (name, descr))
 
     @add_prefix
     def ls(self, directory):
@@ -133,45 +186,7 @@ class OSBase(object):
 
         return meta
 
-    def list_sysprep(self):
-        """List all sysprep actions"""
-
-        is_sysprep = lambda x: x.startswith('sysprep_') and \
-                                                    callable(getattr(self, x))
-        tasks = [getattr(self, x) for x in dir(self) if is_sysprep(x)]
-
-        included = [t for t in tasks if not getattr(t, "excluded", False)]
-        excluded = [t for t in tasks if getattr(t, "excluded", False)]
-
-        return included, excluded
-
-    def list_data_cleanup(self):
-        """List all data_cleanup actions"""
-
-        is_cleanup = lambda x: x.startswith('data_cleanup_') and \
-                                                    callable(getattr(self, x))
-        tasks = [getattr(self, x) for x in dir(self) if is_cleanup(x)]
-
-        included = [t for t in tasks if not getattr(t, "excluded", False)]
-        excluded = [t for t in tasks if getattr(t, "excluded", False)]
-
-        return included, excluded
-
-    def data_cleanup(self):
-        """Cleanup sensitive data out of the OS image."""
-
-        output('Cleaning up sensitive data out of the OS image:')
-
-        tasks, _ = self.list_data_cleanup()
-        size = len(tasks)
-        cnt = 0
-        for task in tasks:
-            cnt += 1
-            output(('(%d/%d)' % (cnt, size)).ljust(7), False)
-            task()
-        output()
-
-    def sysprep(self):
+    def do_sysprep(self):
         """Prepere system for image creation."""
 
         output('Preparing system for image creation:')
@@ -185,50 +200,4 @@ class OSBase(object):
             task()
         output()
 
-    def print_task(self, task):
-        name = task.__name__
-
-        if self.sysprep_regexp.match(name):
-            name = self.sysprep_regexp.sub("", name)
-        elif self.data_cleanup_regexp.match(name):
-            name = self.data_cleanup_regexp.sub("", name)
-        else:
-            raise FatalError("%s is not a task" % name)
-
-        name = name.replace('_', '-')
-
-        output("  %s:\n    %s" % (name, task.__doc__))
-
-    def print_data_cleanup(self):
-        included, excluded = self.list_data_cleanup()
-
-        output("Included data cleanup operations:")
-        if len(included) == 0:
-            ouput("(none)")
-        else:
-            for task in included:
-                self._print_task(task)
-        output("Ommited data cleanup operations:")
-        if len(excluded) == 0:
-            ouput("(none)")
-        else:
-            for task in excluded:
-                self._print_task(task)
-
-    def print_sysprep(self):
-        included, excluded = self.list_sysprep()
-
-        output("Included sysprep operations:")
-        if len(included) == 0:
-            ouput("(none)")
-        else:
-            for task in included:
-                self._print_task(task)
-        output("Ommited sysprep operations:")
-        if len(excluded) == 0:
-            output("(none)")
-        else:
-            for task in excluded:
-                self._print_task(task)
-
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
diff --git a/image_creator/os_type/freebsd.py b/image_creator/os_type/freebsd.py
index c2e1a5c..b5e7491 100644
--- a/image_creator/os_type/freebsd.py
+++ b/image_creator/os_type/freebsd.py
@@ -31,7 +31,7 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.os_type.unix import Unix, exclude_task
+from image_creator.os_type.unix import Unix, sysprep
 
 
 class Freebsd(Unix):
diff --git a/image_creator/os_type/hurd.py b/image_creator/os_type/hurd.py
index bba3e65..34a605d 100644
--- a/image_creator/os_type/hurd.py
+++ b/image_creator/os_type/hurd.py
@@ -31,7 +31,7 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.os_type.unix import Unix, exclude_task
+from image_creator.os_type.unix import Unix, sysprep
 
 
 class Hard(Unix):
diff --git a/image_creator/os_type/linux.py b/image_creator/os_type/linux.py
index 31e651d..de82e7d 100644
--- a/image_creator/os_type/linux.py
+++ b/image_creator/os_type/linux.py
@@ -31,7 +31,7 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.os_type.unix import Unix, exclude_task
+from image_creator.os_type.unix import Unix, sysprep
 from image_creator.util import warn, output
 
 import re
@@ -56,7 +56,8 @@ class Linux(Unix):
                 self._uuid[dev] = attr[1]
                 return attr[1]
 
-    def sysprep_fix_acpid(self, print_header=True):
+    @sysprep()
+    def fix_acpid(self, print_header=True):
         """Replace acpid powerdown action scripts to immediately shutdown the
         system without checking if a GUI is running.
         """
@@ -110,7 +111,8 @@ class Linux(Unix):
                     "event occures" % action)
                 return
 
-    def sysprep_persistent_net_rules(self, print_header=True):
+    @sysprep()
+    def persistent_net_rules(self, print_header=True):
         """Remove udev rules that will keep network interface names persistent
         after hardware changes and reboots. Those rules will be created again
         the next time the image runs.
@@ -123,9 +125,10 @@ class Linux(Unix):
         if self.g.is_file(rule_file):
             self.g.rm(rule_file)
 
-    def sysprep_persistent_devs(self, print_header=True):
-        """Scan fstab and grub configuration files and replace all
-        non-persistent device appearences with UUIDs.
+    @sysprep()
+    def persistent_devs(self, print_header=True):
+        """Scan fstab & grub configuration files and replace all non-persistent
+        device appearences with UUIDs.
         """
 
         if print_header:
diff --git a/image_creator/os_type/netbsd.py b/image_creator/os_type/netbsd.py
index 9973b28..57e50e4 100644
--- a/image_creator/os_type/netbsd.py
+++ b/image_creator/os_type/netbsd.py
@@ -31,7 +31,7 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.os_type.unix import Unix, exclude_task
+from image_creator.os_type.unix import Unix, sysprep
 
 
 class Netbsd(Unix):
diff --git a/image_creator/os_type/slackware.py b/image_creator/os_type/slackware.py
index ac7ded0..c2cbb6c 100644
--- a/image_creator/os_type/slackware.py
+++ b/image_creator/os_type/slackware.py
@@ -31,11 +31,12 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.os_type.linux import Linux, exclude_task
+from image_creator.os_type.linux import Linux, sysprep
 
 
 class Slackware(Linux):
-    def data_cleanup_log(self):
+    @sysprep()
+    def cleanup_log(self):
         # In slackware the metadata about installed packages are
         # stored in /var/log/packages. Clearing all /var/log files
         # will destroy the package management system.
diff --git a/image_creator/os_type/ubuntu.py b/image_creator/os_type/ubuntu.py
index 4dec1e8..f5a82ea 100644
--- a/image_creator/os_type/ubuntu.py
+++ b/image_creator/os_type/ubuntu.py
@@ -31,7 +31,7 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.os_type.linux import Linux, exclude_task
+from image_creator.os_type.linux import Linux, sysprep
 
 
 class Ubuntu(Linux):
diff --git a/image_creator/os_type/unix.py b/image_creator/os_type/unix.py
index 655fd49..3829c3e 100644
--- a/image_creator/os_type/unix.py
+++ b/image_creator/os_type/unix.py
@@ -34,7 +34,7 @@
 import re
 import sys
 
-from image_creator.os_type import OSBase, exclude_task
+from image_creator.os_type import OSBase, sysprep
 from image_creator.util import warn, output
 
 
@@ -70,8 +70,8 @@ class Unix(OSBase):
 
         return users
 
-    @exclude_task
-    def data_cleanup_user_accounts(self, print_header=True):
+    @sysprep(enabled=False)
+    def remove_user_accounts(self, print_header=True):
         """Remove all user account with id more than 1000"""
 
         if print_header:
@@ -113,7 +113,8 @@ class Unix(OSBase):
             if self.g.is_dir(home) and home.startswith('/home/'):
                 self.g.rm_rf(home)
 
-    def data_cleanup_passwords(self, print_header=True):
+    @sysprep()
+    def cleanup_passwords(self, print_header=True):
         """Remove all passwords and lock all user accounts"""
 
         if print_header:
@@ -130,7 +131,8 @@ class Unix(OSBase):
 
         self.g.write('/etc/shadow', "\n".join(shadow) + '\n')
 
-    def data_cleanup_cache(self, print_header=True):
+    @sysprep()
+    def cleanup_cache(self, print_header=True):
         """Remove all regular files under /var/cache"""
 
         if print_header:
@@ -138,7 +140,8 @@ class Unix(OSBase):
 
         self.foreach_file('/var/cache', self.g.rm, ftype='r')
 
-    def data_cleanup_tmp(self, print_header=True):
+    @sysprep()
+    def cleanup_tmp(self, print_header=True):
         """Remove all files under /tmp and /var/tmp"""
 
         if print_header:
@@ -147,7 +150,8 @@ class Unix(OSBase):
         self.foreach_file('/tmp', self.g.rm_rf, maxdepth=1)
         self.foreach_file('/var/tmp', self.g.rm_rf, maxdepth=1)
 
-    def data_cleanup_log(self, print_header=True):
+    @sysprep()
+    def cleanup_log(self, print_header=True):
         """Empty all files under /var/log"""
 
         if print_header:
@@ -155,8 +159,8 @@ class Unix(OSBase):
 
         self.foreach_file('/var/log', self.g.truncate, ftype='r')
 
-    @exclude_task
-    def data_cleanup_mail(self, print_header=True):
+    @sysprep(enabled=False)
+    def cleanup_mail(self, print_header=True):
         """Remove all files under /var/mail and /var/spool/mail"""
 
         if print_header:
@@ -165,7 +169,8 @@ class Unix(OSBase):
         self.foreach_file('/var/spool/mail', self.g.rm_rf, maxdepth=1)
         self.foreach_file('/var/mail', self.g.rm_rf, maxdepth=1)
 
-    def data_cleanup_userdata(self, print_header=True):
+    @sysprep()
+    def cleanup_userdata(self, print_header=True):
         """Delete sensitive userdata"""
 
         homedirs = ['/root'] + self.ls('/home/')
diff --git a/image_creator/os_type/windows.py b/image_creator/os_type/windows.py
index 036e33b..405a7d9 100644
--- a/image_creator/os_type/windows.py
+++ b/image_creator/os_type/windows.py
@@ -31,7 +31,7 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.os_type import OSBase, exclude_task
+from image_creator.os_type import OSBase, sysprep
 
 
 class Windows(OSBase):
diff --git a/image_creator/util.py b/image_creator/util.py
index 588d2b0..1b4994a 100644
--- a/image_creator/util.py
+++ b/image_creator/util.py
@@ -35,6 +35,11 @@ import sys
 import pbs
 from clint.textui import colored, progress as uiprogress
 
+
+class FatalError(Exception):
+    pass
+
+
 silent = False
 
 
@@ -54,7 +59,7 @@ def get_command(command):
 
 def error(msg, new_line=True):
     nl = "\n" if new_line else ''
-    sys.stderr.write('Error: %s' % msg + nl)
+    sys.stderr.write(colored.red('Error: %s' % msg) + nl)
 
 
 def warn(msg, new_line=True):
-- 
1.7.10.4