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