Add Output module for dialog-based output

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

diff --git a/image_creator/disk.py b/image_creator/disk.py
index e0ed6a1..38df2fd 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

@@ -62,12 +62,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))

@@ -98,7 +99,7 @@ class Disk(object):
         instance.
         """
 
         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):

@@ -111,17 +112,16 @@ class Disk(object):
             raise ValueError("Invalid media source. Only block devices, "
                             "regular files and directories are supported.")
         else:
             raise ValueError("Invalid media source. Only block devices, "
                             "regular files and directories are supported.")
         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)
+        self.out.output("Snapshotting media source...", False)

         size = blockdev('--getsize', sourcedev)
         cowfd, cow = tempfile.mkstemp()
         os.close(cowfd)
         self._add_cleanup(os.unlink, cow)
         # Create 1G cow sparse file
         size = blockdev('--getsize', sourcedev)
         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))
+        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))

         cowdev = self._losetup(cow)
 
         snapshot = uuid.uuid4().hex
         cowdev = self._losetup(cow)
 
         snapshot = uuid.uuid4().hex

@@ -138,13 +138,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 +162,31 @@ 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.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,20 +195,17 @@ 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: ", "percent")
-        self.progressbar.max = 100
-        self.progressbar.goto(1)
+        self.progressbar = self.out.Progress(100, "Launching helper VM",
+                                             "percent")

         eh = self.g.set_event_callback(self.progress_callback,
                                                     guestfs.EVENT_PROGRESS)
         self.g.launch()
         self.guestfs_enabled = True
         self.g.delete_event_callback(eh)
         eh = self.g.set_event_callback(self.progress_callback,
                                                     guestfs.EVENT_PROGRESS)
         self.g.launch()
         self.guestfs_enabled = True
         self.g.delete_event_callback(eh)

-        if self.progressbar is not None:
-            output("\rLaunching helper VM...\033[K", False)
-            success("done")
-            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")
         roots = self.g.inspect_os()
         if len(roots) == 0:
             raise FatalError("No operating system found")

@@ -200,9 +213,14 @@ class DiskDevice(object):
             raise FatalError("Multiple operating systems found."
                             "We only support images with one filesystem.")
         self.root = roots[0]
             raise FatalError("Multiple operating systems found."
                             "We only support images with one filesystem.")
         self.root = roots[0]

+        self.guestfs_device = self.g.part_to_dev(self.root)
+        self.meta['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."""

@@ -223,7 +241,7 @@ class DiskDevice(object):
     def mount(self):
         """Mount all disk partitions in a correct order."""
 
     def mount(self):
         """Mount all disk partitions in a correct order."""
 

-        output("Mounting image...", False)
+        self.out.output("Mounting image...", 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.

@@ -240,40 +258,93 @@ class DiskDevice(object):
             try:
                 self.g.mount(dev, mp)
             except RuntimeError as msg:
             try:
                 self.g.mount(dev, mp)
             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)
+
+        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
+
+            self.meta['SIZE'] = last_part['part_end'] + 1
+            break
+
+        if not re.match("ext[234]", fstype):
+            self.out.warn("Don't know how to resize %s partitions." % fstype)
+            return self.meta['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)
 

@@ -283,29 +354,60 @@ 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
+        sector_size = self.g.blockdev_getss(self.guestfs_device)
+        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'])

 
 

-        new_size = (end + 1) * sector_size
-        success("new image size is %dMB" %
-                            ((new_size + 2 ** 20 - 1) // 2 ** 20))
-        return new_size
+            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'])

 
 

-    def size(self):
-        """Returns the "payload" size of the device.
+            if self.meta['PARTITION_TABLE'] == 'msdos':
+                part_set_id(last_part['part_num'], last_part['id'])

 
 

-        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]
+        new_size = (end + 1) * sector_size
+        self.out.success("new size is %dMB" % ((new_size + MB - 1) // MB))
+
+        if self.meta['PARTITION_TABLE'] == 'gpt':
+            ptable = GPTPartitionTable(self.real_device)
+            self.meta['SIZE'] = ptable.shrink(new_size)
+        else:
+            self.meta['SIZE'] = new_size

 
 

-        return last['part_end'] + 1
+        return self.meta['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.

@@ -313,31 +415,23 @@ 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: ", 'mb')
-        progressbar.max = progress_size
-        source = open(self.device, "r")
-        try:
-            dest = open(outfile, "w")
-            try:
+        MB = 2 ** 20
+        blocksize = 4 * MB  # 4MB
+        size = self.meta['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

-                    progressbar.goto((size - left) // 2 ** 20)
-            finally:
-                dest.close()
-        finally:
-            source.close()
-        
-        output("\rDumping image file...\033[K", False)
-        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 :