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