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.util import get_command
  44 from image_creator.util import FatalError
  45
  46 findfs = get_command('findfs')
  47 truncate = get_command('truncate')
  48 dd = get_command('dd')
  49 dmsetup = get_command('dmsetup')
  50 losetup = get_command('losetup')
  51 mount = get_command('mount')
  52 umount = get_command('umount')
  53
  54 MKFS_OPTS = {
  55     'ext2': ['-F'],
  56     'ext3': ['-F'],
  57     'ext4': ['-F'],
  58     'reiserfs': ['-ff'],
  59     'btrfs': [],
  60     'minix': [],
  61     'xfs': ['-f'],
  62     'jfs': ['-f'],
  63     'ntfs': ['-F'],
  64     'msdos': [],
  65     'vfat': []
  66     }
  67
  68
  69 class BundleVolume():
  70
  71     def __init__(self, out, meta):
  72         self.out = out
  73         self.meta = meta
  74
  75         self.out.output('Searching for root device ...', False)
  76         root = self._get_root_partition()
  77
  78         if root.startswith("UUID=") or root.startswith("LABEL="):
  79             root = findfs(root).stdout.strip()
  80
  81         if not re.match('/dev/[hsv]d[a-z][1-9]*$', root):
  82             raise FatalError("Don't know how to handle root device: %s" % root)
  83
  84         out.success(root)
  85
  86         disk_file = re.split('[0-9]', root)[0]
  87         device = parted.Device(disk_file)
  88         self.disk = parted.Disk(device)
  89
  90     def _read_fstable(self, f):
  91
  92         if not os.path.isfile(f):
  93             raise FatalError("Unable to open: `%s'. File is missing." % f)
  94
  95         FileSystemEntry = namedtuple('FileSystemEntry',
  96                                      'dev mpoint fs opts freq passno')
  97         with open(f) as table:
  98             for line in iter(table):
  99                 entry = line.split('#')[0].strip().split()
 100                 if len(entry) != 6:
 101                     continue
 102                 yield FileSystemEntry(*entry)
 103
 104     def _get_root_partition(self):
 105         for entry in self._read_fstable('/etc/fstab'):
 106             if entry.mpoint == '/':
 107                 return entry.dev
 108
 109         raise FatalError("Unable to find root device in /etc/fstab")
 110
 111     def _is_mpoint(self, path):
 112         for entry in self._read_fstable('/proc/mounts'):
 113             if entry.mpoint == path:
 114                 return True
 115         return False
 116
 117     def _get_mount_options(self, device):
 118         for entry in self._read_fstable('/proc/mounts'):
 119             if not entry.dev.startswith('/'):
 120                 continue
 121
 122             if os.path.realpath(entry.dev) == os.path.realpath(device):
 123                 return entry
 124
 125         return None
 126
 127     def _create_partition_table(self, image):
 128
 129         if self.disk.type != 'msdos':
 130             raise FatalError('Only msdos partition tables are supported')
 131
 132         # Copy the MBR and the space between the MBR and the first partition.
 133         # In Grub version 1 Stage 1.5 is located there.
 134         first_sector = self.disk.getPrimaryPartitions()[0].geometry.start
 135
 136         dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 137            'bs=%d' % self.disk.device.sectorSize,
 138            'count=%d' % first_sector, 'conv=notrunc')
 139
 140         # Create the Extended boot records (EBRs) in the image
 141         extended = self.disk.getExtendedPartition()
 142         if not extended:
 143             return
 144
 145         # Extended boot records precede the logical partitions they describe
 146         logical = self.disk.getLogicalPartitions()
 147         start = extended.geometry.start
 148         for i in range(len(logical)):
 149             end = logical[i].geometry.start - 1
 150             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 151                'count=%d' % (end - start + 1), 'conv=notrunc',
 152                'seek=%d' % start, 'skip=%d' % start)
 153             start = logical[i].geometry.end + 1
 154
 155     def _get_partitions(self, disk):
 156         Partition = namedtuple('Partition', 'num start end type fs')
 157
 158         partitions = []
 159         for p in disk.partitions:
 160             num = p.number
 161             start = p.geometry.start
 162             end = p.geometry.end
 163             ptype = p.type
 164             fs = p.fileSystem.type if p.fileSystem is not None else ''
 165             partitions.append(Partition(num, start, end, ptype, fs))
 166
 167         return partitions
 168
 169     def _shrink_partitions(self, image):
 170
 171         new_end = self.disk.device.getLength()
 172
 173         image_dev = parted.Device(image)
 174         image_disk = parted.Disk(image_dev)
 175
 176         is_extended = lambda p: p.type == parted.PARTITION_EXTENDED
 177         is_logical = lambda p: p.type == parted.PARTITION_LOGICAL
 178
 179         partitions = self._get_partitions(self.disk)
 180
 181         last = partitions[-1]
 182         if last.fs == 'linux-swap(v1)':
 183             MB = 2 ** 20
 184             size = (last.end - last.start + 1) * self.disk.device.sectorSize
 185             self.meta['SWAP'] = "%d:%s" % (last.num, (size + MB - 1) // MB)
 186
 187             image_disk.deletePartition(
 188                 image_disk.getPartitionBySector(last.start))
 189             image_disk.commit()
 190
 191             if is_logical(last) and last.num == 5:
 192                 extended = image_disk.getExtendedPartition()
 193                 image_disk.deletePartition(extended)
 194                 image_disk.commit()
 195                 partitions.remove(filter(is_extended, partitions)[0])
 196
 197             partitions.remove(last)
 198             last = partitions[-1]
 199
 200             # Leave 2048 blocks at the end
 201             new_end = last.end + 2048
 202
 203         mount_options = self._get_mount_options(
 204                 self.disk.getPartitionBySector(last.start).path)
 205         if mount_options is not None:
 206             stat = os.statvfs(mount_options.mpoint)
 207             # Shrink the last partition. The new size should be the size of the
 208             # occupied blocks
 209             blcks = stat.f_blocks - stat.f_bavail
 210             new_size = (blcks * stat.f_frsize) // self.disk.device.sectorSize
 211
 212             # Add 10% just to be on the safe side
 213             part_end = last.start + (new_size * 11) // 10
 214             # Alighn to 2048
 215             part_end = ((part_end + 2047) // 2048) * 2048
 216
 217             image_disk.setPartitionGeometry(
 218                 image_disk.getPartitionBySector(last.start),
 219                 parted.Constraint(device=image_disk.device),
 220                 start=last.start, end=last.end)
 221             image_disk.commit()
 222
 223             # Parted may have changed this for better alignment
 224             part_end = image_disk.getPartitionBySector(last.start).geometry.end
 225             last = last._replace(end=part_end)
 226             partitions[-1] = last
 227
 228             # Leave 2048 blocks at the end.
 229             new_end = new_size + 2048
 230
 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 _create_filesystems(self, image):
 282
 283         partitions = self._get_partitions(parted.Disk(parted.Device(image)))
 284         filesystems = {}
 285         for p in self.disk.partitions:
 286             filesystems[p.number] = self._get_mount_options(p.path)
 287
 288         unmounted = filter(lambda p: filesystems[p.num] is None, partitions)
 289         mounted = filter(lambda p: filesystems[p.num] is not None, partitions)
 290
 291         # For partitions that are not mounted right now, we can simply dd them
 292         # into the image.
 293         for p in unmounted:
 294             dd('if=%s' % self.disk.device.path, 'of=%s' % image,
 295                'count=%d' % (p.end - p.start + 1), 'conv=notrunc',
 296                'seek=%d' % p.start, 'skip=%d' % p.start)
 297
 298         loop = str(losetup('-f', '--show', image)).strip()
 299         mapped = {}
 300         try:
 301             for p in mounted:
 302                 i =  p.num
 303                 mapped[i] = self._map_partition(loop, i, p.start, p.end)
 304
 305             # Create the file systems
 306             for i, dev in mapped.iteritems():
 307                 fs = filesystems[i].fs
 308                 self.out.output('Creating %s filesystem on partition %d ... ' %
 309                     (fs, i), False)
 310                 get_command('mkfs.%s' % fs)(*(MKFS_OPTS[fs] + [dev]))
 311                 self.out.success('done')
 312
 313             target = tempfile.mkdtemp()
 314             try:
 315                 absmpoints = self._mount(target,
 316                     [(mapped[i], filesystems[i].mpoint) for i in mapped.keys()]
 317                 )
 318
 319             finally:
 320                 self._umount_all(target)
 321                 os.rmdir(target)
 322         finally:
 323             for dev in mapped.values():
 324                 self._unmap_partition(dev)
 325             losetup('-d', loop)
 326
 327     def create_image(self):
 328
 329         image = '/mnt/%s.diskdump' % uuid.uuid4().hex
 330
 331         disk_size = self.disk.device.getLength() * self.disk.device.sectorSize
 332
 333         # Create sparse file to host the image
 334         truncate("-s", "%d" % disk_size, image)
 335
 336         self._create_partition_table(image)
 337         end_sector = self._shrink_partitions(image)
 338
 339         # Check if the available space is enough to host the image
 340         dirname = os.path.dirname(image)
 341         size = (end_sector + 1) * self.disk.device.sectorSize
 342         self.out.output("Examining available space in %s ..." % dirname, False)
 343         stat = os.statvfs(dirname)
 344         available = stat.f_bavail * stat.f_frsize
 345         if available <= size:
 346             raise FatalError('Not enough space in %s to host the image' %
 347                              dirname)
 348         self.out.success("sufficient")
 349
 350         self._create_filesystems(image)
 351
 352         return image
 353
 354 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :