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.util import try_fail_repeat
  37 from image_creator.util import free_space
  38 from image_creator.gpt import GPTPartitionTable
  39 from image_creator.bundle_volume import BundleVolume
  40
  41 import stat
  42 import os
  43 import tempfile
  44 import uuid
  45 import re
  46 import sys
  47 import guestfs
  48 import shutil
  49 from sendfile import sendfile
  50
  51
  52 dd = get_command('dd')
  53 dmsetup = get_command('dmsetup')
  54 losetup = get_command('losetup')
  55 blockdev = get_command('blockdev')
  56
  57
  58 TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']
  59
  60
  61 class Disk(object):
  62     """This class represents a hard disk hosting an Operating System
  63
  64     A Disk instance never alters the source media it is created from.
  65     Any change is done on a snapshot created by the device-mapper of
  66     the Linux kernel.
  67     """
  68
  69     def __init__(self, source, output, tmp=None):
  70         """Create a new Disk instance out of a source media. The source
  71         media can be an image file, a block device or a directory.
  72         """
  73         self._cleanup_jobs = []
  74         self._devices = []
  75         self.source = source
  76         self.out = output
  77         self.meta = {}
  78         self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.',
  79                                     dir=self._get_tmp_dir(tmp))
  80
  81         self._add_cleanup(shutil.rmtree, self.tmp)
  82
  83     def _get_tmp_dir(self, default=None):
  84         if default is not None:
  85             return default
  86
  87         space = map(free_space, TMP_CANDIDATES)
  88
  89         max_idx = 0
  90         max_val = space[0]
  91         for i, val in zip(range(len(space)), space):
  92             if val > max_val:
  93                 max_val = val
  94                 max_idx = i
  95
  96         # Return the candidate path with more available space
  97         return TMP_CANDIDATES[max_idx]
  98
  99     def _add_cleanup(self, job, *args):
 100         self._cleanup_jobs.append((job, args))
 101
 102     def _losetup(self, fname):
 103         loop = losetup('-f', '--show', fname)
 104         loop = loop.strip()  # remove the new-line char
 105         self._add_cleanup(try_fail_repeat, losetup, '-d', loop)
 106         return loop
 107
 108     def _dir_to_disk(self):
 109         if self.source == '/':
 110             bundle = BundleVolume(self.out, self.meta)
 111             image = '%s/%s.diskdump' % (self.tmp, uuid.uuid4().hex)
 112
 113             def check_unlink(path):
 114                 if os.path.exists(path):
 115                     os.unlink(path)
 116
 117             self._add_cleanup(check_unlink, image)
 118             bundle.create_image(image)
 119             return self._losetup(image)
 120         raise FatalError("Using a directory as media source is supported")
 121
 122     def cleanup(self):
 123         """Cleanup internal data. This needs to be called before the
 124         program ends.
 125         """
 126         try:
 127             while len(self._devices):
 128                 device = self._devices.pop()
 129                 device.destroy()
 130         finally:
 131             # Make sure those are executed even if one of the device.destroy
 132             # methods throws exeptions.
 133             while len(self._cleanup_jobs):
 134                 job, args = self._cleanup_jobs.pop()
 135                 job(*args)
 136
 137     def snapshot(self):
 138         """Creates a snapshot of the original source media of the Disk
 139         instance.
 140         """
 141
 142         self.out.output("Examining source media `%s' ..." % self.source, False)
 143         sourcedev = self.source
 144         mode = os.stat(self.source).st_mode
 145         if stat.S_ISDIR(mode):
 146             self.out.success('looks like a directory')
 147             return self._dir_to_disk()
 148         elif stat.S_ISREG(mode):
 149             self.out.success('looks like an image file')
 150             sourcedev = self._losetup(self.source)
 151         elif not stat.S_ISBLK(mode):
 152             raise ValueError("Invalid media source. Only block devices, "
 153                              "regular files and directories are supported.")
 154         else:
 155             self.out.success('looks like a block device')
 156
 157         # Take a snapshot and return it to the user
 158         self.out.output("Snapshotting media source...", False)
 159         size = blockdev('--getsz', sourcedev)
 160         cowfd, cow = tempfile.mkstemp(dir=self.tmp)
 161         os.close(cowfd)
 162         self._add_cleanup(os.unlink, cow)
 163         # Create cow sparse file
 164         dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
 165         cowdev = self._losetup(cow)
 166
 167         snapshot = uuid.uuid4().hex
 168         tablefd, table = tempfile.mkstemp()
 169         try:
 170             os.write(tablefd, "0 %d snapshot %s %s n 8" %
 171                               (int(size), sourcedev, cowdev))
 172             dmsetup('create', snapshot, table)
 173             self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
 174
 175         finally:
 176             os.unlink(table)
 177         self.out.success('done')
 178         return "/dev/mapper/%s" % snapshot
 179
 180     def get_device(self, media):
 181         """Returns a newly created DiskDevice instance."""
 182
 183         new_device = DiskDevice(media, self.out)
 184         self._devices.append(new_device)
 185         new_device.enable()
 186         return new_device
 187
 188     def destroy_device(self, device):
 189         """Destroys a DiskDevice instance previously created by
 190         get_device method.
 191         """
 192         self._devices.remove(device)
 193         device.destroy()
 194
 195
 196 class DiskDevice(object):
 197     """This class represents a block device hosting an Operating System
 198     as created by the device-mapper.
 199     """
 200
 201     def __init__(self, device, output, bootable=True, meta={}):
 202         """Create a new DiskDevice."""
 203
 204         self.real_device = device
 205         self.out = output
 206         self.bootable = bootable
 207         self.meta = meta
 208         self.progress_bar = None
 209         self.guestfs_device = None
 210         self.size = 0
 211
 212         self.g = guestfs.GuestFS()
 213         self.g.add_drive_opts(self.real_device, readonly=0)
 214
 215         # Before version 1.17.14 the recovery process, which is a fork of the
 216         # original process that called libguestfs, did not close its inherited
 217         # file descriptors. This can cause problems especially if the parent
 218         # process has opened pipes. Since the recovery process is an optional
 219         # feature of libguestfs, it's better to disable it.
 220         self.g.set_recovery_proc(0)
 221         version = self.g.version()
 222         if version['major'] > 1 or \
 223             (version['major'] == 1 and (version['minor'] >= 18 or
 224                                         (version['minor'] == 17 and
 225                                          version['release'] >= 14))):
 226             self.g.set_recovery_proc(1)
 227             self.out.output("Enabling recovery proc")
 228
 229         #self.g.set_trace(1)
 230         #self.g.set_verbose(1)
 231
 232         self.guestfs_enabled = False
 233
 234     def enable(self):
 235         """Enable a newly created DiskDevice"""
 236
 237         self.out.output('Launching helper VM (may take a while) ...', False)
 238         # self.progressbar = self.out.Progress(100, "Launching helper VM",
 239         #                                     "percent")
 240         # eh = self.g.set_event_callback(self.progress_callback,
 241         #                               guestfs.EVENT_PROGRESS)
 242         self.g.launch()
 243         self.guestfs_enabled = True
 244         # self.g.delete_event_callback(eh)
 245         # self.progressbar.success('done')
 246         # self.progressbar = None
 247         self.out.success('done')
 248
 249         self.out.output('Inspecting Operating System ...', False)
 250         roots = self.g.inspect_os()
 251         if len(roots) == 0:
 252             raise FatalError("No operating system found")
 253         if len(roots) > 1:
 254             raise FatalError("Multiple operating systems found."
 255                              "We only support images with one OS.")
 256         self.root = roots[0]
 257         self.guestfs_device = self.g.part_to_dev(self.root)
 258         self.size = self.g.blockdev_getsize64(self.guestfs_device)
 259         self.meta['PARTITION_TABLE'] = \
 260             self.g.part_get_parttype(self.guestfs_device)
 261
 262         self.ostype = self.g.inspect_get_type(self.root)
 263         self.distro = self.g.inspect_get_distro(self.root)
 264         self.out.success('found a(n) %s system' % self.distro)
 265
 266     def destroy(self):
 267         """Destroy this DiskDevice instance."""
 268
 269         # In new guestfs versions, there is a handy shutdown method for this
 270         try:
 271             if self.guestfs_enabled:
 272                 self.g.umount_all()
 273                 self.g.sync()
 274         finally:
 275             # Close the guestfs handler if open
 276             self.g.close()
 277
 278 #    def progress_callback(self, ev, eh, buf, array):
 279 #        position = array[2]
 280 #        total = array[3]
 281 #
 282 #        self.progressbar.goto((position * 100) // total)
 283
 284     def mount(self, readonly=False):
 285         """Mount all disk partitions in a correct order."""
 286
 287         mount = self.g.mount_ro if readonly else self.g.mount
 288         msg = " read-only" if readonly else ""
 289         self.out.output("Mounting the media%s ..." % msg, False)
 290         mps = self.g.inspect_get_mountpoints(self.root)
 291
 292         # Sort the keys to mount the fs in a correct order.
 293         # / should be mounted befor /boot, etc
 294         def compare(a, b):
 295             if len(a[0]) > len(b[0]):
 296                 return 1
 297             elif len(a[0]) == len(b[0]):
 298                 return 0
 299             else:
 300                 return -1
 301         mps.sort(compare)
 302         for mp, dev in mps:
 303             try:
 304                 mount(dev, mp)
 305             except RuntimeError as msg:
 306                 self.out.warn("%s (ignored)" % msg)
 307         self.out.success("done")
 308
 309     def umount(self):
 310         """Umount all mounted filesystems."""
 311         self.g.umount_all()
 312
 313     def _last_partition(self):
 314         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 315             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 316                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 317             raise FatalError(msg)
 318
 319         is_extended = lambda p: \
 320             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 321             in (0x5, 0xf)
 322         is_logical = lambda p: \
 323             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 324
 325         partitions = self.g.part_list(self.guestfs_device)
 326         last_partition = partitions[-1]
 327
 328         if is_logical(last_partition):
 329             # The disk contains extended and logical partitions....
 330             extended = filter(is_extended, partitions)[0]
 331             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 332
 333             # check if extended is the last primary partition
 334             if last_primary['part_num'] > extended['part_num']:
 335                 last_partition = last_primary
 336
 337         return last_partition
 338
 339     def shrink(self):
 340         """Shrink the disk.
 341
 342         This is accomplished by shrinking the last filesystem in the
 343         disk and then updating the partition table. The new disk size
 344         (in bytes) is returned.
 345
 346         ATTENTION: make sure unmount is called before shrink
 347         """
 348         get_fstype = lambda p: \
 349             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 350         is_logical = lambda p: \
 351             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 352         is_extended = lambda p: \
 353             self.meta['PARTITION_TABLE'] == 'msdos' and \
 354             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 355             in (0x5, 0xf)
 356
 357         part_add = lambda ptype, start, stop: \
 358             self.g.part_add(self.guestfs_device, ptype, start, stop)
 359         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 360         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 361         part_set_id = lambda p, id: \
 362             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 363         part_get_bootable = lambda p: \
 364             self.g.part_get_bootable(self.guestfs_device, p)
 365         part_set_bootable = lambda p, bootable: \
 366             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 367
 368         MB = 2 ** 20
 369
 370         self.out.output("Shrinking image (this may take a while) ...", False)
 371
 372         sector_size = self.g.blockdev_getss(self.guestfs_device)
 373
 374         last_part = None
 375         fstype = None
 376         while True:
 377             last_part = self._last_partition()
 378             fstype = get_fstype(last_part)
 379
 380             if fstype == 'swap':
 381                 self.meta['SWAP'] = "%d:%s" % \
 382                     (last_part['part_num'],
 383                      (last_part['part_size'] + MB - 1) // MB)
 384                 part_del(last_part['part_num'])
 385                 continue
 386             elif is_extended(last_part):
 387                 part_del(last_part['part_num'])
 388                 continue
 389
 390             # Most disk manipulation programs leave 2048 sectors after the last
 391             # partition
 392             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 393             self.size = min(self.size, new_size)
 394             break
 395
 396         if not re.match("ext[234]", fstype):
 397             self.out.warn("Don't know how to resize %s partitions." % fstype)
 398             return self.size
 399
 400         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 401         self.g.e2fsck_f(part_dev)
 402         self.g.resize2fs_M(part_dev)
 403
 404         out = self.g.tune2fs_l(part_dev)
 405         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
 406         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
 407
 408         start = last_part['part_start'] / sector_size
 409         end = start + (block_size * block_cnt) / sector_size - 1
 410
 411         if is_logical(last_part):
 412             partitions = self.g.part_list(self.guestfs_device)
 413
 414             logical = []  # logical partitions
 415             for partition in partitions:
 416                 if partition['part_num'] < 4:
 417                     continue
 418                 logical.append({
 419                     'num': partition['part_num'],
 420                     'start': partition['part_start'] / sector_size,
 421                     'end': partition['part_end'] / sector_size,
 422                     'id': part_get_id(partition['part_num']),
 423                     'bootable': part_get_bootable(partition['part_num'])
 424                 })
 425
 426             logical[-1]['end'] = end  # new end after resize
 427
 428             # Recreate the extended partition
 429             extended = filter(is_extended, partitions)[0]
 430             part_del(extended['part_num'])
 431             part_add('e', extended['part_start'] / sector_size, end)
 432
 433             # Create all the logical partitions back
 434             for l in logical:
 435                 part_add('l', l['start'], l['end'])
 436                 part_set_id(l['num'], l['id'])
 437                 part_set_bootable(l['num'], l['bootable'])
 438         else:
 439             # Recreate the last partition
 440             if self.meta['PARTITION_TABLE'] == 'msdos':
 441                 last_part['id'] = part_get_id(last_part['part_num'])
 442
 443             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 444             part_del(last_part['part_num'])
 445             part_add('p', start, end)
 446             part_set_bootable(last_part['part_num'], last_part['bootable'])
 447
 448             if self.meta['PARTITION_TABLE'] == 'msdos':
 449                 part_set_id(last_part['part_num'], last_part['id'])
 450
 451         new_size = (end + 1) * sector_size
 452
 453         assert (new_size <= self.size)
 454
 455         if self.meta['PARTITION_TABLE'] == 'gpt':
 456             ptable = GPTPartitionTable(self.real_device)
 457             self.size = ptable.shrink(new_size, self.size)
 458         else:
 459             self.size = min(new_size + 2048 * sector_size, self.size)
 460
 461         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 462
 463         return self.size
 464
 465     def dump(self, outfile):
 466         """Dumps the content of device into a file.
 467
 468         This method will only dump the actual payload, found by reading the
 469         partition table. Empty space in the end of the device will be ignored.
 470         """
 471         MB = 2 ** 20
 472         blocksize = 4 * MB  # 4MB
 473         size = self.size
 474         progr_size = (size + MB - 1) // MB  # in MB
 475         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 476
 477         with open(self.real_device, 'r') as src:
 478             with open(outfile, "w") as dst:
 479                 left = size
 480                 offset = 0
 481                 progressbar.next()
 482                 while left > 0:
 483                     length = min(left, blocksize)
 484                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 485
 486                     # Workaround for python-sendfile API change. In
 487                     # python-sendfile 1.2.x (py-sendfile) the returning value
 488                     # of sendfile is a tuple, where in version 2.x (pysendfile)
 489                     # it is just a sigle integer.
 490                     if isinstance(sent, tuple):
 491                         sent = sent[1]
 492
 493                     offset += sent
 494                     left -= sent
 495                     progressbar.goto((size - left) // MB)
 496         progressbar.success('image file %s was successfully created' % outfile)
 497
 498 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :