Add "Running System" button in media selection

[snf-image-creator] / image_creator / disk.py

diff --git a/image_creator/disk.py b/image_creator/disk.py
index ef53799..40e1316 100644 (file)
--- a/image_creator/disk.py
+++ b/image_creator/disk.py
@@ -32,8 +32,8 @@
 # or implied, of GRNET S.A.
 
 from image_creator.util import get_command
 # or implied, of GRNET S.A.
 
 from image_creator.util import get_command

-from image_creator.util import warn, progress, success, output, FatalError
-
+from image_creator.util import FatalError
+from image_creator.gpt import GPTPartitionTable

 import stat
 import os
 import tempfile
 import stat
 import os
 import tempfile

@@ -45,9 +45,6 @@ import time
 from sendfile import sendfile
 
 
 from sendfile import sendfile
 
 

-class DiskError(Exception):
-    pass
-

 dd = get_command('dd')
 dmsetup = get_command('dmsetup')
 losetup = get_command('losetup')
 dd = get_command('dd')
 dmsetup = get_command('dmsetup')
 losetup = get_command('losetup')

@@ -62,12 +59,13 @@ class Disk(object):
     the Linux kernel.
     """
 
     the Linux kernel.
     """
 

-    def __init__(self, source):
+    def __init__(self, source, output):

         """Create a new Disk instance out of a source media. The source
         media can be an image file, a block device or a directory."""
         self._cleanup_jobs = []
         self._devices = []
         self.source = source
         """Create a new Disk instance out of a source media. The source
         media can be an image file, a block device or a directory."""
         self._cleanup_jobs = []
         self._devices = []
         self.source = source

+        self.out = output

 
     def _add_cleanup(self, job, *args):
         self._cleanup_jobs.append((job, args))
 
     def _add_cleanup(self, job, *args):
         self._cleanup_jobs.append((job, args))

@@ -79,56 +77,59 @@ class Disk(object):
         return loop
 
     def _dir_to_disk(self):
         return loop
 
     def _dir_to_disk(self):

-        raise NotImplementedError
+        raise FatalError("Using a directory as media source is not supported "
+                         "yet!")

 
     def cleanup(self):
         """Cleanup internal data. This needs to be called before the
         program ends.
         """
 
     def cleanup(self):
         """Cleanup internal data. This needs to be called before the
         program ends.
         """

-        while len(self._devices):
-            device = self._devices.pop()
-            device.destroy()
-
-        while len(self._cleanup_jobs):
-            job, args = self._cleanup_jobs.pop()
-            job(*args)
+        try:
+            while len(self._devices):
+                device = self._devices.pop()
+                device.destroy()
+        finally:
+            # Make sure those are executed even if one of the device.destroy
+            # methods throws exeptions.
+            while len(self._cleanup_jobs):
+                job, args = self._cleanup_jobs.pop()
+                job(*args)

 
     def snapshot(self):
         """Creates a snapshot of the original source media of the Disk
         instance.
         """
 
 
     def snapshot(self):
         """Creates a snapshot of the original source media of the Disk
         instance.
         """
 

-        output("Examining source media `%s'..." % self.source, False)
+        self.out.output("Examining source media `%s'..." % self.source, False)

         sourcedev = self.source
         mode = os.stat(self.source).st_mode
         if stat.S_ISDIR(mode):
         sourcedev = self.source
         mode = os.stat(self.source).st_mode
         if stat.S_ISDIR(mode):

-            success('looks like a directory')
+            self.out.success('looks like a directory')

             return self._losetup(self._dir_to_disk())
         elif stat.S_ISREG(mode):
             return self._losetup(self._dir_to_disk())
         elif stat.S_ISREG(mode):

-            success('looks like an image file')
+            self.out.success('looks like an image file')

             sourcedev = self._losetup(self.source)
         elif not stat.S_ISBLK(mode):
             raise ValueError("Invalid media source. Only block devices, "
             sourcedev = self._losetup(self.source)
         elif not stat.S_ISBLK(mode):
             raise ValueError("Invalid media source. Only block devices, "

-                            "regular files and directories are supported.")
+                             "regular files and directories are supported.")

         else:
         else:

-            success('looks like a block device')
+            self.out.success('looks like a block device')

 
         # Take a snapshot and return it to the user
 
         # Take a snapshot and return it to the user

-        output("Snapshotting media source...", False)
-        size = blockdev('--getsize', sourcedev)
+        self.out.output("Snapshotting media source...", False)
+        size = blockdev('--getsz', sourcedev)

         cowfd, cow = tempfile.mkstemp()
         os.close(cowfd)
         self._add_cleanup(os.unlink, cow)
         cowfd, cow = tempfile.mkstemp()
         os.close(cowfd)
         self._add_cleanup(os.unlink, cow)

-        # Create 1G cow sparse file
-        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
-                                        'seek=%d' % (1024 * 1024))
+        # Create cow sparse file
+        dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))

         cowdev = self._losetup(cow)
 
         snapshot = uuid.uuid4().hex
         tablefd, table = tempfile.mkstemp()
         try:
         cowdev = self._losetup(cow)
 
         snapshot = uuid.uuid4().hex
         tablefd, table = tempfile.mkstemp()
         try:

-            os.write(tablefd, "0 %d snapshot %s %s n 8" % \
-                                        (int(size), sourcedev, cowdev))
+            os.write(tablefd, "0 %d snapshot %s %s n 8" %
+                              (int(size), sourcedev, cowdev))

             dmsetup('create', snapshot, table)
             self._add_cleanup(dmsetup, 'remove', snapshot)
             # Sometimes dmsetup remove fails with Device or resource busy,
             dmsetup('create', snapshot, table)
             self._add_cleanup(dmsetup, 'remove', snapshot)
             # Sometimes dmsetup remove fails with Device or resource busy,

@@ -138,13 +139,13 @@ class Disk(object):
 
         finally:
             os.unlink(table)
 
         finally:
             os.unlink(table)

-        success('done')
+        self.out.success('done')

         return "/dev/mapper/%s" % snapshot
 
     def get_device(self, media):
         """Returns a newly created DiskDevice instance."""
 
         return "/dev/mapper/%s" % snapshot
 
     def get_device(self, media):
         """Returns a newly created DiskDevice instance."""
 

-        new_device = DiskDevice(media)
+        new_device = DiskDevice(media, self.out)

         self._devices.append(new_device)
         new_device.enable()
         return new_device
         self._devices.append(new_device)
         new_device.enable()
         return new_device

@@ -162,15 +163,33 @@ class DiskDevice(object):
     as created by the device-mapper.
     """
 
     as created by the device-mapper.
     """
 

-    def __init__(self, device, bootable=True):
+    def __init__(self, device, output, bootable=True):

         """Create a new DiskDevice."""
 
         """Create a new DiskDevice."""
 

-        self.device = device
+        self.real_device = device
+        self.out = output

         self.bootable = bootable
         self.progress_bar = None
         self.bootable = bootable
         self.progress_bar = None

+        self.guestfs_device = None
+        self.size = 0
+        self.meta = {}

 
         self.g = guestfs.GuestFS()
 
         self.g = guestfs.GuestFS()

-        self.g.add_drive_opts(self.device, readonly=0)
+        self.g.add_drive_opts(self.real_device, readonly=0)
+
+        # Before version 1.17.14 the recovery process, which is a fork of the
+        # original process that called libguestfs, did not close its inherited
+        # file descriptors. This can cause problems especially if the parent
+        # process has opened pipes. Since the recovery process is an optional
+        # feature of libguestfs, it's better to disable it.
+        self.g.set_recovery_proc(0)
+        version = self.g.version()
+        if version['major'] > 1 or \
+            (version['major'] == 1 and (version['minor'] >= 18 or
+                                        (version['minor'] == 17 and
+                                         version['release'] >= 14))):
+            self.g.set_recovery_proc(1)
+            self.out.output("Enabling recovery proc")

 
         #self.g.set_trace(1)
         #self.g.set_verbose(1)
 
         #self.g.set_trace(1)
         #self.g.set_verbose(1)

@@ -179,52 +198,57 @@ class DiskDevice(object):
 
     def enable(self):
         """Enable a newly created DiskDevice"""
 
     def enable(self):
         """Enable a newly created DiskDevice"""

-        self.progressbar = progress("Launching helper VM: ")
-        self.progressbar.next()
+        self.progressbar = self.out.Progress(100, "Launching helper VM",
+                                             "percent")

         eh = self.g.set_event_callback(self.progress_callback,
         eh = self.g.set_event_callback(self.progress_callback,

-                                                    guestfs.EVENT_PROGRESS)
+                                       guestfs.EVENT_PROGRESS)

         self.g.launch()
         self.guestfs_enabled = True
         self.g.delete_event_callback(eh)
         self.g.launch()
         self.guestfs_enabled = True
         self.g.delete_event_callback(eh)

-        if self.progressbar is not None:
-            self.progressbar.send(100)
-            self.progressbar = None
+        self.progressbar.success('done')
+        self.progressbar = None

 
 

-        output('Inspecting Operating System...', False)
+        self.out.output('Inspecting Operating System...', False)

         roots = self.g.inspect_os()
         if len(roots) == 0:
             raise FatalError("No operating system found")
         if len(roots) > 1:
             raise FatalError("Multiple operating systems found."
         roots = self.g.inspect_os()
         if len(roots) == 0:
             raise FatalError("No operating system found")
         if len(roots) > 1:
             raise FatalError("Multiple operating systems found."

-                            "We only support images with one filesystem.")
+                             "We only support images with one OS.")

         self.root = roots[0]
         self.root = roots[0]

+        self.guestfs_device = self.g.part_to_dev(self.root)
+        self.size = self.g.blockdev_getsize64(self.guestfs_device)
+        self.meta['PARTITION_TABLE'] = \
+            self.g.part_get_parttype(self.guestfs_device)
+

         self.ostype = self.g.inspect_get_type(self.root)
         self.distro = self.g.inspect_get_distro(self.root)
         self.ostype = self.g.inspect_get_type(self.root)
         self.distro = self.g.inspect_get_distro(self.root)

-        success('found a %s system' % self.distro)
+        self.out.success('found a(n) %s system' % self.distro)

 
     def destroy(self):
         """Destroy this DiskDevice instance."""
 
 
     def destroy(self):
         """Destroy this DiskDevice instance."""
 

-        if self.guestfs_enabled:
-            self.g.umount_all()
-            self.g.sync()
-
-        # Close the guestfs handler if open
-        self.g.close()
+        # In new guestfs versions, there is a handy shutdown method for this
+        try:
+            if self.guestfs_enabled:
+                self.g.umount_all()
+                self.g.sync()
+        finally:
+            # Close the guestfs handler if open
+            self.g.close()

 
     def progress_callback(self, ev, eh, buf, array):
         position = array[2]
         total = array[3]
 
 
     def progress_callback(self, ev, eh, buf, array):
         position = array[2]
         total = array[3]
 

-        self.progressbar.send((position * 100) // total)
+        self.progressbar.goto((position * 100) // total)

 
 

-        if position == total:
-            self.progressbar = None
-
-    def mount(self):
+    def mount(self, readonly=False):

         """Mount all disk partitions in a correct order."""
 
         """Mount all disk partitions in a correct order."""
 

-        output("Mounting image...", False)
+        mount = self.g.mount_ro if readonly else self.g.mount
+        msg = " read-only" if readonly else ""
+        self.out.output("Mounting the media%s..." % msg, False)

         mps = self.g.inspect_get_mountpoints(self.root)
 
         # Sort the keys to mount the fs in a correct order.
         mps = self.g.inspect_get_mountpoints(self.root)
 
         # Sort the keys to mount the fs in a correct order.

@@ -239,42 +263,101 @@ class DiskDevice(object):
         mps.sort(compare)
         for mp, dev in mps:
             try:
         mps.sort(compare)
         for mp, dev in mps:
             try:

-                self.g.mount(dev, mp)
+                mount(dev, mp)

             except RuntimeError as msg:
             except RuntimeError as msg:

-                warn("%s (ignored)" % msg)
-        success("done")
+                self.out.warn("%s (ignored)" % msg)
+        self.out.success("done")

 
     def umount(self):
         """Umount all mounted filesystems."""
         self.g.umount_all()
 
 
     def umount(self):
         """Umount all mounted filesystems."""
         self.g.umount_all()
 

+    def _last_partition(self):
+        if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
+            msg = "Unsupported partition table: %s. Only msdos and gpt " \
+                "partition tables are supported" % self.meta['PARTITION_TABLE']
+            raise FatalError(msg)
+
+        is_extended = lambda p: \
+            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
+        is_logical = lambda p: \
+            self.meta['PARTITION_TABLE'] != 'msdos' and p['part_num'] > 4
+
+        partitions = self.g.part_list(self.guestfs_device)
+        last_partition = partitions[-1]
+
+        if is_logical(last_partition):
+            # The disk contains extended and logical partitions....
+            extended = [p for p in partitions if is_extended(p)][0]
+            last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
+
+            # check if extended is the last primary partition
+            if last_primary['part_num'] > extended['part_num']:
+                last_partition = last_primary
+
+        return last_partition
+

     def shrink(self):
         """Shrink the disk.
 
         This is accomplished by shrinking the last filesystem in the
         disk and then updating the partition table. The new disk size
         (in bytes) is returned.
     def shrink(self):
         """Shrink the disk.
 
         This is accomplished by shrinking the last filesystem in the
         disk and then updating the partition table. The new disk size
         (in bytes) is returned.

-        """
-        output("Shrinking image (this may take a while)...", False)
-
-        dev = self.g.part_to_dev(self.root)
-        parttype = self.g.part_get_parttype(dev)
-        if parttype != 'msdos':
-            raise FatalError("You have a %s partition table. "
-                "Only msdos partitions are supported" % parttype)
-
-        last_partition = self.g.part_list(dev)[-1]
-
-        if last_partition['part_num'] > 4:
-            raise FatalError("This disk contains logical partitions. "
-                "Only primary partitions are supported.")
-
-        part_dev = "%s%d" % (dev, last_partition['part_num'])
-        fs_type = self.g.vfs_type(part_dev)
-        if not re.match("ext[234]", fs_type):
-            warn("Don't know how to resize %s partitions." % vfs_type)
-            return

 
 

+        ATTENTION: make sure unmount is called before shrink
+        """
+        get_fstype = lambda p: \
+            self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
+        is_logical = lambda p: \
+            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
+        is_extended = lambda p: \
+            self.meta['PARTITION_TABLE'] == 'msdos' and \
+            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
+
+        part_add = lambda ptype, start, stop: \
+            self.g.part_add(self.guestfs_device, ptype, start, stop)
+        part_del = lambda p: self.g.part_del(self.guestfs_device, p)
+        part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
+        part_set_id = lambda p, id: \
+            self.g.part_set_mbr_id(self.guestfs_device, p, id)
+        part_get_bootable = lambda p: \
+            self.g.part_get_bootable(self.guestfs_device, p)
+        part_set_bootable = lambda p, bootable: \
+            self.g.part_set_bootable(self.guestfs_device, p, bootable)
+
+        MB = 2 ** 20
+
+        self.out.output("Shrinking image (this may take a while)...", False)
+
+        sector_size = self.g.blockdev_getss(self.guestfs_device)
+
+        last_part = None
+        fstype = None
+        while True:
+            last_part = self._last_partition()
+            fstype = get_fstype(last_part)
+
+            if fstype == 'swap':
+                self.meta['SWAP'] = "%d:%s" % \
+                    (last_part['part_num'],
+                     (last_part['part_size'] + MB - 1) // MB)
+                part_del(last_part['part_num'])
+                continue
+            elif is_extended(last_part):
+                part_del(last_part['part_num'])
+                continue
+
+            # Most disk manipulation programs leave 2048 sectors after the last
+            # partition
+            new_size = last_part['part_end'] + 1 + 2048 * sector_size
+            self.size = min(self.size, new_size)
+            break
+
+        if not re.match("ext[234]", fstype):
+            self.out.warn("Don't know how to resize %s partitions." % fstype)
+            return self.size
+
+        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])

         self.g.e2fsck_f(part_dev)
         self.g.resize2fs_M(part_dev)
 
         self.g.e2fsck_f(part_dev)
         self.g.resize2fs_M(part_dev)
 

@@ -284,29 +367,62 @@ class DiskDevice(object):
         block_cnt = int(
             filter(lambda x: x[0] == 'Block count', out)[0][1])
 
         block_cnt = int(
             filter(lambda x: x[0] == 'Block count', out)[0][1])
 

-        sector_size = self.g.blockdev_getss(dev)
-
-        start = last_partition['part_start'] / sector_size
+        start = last_part['part_start'] / sector_size

         end = start + (block_size * block_cnt) / sector_size - 1
 
         end = start + (block_size * block_cnt) / sector_size - 1
 

-        self.g.part_del(dev, last_partition['part_num'])
-        self.g.part_add(dev, 'p', start, end)
+        if is_logical(last_part):
+            partitions = self.g.part_list(self.guestfs_device)
+
+            logical = []  # logical partitions
+            for partition in partitions:
+                if partition['part_num'] < 4:
+                    continue
+                logical.append({
+                    'num': partition['part_num'],
+                    'start': partition['part_start'] / sector_size,
+                    'end': partition['part_end'] / sector_size,
+                    'id': part_get_(partition['part_num']),
+                    'bootable': part_get_bootable(partition['part_num'])
+                })
+
+            logical[-1]['end'] = end  # new end after resize
+
+            # Recreate the extended partition
+            extended = [p for p in partitions if self._is_extended(p)][0]
+            part_del(extended['part_num'])
+            part_add('e', extended['part_start'], end)
+
+            # Create all the logical partitions back
+            for l in logical:
+                part_add('l', l['start'], l['end'])
+                part_set_id(l['num'], l['id'])
+                part_set_bootable(l['num'], l['bootable'])
+        else:
+            # Recreate the last partition
+            if self.meta['PARTITION_TABLE'] == 'msdos':
+                last_part['id'] = part_get_id(last_part['part_num'])
+
+            last_part['bootable'] = part_get_bootable(last_part['part_num'])
+            part_del(last_part['part_num'])
+            part_add('p', start, end)
+            part_set_bootable(last_part['part_num'], last_part['bootable'])
+
+            if self.meta['PARTITION_TABLE'] == 'msdos':
+                part_set_id(last_part['part_num'], last_part['id'])

 
         new_size = (end + 1) * sector_size
 
         new_size = (end + 1) * sector_size

-        success("new image size is %dMB" %
-                            ((new_size + 2 ** 20 - 1) // 2 ** 20))
-        return new_size

 
 

-    def size(self):
-        """Returns the "payload" size of the device.
+        assert (new_size <= self.size)

 
 

-        The size returned by this method is the size of the space occupied by
-        the partitions (including the space before the first partition).
-        """
-        dev = self.g.part_to_dev(self.root)
-        last = self.g.part_list(dev)[-1]
+        if self.meta['PARTITION_TABLE'] == 'gpt':
+            ptable = GPTPartitionTable(self.real_device)
+            self.size = ptable.shrink(new_size, self.size)
+        else:
+            self.size = min(new_size + 2048 * sector_size, self.size)

 
 

-        return last['part_end'] + 1
+        self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
+
+        return self.size

 
     def dump(self, outfile):
         """Dumps the content of device into a file.
 
     def dump(self, outfile):
         """Dumps the content of device into a file.

@@ -314,31 +430,24 @@ class DiskDevice(object):
         This method will only dump the actual payload, found by reading the
         partition table. Empty space in the end of the device will be ignored.
         """
         This method will only dump the actual payload, found by reading the
         partition table. Empty space in the end of the device will be ignored.
         """

-        blocksize = 2 ** 22  # 4MB
-        size = self.size()
-        progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
-        progressbar = progress("Dumping image file: ", progress_size)
-
-        source = open(self.device, "r")
-        try:
-            dest = open(outfile, "w")
-            try:
+        MB = 2 ** 20
+        blocksize = 4 * MB  # 4MB
+        size = self.size
+        progr_size = (size + MB - 1) // MB  # in MB
+        progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
+
+        with open(self.real_device, 'r') as src:
+            with open(outfile, "w") as dst:

                 left = size
                 offset = 0
                 progressbar.next()
                 while left > 0:
                     length = min(left, blocksize)
                 left = size
                 offset = 0
                 progressbar.next()
                 while left > 0:
                     length = min(left, blocksize)

-                    sent = sendfile(dest.fileno(), source.fileno(), offset,
-                                                                        length)
+                    _, sent = sendfile(dst.fileno(), src.fileno(), offset,
+                        length)

                     offset += sent
                     left -= sent
                     offset += sent
                     left -= sent

-                    for i in range((length + 2 ** 20 - 1) // 2 ** 20):
-                        progressbar.next()
-            finally:
-                dest.close()
-        finally:
-            source.close()
-
-        success('Image file %s was successfully created' % outfile)
+                    progressbar.goto((size - left) // MB)
+        progressbar.success('image file %s was successfully created' % outfile)

 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :