In bundle volume copy host files with rsync
[snf-image-creator] / 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 :