From 71b0ab283eeef37c3d3ecd98fbafa6eaa50b58d8 Mon Sep 17 00:00:00 2001
From: Nikos Skalkotos <skalkoto@grnet.gr>
Date: Fri, 7 Jun 2013 15:24:51 +0300
Subject: [PATCH] Move mount/umount from image to os_type module

---
 image_creator/dialog_menu.py      |   35 ++++++----------
 image_creator/dialog_wizard.py    |    8 ----
 image_creator/image.py            |   81 ++-----------------------------------
 image_creator/main.py             |   41 +++++++------------
 image_creator/os_type/__init__.py |   76 ++++++++++++++++++++++++++++------
 image_creator/os_type/freebsd.py  |   31 ++++++++++++++
 image_creator/os_type/linux.py    |    4 ++
 image_creator/os_type/ubuntu.py   |    4 ++
 image_creator/os_type/unix.py     |   37 ++++++++++++++++-
 image_creator/version.py          |    9 +++--
 10 files changed, 171 insertions(+), 155 deletions(-)

diff --git a/image_creator/dialog_menu.py b/image_creator/dialog_menu.py
index b6d3162..448060b 100644
--- a/image_creator/dialog_menu.py
+++ b/image_creator/dialog_menu.py
@@ -38,7 +38,7 @@ import textwrap
 import StringIO
 
 from image_creator import __version__ as version
-from image_creator.util import MD5
+from image_creator.util import MD5, FatalError
 from image_creator.output.dialog import GaugeOutput, InfoBoxOutput
 from image_creator.kamaki_wrapper import Kamaki, ClientError
 from image_creator.help import get_help_file
@@ -568,30 +568,19 @@ def sysprep(session):
             try:
                 image.out.add(infobox)
                 try:
-                    image.mount(readonly=False)
-                    try:
-                        err = "Unable to execute the system preparation " \
-                            "tasks. Couldn't mount the media%s."
-                        title = "System Preparation"
-                        if not image.mounted:
-                            d.msgbox(err % "", title=title, width=SMALL_WIDTH)
-                            return
-                        elif image.mounted_ro:
-                            d.msgbox(err % " read-write", title=title,
-                                     width=SMALL_WIDTH)
-                            return
-
-                        # The checksum is invalid. We have mounted the image rw
-                        if 'checksum' in session:
-                            del session['checksum']
-
-                        # Monitor the metadata changes during syspreps
-                        with MetadataMonitor(session, image.os.meta):
+                    # The checksum is invalid. We have mounted the image rw
+                    if 'checksum' in session:
+                        del session['checksum']
+
+                    # Monitor the metadata changes during syspreps
+                    with MetadataMonitor(session, image.os.meta):
+                        try:
                             image.os.do_sysprep()
                             infobox.finalize()
-
-                    finally:
-                        image.umount()
+                        except FatalError as e:
+                            title = "System Preparation"
+                            d.msgbox("System Preparation failed: %s" % e,
+                                     title=title, width=SMALL_WIDTH)
                 finally:
                     image.out.remove(infobox)
             finally:
diff --git a/image_creator/dialog_wizard.py b/image_creator/dialog_wizard.py
index 30dc7c6..79d18db 100644
--- a/image_creator/dialog_wizard.py
+++ b/image_creator/dialog_wizard.py
@@ -266,16 +266,8 @@ def create_image(session):
         out.clear()
 
         #Sysprep
-        image.mount(False)
-        err_msg = "Unable to execute the system preparation tasks."
-        if not image.mounted:
-            raise FatalError("%s Couldn't mount the media." % err_msg)
-        elif image.mounted_ro:
-            raise FatalError("%s Couldn't mount the media read-write."
-                             % err_msg)
         image.os.do_sysprep()
         metadata = image.os.meta
-        image.umount()
 
         #Shrink
         size = image.shrink()
diff --git a/image_creator/image.py b/image_creator/image.py
index de28756..1ad915d 100644
--- a/image_creator/image.py
+++ b/image_creator/image.py
@@ -53,8 +53,6 @@ class Image(object):
         self.progress_bar = None
         self.guestfs_device = None
         self.size = 0
-        self.mounted = False
-        self.mounted_ro = False
 
         self.g = guestfs.GuestFS()
         self.g.add_drive_opts(self.device, readonly=0, format="raw")
@@ -120,19 +118,10 @@ class Image(object):
         if not self.guestfs_enabled:
             self.enable()
 
-        if not self.mounted:
-            do_unmount = True
-            self.mount(readonly=True)
-        else:
-            do_unmount = False
-
-        try:
-            cls = os_cls(self.distro, self.ostype)
-            self._os = cls(self.root, self.g, self.out)
+        cls = os_cls(self.distro, self.ostype)
+        self._os = cls(self.root, self.g, self.out)
 
-        finally:
-            if do_unmount:
-                self.umount()
+        self._os.collect_metadata()
 
         return self._os
 
@@ -156,70 +145,6 @@ class Image(object):
 #
 #        self.progressbar.goto((position * 100) // total)
 
-    def mount(self, readonly=False):
-        """Mount all disk partitions in a correct order."""
-
-        msg = "Mounting the media%s ..." % (" read-only" if readonly else "")
-        self.out.output(msg, False)
-
-        #If something goes wrong when mounting rw, remount the filesystem ro
-        remount_ro = False
-        rw_mpoints = ('/', '/etc', '/root', '/home', '/var')
-
-        # Sort the keys to mount the fs in a correct order.
-        # / should be mounted befor /boot, etc
-        def compare(a, b):
-            if len(a[0]) > len(b[0]):
-                return 1
-            elif len(a[0]) == len(b[0]):
-                return 0
-            else:
-                return -1
-        mps = self.g.inspect_get_mountpoints(self.root)
-        mps.sort(compare)
-
-        mopts = 'ro' if readonly else 'rw'
-        for mp, dev in mps:
-            if self.ostype == 'freebsd':
-                # libguestfs can't handle correct freebsd partitions on GUID
-                # Partition Table. We have to do the translation to linux
-                # device names ourselves
-                m = re.match('^/dev/((?:ada)|(?:vtbd))(\d+)p(\d+)$', dev)
-                if m:
-                    m2 = int(m.group(2))
-                    m3 = int(m.group(3))
-                    dev = '/dev/sd%c%d' % (chr(ord('a') + m2), m3)
-            try:
-                self.g.mount_options(mopts, dev, mp)
-            except RuntimeError as msg:
-                if self.ostype == 'freebsd':
-                    freebsd_mopts = "ufstype=ufs2,%s" % mopts
-                    try:
-                        self.g.mount_vfs(freebsd_mopts, 'ufs', dev, mp)
-                    except RuntimeError as msg:
-                        if readonly is False and mp in rw_mpoints:
-                            remount_ro = True
-                            break
-                elif readonly is False and mp in rw_mpoints:
-                    remount_ro = True
-                    break
-                else:
-                    self.out.warn("%s (ignored)" % msg)
-        if remount_ro:
-            self.out.warn("Unable to mount %s read-write. "
-                          "Remounting everything read-only..." % mp)
-            self.umount()
-            self.mount(True)
-        else:
-            self.mounted = True
-            self.mounted_ro = readonly
-            self.out.success("done")
-
-    def umount(self):
-        """Umount all mounted filesystems."""
-        self.g.umount_all()
-        self.mounted = False
-
     def _last_partition(self):
         """Return the last partition of the image disk"""
         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
diff --git a/image_creator/main.py b/image_creator/main.py
index 10fc0b5..a32dbd0 100644
--- a/image_creator/main.py
+++ b/image_creator/main.py
@@ -221,36 +221,23 @@ def image_creator():
 
         image = disk.get_image(snapshot)
 
-        # If no customization is to be done, the image should be mounted ro
-        ro = (not (options.sysprep or options.shrink) or options.print_sysprep)
-        image.mount(ro)
-        try:
-            for sysprep in options.disabled_syspreps:
-                image.os.disable_sysprep(image.os.get_sysprep_by_name(sysprep))
+        for sysprep in options.disabled_syspreps:
+            image.os.disable_sysprep(image.os.get_sysprep_by_name(sysprep))
 
-            for sysprep in options.enabled_syspreps:
-                image.os.enable_sysprep(image.os.get_sysprep_by_name(sysprep))
+        for sysprep in options.enabled_syspreps:
+            image.os.enable_sysprep(image.os.get_sysprep_by_name(sysprep))
 
-            if options.print_sysprep:
-                image.os.print_syspreps()
-                out.output()
+        if options.print_sysprep:
+            image.os.print_syspreps()
+            out.output()
+
+        if options.outfile is None and not options.upload:
+            return 0
+
+        if options.sysprep:
+            image.os.do_sysprep()
 
-            if options.outfile is None and not options.upload:
-                return 0
-
-            if options.sysprep:
-                err_msg = "Unable to perform the system preparation tasks. " \
-                    "Couldn't mount the media%s. Use --no-sysprep if you " \
-                    "don't won't to perform any system preparation task."
-                if not image.mounted:
-                    raise FatalError(err_msg % "")
-                elif image.mounted_ro:
-                    raise FatalError(err_msg % " read-write")
-                image.os.do_sysprep()
-
-            metadata = image.os.meta
-        finally:
-            image.umount()
+        metadata = image.os.meta
 
         size = options.shrink and image.shrink() or image.size
         metadata.update(image.meta)
diff --git a/image_creator/os_type/__init__.py b/image_creator/os_type/__init__.py
index cea08a3..d093172 100644
--- a/image_creator/os_type/__init__.py
+++ b/image_creator/os_type/__init__.py
@@ -77,9 +77,23 @@ class OSBase(object):
         self.root = rootdev
         self.g = ghandler
         self.out = output
-
-        # Collect metadata about the OS
         self.meta = {}
+
+    def collect_metadata(self):
+        """Collect metadata about the OS"""
+
+        try:
+            if not self.mount(readonly=True):
+                raise FatalError("Unable to mount the media read-only")
+
+            self.out.output('Collecting image metadata ...', False)
+            self._do_collect_metadata()
+            self.out.success('done')
+        finally:
+            self.umount()
+
+    def _do_collect_metadata(self):
+
         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)
@@ -212,18 +226,54 @@ class OSBase(object):
     def do_sysprep(self):
         """Prepere system for image creation."""
 
-        self.out.output('Preparing system for image creation:')
+        try:
+            if not self.mount(readonly=False):
+                raise FatalError("Unable to mount the media read-write")
+
+            self.out.output('Preparing system for image creation:')
+
+            tasks = self.list_syspreps()
+            enabled = filter(lambda x: x.enabled, tasks)
+
+            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)
+            self.out.output()
+        finally:
+            self.umount()
+
+    def _do_mount(self, readonly):
+        try:
+            self.g.mount_options('ro' if readonly else 'rw', self.root, '/')
+        except RuntimeError as msg:
+            self.out.warn("unable to mount the root partition: %s" % msg)
+            return False
+
+        return True
+
+    def mount(self, readonly=False):
+        """Mount image."""
+
+        if getattr(self, "mounted", False):
+            return True
+
+        mount_type = 'read-only' if readonly else 'read-write'
+        self.out.output("Mount the media %s ..." % mount_type, False)
+
+        if not self._do_mount(readonly):
+            return False
 
-        tasks = self.list_syspreps()
-        enabled = filter(lambda x: x.enabled, tasks)
+        self.mounted = True
+        self.out.success('done')
+        return True
 
-        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)
-        self.out.output()
+    def umount(self):
+        """Umount all mounted filesystems."""
+        self.g.umount_all()
+        self.mounted = False
 
 # 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 d0ed78b..28d58bd 100644
--- a/image_creator/os_type/freebsd.py
+++ b/image_creator/os_type/freebsd.py
@@ -41,6 +41,9 @@ class Freebsd(Unix):
     def __init__(self, rootdev, ghandler, output):
         super(Freebsd, self).__init__(rootdev, ghandler, output)
 
+    def _do_collect_metadata(self):
+
+        super(Freebsd, self)._do_collect_metadata()
         self.meta["USERS"] = " ".join(self._get_passworded_users())
 
         #The original product name key is long and ugly
@@ -72,6 +75,34 @@ class Freebsd(Unix):
 
         return users
 
+    def _do_mount(self, readonly):
+        """Mount partitions in the correct order"""
+
+        critical_mpoints = ('/', '/etc', '/root', '/home', '/var')
+
+        # libguestfs can't handle correct freebsd partitions on GUID
+        # Partition Table. We have to do the translation to linux
+        # device names ourselves
+        guid_device = re.compile('^/dev/((?:ada)|(?:vtbd))(\d+)p(\d+)$')
+
+        mopts = "ufstype=ufs2,%s" % ('ro' if readonly else 'rw')
+        for mp, dev in self._mountpoints():
+            match = guid_device.match(dev)
+            if match:
+                m2 = int(match.group(2))
+                m3 = int(match.group(3))
+                dev = '/dev/sd%c%d' % (chr(ord('a') + m2), m3)
+            try:
+                self.g.mount_vfs(mopts, 'ufs', dev, mp)
+            except RuntimeError as msg:
+                if mp in critical_mpoints:
+                    self.out.warn('unable to mount %s. Reason: %s' % (mp, msg))
+                    return False
+                else:
+                    self.out.warn('%s (ignored)' % msg)
+
+        return True
+
     @sysprep()
     def cleanup_password(self, print_header=True):
         """Remove all passwords and lock all user accounts"""
diff --git a/image_creator/os_type/linux.py b/image_creator/os_type/linux.py
index 86638c0..5a3d29d 100644
--- a/image_creator/os_type/linux.py
+++ b/image_creator/os_type/linux.py
@@ -44,6 +44,10 @@ class Linux(Unix):
         self._uuid = dict()
         self._persistent = re.compile('/dev/[hsv]d[a-z][1-9]*')
 
+    def _do_collect_metadata(self):
+        """Collect metadata about the OS"""
+
+        super(Linux, self)._do_collect_metadata()
         self.meta["USERS"] = " ".join(self._get_passworded_users())
 
         # Delete the USERS metadata if empty
diff --git a/image_creator/os_type/ubuntu.py b/image_creator/os_type/ubuntu.py
index c0ed913..8bfe04d 100644
--- a/image_creator/os_type/ubuntu.py
+++ b/image_creator/os_type/ubuntu.py
@@ -39,6 +39,10 @@ class Ubuntu(Linux):
     def __init__(self, rootdev, ghandler, output):
         super(Ubuntu, self).__init__(rootdev, ghandler, output)
 
+    def _do_collect_metadata(self):
+        """Collect metadata about the OS"""
+
+        super(Ubuntu, self)._do_collect_metadata()
         apps = self.g.inspect_list_applications(self.root)
         for app in apps:
             if app['app_name'] == 'kubuntu-desktop':
diff --git a/image_creator/os_type/unix.py b/image_creator/os_type/unix.py
index b0b3ae6..e368231 100644
--- a/image_creator/os_type/unix.py
+++ b/image_creator/os_type/unix.py
@@ -47,8 +47,41 @@ class Unix(OSBase):
         '.kamaki.history'
     ]
 
-    def __init__(self, rootdev, ghandler, output):
-        super(Unix, self).__init__(rootdev, ghandler, output)
+    def _mountpoints(self):
+        """Return mountpoints in the correct order.
+        / should be mounted before /boot or /usr, /usr befor /usr/bin ...
+        """
+        mps = self.g.inspect_get_mountpoints(self.root)
+
+        def compare(a, b):
+            if len(a[0]) > len(b[0]):
+                return 1
+            elif len(a[0]) == len(b[0]):
+                return 0
+            else:
+                return -1
+        mps.sort(compare)
+
+        for mp in mps:
+            yield mp
+
+    def _do_mount(self, readonly):
+        """Mount partitions in the correct order"""
+
+        critical_mpoints = ('/', '/etc', '/root', '/home', '/var')
+
+        mopts = 'ro' if readonly else 'rw'
+        for mp, dev in self._mountpoints():
+            try:
+                self.g.mount_options(mopts, dev, mp)
+            except RuntimeError as msg:
+                if mp in critical_mpoint:
+                    self.out.warn('unable to mount %s. Reason: %s' % (mp, msg))
+                    return False
+                else:
+                    self.out.warn('%s (ignored)' % msg)
+
+        return True
 
     @sysprep()
     def cleanup_cache(self, print_header=True):
diff --git a/image_creator/version.py b/image_creator/version.py
index e7caa11..249fd43 100644
--- a/image_creator/version.py
+++ b/image_creator/version.py
@@ -1,7 +1,8 @@
 __version__ = "0.3next"
 __version_info__ = ['0', '3next']
-__version_vcs_info__ = {'branch': 'develop',
- 'revid': '9c060ab',
- 'revno': 297,
- 'toplevel': '/home/skalkoto/src/snf-image-creator'}
+__version_vcs_info__ = {
+    'branch': 'develop',
+    'revid': '9c060ab',
+    'revno': 297,
+    'toplevel': '/home/skalkoto/src/snf-image-creator'}
 __version_user_info__ = "skalkoto@darkstar.admin.grnet.gr"
-- 
1.7.10.4