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