Reimplement the sysprep selection mechanism

Change the sysprep enabling and disabling mechanism to always work
correct. The previous implementation was buggy and would erroneously
"remember" the executed syspreps on a reset.

image_creator/dialog_menu.py

@@ -826,7 +826,7 @@ def sysprep(session):
             sysprep_help += "%s\n" % display_name
             sysprep_help += "%s\n" % ('-' * len(display_name))
             sysprep_help += "%s\n\n" % wrapper.fill(" ".join(descr.split()))
-            enabled = 1 if task.enabled else 0
+            enabled = 1 if image.os.sysprep_enabled(task) else 0
             choices.append((str(index + 1), display_name, enabled))
             index += 1
 
@@ -854,7 +854,8 @@ def sysprep(session):
                 else:
                     image.os.disable_sysprep(syspreps[i])
 
-            if len([s for s in image.os.list_syspreps() if s.enabled]) == 0:
+            if len([s for s in image.os.list_syspreps()
+                    if image.os.sysprep_enabled(s)]) == 0:
                 d.msgbox("No system preparation task is selected!",
                          title="System Preparation", width=SMALL_WIDTH)
                 continue

image_creator/os_type/__init__.py

@@ -64,12 +64,14 @@ def add_prefix(target):
 def sysprep(message, enabled=True, **kwargs):
     """Decorator for system preparation tasks"""
     def wrapper(method):
-        method.sysprep = True
-        method.enabled = enabled
-        method.executed = False
+        assert method.__name__.startswith('_'), \
+            "Invalid sysprep name:` %s'. Should start with _" % method.__name__
+
+        method._sysprep = True
+        method._sysprep_enabled = enabled
 
         for key, val in kwargs.items():
-            setattr(method, key, val)
+            setattr(method, "_sysprep_%s" % key, val)
 
         @wraps(method)
         def inner(self, print_message=True):
@@ -227,6 +229,14 @@ class OSBase(object):
         except RuntimeError:
             self._scrub_support = False
 
+        # Create a list of available syspreps
+        self._sysprep_tasks = {}
+        for name in dir(self):
+            obj = getattr(self, name)
+            if not hasattr(obj, '_sysprep'):
+                continue
+            self._sysprep_tasks[name] = obj._sysprep_enabled
+
         self._cleanup_jobs = {}
 
     def _add_cleanup(self, namespace, job, *args):
@@ -274,51 +284,61 @@ class OSBase(object):
 
     def list_syspreps(self):
         """Returns a list of sysprep objects"""
-        objs = [getattr(self, name) for name in dir(self)
-                if not name.startswith('_')]
-
-        return [x for x in objs if self._is_sysprep(x) and x.executed is False]
+        return [getattr(self, name) for name in self._sysprep_tasks]
 
     def sysprep_info(self, obj):
         """Returns information about a sysprep object"""
-        assert self._is_sysprep(obj), "Object is not a sysprep"
+        assert hasattr(obj, '_sysprep'), "Object is not a sysprep"
 
         SysprepInfo = namedtuple("SysprepInfo", "name description")
 
-        return SysprepInfo(obj.__name__.replace('_', '-'),
-                           textwrap.dedent(obj.__doc__))
+        name = obj.__name__.replace('_', '-')[1:]
+        description = textwrap.dedent(obj.__doc__)
+
+        return SysprepInfo(name, description)
 
     def get_sysprep_by_name(self, name):
         """Returns the sysprep object with the given name"""
         error_msg = "Syprep operation %s does not exist for %s" % \
                     (name, self.__class__.__name__)
 
-        method_name = name.replace('-', '_')
-        method = None
-        try:
+        method_name = '_' + name.replace('-', '_')
+
+        if hasattr(self, method_name):
             method = getattr(self, method_name)
-        except AttributeError:
-            raise FatalError(error_msg)
 
-        if not self._is_sysprep(method):
-            raise FatalError(error_msg)
+            if hasattr(method, '_sysprep'):
+                return method
 
-        return method
+        raise FatalError(error_msg)
 
     def enable_sysprep(self, obj):
         """Enable a system preparation operation"""
-        setattr(obj.im_func, 'enabled', True)
+        assert hasattr(obj, '_sysprep'), "Object is not a sysprep"
+        assert obj.__name__ in self._sysprep_tasks, "Sysprep already executed"
+
+        self._sysprep_tasks[obj.__name__] = True
 
     def disable_sysprep(self, obj):
         """Disable a system preparation operation"""
-        setattr(obj.im_func, 'enabled', False)
+        assert hasattr(obj, '_sysprep'), "Object is not a sysprep"
+        assert obj.__name__ in self._sysprep_tasks, "Sysprep already executed"
+
+        self._sysprep_tasks[obj.__name__] = False
+
+    def sysprep_enabled(self, obj):
+        """Returns True if this system praparation operation is enabled"""
+        assert hasattr(obj, '_sysprep'), "Object is not a sysprep"
+        assert obj.__name__ in self._sysprep_tasks, "Sysprep already executed"
+
+        return self._sysprep_tasks[obj.__name__]
 
     def print_syspreps(self):
         """Print enabled and disabled system preparation operations."""
 
         syspreps = self.list_syspreps()
-        enabled = [sysprep for sysprep in syspreps if sysprep.enabled]
-        disabled = [sysprep for sysprep in syspreps if not sysprep.enabled]
+        enabled = [s for s in syspreps if self.sysprep_enabled(s)]
+        disabled = [s for s in syspreps if not self.sysprep_enabled(s)]
 
         wrapper = textwrap.TextWrapper()
         wrapper.subsequent_indent = '\t'
@@ -330,7 +350,7 @@ class OSBase(object):
             self.out.output("(none)")
         else:
             for sysprep in enabled:
-                name = sysprep.__name__.replace('_', '-')
+                name = sysprep.__name__.replace('_', '-')[1:]
                 descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
                 self.out.output('    %s:\n%s\n' % (name, descr))
 
@@ -339,7 +359,7 @@ class OSBase(object):
             self.out.output("(none)")
         else:
             for sysprep in disabled:
-                name = sysprep.__name__.replace('_', '-')
+                name = sysprep.__name__.replace('_', '-')[1:]
                 descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
                 self.out.output('    %s:\n%s\n' % (name, descr))
 
@@ -378,16 +398,15 @@ class OSBase(object):
                 "System preparation is disabled for unsupported media")
             return
 
+        enabled = [s for s in self.list_syspreps() if self.sysprep_enabled(s)]
+        size = len(enabled)
         with self.mount():
-            enabled = [task for task in self.list_syspreps() if task.enabled]
-
-            size = len(enabled)
             cnt = 0
             for task in enabled:
                 cnt += 1
                 self.out.output(('(%d/%d)' % (cnt, size)).ljust(7), False)
                 task()
-                setattr(task.im_func, 'executed', True)
+                del self._sysprep_tasks[task.__name__]
 
         self.out.output()
 
@@ -457,10 +476,6 @@ class OSBase(object):
 
         return 0
 
-    def _is_sysprep(self, obj):
-        """Checks if an object is a sysprep"""
-        return getattr(obj, 'sysprep', False) and callable(obj)
-
     @add_prefix
     def _ls(self, directory):
         """List the name of all files under a directory"""

image_creator/os_type/freebsd.py

@@ -26,7 +26,7 @@ class Freebsd(Unix):
     """OS class for FreeBSD Unix-like operating system"""
 
     @sysprep("Cleaning up passwords & locking all user accounts")
-    def cleanup_password(self):
+    def _cleanup_password(self):
         """Remove all passwords and lock all user accounts"""
 
         master_passwd = []

image_creator/os_type/linux.py

@@ -31,7 +31,7 @@ class Linux(Unix):
         self._persistent = re.compile('/dev/[hsv]d[a-z][1-9]*')
 
     @sysprep('Removing user accounts with id greater that 1000', enabled=False)
-    def remove_user_accounts(self):
+    def _remove_user_accounts(self):
         """Remove all user accounts with id greater than 1000"""
 
         if 'USERS' not in self.meta:
@@ -84,7 +84,7 @@ class Linux(Unix):
                 self.image.g.rm_rf(home)
 
     @sysprep('Cleaning up password & locking all user accounts')
-    def cleanup_passwords(self):
+    def _cleanup_passwords(self):
         """Remove all passwords and lock all user accounts"""
 
         shadow = []
@@ -102,7 +102,7 @@ class Linux(Unix):
         self.image.g.rm_rf('/etc/shadow-')
 
     @sysprep('Fixing acpid powerdown action')
-    def fix_acpid(self):
+    def _fix_acpid(self):
         """Replace acpid powerdown action scripts to immediately shutdown the
         system without checking if a GUI is running.
         """
@@ -160,7 +160,7 @@ class Linux(Unix):
         self.out.warn("No acpi power button event found!")
 
     @sysprep('Removing persistent network interface names')
-    def remove_persistent_net_rules(self):
+    def _remove_persistent_net_rules(self):
         """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.
@@ -171,7 +171,7 @@ class Linux(Unix):
             self.image.g.rm(rule_file)
 
     @sysprep('Removing swap entry from fstab')
-    def remove_swap_entry(self):
+    def _remove_swap_entry(self):
         """Remove swap entry from /etc/fstab. If swap is the last partition
         then the partition will be removed when shrinking is performed. If the
         swap partition is not the last partition in the disk or if you are not
@@ -191,7 +191,7 @@ class Linux(Unix):
         self.image.g.write('/etc/fstab', new_fstab)
 
     @sysprep('Replacing fstab & grub non-persistent device references')
-    def use_persistent_block_device_names(self):
+    def _use_persistent_block_device_names(self):
         """Scan fstab & grub configuration files and replace all non-persistent
         device references with UUIDs.
         """

image_creator/os_type/slackware.py

@@ -23,7 +23,7 @@ from image_creator.os_type.linux import Linux, sysprep
 class Slackware(Linux):
     """OS class for Slackware Linux"""
     @sysprep("Emptying all files under /var/log")
-    def cleanup_log(self):
+    def _cleanup_log(self):
         """Empty all files under /var/log"""
 
         # In Slackware the metadata about installed packages are

image_creator/os_type/unix.py

@@ -71,26 +71,26 @@ class Unix(OSBase):
         return True
 
     @sysprep('Removing files under /var/cache')
-    def cleanup_cache(self):
+    def _cleanup_cache(self):
         """Remove all regular files under /var/cache"""
 
         self._foreach_file('/var/cache', self.image.g.rm, ftype='r')
 
     @sysprep('Removing files under /tmp and /var/tmp')
-    def cleanup_tmp(self):
+    def _cleanup_tmp(self):
         """Remove all files under /tmp and /var/tmp"""
 
         self._foreach_file('/tmp', self.image.g.rm_rf, maxdepth=1)
         self._foreach_file('/var/tmp', self.image.g.rm_rf, maxdepth=1)
 
     @sysprep('Emptying all files under /var/log')
-    def cleanup_log(self):
+    def _cleanup_log(self):
         """Empty all files under /var/log"""
 
         self._foreach_file('/var/log', self.image.g.truncate, ftype='r')
 
     @sysprep('Removing files under /var/mail & /var/spool/mail', enabled=False)
-    def cleanup_mail(self):
+    def _cleanup_mail(self):
         """Remove all files under /var/mail and /var/spool/mail"""
 
         self._foreach_file('/var/spool/mail', self.image.g.rm_rf, maxdepth=1)
@@ -98,7 +98,7 @@ class Unix(OSBase):
         self._foreach_file('/var/mail', self.image.g.rm_rf, maxdepth=1)
 
     @sysprep('Removing sensitive user data')
-    def cleanup_userdata(self):
+    def _cleanup_userdata(self):
         """Delete sensitive user data"""
 
         homedirs = ['/root']

image_creator/os_type/windows/__init__.py

@@ -313,32 +313,32 @@ class Windows(OSBase):
             namedtuple('User', 'rid name')(admin, self.usernames[admin]))
 
     @sysprep('Disabling IPv6 privacy extensions')
-    def disable_ipv6_privacy_extensions(self):
+    def _disable_ipv6_privacy_extensions(self):
         """Disable IPv6 privacy extensions"""
 
         self.vm.rexec('netsh interface ipv6 set global '
                       'randomizeidentifiers=disabled store=persistent')
 
     @sysprep('Disabling Teredo interface')
-    def disable_teredo(self):
+    def _disable_teredo(self):
         """Disable Teredo interface"""
 
         self.vm.rexec('netsh interface teredo set state disabled')
 
     @sysprep('Disabling ISATAP Adapters')
-    def disable_isatap(self):
+    def _disable_isatap(self):
         """Disable ISATAP Adapters"""
 
         self.vm.rexec('netsh interface isa set state disabled')
 
     @sysprep('Enabling ping responses')
-    def enable_pings(self):
+    def _enable_pings(self):
         """Enable ping responses"""
 
         self.vm.rexec('netsh firewall set icmpsetting 8')
 
     @sysprep('Setting the system clock to UTC')
-    def utc(self):
+    def _utc(self):
         """Set the hardware clock to UTC"""
 
         path = r'HKLM\SYSTEM\CurrentControlSet\Control\TimeZoneInformation'
@@ -346,7 +346,7 @@ class Windows(OSBase):
             r'REG ADD %s /v RealTimeIsUniversal /t REG_DWORD /d 1 /f' % path)
 
     @sysprep('Clearing the event logs')
-    def clear_logs(self):
+    def _clear_logs(self):
         """Clear all the event logs"""
 
         self.vm.rexec(
@@ -354,7 +354,7 @@ class Windows(OSBase):
             "wevtutil cl \"%l\"")
 
     @sysprep('Executing Sysprep on the image (may take more that 10 min)')
-    def microsoft_sysprep(self):
+    def _microsoft_sysprep(self):
         """Run the Microsoft System Preparation Tool. This will remove
         system-specific data and will make the image ready to be deployed.
         After this no other task may run.
@@ -365,7 +365,7 @@ class Windows(OSBase):
         self.sysprepped = True
 
     @sysprep('Converting the image into a KMS client', enabled=False)
-    def kms_client_setup(self):
+    def _kms_client_setup(self):
         """Install the appropriate KMS client setup key to the image to convert
         it to a KMS client. Computers that are running volume licensing
         editions of Windows 8, Windows Server 2012, Windows 7, Windows Server
@@ -384,7 +384,7 @@ class Windows(OSBase):
             r"cscript \Windows\system32\slmgr.vbs /ipk %s" % setup_key)
 
     @sysprep('Shrinking the last filesystem')
-    def shrink(self):
+    def _shrink(self):
         """Shrink the last filesystem. Make sure the filesystem is defragged"""
 
         # Query for the maximum number of reclaimable bytes
@@ -591,28 +591,27 @@ class Windows(OSBase):
         sysprep tasks. At the end of this method the VM is shut down if needed.
         """
         tasks = self.list_syspreps()
-        enabled = [task for task in tasks if task.enabled]
+        enabled = [t for t in tasks if self.sysprep_enabled(t)]
         size = len(enabled)
 
         # Make sure shrink runs in the end, before ms sysprep
-        enabled = [task for task in enabled
-                   if self.sysprep_info(task).name != 'shrink']
+        enabled = [t for t in enabled if self.sysprep_info(t).name != 'shrink']
         if len(enabled) != size:
-            enabled.append(self.shrink)
+            enabled.append(self._shrink)
 
         # Make sure the ms sysprep is the last task to run if it is enabled
-        enabled = [task for task in enabled
-                   if self.sysprep_info(task).name != 'microsoft-sysprep']
+        enabled = [t for t in enabled
+                   if self.sysprep_info(t).name != 'microsoft-sysprep']
 
         if len(enabled) != size:
-            enabled.append(self.microsoft_sysprep)
+            enabled.append(self._microsoft_sysprep)
 
         cnt = 0
         for task in enabled:
             cnt += 1
             self.out.output(('(%d/%d)' % (cnt, size)).ljust(7), False)
             task()
-            setattr(task.im_func, 'executed', True)
+            del self._sysprep_tasks[task.__name__]
 
         self.out.output("Sending shut down command ...", False)
         if not self.sysprepped: