Fix bugs in bundle_volume module
[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 from collections import namedtuple
39
40 import parted
41
42 from image_creator.util import get_command
43 from image_creator.util import FatalError
44
45 findfs = get_command('findfs')
46 truncate = get_command('truncate')
47 dd = get_command('dd')
48 dmsetup = get_command('dmsetup')
49
50
51 class BundleVolume():
52     _FileSystemEntry = namedtuple('FileSystemEntry',
53                                   'dev mpoint fs opts freq passno')
54
55     _Partition = namedtuple('Partition', 'num start end type fs mopts')
56
57     def __init__(self, out, meta):
58         self.out = out
59         self.meta = meta
60
61         self.out.output('Searching for root device ...', False)
62         root = self._get_root_partition()
63
64         if root.startswith("UUID=") or root.startswith("LABEL="):
65             self.root = findfs(root).stdout.strip()
66         else:
67             self.root = root
68
69         if not re.match('/dev/[hsv]d[a-z][1-9]*$', self.root):
70             raise FatalError("Don't know how to handle root device: %s" % \
71                              self.root)
72
73         self.disk = re.split('[0-9]', self.root)[0]
74
75         out.success('%s' % self.root)
76
77     def _read_fstable(self, f):
78         if not os.path.isfile(f):
79             raise FatalError("Unable to open: `%s'. File is missing." % f)
80
81         with open(f) as table:
82             for line in iter(table):
83                 entry = line.split('#')[0].strip().split()
84                 if len(entry) != 6:
85                     continue
86                 yield self._FileSystemEntry(*entry)
87
88     def _get_root_partition(self):
89         for entry in self._read_fstable('/etc/fstab'):
90             if entry.mpoint == '/':
91                 return entry.dev
92
93         raise FatalError("Unable to find root device in /etc/fstab")
94
95     def _is_mpoint(self, path):
96         for entry in self._read_fstable('/proc/mounts'):
97             if entry.mpoint == path:
98                 return True
99         return False
100
101     def _mount_options(self, device):
102         for entry in self._read_fstable('/proc/mounts'):
103             if not entry.dev.startswith('/'):
104                 continue
105
106             if os.path.realpath(entry.dev) == os.path.realpath(device):
107                 return entry
108
109         return
110
111     def _create_partition_table(self, src_disk, dest_file):
112
113         if src_disk.type != 'msdos':
114             raise FatalError('Only msdos partition tables are supported')
115
116         first_sector = src_disk.getPrimaryPartitions()[0].geometry.start
117
118         # Copy the MBR and the space between the MBR and the first partition.
119         # In Grub version 1 Stage 1.5 is located there.
120         first_sector = src_disk.getPrimaryPartitions()[0].geometry.start
121
122         dd('if=%s' % src_disk.device.path, 'of=%s' % dest_file,
123            'bs=%d' % src_disk.device.sectorSize,
124            'count=%d' % first_sector, 'conv=notrunc')
125
126         # Create the Extended boot records (EBRs) in the image
127         extended = src_disk.getExtendedPartition()
128         if not extended:
129             return
130
131         # Extended boot records precede the logical partitions they describe
132         logical = src_disk.getLogicalPartitions()
133         start = extended.geometry.start
134         for i in range(len(logical)):
135             end = logical[i].geometry.start - 1
136             dd('if=%s' % src_disk.device.path, 'of=%s' % dest_file,
137                'count=%d' % (end - start + 1), 'conv=notrunc',
138                'seek=%d' % start, 'skip=%d' % start)
139             start = logical[i].geometry.end + 1
140
141     def _shrink_partitions(self, src_disk, image_file):
142
143         partitions = []
144         new_end = 0
145
146         image_dev = parted.Device(image_file)
147         image_disk = parted.Disk(image_dev)
148
149         is_extended = lambda p: p.type == parted.PARTITION_EXTENDED
150         is_logical = lambda p: p.type == parted.PARTITION_LOGICAL
151
152         partitions = []
153         for p in src_disk.partitions:
154             g = p.geometry
155             f = p.fileSystem
156             partitions.append(self._Partition(p.number, g.start, g.end,
157                               p.type, f.type if f is not None else '',
158                               self._mount_options(p.path)))
159
160         last = partitions[-1]
161         new_end = src_disk.device.getLength()
162         if last.fs == 'linux-swap(v1)':
163             MB = 2 ** 20
164             size = (last.end - last.start + 1) * src_disk.device.sectorSize
165             self.meta['SWAP'] = "%d:%s" % (last.num, (size + MB - 1) // MB)
166
167             image_disk.deletePartition(
168                 image_disk.getPartitionBySector(last.start))
169             image_disk.commit()
170
171             if is_logical(last) and last.num == 5:
172                 extended = image_disk.getExtendedPartition()
173                 image_disk.deletePartition(extended)
174                 image_disk.commit()
175                 partitions.remove(filter(is_extended, partitions)[0])
176
177             partitions.remove(last)
178             last = partitions[-1]
179
180             # Leave 2048 blocks at the end
181             new_end = last.end + 2048
182
183         if last.mopts.mpoint:
184             stat = os.statvfs(last.mopts.mpoint)
185             # Shrink the last partition. The new size should be the
186             # size of the occupied blocks
187             blcks = stat.f_blocks - stat.f_bavail
188             new_size = (blcks * stat.f_frsize) // src_disk.device.sectorSize
189
190             # Add 10% just to be on the safe side
191             part_end = last.start + (new_size * 11) // 10
192             # Alighn to 2048
193             part_end = ((part_end + 2047) // 2048) * 2048
194             last = last._replace(end=part_end)
195             partitions[-1] = last
196
197             # Leave 2048 blocks at the end.
198             new_end = new_size + 2048
199
200             image_disk.setPartitionGeometry(
201                 image_disk.getPartitionBySector(last.start),
202                 parted.Constraint(device=image_disk.device),
203                 start=last.start, end=last.end)
204             image_disk.commit()
205
206             if last.type == parted.PARTITION_LOGICAL:
207                 # Fix the extended partition
208                 extended = disk.getExtendedPartition()
209
210                 image_disk.setPartitionGeometry(extended,
211                     parted.Constraint(device=img_dev),
212                     ext.geometry.start, end=last.end)
213                 image_disk.commit()
214
215         # Check if the available space is enough to host the image
216         location = os.path.dirname(image_file)
217         size = (new_end + 1) * src_disk.device.sectorSize
218         self.out.output("Examining available space in %s..." % location, False)
219         stat = os.statvfs(location)
220         available = stat.f_bavail * stat.f_frsize
221         if available <= size:
222             raise FatalError('Not enough space in %s to host the image' % \
223                              location)
224         self.out.success("sufficient")
225
226         return partitions
227
228     def _fill_partitions(self, src_disk, image, partitions):
229         pass
230
231     def create_image(self):
232
233         image_file = '/mnt/%s.diskdump' % uuid.uuid4().hex
234
235         src_dev = parted.Device(self.disk)
236
237         size = src_dev.getLength() * src_dev.sectorSize
238
239         # Create sparse file to host the image
240         truncate("-s", "%d" % size, image_file)
241
242         src_disk = parted.Disk(src_dev)
243         self._create_partition_table(src_disk, image_file)
244         partitions = self._shrink_partitions(src_disk, image_file)
245         self._fill_partitions(src_disk, image_file, partitions)
246
247         return image_file
248
249 #       unmounted = filter(lambda p: not p.mopts.mpoint, partitions)
250 #        mounted = filter(lambda p: p.mopts.mpoint, partitions)
251 #
252 #        for p in unmounted:
253 #            dd('if=%s' % src_dev.path, 'of=%s' % img_dev.path,
254 #               'count=%d' % (p.end - p.start + 1), 'conv=notrunc',
255 #                'seek=%d' % p.start, 'skip=%d' % p.start)
256 #
257 #        partition_devices = create_devices(dest, partitions)
258 #
259 #        mounted.sort(key=lambda p: p.mopts.mpoint)
260 #
261 #        return img
262
263 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :