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