code.grnet.gr Git - snf-image-creator/blob - image_creator/disk.py

   1 # Copyright 2012 GRNET S.A. All rights reserved.
   2 #
   3 # Redistribution and use in source and binary forms, with or
   4 # without modification, are permitted provided that the following
   5 # conditions are met:
   6 #
   7 #   1. Redistributions of source code must retain the above
   8 #      copyright notice, this list of conditions and the following
   9 #      disclaimer.
  10 #
  11 #   2. Redistributions in binary form must reproduce the above
  12 #      copyright notice, this list of conditions and the following
  13 #      disclaimer in the documentation and/or other materials
  14 #      provided with the distribution.
  15 #
  16 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
  17 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  18 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  19 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
  20 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  21 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  22 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  23 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
  24 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  25 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  26 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  27 # POSSIBILITY OF SUCH DAMAGE.
  28 #
  29 # The views and conclusions contained in the software and
  30 # documentation are those of the authors and should not be
  31 # interpreted as representing official policies, either expressed
  32 # or implied, of GRNET S.A.
  33
  34 from image_creator.util import get_command
  35 from image_creator.util import FatalError
  36 from image_creator.gpt import GPTPartitionTable
  37
  38 import image_creator.bundle_volume
  39
  40 import stat
  41 import os
  42 import tempfile
  43 import uuid
  44 import re
  45 import sys
  46 import guestfs
  47 import time
  48 from sendfile import sendfile
  49
  50
  51 dd = get_command('dd')
  52 dmsetup = get_command('dmsetup')
  53 losetup = get_command('losetup')
  54 blockdev = get_command('blockdev')
  55
  56
  57 class Disk(object):
  58     """This class represents a hard disk hosting an Operating System
  59
  60     A Disk instance never alters the source media it is created from.
  61     Any change is done on a snapshot created by the device-mapper of
  62     the Linux kernel.
  63     """
  64
  65     def __init__(self, source, output):
  66         """Create a new Disk instance out of a source media. The source
  67         media can be an image file, a block device or a directory."""
  68         self._cleanup_jobs = []
  69         self._devices = []
  70         self.source = source
  71         self.out = output
  72         self.meta = {}
  73
  74     def _add_cleanup(self, job, *args):
  75         self._cleanup_jobs.append((job, args))
  76
  77     def _losetup(self, fname):
  78         loop = losetup('-f', '--show', fname)
  79         loop = loop.strip()  # remove the new-line char
  80         self._add_cleanup(losetup, '-d', loop)
  81         return loop
  82
  83     def _map_partition(self, dev, index, start, end):
  84         name = "%sp%d" % (os.path.basename(dev), index)
  85         tablefd, table = tempfile.mkstemp()
  86         try:
  87             size = end - start + 1
  88             os.write(tablefd, "0 %d linear %s %d" % (start, dev, size))
  89             dmsetup('create', name, table)
  90         finally:
  91             os.unlink(table)
  92
  93     def _unmap_partition(self, dev, index):
  94         name = "%sp%d" % (os.path.basename(dev), index)
  95         if not os.path.exists("/dev/mapper/%s" % name):
  96             return
  97
  98         dmsetup('remove', name)
  99         time.sleep(0.5)
 100
 101     def _dir_to_disk(self):
 102         if self.source == '/':
 103             bundle = BundleVolume(self.out, self.meta)
 104             return _losetup(bundle.create_image())
 105         raise FatalError("Using a directory as media source is supported")
 106
 107     def cleanup(self):
 108         """Cleanup internal data. This needs to be called before the
 109         program ends.
 110         """
 111         try:
 112             while len(self._devices):
 113                 device = self._devices.pop()
 114                 device.destroy()
 115         finally:
 116             # Make sure those are executed even if one of the device.destroy
 117             # methods throws exeptions.
 118             while len(self._cleanup_jobs):
 119                 job, args = self._cleanup_jobs.pop()
 120                 job(*args)
 121
 122     def snapshot(self):
 123         """Creates a snapshot of the original source media of the Disk
 124         instance.
 125         """
 126
 127         self.out.output("Examining source media `%s'..." % self.source, False)
 128         sourcedev = self.source
 129         mode = os.stat(self.source).st_mode
 130         if stat.S_ISDIR(mode):
 131             self.out.success('looks like a directory')
 132             return self._dir_to_disk()
 133         elif stat.S_ISREG(mode):
 134             self.out.success('looks like an image file')
 135             sourcedev = self._losetup(self.source)
 136         elif not stat.S_ISBLK(mode):
 137             raise ValueError("Invalid media source. Only block devices, "
 138                              "regular files and directories are supported.")
 139         else:
 140             self.out.success('looks like a block device')
 141
 142         # Take a snapshot and return it to the user
 143         self.out.output("Snapshotting media source...", False)
 144         size = blockdev('--getsz', sourcedev)
 145         cowfd, cow = tempfile.mkstemp()
 146         os.close(cowfd)
 147         self._add_cleanup(os.unlink, cow)
 148         # Create cow sparse file
 149         dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
 150         cowdev = self._losetup(cow)
 151
 152         snapshot = uuid.uuid4().hex
 153         tablefd, table = tempfile.mkstemp()
 154         try:
 155             os.write(tablefd, "0 %d snapshot %s %s n 8" %
 156                               (int(size), sourcedev, cowdev))
 157             dmsetup('create', snapshot, table)
 158             self._add_cleanup(dmsetup, 'remove', snapshot)
 159             # Sometimes dmsetup remove fails with Device or resource busy,
 160             # although everything is cleaned up and the snapshot is not
 161             # used by anyone. Add a 2 seconds delay to be on the safe side.
 162             self._add_cleanup(time.sleep, 2)
 163
 164         finally:
 165             os.unlink(table)
 166         self.out.success('done')
 167         return "/dev/mapper/%s" % snapshot
 168
 169     def get_device(self, media):
 170         """Returns a newly created DiskDevice instance."""
 171
 172         new_device = DiskDevice(media, self.out)
 173         self._devices.append(new_device)
 174         new_device.enable()
 175         return new_device
 176
 177     def destroy_device(self, device):
 178         """Destroys a DiskDevice instance previously created by
 179         get_device method.
 180         """
 181         self._devices.remove(device)
 182         device.destroy()
 183
 184
 185 class DiskDevice(object):
 186     """This class represents a block device hosting an Operating System
 187     as created by the device-mapper.
 188     """
 189
 190     def __init__(self, device, output, bootable=True, meta={}):
 191         """Create a new DiskDevice."""
 192
 193         self.real_device = device
 194         self.out = output
 195         self.bootable = bootable
 196         self.meta = meta
 197         self.progress_bar = None
 198         self.guestfs_device = None
 199         self.size = 0
 200
 201         self.g = guestfs.GuestFS()
 202         self.g.add_drive_opts(self.real_device, readonly=0)
 203
 204         # Before version 1.17.14 the recovery process, which is a fork of the
 205         # original process that called libguestfs, did not close its inherited
 206         # file descriptors. This can cause problems especially if the parent
 207         # process has opened pipes. Since the recovery process is an optional
 208         # feature of libguestfs, it's better to disable it.
 209         self.g.set_recovery_proc(0)
 210         version = self.g.version()
 211         if version['major'] > 1 or \
 212             (version['major'] == 1 and (version['minor'] >= 18 or
 213                                         (version['minor'] == 17 and
 214                                          version['release'] >= 14))):
 215             self.g.set_recovery_proc(1)
 216             self.out.output("Enabling recovery proc")
 217
 218         #self.g.set_trace(1)
 219         #self.g.set_verbose(1)
 220
 221         self.guestfs_enabled = False
 222
 223     def enable(self):
 224         """Enable a newly created DiskDevice"""
 225         self.progressbar = self.out.Progress(100, "Launching helper VM",
 226                                              "percent")
 227         eh = self.g.set_event_callback(self.progress_callback,
 228                                        guestfs.EVENT_PROGRESS)
 229         self.g.launch()
 230         self.guestfs_enabled = True
 231         self.g.delete_event_callback(eh)
 232         self.progressbar.success('done')
 233         self.progressbar = None
 234
 235         self.out.output('Inspecting Operating System...', False)
 236         roots = self.g.inspect_os()
 237         if len(roots) == 0:
 238             raise FatalError("No operating system found")
 239         if len(roots) > 1:
 240             raise FatalError("Multiple operating systems found."
 241                              "We only support images with one OS.")
 242         self.root = roots[0]
 243         self.guestfs_device = self.g.part_to_dev(self.root)
 244         self.size = self.g.blockdev_getsize64(self.guestfs_device)
 245         self.meta['PARTITION_TABLE'] = \
 246             self.g.part_get_parttype(self.guestfs_device)
 247
 248         self.ostype = self.g.inspect_get_type(self.root)
 249         self.distro = self.g.inspect_get_distro(self.root)
 250         self.out.success('found a(n) %s system' % self.distro)
 251
 252     def destroy(self):
 253         """Destroy this DiskDevice instance."""
 254
 255         # In new guestfs versions, there is a handy shutdown method for this
 256         try:
 257             if self.guestfs_enabled:
 258                 self.g.umount_all()
 259                 self.g.sync()
 260         finally:
 261             # Close the guestfs handler if open
 262             self.g.close()
 263
 264     def progress_callback(self, ev, eh, buf, array):
 265         position = array[2]
 266         total = array[3]
 267
 268         self.progressbar.goto((position * 100) // total)
 269
 270     def mount(self, readonly=False):
 271         """Mount all disk partitions in a correct order."""
 272
 273         mount = self.g.mount_ro if readonly else self.g.mount
 274         msg = " read-only" if readonly else ""
 275         self.out.output("Mounting the media%s..." % msg, False)
 276         mps = self.g.inspect_get_mountpoints(self.root)
 277
 278         # Sort the keys to mount the fs in a correct order.
 279         # / should be mounted befor /boot, etc
 280         def compare(a, b):
 281             if len(a[0]) > len(b[0]):
 282                 return 1
 283             elif len(a[0]) == len(b[0]):
 284                 return 0
 285             else:
 286                 return -1
 287         mps.sort(compare)
 288         for mp, dev in mps:
 289             try:
 290                 mount(dev, mp)
 291             except RuntimeError as msg:
 292                 self.out.warn("%s (ignored)" % msg)
 293         self.out.success("done")
 294
 295     def umount(self):
 296         """Umount all mounted filesystems."""
 297         self.g.umount_all()
 298
 299     def _last_partition(self):
 300         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 301             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 302                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 303             raise FatalError(msg)
 304
 305         is_extended = lambda p: \
 306             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
 307         is_logical = lambda p: \
 308             self.meta['PARTITION_TABLE'] != 'msdos' and p['part_num'] > 4
 309
 310         partitions = self.g.part_list(self.guestfs_device)
 311         last_partition = partitions[-1]
 312
 313         if is_logical(last_partition):
 314             # The disk contains extended and logical partitions....
 315             extended = [p for p in partitions if is_extended(p)][0]
 316             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 317
 318             # check if extended is the last primary partition
 319             if last_primary['part_num'] > extended['part_num']:
 320                 last_partition = last_primary
 321
 322         return last_partition
 323
 324     def shrink(self):
 325         """Shrink the disk.
 326
 327         This is accomplished by shrinking the last filesystem in the
 328         disk and then updating the partition table. The new disk size
 329         (in bytes) is returned.
 330
 331         ATTENTION: make sure unmount is called before shrink
 332         """
 333         get_fstype = lambda p: \
 334             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 335         is_logical = lambda p: \
 336             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 337         is_extended = lambda p: \
 338             self.meta['PARTITION_TABLE'] == 'msdos' and \
 339             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
 340
 341         part_add = lambda ptype, start, stop: \
 342             self.g.part_add(self.guestfs_device, ptype, start, stop)
 343         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 344         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 345         part_set_id = lambda p, id: \
 346             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 347         part_get_bootable = lambda p: \
 348             self.g.part_get_bootable(self.guestfs_device, p)
 349         part_set_bootable = lambda p, bootable: \
 350             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 351
 352         MB = 2 ** 20
 353
 354         self.out.output("Shrinking image (this may take a while)...", False)
 355
 356         sector_size = self.g.blockdev_getss(self.guestfs_device)
 357
 358         last_part = None
 359         fstype = None
 360         while True:
 361             last_part = self._last_partition()
 362             fstype = get_fstype(last_part)
 363
 364             if fstype == 'swap':
 365                 self.meta['SWAP'] = "%d:%s" % \
 366                     (last_part['part_num'],
 367                      (last_part['part_size'] + MB - 1) // MB)
 368                 part_del(last_part['part_num'])
 369                 continue
 370             elif is_extended(last_part):
 371                 part_del(last_part['part_num'])
 372                 continue
 373
 374             # Most disk manipulation programs leave 2048 sectors after the last
 375             # partition
 376             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 377             self.size = min(self.size, new_size)
 378             break
 379
 380         if not re.match("ext[234]", fstype):
 381             self.out.warn("Don't know how to resize %s partitions." % fstype)
 382             return self.size
 383
 384         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 385         self.g.e2fsck_f(part_dev)
 386         self.g.resize2fs_M(part_dev)
 387
 388         out = self.g.tune2fs_l(part_dev)
 389         block_size = int(
 390             filter(lambda x: x[0] == 'Block size', out)[0][1])
 391         block_cnt = int(
 392             filter(lambda x: x[0] == 'Block count', out)[0][1])
 393
 394         start = last_part['part_start'] / sector_size
 395         end = start + (block_size * block_cnt) / sector_size - 1
 396
 397         if is_logical(last_part):
 398             partitions = self.g.part_list(self.guestfs_device)
 399
 400             logical = []  # logical partitions
 401             for partition in partitions:
 402                 if partition['part_num'] < 4:
 403                     continue
 404                 logical.append({
 405                     'num': partition['part_num'],
 406                     'start': partition['part_start'] / sector_size,
 407                     'end': partition['part_end'] / sector_size,
 408                     'id': part_get_(partition['part_num']),
 409                     'bootable': part_get_bootable(partition['part_num'])
 410                 })
 411
 412             logical[-1]['end'] = end  # new end after resize
 413
 414             # Recreate the extended partition
 415             extended = [p for p in partitions if self._is_extended(p)][0]
 416             part_del(extended['part_num'])
 417             part_add('e', extended['part_start'], end)
 418
 419             # Create all the logical partitions back
 420             for l in logical:
 421                 part_add('l', l['start'], l['end'])
 422                 part_set_id(l['num'], l['id'])
 423                 part_set_bootable(l['num'], l['bootable'])
 424         else:
 425             # Recreate the last partition
 426             if self.meta['PARTITION_TABLE'] == 'msdos':
 427                 last_part['id'] = part_get_id(last_part['part_num'])
 428
 429             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 430             part_del(last_part['part_num'])
 431             part_add('p', start, end)
 432             part_set_bootable(last_part['part_num'], last_part['bootable'])
 433
 434             if self.meta['PARTITION_TABLE'] == 'msdos':
 435                 part_set_id(last_part['part_num'], last_part['id'])
 436
 437         new_size = (end + 1) * sector_size
 438
 439         assert (new_size <= self.size)
 440
 441         if self.meta['PARTITION_TABLE'] == 'gpt':
 442             ptable = GPTPartitionTable(self.real_device)
 443             self.size = ptable.shrink(new_size, self.size)
 444         else:
 445             self.size = min(new_size + 2048 * sector_size, self.size)
 446
 447         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 448
 449         return self.size
 450
 451     def dump(self, outfile):
 452         """Dumps the content of device into a file.
 453
 454         This method will only dump the actual payload, found by reading the
 455         partition table. Empty space in the end of the device will be ignored.
 456         """
 457         MB = 2 ** 20
 458         blocksize = 4 * MB  # 4MB
 459         size = self.size
 460         progr_size = (size + MB - 1) // MB  # in MB
 461         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 462
 463         with open(self.real_device, 'r') as src:
 464             with open(outfile, "w") as dst:
 465                 left = size
 466                 offset = 0
 467                 progressbar.next()
 468                 while left > 0:
 469                     length = min(left, blocksize)
 470                     _, sent = sendfile(dst.fileno(), src.fileno(), offset,
 471                         length)
 472                     offset += sent
 473                     left -= sent
 474                     progressbar.goto((size - left) // MB)
 475         progressbar.success('image file %s was successfully created' % outfile)
 476
 477 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :