code.grnet.gr Git - snf-image-creator/blob - image_creator/bundle_volume.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 import os
  35 import re
  36 import tempfile
  37 import uuid
  38 from collections import namedtuple
  39
  40 import parted
  41
  42 from image_creator.rsync import Rsync
  43 from image_creator.util import get_command
  44 from image_creator.util import FatalError
  45 from image_creator.util import try_fail_repeat
  46 from image_creator.util import free_space
  47 from image_creator.gpt import GPTPartitionTable
  48
  49 findfs = get_command('findfs')
  50 dd = get_command('dd')
  51 dmsetup = get_command('dmsetup')
  52 losetup = get_command('losetup')
  53 mount = get_command('mount')
  54 umount = get_command('umount')
  55 blkid = get_command('blkid')
  56
  57 MKFS_OPTS = {'ext2': ['-F'],
  58              'ext3': ['-F'],
  59              'ext4': ['-F'],
  60              'reiserfs': ['-ff'],
  61              'btrfs': [],
  62              'minix': [],
  63              'xfs': ['-f'],
  64              'jfs': ['-f'],
  65              'ntfs': ['-F'],
  66              'msdos': [],
  67              'vfat': []}
  68
  69
  70 class BundleVolume(object):
  71     """This class can be used to create an image out of the running system"""
  72
  73     def __init__(self, out, meta, tmp=None):
  74         """Create an instance of the BundleVolume class."""
  75         self.out = out
  76         self.meta = meta
  77         self.tmp = tmp
  78
  79         self.out.output('Searching for root device ...', False)
  80         root = self._get_root_partition()
  81
  82         if root.startswith("UUID=") or root.startswith("LABEL="):
  83             root = findfs(root).stdout.strip()
  84
  85         if not re.match('/dev/[hsv]d[a-z][1-9]*$', root):
  86             raise FatalError("Don't know how to handle root device: %s" % root)
  87
  88         out.success(root)
  89
  90         disk_file = re.split('[0-9]', root)[0]
  91         device = parted.Device(disk_file)
  92         self.disk = parted.Disk(device)
  93
  94     def _read_fstable(self, f):
  95         """Use this generator to iterate over the lines of and fstab file"""
  96
  97         if not os.path.isfile(f):
  98             raise FatalError("Unable to open: `%s'. File is missing." % f)
  99
 100         FileSystemTableEntry = namedtuple('FileSystemTableEntry',
 101                                           'dev mpoint fs opts freq passno')
 102         with open(f) as table:
 103             for line in iter(table):
 104                 entry = line.split('#')[0].strip().split()
 105                 if len(entry) != 6:
 106                     continue
 107                 yield FileSystemTableEntry(*entry)
 108
 109     def _get_root_partition(self):
 110         """Return the fstab entry accosiated with the root filesystem"""
 111         for entry in self._read_fstable('/etc/fstab'):
 112             if entry.mpoint == '/':
 113                 return entry.dev
 114
 115         raise FatalError("Unable to find root device in /etc/fstab")
 116
 117     def _is_mpoint(self, path):
 118         """Check if a directory is currently a mount point"""
 119         for entry in self._read_fstable('/proc/mounts'):
 120             if entry.mpoint == path:
 121                 return True
 122         return False
 123
 124     def _get_mount_options(self, device):
 125         """Return the mount entry associated with a mounted device"""
 126         for entry in self._read_fstable('/proc/mounts'):
 127             if not entry.dev.startswith('/'):
 128                 continue
 129
 130             if os.path.realpath(entry.dev) == os.path.realpath(device):
 131                 return entry
 132
 133         return None
 134
 135     def _create_partition_table(self, image):
 136         """Copy the partition table of the host system into the image"""
 137
 138         # Copy the MBR and the space between the MBR and the first partition.
 139         # In msdos partition tables Grub Stage 1.5 is located there.
 140         # In gpt partition tables the Primary GPT Header is there.
 141         first_sector = self.disk.getPrimaryPartitions()[0].geometry.start
 142
 143         dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 144            'bs=%d' % self.disk.device.sectorSize,
 145            'count=%d' % first_sector, 'conv=notrunc')
 146
 147         if self.disk.type == 'gpt':
 148             # Copy the Secondary GPT Header
 149             table = GPTPartitionTable(self.disk.device.path)
 150             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 151                'bs=%d' % self.disk.device.sectorSize, 'conv=notrunc',
 152                'seek=%d' % table.primary.last_usable_lba,
 153                'skip=%d' % table.primary.last_usable_lba)
 154
 155         # Create the Extended boot records (EBRs) in the image
 156         extended = self.disk.getExtendedPartition()
 157         if not extended:
 158             return
 159
 160         # Extended boot records precede the logical partitions they describe
 161         logical = self.disk.getLogicalPartitions()
 162         start = extended.geometry.start
 163         for i in range(len(logical)):
 164             end = logical[i].geometry.start - 1
 165             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 166                'count=%d' % (end - start + 1), 'conv=notrunc',
 167                'seek=%d' % start, 'skip=%d' % start)
 168             start = logical[i].geometry.end + 1
 169
 170     def _get_partitions(self, disk):
 171         """Returns a list with the partitions of the provided disk"""
 172         Partition = namedtuple('Partition', 'num start end type fs')
 173
 174         partitions = []
 175         for p in disk.partitions:
 176             num = p.number
 177             start = p.geometry.start
 178             end = p.geometry.end
 179             ptype = p.type
 180             fs = p.fileSystem.type if p.fileSystem is not None else ''
 181             partitions.append(Partition(num, start, end, ptype, fs))
 182
 183         return partitions
 184
 185     def _shrink_partitions(self, image):
 186         """Remove the last partition of the image if it is a swap partition and
 187         shrink the partition before that. Make sure it can still host all the
 188         files the corresponding host file system hosts
 189         """
 190         new_end = self.disk.device.length
 191
 192         image_disk = parted.Disk(parted.Device(image))
 193
 194         is_extended = lambda p: p.type == parted.PARTITION_EXTENDED
 195         is_logical = lambda p: p.type == parted.PARTITION_LOGICAL
 196
 197         partitions = self._get_partitions(self.disk)
 198
 199         last = partitions[-1]
 200         if last.fs == 'linux-swap(v1)':
 201             MB = 2 ** 20
 202             size = (last.end - last.start + 1) * self.disk.device.sectorSize
 203             self.meta['SWAP'] = "%d:%s" % (last.num, (size + MB - 1) // MB)
 204
 205             image_disk.deletePartition(
 206                 image_disk.getPartitionBySector(last.start))
 207             image_disk.commitToDevice()
 208
 209             if is_logical(last) and last.num == 5:
 210                 extended = image_disk.getExtendedPartition()
 211                 image_disk.deletePartition(extended)
 212                 image_disk.commitToDevice()
 213                 partitions.remove(filter(is_extended, partitions)[0])
 214
 215             partitions.remove(last)
 216             last = partitions[-1]
 217
 218             new_end = last.end
 219
 220         mount_options = self._get_mount_options(
 221             self.disk.getPartitionBySector(last.start).path)
 222         if mount_options is not None:
 223             stat = os.statvfs(mount_options.mpoint)
 224             # Shrink the last partition. The new size should be the size of the
 225             # occupied blocks
 226             blcks = stat.f_blocks - stat.f_bavail
 227             new_size = (blcks * stat.f_frsize) // self.disk.device.sectorSize
 228
 229             # Add 10% just to be on the safe side
 230             part_end = last.start + (new_size * 11) // 10
 231             # Align to 2048
 232             part_end = ((part_end + 2047) // 2048) * 2048
 233
 234             # Make sure the partition starts where the old partition started.
 235             constraint = parted.Constraint(device=image_disk.device)
 236             constraint.startRange = parted.Geometry(device=image_disk.device,
 237                                                     start=last.start, length=1)
 238
 239             image_disk.setPartitionGeometry(
 240                 image_disk.getPartitionBySector(last.start), constraint,
 241                 start=last.start, end=part_end)
 242             image_disk.commitToDevice()
 243
 244             # Parted may have changed this for better alignment
 245             part_end = image_disk.getPartitionBySector(last.start).geometry.end
 246             last = last._replace(end=part_end)
 247             partitions[-1] = last
 248
 249             new_end = part_end
 250
 251             if last.type == parted.PARTITION_LOGICAL:
 252                 # Fix the extended partition
 253                 image_disk.minimizeExtendedPartition()
 254
 255         return (new_end, self._get_partitions(image_disk))
 256
 257     def _map_partition(self, dev, num, start, end):
 258         """Map a partition into a block device using the device mapper"""
 259         name = os.path.basename(dev) + "_" + uuid.uuid4().hex
 260         tablefd, table = tempfile.mkstemp()
 261         try:
 262             size = end - start + 1
 263             os.write(tablefd, "0 %d linear %s %d" % (size, dev, start))
 264             dmsetup('create', "%sp%d" % (name, num), table)
 265         finally:
 266             os.unlink(table)
 267
 268         return "/dev/mapper/%sp%d" % (name, num)
 269
 270     def _unmap_partition(self, dev):
 271         """Unmap a previously mapped partition"""
 272         if not os.path.exists(dev):
 273             return
 274
 275         try_fail_repeat(dmsetup, 'remove', dev.split('/dev/mapper/')[1])
 276
 277     def _mount(self, target, devs):
 278         """Mount a list of filesystems in mountpoints relative to target"""
 279         devs.sort(key=lambda d: d[1])
 280         for dev, mpoint in devs:
 281             absmpoint = os.path.abspath(target + mpoint)
 282             if not os.path.exists(absmpoint):
 283                 os.makedirs(absmpoint)
 284             mount(dev, absmpoint)
 285
 286     def _umount_all(self, target):
 287         """Unmount all filesystems that are mounted under the directory target
 288         """
 289         mpoints = []
 290         for entry in self._read_fstable('/proc/mounts'):
 291             if entry.mpoint.startswith(os.path.abspath(target)):
 292                     mpoints.append(entry.mpoint)
 293
 294         mpoints.sort()
 295         for mpoint in reversed(mpoints):
 296             try_fail_repeat(umount, mpoint)
 297
 298     def _to_exclude(self):
 299         """Find which directories to exclude during the image copy. This is
 300         accompliced by checking which directories serve as mount points for
 301         virtual file systems
 302         """
 303         excluded = ['/tmp', '/var/tmp']
 304         if self.tmp is not None:
 305             excluded.append(self.tmp)
 306         local_filesystems = MKFS_OPTS.keys() + ['rootfs']
 307         for entry in self._read_fstable('/proc/mounts'):
 308             if entry.fs in local_filesystems:
 309                 continue
 310
 311             mpoint = entry.mpoint
 312             if mpoint in excluded:
 313                 continue
 314
 315             descendants = filter(
 316                 lambda p: p.startswith(mpoint + '/'), excluded)
 317             if len(descendants):
 318                 for d in descendants:
 319                     excluded.remove(d)
 320                 excluded.append(mpoint)
 321                 continue
 322
 323             dirname = mpoint
 324             basename = ''
 325             found_ancestor = False
 326             while dirname != '/':
 327                 (dirname, basename) = os.path.split(dirname)
 328                 if dirname in excluded:
 329                     found_ancestor = True
 330                     break
 331
 332             if not found_ancestor:
 333                 excluded.append(mpoint)
 334
 335         return excluded
 336
 337     def _replace_uuids(self, target, new_uuid):
 338         """Replace UUID references in various files. This is needed after
 339         copying system files of the host into a new filesystem
 340         """
 341
 342         files = ['/etc/fstab',
 343                  '/boot/grub/grub.cfg',
 344                  '/boot/grub/menu.lst',
 345                  '/boot/grub/grub.conf']
 346
 347         orig = {}
 348         for p in self.disk.partitions:
 349             if p.number in new_uuid.keys():
 350                 orig[p.number] = \
 351                     blkid('-s', 'UUID', '-o', 'value', p.path).stdout.strip()
 352
 353         for f in map(lambda f: target + f, files):
 354             if not os.path.exists(f):
 355                 continue
 356
 357             with open(f, 'r') as src:
 358                 lines = src.readlines()
 359             with open(f, 'w') as dest:
 360                 for line in lines:
 361                     for i, uuid in new_uuid.items():
 362                         line = re.sub(orig[i], uuid, line)
 363                     dest.write(line)
 364
 365     def _create_filesystems(self, image, partitions):
 366         """Fill the image with data. Host file systems that are not currently
 367         mounted are binary copied into the image. For mounted file systems, a
 368         file system level copy is performed.
 369         """
 370
 371         filesystem = {}
 372         for p in self.disk.partitions:
 373             filesystem[p.number] = self._get_mount_options(p.path)
 374
 375         unmounted = filter(lambda p: filesystem[p.num] is None, partitions)
 376         mounted = filter(lambda p: filesystem[p.num] is not None, partitions)
 377
 378         # For partitions that are not mounted right now, we can simply dd them
 379         # into the image.
 380         for p in unmounted:
 381             self.out.output('Cloning partition %d ... ' % p.num, False)
 382             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 383                'count=%d' % (p.end - p.start + 1), 'conv=notrunc',
 384                'seek=%d' % p.start, 'skip=%d' % p.start)
 385             self.out.success("done")
 386
 387         loop = str(losetup('-f', '--show', image)).strip()
 388         mapped = {}
 389         try:
 390             for p in mounted:
 391                 i = p.num
 392                 mapped[i] = self._map_partition(loop, i, p.start, p.end)
 393
 394             new_uuid = {}
 395             # Create the file systems
 396             for i, dev in mapped.iteritems():
 397                 fs = filesystem[i].fs
 398                 self.out.output('Creating %s filesystem on partition %d ... ' %
 399                                 (fs, i), False)
 400                 get_command('mkfs.%s' % fs)(*(MKFS_OPTS[fs] + [dev]))
 401                 self.out.success('done')
 402                 new_uuid[i] = blkid(
 403                     '-s', 'UUID', '-o', 'value', dev).stdout.strip()
 404
 405             target = tempfile.mkdtemp()
 406             try:
 407                 self._mount(
 408                     target,
 409                     [(mapped[i], filesystem[i].mpoint) for i in mapped.keys()])
 410
 411                 excluded = self._to_exclude()
 412
 413                 rsync = Rsync(self.out)
 414
 415                 for excl in excluded + [image]:
 416                     rsync.exclude(excl)
 417
 418                 rsync.archive().hard_links().xattrs().sparse().acls()
 419                 rsync.run('/', target, 'host', 'temporary image')
 420
 421                 # Create missing mountpoints. Since they are mountpoints, we
 422                 # cannot determine the ownership and the mode of the real
 423                 # directory. Make them inherit those properties from their
 424                 # parent dir
 425                 for excl in excluded:
 426                     dirname = os.path.dirname(excl)
 427                     stat = os.stat(dirname)
 428                     os.mkdir(target + excl)
 429                     os.chmod(target + excl, stat.st_mode)
 430                     os.chown(target + excl, stat.st_uid, stat.st_gid)
 431
 432                 # /tmp and /var/tmp are special cases. We exclude then even if
 433                 # they aren't mountpoints. Restore their permissions.
 434                 for excl in ('/tmp', '/var/tmp'):
 435                     if self._is_mpoint(excl):
 436                         os.chmod(target + excl, 041777)
 437                         os.chown(target + excl, 0, 0)
 438                     else:
 439                         stat = os.stat(excl)
 440                         os.chmod(target + excl, stat.st_mode)
 441                         os.chown(target + excl, stat.st_uid, stat.st_gid)
 442
 443                 # We need to replace the old UUID referencies with the new
 444                 # ones in grub configuration files and /etc/fstab for file
 445                 # systems that have been recreated.
 446                 self._replace_uuids(target, new_uuid)
 447
 448             finally:
 449                 self._umount_all(target)
 450                 os.rmdir(target)
 451         finally:
 452             for dev in mapped.values():
 453                 self._unmap_partition(dev)
 454             losetup('-d', loop)
 455
 456     def create_image(self, image):
 457         """Given an image filename, this method will create an image out of the
 458         running system.
 459         """
 460
 461         size = self.disk.device.length * self.disk.device.sectorSize
 462
 463         # Create sparse file to host the image
 464         fd = os.open(image, os.O_WRONLY | os.O_CREAT)
 465         try:
 466             os.ftruncate(fd, size)
 467         finally:
 468             os.close(fd)
 469
 470         self._create_partition_table(image)
 471         end_sector, partitions = self._shrink_partitions(image)
 472
 473         if self.disk.type == 'gpt':
 474             old_size = size
 475             size = (end_sector + 1) * self.disk.device.sectorSize
 476             ptable = GPTPartitionTable(image)
 477             size = ptable.shrink(size, old_size)
 478         else:
 479             # Alighn to 2048
 480             end_sector = ((end_sector + 2047) // 2048) * 2048
 481             size = (end_sector + 1) * self.disk.device.sectorSize
 482
 483         # Truncate image to the new size.
 484         fd = os.open(image, os.O_RDWR)
 485         try:
 486             os.ftruncate(fd, size)
 487         finally:
 488             os.close(fd)
 489
 490         # Check if the available space is enough to host the image
 491         dirname = os.path.dirname(image)
 492         self.out.output("Examining available space ...", False)
 493         if free_space(dirname) <= size:
 494             raise FatalError("Not enough space under %s to host the temporary "
 495                              "image" % dirname)
 496         self.out.success("sufficient")
 497
 498         self._create_filesystems(image, partitions)
 499
 500         return image
 501
 502 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :