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