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=part_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 = part_end + 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         image_dev.destroy()
 243         return new_end
 244
 245     def _map_partition(self, dev, num, start, end):
 246         name = os.path.basename(dev)
 247         tablefd, table = tempfile.mkstemp()
 248         try:
 249             size = end - start + 1
 250             os.write(tablefd, "0 %d linear %s %d" % (size, dev, start))
 251             dmsetup('create', "%sp%d" % (name, num), table)
 252         finally:
 253             os.unlink(table)
 254
 255         return "/dev/mapper/%sp%d" % (name, num)
 256
 257     def _unmap_partition(self, dev):
 258         if not os.path.exists(dev):
 259             return
 260
 261         dmsetup('remove', dev.split('/dev/mapper/')[1])
 262         time.sleep(0.1)
 263
 264     def _mount(self, target, devs):
 265
 266         devs.sort(key=lambda d: d[1])
 267         for dev, mpoint in devs:
 268             absmpoint = os.path.abspath(target + mpoint)
 269             if not os.path.exists(absmpoint):
 270                 os.makedirs(absmpoint)
 271             mount(dev, absmpoint)
 272
 273     def _umount_all(self, target):
 274         mpoints = []
 275         for entry in self._read_fstable('/proc/mounts'):
 276             if entry.mpoint.startswith(os.path.abspath(target)):
 277                     mpoints.append(entry.mpoint)
 278
 279         mpoints.sort()
 280         for mpoint in reversed(mpoints):
 281             umount(mpoint)
 282
 283     def _to_exclude(self):
 284         excluded = ['/tmp', '/var/tmp']
 285         local_filesystems = MKFS_OPTS.keys() + ['rootfs']
 286         for entry in self._read_fstable('/proc/mounts'):
 287             if entry.fs in local_filesystems:
 288                 continue
 289
 290             mpoint = entry.mpoint
 291             if mpoint in excluded:
 292                 continue
 293
 294             descendants = filter(lambda p: p.startswith(mpoint + '/'),
 295                     excluded)
 296             if len(descendants):
 297                 for d in descendants:
 298                     excluded.remove(d)
 299                 excluded.append(mpoint)
 300                 continue
 301
 302             dirname = mpoint
 303             basename = ''
 304             found_ancestor = False
 305             while dirname != '/':
 306                 (dirname, basename) = os.path.split(dirname)
 307                 if dirname in excluded:
 308                     found_ancestor = True
 309                     break
 310
 311             if not found_ancestor:
 312                 excluded.append(mpoint)
 313
 314         return map(lambda d: d + "/*", excluded)
 315
 316     def _replace_uuids(self, target, new_uuid):
 317
 318         files = ['/etc/fstab',
 319                  '/boot/grub/grub.cfg',
 320                  '/boot/grub/menu.lst',
 321                  '/boot/grub/grub.conf']
 322
 323         orig = dict(map(lambda p: (p.number, blkid( '-s', 'UUID', '-o',
 324             'value', p.path).stdout.strip()), self.disk.partitions))
 325
 326         for f in map(lambda f: target + f, files):
 327
 328             if not os.path.exists(f):
 329                 continue
 330
 331             with open(f, 'r') as src:
 332                 lines = src.readlines()
 333             with open(f, 'w') as dest:
 334                 for line in lines:
 335                     for i, uuid in new_uuid.items():
 336                         line = re.sub(orig[i], uuid, line)
 337                     dest.write(line)
 338
 339     def _create_filesystems(self, image):
 340
 341         filesystem = {}
 342         for p in self.disk.partitions:
 343             filesystem[p.number] = self._get_mount_options(p.path)
 344
 345         partitions = self._get_partitions(parted.Disk(parted.Device(image)))
 346         unmounted = filter(lambda p: filesystem[p.num] is None, partitions)
 347         mounted = filter(lambda p: filesystem[p.num] is not None, partitions)
 348
 349         # For partitions that are not mounted right now, we can simply dd them
 350         # into the image.
 351         for p in unmounted:
 352             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 353                'count=%d' % (p.end - p.start + 1), 'conv=notrunc',
 354                'seek=%d' % p.start, 'skip=%d' % p.start)
 355
 356         loop = str(losetup('-f', '--show', image)).strip()
 357         mapped = {}
 358         try:
 359             for p in mounted:
 360                 i = p.num
 361                 mapped[i] = self._map_partition(loop, i, p.start, p.end)
 362
 363             new_uuid = {}
 364             # Create the file systems
 365             for i, dev in mapped.iteritems():
 366                 fs = filesystem[i].fs
 367                 self.out.output('Creating %s filesystem on partition %d ... ' %
 368                     (fs, i), False)
 369                 get_command('mkfs.%s' % fs)(*(MKFS_OPTS[fs] + [dev]))
 370                 self.out.success('done')
 371                 new_uuid[i] = blkid('-s', 'UUID', '-o', 'value', dev
 372                     ).stdout.strip()
 373
 374             target = tempfile.mkdtemp()
 375             try:
 376                 absmpoints = self._mount(target,
 377                     [(mapped[i], filesystem[i].mpoint) for i in mapped.keys()]
 378                 )
 379                 exclude = self._to_exclude() + [image]
 380                 rsync = Rsync('/', target,
 381                               map(lambda p: os.path.relpath(p, '/'), exclude))
 382                 msg = "Copying host files into the image"
 383                 rsync.archive().run(self.out, msg)
 384
 385                 self._replace_uuids(target, new_uuid)
 386
 387             finally:
 388                 self._umount_all(target)
 389                 os.rmdir(target)
 390         finally:
 391             for dev in mapped.values():
 392                 self._unmap_partition(dev)
 393             losetup('-d', loop)
 394
 395     def create_image(self):
 396
 397         image = '/mnt/%s.diskdump' % uuid.uuid4().hex
 398
 399         disk_size = self.disk.device.getLength() * self.disk.device.sectorSize
 400
 401         # Create sparse file to host the image
 402         truncate("-s", "%d" % disk_size, image)
 403
 404         self._create_partition_table(image)
 405
 406         end_sector = self._shrink_partitions(image)
 407
 408         # Check if the available space is enough to host the image
 409         dirname = os.path.dirname(image)
 410         size = (end_sector + 1) * self.disk.device.sectorSize
 411         self.out.output("Examining available space in %s ..." % dirname, False)
 412         stat = os.statvfs(dirname)
 413         available = stat.f_bavail * stat.f_frsize
 414         if available <= size:
 415             raise FatalError('Not enough space in %s to host the image' %
 416                              dirname)
 417         self.out.success("sufficient")
 418
 419         self._create_filesystems(image)
 420
 421         # Truncate image to the new size. I counldn't find a better way to do
 422         # this. It seems that python's high level functions work in a different
 423         # way.
 424         fd = os.open(image, os.O_RDWR)
 425         try:
 426             os.ftruncate(fd, size)
 427         finally:
 428             os.close(fd)
 429
 430         return image
 431
 432 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :