Add remove_swap_entry sysprep in linux

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

diff --git a/image_creator/disk.py b/image_creator/disk.py
index be728a2..6632b3e 100644 (file)
--- a/image_creator/disk.py
+++ b/image_creator/disk.py
@@ -33,7 +33,7 @@
 
 from image_creator.util import get_command
 from image_creator.util import warn, progress, success, output, FatalError
 
 from image_creator.util import get_command
 from image_creator.util import warn, progress, success, output, FatalError

-
+from image_creator.gpt import GPTPartitionTable

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

@@ -93,11 +93,9 @@ class Disk(object):
             job, args = self._cleanup_jobs.pop()
             job(*args)
 
             job, args = self._cleanup_jobs.pop()
             job(*args)
 

-    def get_device(self):
-        """Returns a newly created DiskDevice instance.
-
-        This instance is 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)
         """
 
         output("Examining source media `%s'..." % self.source, False)

@@ -122,8 +120,7 @@ class Disk(object):
         os.close(cowfd)
         self._add_cleanup(os.unlink, cow)
         # Create 1G cow sparse file
         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

@@ -141,7 +138,12 @@ class Disk(object):
         finally:
             os.unlink(table)
         success('done')
         finally:
             os.unlink(table)
         success('done')

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

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

@@ -162,12 +164,14 @@ class DiskDevice(object):
     def __init__(self, device, bootable=True):
         """Create a new DiskDevice."""
 
     def __init__(self, device, bootable=True):
         """Create a new DiskDevice."""
 

-        self.device = device
+        self.real_device = device

         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)

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

@@ -176,15 +180,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: ")
-        self.progressbar.next()
+        self.progressbar = progress("Launching helper VM: ", "percent")
+        self.progressbar.max = 100
+        self.progressbar.goto(1)

         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:
         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:

-            self.progressbar.send(100)
+            output("\rLaunching helper VM...\033[K", False)
+            success("done")

             self.progressbar = None
 
         output('Inspecting Operating System...', False)
             self.progressbar = None
 
         output('Inspecting Operating System...', False)

@@ -195,9 +201,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)
+        success('found a(n) %s system' % self.distro)

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

@@ -213,10 +224,7 @@ class DiskDevice(object):
         position = array[2]
         total = array[3]
 
         position = array[2]
         total = array[3]
 

-        self.progressbar.send((position * 100) // total)
-
-        if position == total:
-            self.progressbar = None
+        self.progressbar.goto((position * 100) // total)

 
     def mount(self):
         """Mount all disk partitions in a correct order."""
 
     def mount(self):
         """Mount all disk partitions in a correct order."""

@@ -245,65 +253,149 @@ class DiskDevice(object):
         """Umount all mounted filesystems."""
         self.g.umount_all()
 
         """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.")
+        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

 
 

-        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
+        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):
+            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)
 

-        output = self.g.tune2fs_l(part_dev)
+        out = self.g.tune2fs_l(part_dev)

         block_size = int(
         block_size = int(

-            filter(lambda x: x[0] == 'Block size', output)[0][1])
+            filter(lambda x: x[0] == 'Block size', out)[0][1])

         block_cnt = int(
         block_cnt = int(

-            filter(lambda x: x[0] == 'Block count', output)[0][1])
-
-        sector_size = self.g.blockdev_getss(dev)
+            filter(lambda x: x[0] == 'Block count', out)[0][1])

 
 

-        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
+        success("new size is %dMB" % ((new_size + MB - 1) // MB))

 
 

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

@@ -311,31 +403,25 @@ 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.meta['SIZE']
+        progress_size = (size + MB - 1) // MB  # in MB
+        progressbar = progress("Dumping image file: ", 'mb')
+        progressbar.max = progress_size
+
+        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)
+        output("\rDumping image file...\033[K", False)
+        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 :