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