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