Enforce raw img format in libguestfs
[snf-image-creator] / image_creator / disk.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 from image_creator.util import get_command
35 from image_creator.util import FatalError
36 from image_creator.util import try_fail_repeat
37 from image_creator.util import free_space
38 from image_creator.gpt import GPTPartitionTable
39 from image_creator.bundle_volume import BundleVolume
40
41 import stat
42 import os
43 import tempfile
44 import uuid
45 import re
46 import sys
47 import guestfs
48 import shutil
49 from sendfile import sendfile
50
51
52 dd = get_command('dd')
53 dmsetup = get_command('dmsetup')
54 losetup = get_command('losetup')
55 blockdev = get_command('blockdev')
56
57
58 TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']
59
60
61 class Disk(object):
62     """This class represents a hard disk hosting an Operating System
63
64     A Disk instance never alters the source media it is created from.
65     Any change is done on a snapshot created by the device-mapper of
66     the Linux kernel.
67     """
68
69     def __init__(self, source, output, tmp=None):
70         """Create a new Disk instance out of a source media. The source
71         media can be an image file, a block device or a directory.
72         """
73         self._cleanup_jobs = []
74         self._devices = []
75         self.source = source
76         self.out = output
77         self.meta = {}
78         self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.',
79                                     dir=self._get_tmp_dir(tmp))
80
81         self._add_cleanup(shutil.rmtree, self.tmp)
82
83     def _get_tmp_dir(self, default=None):
84         if default is not None:
85             return default
86
87         space = map(free_space, TMP_CANDIDATES)
88
89         max_idx = 0
90         max_val = space[0]
91         for i, val in zip(range(len(space)), space):
92             if val > max_val:
93                 max_val = val
94                 max_idx = i
95
96         # Return the candidate path with more available space
97         return TMP_CANDIDATES[max_idx]
98
99     def _add_cleanup(self, job, *args):
100         self._cleanup_jobs.append((job, args))
101
102     def _losetup(self, fname):
103         loop = losetup('-f', '--show', fname)
104         loop = loop.strip()  # remove the new-line char
105         self._add_cleanup(try_fail_repeat, losetup, '-d', loop)
106         return loop
107
108     def _dir_to_disk(self):
109         if self.source == '/':
110             bundle = BundleVolume(self.out, self.meta)
111             image = '%s/%s.diskdump' % (self.tmp, uuid.uuid4().hex)
112
113             def check_unlink(path):
114                 if os.path.exists(path):
115                     os.unlink(path)
116
117             self._add_cleanup(check_unlink, image)
118             bundle.create_image(image)
119             return self._losetup(image)
120         raise FatalError("Using a directory as media source is supported")
121
122     def cleanup(self):
123         """Cleanup internal data. This needs to be called before the
124         program ends.
125         """
126         try:
127             while len(self._devices):
128                 device = self._devices.pop()
129                 device.destroy()
130         finally:
131             # Make sure those are executed even if one of the device.destroy
132             # methods throws exeptions.
133             while len(self._cleanup_jobs):
134                 job, args = self._cleanup_jobs.pop()
135                 job(*args)
136
137     def snapshot(self):
138         """Creates a snapshot of the original source media of the Disk
139         instance.
140         """
141
142         self.out.output("Examining source media `%s' ..." % self.source, False)
143         sourcedev = self.source
144         mode = os.stat(self.source).st_mode
145         if stat.S_ISDIR(mode):
146             self.out.success('looks like a directory')
147             return self._dir_to_disk()
148         elif stat.S_ISREG(mode):
149             self.out.success('looks like an image file')
150             sourcedev = self._losetup(self.source)
151         elif not stat.S_ISBLK(mode):
152             raise ValueError("Invalid media source. Only block devices, "
153                              "regular files and directories are supported.")
154         else:
155             self.out.success('looks like a block device')
156
157         # Take a snapshot and return it to the user
158         self.out.output("Snapshotting media source...", False)
159         size = blockdev('--getsz', sourcedev)
160         cowfd, cow = tempfile.mkstemp(dir=self.tmp)
161         os.close(cowfd)
162         self._add_cleanup(os.unlink, cow)
163         # Create cow sparse file
164         dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
165         cowdev = self._losetup(cow)
166
167         snapshot = uuid.uuid4().hex
168         tablefd, table = tempfile.mkstemp()
169         try:
170             os.write(tablefd, "0 %d snapshot %s %s n 8" %
171                               (int(size), sourcedev, cowdev))
172             dmsetup('create', snapshot, table)
173             self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
174
175         finally:
176             os.unlink(table)
177         self.out.success('done')
178         return "/dev/mapper/%s" % snapshot
179
180     def get_device(self, media):
181         """Returns a newly created DiskDevice instance."""
182
183         new_device = DiskDevice(media, self.out)
184         self._devices.append(new_device)
185         new_device.enable()
186         return new_device
187
188     def destroy_device(self, device):
189         """Destroys a DiskDevice instance previously created by
190         get_device method.
191         """
192         self._devices.remove(device)
193         device.destroy()
194
195
196 class DiskDevice(object):
197     """This class represents a block device hosting an Operating System
198     as created by the device-mapper.
199     """
200
201     def __init__(self, device, output, bootable=True, meta={}):
202         """Create a new DiskDevice."""
203
204         self.real_device = device
205         self.out = output
206         self.bootable = bootable
207         self.meta = meta
208         self.progress_bar = None
209         self.guestfs_device = None
210         self.size = 0
211
212         self.g = guestfs.GuestFS()
213         self.g.add_drive_opts(self.real_device, readonly=0, format="raw")
214
215         # Before version 1.17.14 the recovery process, which is a fork of the
216         # original process that called libguestfs, did not close its inherited
217         # file descriptors. This can cause problems especially if the parent
218         # process has opened pipes. Since the recovery process is an optional
219         # feature of libguestfs, it's better to disable it.
220         self.g.set_recovery_proc(0)
221         version = self.g.version()
222         if version['major'] > 1 or \
223             (version['major'] == 1 and (version['minor'] >= 18 or
224                                         (version['minor'] == 17 and
225                                          version['release'] >= 14))):
226             self.g.set_recovery_proc(1)
227             self.out.output("Enabling recovery proc")
228
229         #self.g.set_trace(1)
230         #self.g.set_verbose(1)
231
232         self.guestfs_enabled = False
233
234     def enable(self):
235         """Enable a newly created DiskDevice"""
236
237         self.out.output('Launching helper VM (may take a while) ...', False)
238         # self.progressbar = self.out.Progress(100, "Launching helper VM",
239         #                                     "percent")
240         # eh = self.g.set_event_callback(self.progress_callback,
241         #                               guestfs.EVENT_PROGRESS)
242         self.g.launch()
243         self.guestfs_enabled = True
244         # self.g.delete_event_callback(eh)
245         # self.progressbar.success('done')
246         # self.progressbar = None
247         self.out.success('done')
248
249         self.out.output('Inspecting Operating System ...', False)
250         roots = self.g.inspect_os()
251         if len(roots) == 0:
252             raise FatalError("No operating system found")
253         if len(roots) > 1:
254             raise FatalError("Multiple operating systems found."
255                              "We only support images with one OS.")
256         self.root = roots[0]
257         self.guestfs_device = self.g.part_to_dev(self.root)
258         self.size = self.g.blockdev_getsize64(self.guestfs_device)
259         self.meta['PARTITION_TABLE'] = \
260             self.g.part_get_parttype(self.guestfs_device)
261
262         self.ostype = self.g.inspect_get_type(self.root)
263         self.distro = self.g.inspect_get_distro(self.root)
264         self.out.success('found a(n) %s system' % self.distro)
265
266     def destroy(self):
267         """Destroy this DiskDevice instance."""
268
269         # In new guestfs versions, there is a handy shutdown method for this
270         try:
271             if self.guestfs_enabled:
272                 self.g.umount_all()
273                 self.g.sync()
274         finally:
275             # Close the guestfs handler if open
276             self.g.close()
277
278 #    def progress_callback(self, ev, eh, buf, array):
279 #        position = array[2]
280 #        total = array[3]
281 #
282 #        self.progressbar.goto((position * 100) // total)
283
284     def mount(self, readonly=False):
285         """Mount all disk partitions in a correct order."""
286
287         mount = self.g.mount_ro if readonly else self.g.mount
288         msg = " read-only" if readonly else ""
289         self.out.output("Mounting the media%s ..." % msg, False)
290         mps = self.g.inspect_get_mountpoints(self.root)
291
292         # Sort the keys to mount the fs in a correct order.
293         # / should be mounted befor /boot, etc
294         def compare(a, b):
295             if len(a[0]) > len(b[0]):
296                 return 1
297             elif len(a[0]) == len(b[0]):
298                 return 0
299             else:
300                 return -1
301         mps.sort(compare)
302         for mp, dev in mps:
303             try:
304                 mount(dev, mp)
305             except RuntimeError as msg:
306                 self.out.warn("%s (ignored)" % msg)
307         self.out.success("done")
308
309     def umount(self):
310         """Umount all mounted filesystems."""
311         self.g.umount_all()
312
313     def _last_partition(self):
314         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
315             msg = "Unsupported partition table: %s. Only msdos and gpt " \
316                 "partition tables are supported" % self.meta['PARTITION_TABLE']
317             raise FatalError(msg)
318
319         is_extended = lambda p: \
320             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
321             in (0x5, 0xf)
322         is_logical = lambda p: \
323             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
324
325         partitions = self.g.part_list(self.guestfs_device)
326         last_partition = partitions[-1]
327
328         if is_logical(last_partition):
329             # The disk contains extended and logical partitions....
330             extended = filter(is_extended, partitions)[0]
331             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
332
333             # check if extended is the last primary partition
334             if last_primary['part_num'] > extended['part_num']:
335                 last_partition = last_primary
336
337         return last_partition
338
339     def shrink(self):
340         """Shrink the disk.
341
342         This is accomplished by shrinking the last filesystem in the
343         disk and then updating the partition table. The new disk size
344         (in bytes) is returned.
345
346         ATTENTION: make sure unmount is called before shrink
347         """
348         get_fstype = lambda p: \
349             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
350         is_logical = lambda p: \
351             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
352         is_extended = lambda p: \
353             self.meta['PARTITION_TABLE'] == 'msdos' and \
354             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
355             in (0x5, 0xf)
356
357         part_add = lambda ptype, start, stop: \
358             self.g.part_add(self.guestfs_device, ptype, start, stop)
359         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
360         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
361         part_set_id = lambda p, id: \
362             self.g.part_set_mbr_id(self.guestfs_device, p, id)
363         part_get_bootable = lambda p: \
364             self.g.part_get_bootable(self.guestfs_device, p)
365         part_set_bootable = lambda p, bootable: \
366             self.g.part_set_bootable(self.guestfs_device, p, bootable)
367
368         MB = 2 ** 20
369
370         self.out.output("Shrinking image (this may take a while) ...", False)
371
372         sector_size = self.g.blockdev_getss(self.guestfs_device)
373
374         last_part = None
375         fstype = None
376         while True:
377             last_part = self._last_partition()
378             fstype = get_fstype(last_part)
379
380             if fstype == 'swap':
381                 self.meta['SWAP'] = "%d:%s" % \
382                     (last_part['part_num'],
383                      (last_part['part_size'] + MB - 1) // MB)
384                 part_del(last_part['part_num'])
385                 continue
386             elif is_extended(last_part):
387                 part_del(last_part['part_num'])
388                 continue
389
390             # Most disk manipulation programs leave 2048 sectors after the last
391             # partition
392             new_size = last_part['part_end'] + 1 + 2048 * sector_size
393             self.size = min(self.size, new_size)
394             break
395
396         if not re.match("ext[234]", fstype):
397             self.out.warn("Don't know how to resize %s partitions." % fstype)
398             return self.size
399
400         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
401         self.g.e2fsck_f(part_dev)
402         self.g.resize2fs_M(part_dev)
403
404         out = self.g.tune2fs_l(part_dev)
405         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
406         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
407
408         start = last_part['part_start'] / sector_size
409         end = start + (block_size * block_cnt) / sector_size - 1
410
411         if is_logical(last_part):
412             partitions = self.g.part_list(self.guestfs_device)
413
414             logical = []  # logical partitions
415             for partition in partitions:
416                 if partition['part_num'] < 4:
417                     continue
418                 logical.append({
419                     'num': partition['part_num'],
420                     'start': partition['part_start'] / sector_size,
421                     'end': partition['part_end'] / sector_size,
422                     'id': part_get_id(partition['part_num']),
423                     'bootable': part_get_bootable(partition['part_num'])
424                 })
425
426             logical[-1]['end'] = end  # new end after resize
427
428             # Recreate the extended partition
429             extended = filter(is_extended, partitions)[0]
430             part_del(extended['part_num'])
431             part_add('e', extended['part_start'] / sector_size, end)
432
433             # Create all the logical partitions back
434             for l in logical:
435                 part_add('l', l['start'], l['end'])
436                 part_set_id(l['num'], l['id'])
437                 part_set_bootable(l['num'], l['bootable'])
438         else:
439             # Recreate the last partition
440             if self.meta['PARTITION_TABLE'] == 'msdos':
441                 last_part['id'] = part_get_id(last_part['part_num'])
442
443             last_part['bootable'] = part_get_bootable(last_part['part_num'])
444             part_del(last_part['part_num'])
445             part_add('p', start, end)
446             part_set_bootable(last_part['part_num'], last_part['bootable'])
447
448             if self.meta['PARTITION_TABLE'] == 'msdos':
449                 part_set_id(last_part['part_num'], last_part['id'])
450
451         new_size = (end + 1) * sector_size
452
453         assert (new_size <= self.size)
454
455         if self.meta['PARTITION_TABLE'] == 'gpt':
456             ptable = GPTPartitionTable(self.real_device)
457             self.size = ptable.shrink(new_size, self.size)
458         else:
459             self.size = min(new_size + 2048 * sector_size, self.size)
460
461         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
462
463         return self.size
464
465     def dump(self, outfile):
466         """Dumps the content of device into a file.
467
468         This method will only dump the actual payload, found by reading the
469         partition table. Empty space in the end of the device will be ignored.
470         """
471         MB = 2 ** 20
472         blocksize = 4 * MB  # 4MB
473         size = self.size
474         progr_size = (size + MB - 1) // MB  # in MB
475         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
476
477         with open(self.real_device, 'r') as src:
478             with open(outfile, "w") as dst:
479                 left = size
480                 offset = 0
481                 progressbar.next()
482                 while left > 0:
483                     length = min(left, blocksize)
484                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
485
486                     # Workaround for python-sendfile API change. In
487                     # python-sendfile 1.2.x (py-sendfile) the returning value
488                     # of sendfile is a tuple, where in version 2.x (pysendfile)
489                     # it is just a sigle integer.
490                     if isinstance(sent, tuple):
491                         sent = sent[1]
492
493                     offset += sent
494                     left -= sent
495                     progressbar.goto((size - left) // MB)
496         progressbar.success('image file %s was successfully created' % outfile)
497
498 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :