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