Bump version to 0.2.8
[snf-image-creator] / image_creator / image.py
1 # Copyright 2013 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 from image_creator.util import FatalError
35 from image_creator.gpt import GPTPartitionTable
36 from image_creator.os_type import os_cls
37
38 import re
39 import guestfs
40 from sendfile import sendfile
41
42
43 class Image(object):
44     """The instances of this class can create images out of block devices."""
45
46     def __init__(self, device, output, bootable=True, meta={}):
47         """Create a new Image instance"""
48
49         self.device = device
50         self.out = output
51         self.bootable = bootable
52         self.meta = meta
53         self.progress_bar = None
54         self.guestfs_device = None
55         self.size = 0
56         self.mounted = False
57
58         self.g = guestfs.GuestFS()
59         self.g.add_drive_opts(self.device, readonly=0, format="raw")
60
61         # Before version 1.17.14 the recovery process, which is a fork of the
62         # original process that called libguestfs, did not close its inherited
63         # file descriptors. This can cause problems especially if the parent
64         # process has opened pipes. Since the recovery process is an optional
65         # feature of libguestfs, it's better to disable it.
66         self.g.set_recovery_proc(0)
67         version = self.g.version()
68         if version['major'] > 1 or \
69             (version['major'] == 1 and (version['minor'] >= 18 or
70                                         (version['minor'] == 17 and
71                                          version['release'] >= 14))):
72             self.g.set_recovery_proc(1)
73             self.out.output("Enabling recovery proc")
74
75         #self.g.set_trace(1)
76         #self.g.set_verbose(1)
77
78         self.guestfs_enabled = False
79
80     def enable(self):
81         """Enable a newly created Image instance"""
82
83         self.out.output('Launching helper VM (may take a while) ...', False)
84         # self.progressbar = self.out.Progress(100, "Launching helper VM",
85         #                                     "percent")
86         # eh = self.g.set_event_callback(self.progress_callback,
87         #                               guestfs.EVENT_PROGRESS)
88         self.g.launch()
89         self.guestfs_enabled = True
90         # self.g.delete_event_callback(eh)
91         # self.progressbar.success('done')
92         # self.progressbar = None
93         self.out.success('done')
94
95         self.out.output('Inspecting Operating System ...', False)
96         roots = self.g.inspect_os()
97         if len(roots) == 0:
98             raise FatalError("No operating system found")
99         if len(roots) > 1:
100             raise FatalError("Multiple operating systems found."
101                              "We only support images with one OS.")
102         self.root = roots[0]
103         self.guestfs_device = self.g.part_to_dev(self.root)
104         self.size = self.g.blockdev_getsize64(self.guestfs_device)
105         self.meta['PARTITION_TABLE'] = \
106             self.g.part_get_parttype(self.guestfs_device)
107
108         self.ostype = self.g.inspect_get_type(self.root)
109         self.distro = self.g.inspect_get_distro(self.root)
110         self.out.success('found a(n) %s system' % self.distro)
111
112     def _get_os(self):
113         """Return an OS class instance for this image"""
114         if hasattr(self, "_os"):
115             return self._os
116
117         if not self.guestfs_enabled:
118             self.enable()
119
120         if not self.mounted:
121             do_unmount = True
122             self.mount(readonly=True)
123         else:
124             do_unmount = False
125
126         try:
127             cls = os_cls(self.distro, self.ostype)
128             self._os = cls(self.root, self.g, self.out)
129
130         finally:
131             if do_unmount:
132                 self.umount()
133
134         return self._os
135
136     os = property(_get_os)
137
138     def destroy(self):
139         """Destroy this Image instance."""
140
141         # In new guestfs versions, there is a handy shutdown method for this
142         try:
143             if self.guestfs_enabled:
144                 self.g.umount_all()
145                 self.g.sync()
146         finally:
147             # Close the guestfs handler if open
148             self.g.close()
149
150 #    def progress_callback(self, ev, eh, buf, array):
151 #        position = array[2]
152 #        total = array[3]
153 #
154 #        self.progressbar.goto((position * 100) // total)
155
156     def mount(self, readonly=False):
157         """Mount all disk partitions in a correct order."""
158
159         mount = self.g.mount_ro if readonly else self.g.mount
160         msg = " read-only" if readonly else ""
161         self.out.output("Mounting the media%s ..." % msg, False)
162         mps = self.g.inspect_get_mountpoints(self.root)
163
164         # Sort the keys to mount the fs in a correct order.
165         # / should be mounted befor /boot, etc
166         def compare(a, b):
167             if len(a[0]) > len(b[0]):
168                 return 1
169             elif len(a[0]) == len(b[0]):
170                 return 0
171             else:
172                 return -1
173         mps.sort(compare)
174         for mp, dev in mps:
175             try:
176                 mount(dev, mp)
177             except RuntimeError as msg:
178                 self.out.warn("%s (ignored)" % msg)
179
180         self.mounted = True
181         self.out.success("done")
182
183     def umount(self):
184         """Umount all mounted filesystems."""
185         self.g.umount_all()
186         self.mounted = False
187
188     def _last_partition(self):
189         """Return the last partition of the image disk"""
190         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
191             msg = "Unsupported partition table: %s. Only msdos and gpt " \
192                 "partition tables are supported" % self.meta['PARTITION_TABLE']
193             raise FatalError(msg)
194
195         is_extended = lambda p: \
196             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
197             in (0x5, 0xf)
198         is_logical = lambda p: \
199             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
200
201         partitions = self.g.part_list(self.guestfs_device)
202         last_partition = partitions[-1]
203
204         if is_logical(last_partition):
205             # The disk contains extended and logical partitions....
206             extended = filter(is_extended, partitions)[0]
207             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
208
209             # check if extended is the last primary partition
210             if last_primary['part_num'] > extended['part_num']:
211                 last_partition = last_primary
212
213         return last_partition
214
215     def shrink(self):
216         """Shrink the image.
217
218         This is accomplished by shrinking the last file system of the
219         image and then updating the partition table. The new disk size
220         (in bytes) is returned.
221
222         ATTENTION: make sure unmount is called before shrink
223         """
224         get_fstype = lambda p: \
225             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
226         is_logical = lambda p: \
227             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
228         is_extended = lambda p: \
229             self.meta['PARTITION_TABLE'] == 'msdos' and \
230             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
231             in (0x5, 0xf)
232
233         part_add = lambda ptype, start, stop: \
234             self.g.part_add(self.guestfs_device, ptype, start, stop)
235         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
236         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
237         part_set_id = lambda p, id: \
238             self.g.part_set_mbr_id(self.guestfs_device, p, id)
239         part_get_bootable = lambda p: \
240             self.g.part_get_bootable(self.guestfs_device, p)
241         part_set_bootable = lambda p, bootable: \
242             self.g.part_set_bootable(self.guestfs_device, p, bootable)
243
244         MB = 2 ** 20
245
246         self.out.output("Shrinking image (this may take a while) ...", False)
247
248         sector_size = self.g.blockdev_getss(self.guestfs_device)
249
250         last_part = None
251         fstype = None
252         while True:
253             last_part = self._last_partition()
254             fstype = get_fstype(last_part)
255
256             if fstype == 'swap':
257                 self.meta['SWAP'] = "%d:%s" % \
258                     (last_part['part_num'],
259                      (last_part['part_size'] + MB - 1) // MB)
260                 part_del(last_part['part_num'])
261                 continue
262             elif is_extended(last_part):
263                 part_del(last_part['part_num'])
264                 continue
265
266             # Most disk manipulation programs leave 2048 sectors after the last
267             # partition
268             new_size = last_part['part_end'] + 1 + 2048 * sector_size
269             self.size = min(self.size, new_size)
270             break
271
272         if not re.match("ext[234]", fstype):
273             self.out.warn("Don't know how to resize %s partitions." % fstype)
274             return self.size
275
276         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
277         self.g.e2fsck_f(part_dev)
278         self.g.resize2fs_M(part_dev)
279
280         out = self.g.tune2fs_l(part_dev)
281         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
282         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
283
284         start = last_part['part_start'] / sector_size
285         end = start + (block_size * block_cnt) / sector_size - 1
286
287         if is_logical(last_part):
288             partitions = self.g.part_list(self.guestfs_device)
289
290             logical = []  # logical partitions
291             for partition in partitions:
292                 if partition['part_num'] < 4:
293                     continue
294                 logical.append({
295                     'num': partition['part_num'],
296                     'start': partition['part_start'] / sector_size,
297                     'end': partition['part_end'] / sector_size,
298                     'id': part_get_id(partition['part_num']),
299                     'bootable': part_get_bootable(partition['part_num'])
300                 })
301
302             logical[-1]['end'] = end  # new end after resize
303
304             # Recreate the extended partition
305             extended = filter(is_extended, partitions)[0]
306             part_del(extended['part_num'])
307             part_add('e', extended['part_start'] / sector_size, end)
308
309             # Create all the logical partitions back
310             for l in logical:
311                 part_add('l', l['start'], l['end'])
312                 part_set_id(l['num'], l['id'])
313                 part_set_bootable(l['num'], l['bootable'])
314         else:
315             # Recreate the last partition
316             if self.meta['PARTITION_TABLE'] == 'msdos':
317                 last_part['id'] = part_get_id(last_part['part_num'])
318
319             last_part['bootable'] = part_get_bootable(last_part['part_num'])
320             part_del(last_part['part_num'])
321             part_add('p', start, end)
322             part_set_bootable(last_part['part_num'], last_part['bootable'])
323
324             if self.meta['PARTITION_TABLE'] == 'msdos':
325                 part_set_id(last_part['part_num'], last_part['id'])
326
327         new_size = (end + 1) * sector_size
328
329         assert (new_size <= self.size)
330
331         if self.meta['PARTITION_TABLE'] == 'gpt':
332             ptable = GPTPartitionTable(self.device)
333             self.size = ptable.shrink(new_size, self.size)
334         else:
335             self.size = min(new_size + 2048 * sector_size, self.size)
336
337         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
338
339         return self.size
340
341     def dump(self, outfile):
342         """Dumps the content of the image into a file.
343
344         This method will only dump the actual payload, found by reading the
345         partition table. Empty space in the end of the device will be ignored.
346         """
347         MB = 2 ** 20
348         blocksize = 4 * MB  # 4MB
349         size = self.size
350         progr_size = (size + MB - 1) // MB  # in MB
351         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
352
353         with open(self.device, 'r') as src:
354             with open(outfile, "w") as dst:
355                 left = size
356                 offset = 0
357                 progressbar.next()
358                 while left > 0:
359                     length = min(left, blocksize)
360                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
361
362                     # Workaround for python-sendfile API change. In
363                     # python-sendfile 1.2.x (py-sendfile) the returning value
364                     # of sendfile is a tuple, where in version 2.x (pysendfile)
365                     # it is just a sigle integer.
366                     if isinstance(sent, tuple):
367                         sent = sent[1]
368
369                     offset += sent
370                     left -= sent
371                     progressbar.goto((size - left) // MB)
372         progressbar.success('image file %s was successfully created' % outfile)
373
374 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :