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