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