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