Move get_os_class from image_creator to os_type

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

diff --git a/image_creator/disk.py b/image_creator/disk.py
index 1c9d04e..0a4f222 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

@@ -120,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

@@ -165,12 +164,15 @@ 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.size = None
+        self.parttype = None

 
         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)

@@ -179,16 +181,17 @@ class DiskDevice(object):
 
     def enable(self):
         """Enable a newly created DiskDevice"""
 
     def enable(self):
         """Enable a newly created DiskDevice"""

-        new_progress = progress("Launching helper VM: ")
-        self.progressbar = new_progress()
-        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)

@@ -199,9 +202,13 @@ 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.size = self.g.blockdev_getsize64(self.guestfs_device)
+        self.parttype = 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."""

@@ -217,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."""

@@ -255,26 +259,26 @@ class DiskDevice(object):
         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.
         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.

+
+        ATTENTION: make sure unmount is called before shrink

         """
         output("Shrinking image (this may take a while)...", False)
 
         """
         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':
+        if self.parttype not in 'msdos' 'gpt':

             raise FatalError("You have a %s partition table. "
             raise FatalError("You have a %s partition table. "

-                "Only msdos partitions are supported" % parttype)
+                "Only msdos and gpt partitions are supported" % self.parttype)

 
 

-        last_partition = self.g.part_list(dev)[-1]
+        last_partition = self.g.part_list(self.guestfs_device)[-1]

 
 

-        if last_partition['part_num'] > 4:
+        if self.parttype == 'msdos' and last_partition['part_num'] > 4:

             raise FatalError("This disk contains logical partitions. "
             raise FatalError("This disk contains logical partitions. "

-                "Only primary partitions are supported.")
+                                    "Only primary partitions are supported.")

 
 

-        part_dev = "%s%d" % (dev, last_partition['part_num'])
+        part_dev = "%s%d" % (self.guestfs_device, last_partition['part_num'])

         fs_type = self.g.vfs_type(part_dev)
         if not re.match("ext[234]", fs_type):
         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
+            warn("Don't know how to resize %s partitions." % fs_type)
+            return self.size

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

@@ -285,29 +289,22 @@ 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)
+        sector_size = self.g.blockdev_getss(self.guestfs_device)

 
         start = last_partition['part_start'] / sector_size
         end = start + (block_size * block_cnt) / sector_size - 1
 
 
         start = last_partition['part_start'] / sector_size
         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)
-
-        new_size = (end + 1) * sector_size
-        success("new image size is %dMB" %
-                            ((new_size + 2 ** 20 - 1) // 2 ** 20))
-        return new_size
+        self.g.part_del(self.guestfs_device, last_partition['part_num'])
+        self.g.part_add(self.guestfs_device, 'p', start, end)

 
 

-    def size(self):
-        """Returns the "payload" size of the device.
+        self.size = (end + 1) * sector_size
+        success("new size is %dMB" % ((self.size + 2 ** 20 - 1) // 2 ** 20))

 
 

-        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.parttype == 'gpt':
+            ptable = GPTPartitionTable(self.real_device)
+            self.size = ptable.shrink(self.size)

 
 

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

@@ -316,30 +313,22 @@ class DiskDevice(object):
         partition table. Empty space in the end of the device will be ignored.
         """
         blocksize = 2 ** 22  # 4MB
         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
-        new_progress = progress("Dumping image file: ")
-        progressbar = new_progress(progress_size)
-        source = open(self.device, "r")
-        try:
-            dest = open(outfile, "w")
-            try:
-                left = size
+        progress_size = (self.size + 2 ** 20 - 1) // 2 ** 20  # 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 = self.size

                 offset = 0
                 progressbar.next()
                 while left > 0:
                     length = min(left, blocksize)
                 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((self.size - left) // 2 ** 20)
+        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 :