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