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('--getsize', sourcedev)
 118         cowfd, cow = tempfile.mkstemp()
 119         os.close(cowfd)
 120         self._add_cleanup(os.unlink, cow)
 121         # Create 1G cow sparse file
 122         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
 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.meta = {}
 172
 173         self.g = guestfs.GuestFS()
 174         self.g.add_drive_opts(self.real_device, readonly=0)
 175
 176         # Before version 1.17.14 the recovery process, which is a fork of the
 177         # original process that called libguestfs, did not close its inherited
 178         # file descriptors. This can cause problems especially if the parent
 179         # process has opened pipes. Since the recovery process is an optional
 180         # feature of libguestfs, it's better to disable it.
 181         self.g.set_recovery_proc(0)
 182         version = self.g.version()
 183         if version['major'] > 1 or (version['major'] == 1 and
 184             (version['minor'] >= 18 or \
 185             (version['minor'] == 17 and version['release'] >= 14))):
 186             self.g.set_recovery_proc(1)
 187             self.out.output("Enabling recovery proc")
 188
 189         #self.g.set_trace(1)
 190         #self.g.set_verbose(1)
 191
 192         self.guestfs_enabled = False
 193
 194     def enable(self):
 195         """Enable a newly created DiskDevice"""
 196         self.progressbar = self.out.Progress(100, "Launching helper VM",
 197                                              "percent")
 198         eh = self.g.set_event_callback(self.progress_callback,
 199                                                     guestfs.EVENT_PROGRESS)
 200         self.g.launch()
 201         self.guestfs_enabled = True
 202         self.g.delete_event_callback(eh)
 203         self.progressbar.success('done')
 204         self.progressbar = None
 205
 206         self.out.output('Inspecting Operating System...', False)
 207         roots = self.g.inspect_os()
 208         if len(roots) == 0:
 209             raise FatalError("No operating system found")
 210         if len(roots) > 1:
 211             raise FatalError("Multiple operating systems found."
 212                             "We only support images with one filesystem.")
 213         self.root = roots[0]
 214         self.guestfs_device = self.g.part_to_dev(self.root)
 215         self.meta['SIZE'] = self.g.blockdev_getsize64(self.guestfs_device)
 216         self.meta['PARTITION_TABLE'] = \
 217                                 self.g.part_get_parttype(self.guestfs_device)
 218
 219         self.ostype = self.g.inspect_get_type(self.root)
 220         self.distro = self.g.inspect_get_distro(self.root)
 221         self.out.success('found a(n) %s system' % self.distro)
 222
 223     def destroy(self):
 224         """Destroy this DiskDevice instance."""
 225
 226         if self.guestfs_enabled:
 227             self.g.umount_all()
 228             self.g.sync()
 229
 230         # Close the guestfs handler if open
 231         self.g.close()
 232
 233     def progress_callback(self, ev, eh, buf, array):
 234         position = array[2]
 235         total = array[3]
 236
 237         self.progressbar.goto((position * 100) // total)
 238
 239     def mount(self, readonly=False):
 240         """Mount all disk partitions in a correct order."""
 241
 242         mount = self.g.mount_ro if readonly else self.g.mount
 243         self.out.output("Mounting image...", False)
 244         mps = self.g.inspect_get_mountpoints(self.root)
 245
 246         # Sort the keys to mount the fs in a correct order.
 247         # / should be mounted befor /boot, etc
 248         def compare(a, b):
 249             if len(a[0]) > len(b[0]):
 250                 return 1
 251             elif len(a[0]) == len(b[0]):
 252                 return 0
 253             else:
 254                 return -1
 255         mps.sort(compare)
 256         for mp, dev in mps:
 257             try:
 258                 mount(dev, mp)
 259             except RuntimeError as msg:
 260                 self.out.warn("%s (ignored)" % msg)
 261         self.out.success("done")
 262
 263     def umount(self):
 264         """Umount all mounted filesystems."""
 265         self.g.umount_all()
 266
 267     def _last_partition(self):
 268         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 269             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 270             "partition tables are supported" % self.meta['PARTITION_TABLE']
 271             raise FatalError(msg)
 272
 273         is_extended = lambda p: self.g.part_get_mbr_id(
 274                                     self.guestfs_device, p['part_num']) == 5
 275         is_logical = lambda p: self.meta['PARTITION_TABLE'] != 'msdos' and \
 276                                                             p['part_num'] > 4
 277
 278         partitions = self.g.part_list(self.guestfs_device)
 279         last_partition = partitions[-1]
 280
 281         if is_logical(last_partition):
 282             # The disk contains extended and logical partitions....
 283             extended = [p for p in partitions if is_extended(p)][0]
 284             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 285
 286             # check if extended is the last primary partition
 287             if last_primary['part_num'] > extended['part_num']:
 288                 last_partition = last_primary
 289
 290         return last_partition
 291
 292     def shrink(self):
 293         """Shrink the disk.
 294
 295         This is accomplished by shrinking the last filesystem in the
 296         disk and then updating the partition table. The new disk size
 297         (in bytes) is returned.
 298
 299         ATTENTION: make sure unmount is called before shrink
 300         """
 301         get_fstype = lambda p: self.g.vfs_type("%s%d" % \
 302                                         (self.guestfs_device, p['part_num']))
 303         is_logical = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
 304                                                             p['part_num'] > 4
 305         is_extended = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
 306                 self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
 307
 308         part_add = lambda ptype, start, stop: \
 309                     self.g.part_add(self.guestfs_device, ptype, start, stop)
 310         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 311         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 312         part_set_id = lambda p, id: self.g.part_set_mbr_id(
 313                                                     self.guestfs_device, p, id)
 314         part_get_bootable = lambda p: self.g.part_get_bootable(
 315                                                         self.guestfs_device, p)
 316         part_set_bootable = lambda p, bootable: self.g.part_set_bootable(
 317                                             self.guestfs_device, p, bootable)
 318
 319         MB = 2 ** 20
 320
 321         self.out.output("Shrinking image (this may take a while)...", False)
 322
 323         last_part = None
 324         fstype = None
 325         while True:
 326             last_part = self._last_partition()
 327             fstype = get_fstype(last_part)
 328
 329             if fstype == 'swap':
 330                 self.meta['SWAP'] = "%d:%s" % \
 331                         (last_part['part_num'],
 332                         (last_part['part_size'] + MB - 1) // MB)
 333                 part_del(last_part['part_num'])
 334                 continue
 335             elif is_extended(last_part):
 336                 part_del(last_part['part_num'])
 337                 continue
 338
 339             self.meta['SIZE'] = last_part['part_end'] + 1
 340             break
 341
 342         if not re.match("ext[234]", fstype):
 343             self.out.warn("Don't know how to resize %s partitions." % fstype)
 344             return self.meta['SIZE']
 345
 346         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 347         self.g.e2fsck_f(part_dev)
 348         self.g.resize2fs_M(part_dev)
 349
 350         out = self.g.tune2fs_l(part_dev)
 351         block_size = int(
 352             filter(lambda x: x[0] == 'Block size', out)[0][1])
 353         block_cnt = int(
 354             filter(lambda x: x[0] == 'Block count', out)[0][1])
 355
 356         sector_size = self.g.blockdev_getss(self.guestfs_device)
 357         start = last_part['part_start'] / sector_size
 358         end = start + (block_size * block_cnt) / sector_size - 1
 359
 360         if is_logical(last_part):
 361             partitions = self.g.part_list(self.guestfs_device)
 362
 363             logical = []  # logical partitions
 364             for partition in partitions:
 365                 if partition['part_num'] < 4:
 366                     continue
 367                 logical.append({
 368                     'num': partition['part_num'],
 369                     'start': partition['part_start'] / sector_size,
 370                     'end': partition['part_end'] / sector_size,
 371                     'id': part_get_(partition['part_num']),
 372                     'bootable': part_get_bootable(partition['part_num'])
 373                 })
 374
 375             logical[-1]['end'] = end  # new end after resize
 376
 377             # Recreate the extended partition
 378             extended = [p for p in partitions if self._is_extended(p)][0]
 379             part_del(extended['part_num'])
 380             part_add('e', extended['part_start'], end)
 381
 382             # Create all the logical partitions back
 383             for l in logical:
 384                 part_add('l', l['start'], l['end'])
 385                 part_set_id(l['num'], l['id'])
 386                 part_set_bootable(l['num'], l['bootable'])
 387         else:
 388             # Recreate the last partition
 389             if self.meta['PARTITION_TABLE'] == 'msdos':
 390                 last_part['id'] = part_get_id(last_part['part_num'])
 391
 392             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 393             part_del(last_part['part_num'])
 394             part_add('p', start, end)
 395             part_set_bootable(last_part['part_num'], last_part['bootable'])
 396
 397             if self.meta['PARTITION_TABLE'] == 'msdos':
 398                 part_set_id(last_part['part_num'], last_part['id'])
 399
 400         new_size = (end + 1) * sector_size
 401         self.out.success("new size is %dMB" % ((new_size + MB - 1) // MB))
 402
 403         if self.meta['PARTITION_TABLE'] == 'gpt':
 404             ptable = GPTPartitionTable(self.real_device)
 405             self.meta['SIZE'] = ptable.shrink(new_size)
 406         else:
 407             self.meta['SIZE'] = new_size
 408
 409         return self.meta['SIZE']
 410
 411     def dump(self, outfile):
 412         """Dumps the content of device into a file.
 413
 414         This method will only dump the actual payload, found by reading the
 415         partition table. Empty space in the end of the device will be ignored.
 416         """
 417         MB = 2 ** 20
 418         blocksize = 4 * MB  # 4MB
 419         size = self.meta['SIZE']
 420         progr_size = (size + MB - 1) // MB  # in MB
 421         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 422
 423         with open(self.real_device, 'r') as src:
 424             with open(outfile, "w") as dst:
 425                 left = size
 426                 offset = 0
 427                 progressbar.next()
 428                 while left > 0:
 429                     length = min(left, blocksize)
 430                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 431                     offset += sent
 432                     left -= sent
 433                     progressbar.goto((size - left) // MB)
 434         progressbar.success('image file %s was successfully created' % outfile)
 435
 436 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :