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