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 uuid
  37 import tempfile
  38 import time
  39 from collections import namedtuple
  40
  41 import parted
  42
  43 from image_creator.rsync import Rsync
  44 from image_creator.util import get_command
  45 from image_creator.util import FatalError
  46
  47 findfs = get_command('findfs')
  48 truncate = get_command('truncate')
  49 dd = get_command('dd')
  50 dmsetup = get_command('dmsetup')
  51 losetup = get_command('losetup')
  52 mount = get_command('mount')
  53 umount = get_command('umount')
  54 blkid = get_command('blkid')
  55
  56 MKFS_OPTS = {
  57     '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
  71 class BundleVolume():
  72
  73     def __init__(self, out, meta):
  74         self.out = out
  75         self.meta = meta
  76
  77         self.out.output('Searching for root device ...', False)
  78         root = self._get_root_partition()
  79
  80         if root.startswith("UUID=") or root.startswith("LABEL="):
  81             root = findfs(root).stdout.strip()
  82
  83         if not re.match('/dev/[hsv]d[a-z][1-9]*$', root):
  84             raise FatalError("Don't know how to handle root device: %s" % root)
  85
  86         out.success(root)
  87
  88         disk_file = re.split('[0-9]', root)[0]
  89         device = parted.Device(disk_file)
  90         self.disk = parted.Disk(device)
  91
  92     def _read_fstable(self, f):
  93
  94         if not os.path.isfile(f):
  95             raise FatalError("Unable to open: `%s'. File is missing." % f)
  96
  97         FileSystemTableEntry = namedtuple('FileSystemTableEntry',
  98                                      'dev mpoint fs opts freq passno')
  99         with open(f) as table:
 100             for line in iter(table):
 101                 entry = line.split('#')[0].strip().split()
 102                 if len(entry) != 6:
 103                     continue
 104                 yield FileSystemTableEntry(*entry)
 105
 106     def _get_root_partition(self):
 107         for entry in self._read_fstable('/etc/fstab'):
 108             if entry.mpoint == '/':
 109                 return entry.dev
 110
 111         raise FatalError("Unable to find root device in /etc/fstab")
 112
 113     def _is_mpoint(self, path):
 114         for entry in self._read_fstable('/proc/mounts'):
 115             if entry.mpoint == path:
 116                 return True
 117         return False
 118
 119     def _get_mount_options(self, device):
 120         for entry in self._read_fstable('/proc/mounts'):
 121             if not entry.dev.startswith('/'):
 122                 continue
 123
 124             if os.path.realpath(entry.dev) == os.path.realpath(device):
 125                 return entry
 126
 127         return None
 128
 129     def _create_partition_table(self, image):
 130
 131         if self.disk.type != 'msdos':
 132             raise FatalError('Only msdos partition tables are supported')
 133
 134         # Copy the MBR and the space between the MBR and the first partition.
 135         # In Grub version 1 Stage 1.5 is located there.
 136         first_sector = self.disk.getPrimaryPartitions()[0].geometry.start
 137
 138         dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 139            'bs=%d' % self.disk.device.sectorSize,
 140            'count=%d' % first_sector, 'conv=notrunc')
 141
 142         # Create the Extended boot records (EBRs) in the image
 143         extended = self.disk.getExtendedPartition()
 144         if not extended:
 145             return
 146
 147         # Extended boot records precede the logical partitions they describe
 148         logical = self.disk.getLogicalPartitions()
 149         start = extended.geometry.start
 150         for i in range(len(logical)):
 151             end = logical[i].geometry.start - 1
 152             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 153                'count=%d' % (end - start + 1), 'conv=notrunc',
 154                'seek=%d' % start, 'skip=%d' % start)
 155             start = logical[i].geometry.end + 1
 156
 157     def _get_partitions(self, disk):
 158         Partition = namedtuple('Partition', 'num start end type fs')
 159
 160         partitions = []
 161         for p in disk.partitions:
 162             num = p.number
 163             start = p.geometry.start
 164             end = p.geometry.end
 165             ptype = p.type
 166             fs = p.fileSystem.type if p.fileSystem is not None else ''
 167             partitions.append(Partition(num, start, end, ptype, fs))
 168
 169         return partitions
 170
 171     def _shrink_partitions(self, image):
 172
 173         new_end = self.disk.device.getLength()
 174
 175         image_dev = parted.Device(image)
 176         image_disk = parted.Disk(image_dev)
 177
 178         is_extended = lambda p: p.type == parted.PARTITION_EXTENDED
 179         is_logical = lambda p: p.type == parted.PARTITION_LOGICAL
 180
 181         partitions = self._get_partitions(self.disk)
 182
 183         last = partitions[-1]
 184         if last.fs == 'linux-swap(v1)':
 185             MB = 2 ** 20
 186             size = (last.end - last.start + 1) * self.disk.device.sectorSize
 187             self.meta['SWAP'] = "%d:%s" % (last.num, (size + MB - 1) // MB)
 188
 189             image_disk.deletePartition(
 190                 image_disk.getPartitionBySector(last.start))
 191             image_disk.commit()
 192
 193             if is_logical(last) and last.num == 5:
 194                 extended = image_disk.getExtendedPartition()
 195                 image_disk.deletePartition(extended)
 196                 image_disk.commit()
 197                 partitions.remove(filter(is_extended, partitions)[0])
 198
 199             partitions.remove(last)
 200             last = partitions[-1]
 201
 202             # Leave 2048 blocks at the end
 203             new_end = last.end + 2048
 204
 205         mount_options = self._get_mount_options(
 206                 self.disk.getPartitionBySector(last.start).path)
 207         if mount_options is not None:
 208             stat = os.statvfs(mount_options.mpoint)
 209             # Shrink the last partition. The new size should be the size of the
 210             # occupied blocks
 211             blcks = stat.f_blocks - stat.f_bavail
 212             new_size = (blcks * stat.f_frsize) // self.disk.device.sectorSize
 213
 214             # Add 10% just to be on the safe side
 215             part_end = last.start + (new_size * 11) // 10
 216             # Alighn to 2048
 217             part_end = ((part_end + 2047) // 2048) * 2048
 218
 219             image_disk.setPartitionGeometry(
 220                 image_disk.getPartitionBySector(last.start),
 221                 parted.Constraint(device=image_disk.device),
 222                 start=last.start, end=last.end)
 223             image_disk.commit()
 224
 225             # Parted may have changed this for better alignment
 226             part_end = image_disk.getPartitionBySector(last.start).geometry.end
 227             last = last._replace(end=part_end)
 228             partitions[-1] = last
 229
 230             # Leave 2048 blocks at the end.
 231             new_end = new_size + 2048
 232
 233             if last.type == parted.PARTITION_LOGICAL:
 234                 # Fix the extended partition
 235                 extended = disk.getExtendedPartition()
 236
 237                 image_disk.setPartitionGeometry(extended,
 238                     parted.Constraint(device=img_dev),
 239                     ext.geometry.start, end=last.end)
 240                 image_disk.commit()
 241
 242         return new_end
 243
 244     def _map_partition(self, dev, num, start, end):
 245         name = os.path.basename(dev)
 246         tablefd, table = tempfile.mkstemp()
 247         try:
 248             size = end - start + 1
 249             os.write(tablefd, "0 %d linear %s %d" % (size, dev, start))
 250             dmsetup('create', "%sp%d" % (name, num), table)
 251         finally:
 252             os.unlink(table)
 253
 254         return "/dev/mapper/%sp%d" % (name, num)
 255
 256     def _unmap_partition(self, dev):
 257         if not os.path.exists(dev):
 258             return
 259
 260         dmsetup('remove', dev.split('/dev/mapper/')[1])
 261         time.sleep(0.1)
 262
 263     def _mount(self, target, devs):
 264
 265         devs.sort(key=lambda d: d[1])
 266         for dev, mpoint in devs:
 267             absmpoint = os.path.abspath(target + mpoint)
 268             if not os.path.exists(absmpoint):
 269                 os.makedirs(absmpoint)
 270             mount(dev, absmpoint)
 271
 272     def _umount_all(self, target):
 273         mpoints = []
 274         for entry in self._read_fstable('/proc/mounts'):
 275             if entry.mpoint.startswith(os.path.abspath(target)):
 276                     mpoints.append(entry.mpoint)
 277
 278         mpoints.sort()
 279         for mpoint in reversed(mpoints):
 280             umount(mpoint)
 281
 282     def _to_exclude(self):
 283         excluded = ['/tmp', '/var/tmp']
 284         local_filesystems = MKFS_OPTS.keys() + ['rootfs']
 285         for entry in self._read_fstable('/proc/mounts'):
 286             if entry.fs in local_filesystems:
 287                 continue
 288
 289             mpoint = entry.mpoint
 290             if mpoint in excluded:
 291                 continue
 292
 293             descendants = filter(lambda p: p.startswith(mpoint + '/'),
 294                     excluded)
 295             if len(descendants):
 296                 for d in descendants:
 297                     excluded.remove(d)
 298                 excluded.append(mpoint)
 299                 continue
 300
 301             dirname = mpoint
 302             basename = ''
 303             found_ancestor = False
 304             while dirname != '/':
 305                 (dirname, basename) = os.path.split(dirname)
 306                 if dirname in excluded:
 307                     found_ancestor = True
 308                     break
 309
 310             if not found_ancestor:
 311                 excluded.append(mpoint)
 312
 313         return map(lambda d: d + "/*", excluded)
 314
 315     def _replace_uuids(self, target, new_uuid):
 316
 317         files = ['/etc/fstab',
 318                  '/boot/grub/grub.cfg',
 319                  '/boot/grub/menu.lst',
 320                  '/boot/grub/grub.conf']
 321
 322         orig = dict(map(lambda p: (p.number, blkid( '-s', 'UUID', '-o',
 323             'value', p.path).stdout.strip()), self.disk.partitions))
 324
 325         for f in map(lambda f: target + f, files):
 326
 327             if not os.path.exists(f):
 328                 continue
 329
 330             with open(f, 'r') as src:
 331                 lines = src.readlines()
 332             with open(f, 'w') as dest:
 333                 for line in lines:
 334                     for i, uuid in new_uuid.items():
 335                         line = re.sub(orig[i], uuid, line)
 336                     dest.write(line)
 337
 338     def _create_filesystems(self, image):
 339
 340         filesystem = {}
 341         for p in self.disk.partitions:
 342             filesystem[p.number] = self._get_mount_options(p.path)
 343
 344         partitions = self._get_partitions(parted.Disk(parted.Device(image)))
 345         unmounted = filter(lambda p: filesystem[p.num] is None, partitions)
 346         mounted = filter(lambda p: filesystem[p.num] is not None, partitions)
 347
 348         # For partitions that are not mounted right now, we can simply dd them
 349         # into the image.
 350         for p in unmounted:
 351             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 352                'count=%d' % (p.end - p.start + 1), 'conv=notrunc',
 353                'seek=%d' % p.start, 'skip=%d' % p.start)
 354
 355         loop = str(losetup('-f', '--show', image)).strip()
 356         mapped = {}
 357         try:
 358             for p in mounted:
 359                 i = p.num
 360                 mapped[i] = self._map_partition(loop, i, p.start, p.end)
 361
 362             new_uuid = {}
 363             # Create the file systems
 364             for i, dev in mapped.iteritems():
 365                 fs = filesystem[i].fs
 366                 self.out.output('Creating %s filesystem on partition %d ... ' %
 367                     (fs, i), False)
 368                 get_command('mkfs.%s' % fs)(*(MKFS_OPTS[fs] + [dev]))
 369                 self.out.success('done')
 370                 new_uuid[i] = blkid('-s', 'UUID', '-o', 'value', dev
 371                     ).stdout.strip()
 372
 373             target = tempfile.mkdtemp()
 374             try:
 375                 absmpoints = self._mount(target,
 376                     [(mapped[i], filesystem[i].mpoint) for i in mapped.keys()]
 377                 )
 378                 exclude = self._to_exclude() + [image]
 379                 rsync = Rsync('/', target,
 380                               map(lambda p: os.path.relpath(p, '/'), exclude))
 381                 msg = "Copying host files into the image"
 382                 rsync.archive().run(self.out, msg)
 383
 384                 self._replace_uuids(target, new_uuid)
 385
 386             finally:
 387                 self._umount_all(target)
 388                 os.rmdir(target)
 389         finally:
 390             for dev in mapped.values():
 391                 self._unmap_partition(dev)
 392             losetup('-d', loop)
 393
 394     def create_image(self):
 395
 396         image = '/mnt/%s.diskdump' % uuid.uuid4().hex
 397
 398         disk_size = self.disk.device.getLength() * self.disk.device.sectorSize
 399
 400         # Create sparse file to host the image
 401         truncate("-s", "%d" % disk_size, image)
 402
 403         self._create_partition_table(image)
 404
 405         end_sector = self._shrink_partitions(image)
 406
 407         # Check if the available space is enough to host the image
 408         dirname = os.path.dirname(image)
 409         size = (end_sector + 1) * self.disk.device.sectorSize
 410         self.out.output("Examining available space in %s ..." % dirname, False)
 411         stat = os.statvfs(dirname)
 412         available = stat.f_bavail * stat.f_frsize
 413         if available <= size:
 414             raise FatalError('Not enough space in %s to host the image' %
 415                              dirname)
 416         self.out.success("sufficient")
 417
 418         self._create_filesystems(image)
 419
 420         return image
 421
 422 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :