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