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