Support non-raw images

Use qemu-nbd to export a QEMU image as raw where needed

image_creator/dialog_menu.py

@@ -153,11 +153,12 @@ def upload_image(session):
 
             try:
                 # Upload image file
-                with open(image.device, 'rb') as f:
-                    session["pithos_uri"] = \
-                        kamaki.upload(f, image.size, filename,
-                                      "Calculating block hashes",
-                                      "Uploading missing blocks")
+                with image.raw_device() as raw:
+                    with open(raw, 'rb') as f:
+                        session["pithos_uri"] = \
+                            kamaki.upload(f, image.size, filename,
+                                          "Calculating block hashes",
+                                          "Uploading missing blocks")
                 # Upload md5sum file
                 out.output("Uploading md5sum file ...")
                 md5str = "%s %s\n" % (session['checksum'], filename)

image_creator/dialog_wizard.py

@@ -472,10 +472,11 @@ def create_image(session, answers):
 
             name = "%s-%s.diskdump" % (answers['ImageName'],
                                        time.strftime("%Y%m%d%H%M"))
-            with open(image.device, 'rb') as device:
-                remote = kamaki.upload(device, image.size, name,
-                                       "(1/3)  Calculating block hashes",
-                                       "(2/3)  Uploading image blocks")
+            with image.raw_device() as raw:
+                with open(raw, 'rb') as device:
+                    remote = kamaki.upload(device, image.size, name,
+                                           "(1/3)  Calculating block hashes",
+                                           "(2/3)  Uploading image blocks")
 
             image.out.output("(3/3)  Uploading md5sum file ...", False)
             md5sumstr = '%s %s\n' % (session['checksum'], name)

image_creator/image.py

@@ -15,7 +15,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from image_creator.util import FatalError
+from image_creator.util import FatalError, QemuNBD, image_info
 from image_creator.gpt import GPTPartitionTable
 from image_creator.os_type import os_cls
 
@@ -33,6 +33,7 @@ class Image(object):
 
         self.device = device
         self.out = output
+        self.info = image_info(device)
 
         self.meta = kargs['meta'] if 'meta' in kargs else {}
         self.sysprep_params = \
@@ -46,6 +47,9 @@ class Image(object):
         self.guestfs_enabled = False
         self.guestfs_version = self.g.version()
 
+        # This is needed if the image format is not raw
+        self.nbd = QemuNBD(device)
+
     def check_guestfs_version(self, major, minor, release):
         """Checks if the version of the used libguestfs is smaller, equal or
         greater than the one specified by the major, minor and release triplet
@@ -125,7 +129,7 @@ class Image(object):
         if self.check_guestfs_version(1, 18, 4) < 0:
             self.g = guestfs.GuestFS()
 
-        self.g.add_drive_opts(self.device, readonly=0, format="raw")
+        self.g.add_drive_opts(self.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
@@ -202,6 +206,30 @@ class Image(object):
 
         return self._os
 
+    def raw_device(self, readonly=True):
+        """Returns a context manager that exports the raw image device. If
+        readonly is true, the block device that is returned is read only.
+        """
+
+        if self.guestfs_enabled:
+            self.g.umount_all()
+            self.g.sync()
+            self.g.drop_caches(3)  # drop everything
+
+        # Self gets overwritten
+        img = self
+
+        class RawImage:
+            """The RawImage context manager"""
+            def __enter__(self):
+                return img.device if img.info['format'] == 'raw' else \
+                    img.nbd.connect(readonly)
+            def __exit__(self, exc_type, exc_value, traceback):
+                if img.info['format'] != 'raw':
+                    img.nbd.disconnect()
+
+        return RawImage()
+
     def destroy(self):
         """Destroy this Image instance."""
 
@@ -374,8 +402,9 @@ class Image(object):
         assert (new_size <= self.size)
 
         if self.meta['PARTITION_TABLE'] == 'gpt':
-            ptable = GPTPartitionTable(self.device)
-            self.size = ptable.shrink(new_size, self.size)
+            with self.raw_device(readonly=False) as raw:
+                ptable = GPTPartitionTable(raw)
+                self.size = ptable.shrink(new_size, self.size)
         else:
             self.size = min(new_size + 2048 * sector_size, self.size)
 
@@ -396,25 +425,28 @@ class Image(object):
         progr_size = (self.size + MB - 1) // MB  # in MB
         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 
-        with open(self.device, 'rb') as src:
-            with open(outfile, "wb") as dst:
-                left = self.size
-                offset = 0
-                progressbar.next()
-                while left > 0:
-                    length = min(left, blocksize)
-                    sent = sendfile(dst.fileno(), src.fileno(), offset, length)
+        with self.raw_device() as raw:
+            with open(raw, 'rb') as src:
+                with open(outfile, "wb") as dst:
+                    left = self.size
+                    offset = 0
+                    progressbar.next()
+                    while left > 0:
+                        length = min(left, blocksize)
+                        sent = sendfile(dst.fileno(), src.fileno(), offset,
+                                        length)
+
+                        # Workaround for python-sendfile API change. In
+                        # python-sendfile 1.2.x (py-sendfile) the returning
+                        # value of sendfile is a tuple, where in version 2.x
+                        # (pysendfile) it is just a single integer.
+                        if isinstance(sent, tuple):
+                            sent = sent[1]
+
+                        offset += sent
+                        left -= sent
+                        progressbar.goto((self.size - left) // MB)
 
-                    # Workaround for python-sendfile API change. In
-                    # python-sendfile 1.2.x (py-sendfile) the returning value
-                    # of sendfile is a tuple, where in version 2.x (pysendfile)
-                    # it is just a single integer.
-                    if isinstance(sent, tuple):
-                        sent = sent[1]
-
-                    offset += sent
-                    left -= sent
-                    progressbar.goto((self.size - left) // MB)
         progressbar.success('image file %s was successfully created' % outfile)
 
     def md5(self):
@@ -426,14 +458,15 @@ class Image(object):
         progressbar = self.out.Progress(progr_size, "Calculating md5sum", 'mb')
         md5 = hashlib.md5()
 
-        with open(self.device, "rb") as src:
-            left = self.size
-            while left > 0:
-                length = min(left, blocksize)
-                data = src.read(length)
-                md5.update(data)
-                left -= length
-                progressbar.goto((self.size - left) // MB)
+        with self.raw_device() as raw:
+            with open(raw, "rb") as src:
+                left = self.size
+                while left > 0:
+                    length = min(left, blocksize)
+                    data = src.read(length)
+                    md5.update(data)
+                    left -= length
+                    progressbar.goto((self.size - left) // MB)
 
         checksum = md5.hexdigest()
         progressbar.success(checksum)

image_creator/main.py

@@ -334,19 +334,17 @@ def image_creator():
                                      os.path.basename(options.outfile)))
             out.success('done')
 
-        # Destroy the image instance. We only need the disk device from now on
-        disk.destroy_image(image)
-
         out.output()
         try:
-            uploaded_obj = ""
             if options.upload:
                 out.output("Uploading image to the storage service:")
-                with open(device, 'rb') as f:
-                    uploaded_obj = kamaki.upload(
-                        f, image.size, options.upload,
-                        "(1/3)  Calculating block hashes",
-                        "(2/3)  Uploading missing blocks")
+                with image.raw_device() as raw:
+                    with open(raw, 'rb') as f:
+                        remote = kamaki.upload(
+                            f, image.size, options.upload,
+                            "(1/3)  Calculating block hashes",
+                            "(2/3)  Uploading missing blocks")
+
                 out.output("(3/3)  Uploading md5sum file ...", False)
                 md5sumstr = '%s %s\n' % (checksum,
                                          os.path.basename(options.upload))
@@ -360,7 +358,7 @@ def image_creator():
                 img_type = 'public' if options.public else 'private'
                 out.output('Registering %s image with the compute service ...'
                            % img_type, False)
-                result = kamaki.register(options.register, uploaded_obj,
+                result = kamaki.register(options.register, remote,
                                          metadata, options.public)
                 out.success('done')
                 out.output("Uploading metadata file ...", False)

image_creator/os_type/__init__.py

@@ -420,7 +420,7 @@ class OSBase(object):
 
         self.out.output()
 
-    @sysprep('Shrinking image', nomount=True)
+    @sysprep('Shrinking image (may take a while)', nomount=True)
     def _shrink(self):
         """Shrink the last file system and update the partition table"""
         self.image.shrink()

image_creator/util.py

@@ -25,6 +25,8 @@ import os
 import re
 import json
 from sh import qemu_img
+from sh import qemu_nbd
+from sh import modprobe
 
 
 class FatalError(Exception):
@@ -33,6 +35,7 @@ class FatalError(Exception):
 
 
 def image_info(image):
+    """Returns information about an image file"""
     info = qemu_img('info', '--output', 'json', image)
     return json.loads(str(info))
 
@@ -108,4 +111,59 @@ def virtio_versions(virtio_state):
 
     return ret
 
+
+class QemuNBD(object):
+    """Wrapper class for the qemu-nbd tool"""
+
+    def __init__(self, image):
+        """Initialize an instance"""
+        self.image = image
+        self.device = None
+        self.pattern = re.compile('^nbd\d+$')
+
+    def _list_devices(self):
+        """Returns all the NBD block devices"""
+        return set([d for d in os.listdir('/dev/') if self.pattern.match(d)])
+
+    def connect(self, ro=True):
+        """Connect the image to a free NBD device"""
+        devs = self._list_devices()
+
+        if len(devs) == 0:  # Is nbd module loaded?
+            modprobe('nbd', 'max_part=16')
+            # Wait a second for /dev to be populated
+            time.sleep(1)
+            devs = self._list_devices()
+            if len(devs) == 0:
+                raise FatalError("/dev/nbd* devices not present!")
+
+        # Ignore the nbd block devices that are in use
+        with open('/proc/partitions') as partitions:
+            for line in iter(partitions):
+                entry = line.split()
+                if len(entry) != 4:
+                    continue
+                if entry[3] in devs:
+                    devs.remove(entry[3])
+
+        if len(devs) == 0:
+            raise FatalError("All NBD block devices are busy!")
+
+        device = '/dev/%s' % devs.pop()
+        args = ['-c', device]
+        if ro:
+            args.append('-r')
+        args.append(self.image)
+
+        qemu_nbd(*args)
+        self.device = device
+        return device
+
+    def disconnect(self):
+        """Disconnect the image from the connected device"""
+        assert self.device is not None, "No device connected"
+
+        qemu_nbd('-d', self.device)
+        self.device = None
+
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :