Remove temporary bundeled image during cleanup
[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         self.progressbar = self.out.Progress(100, "Launching helper VM",
209                                              "percent")
210         eh = self.g.set_event_callback(self.progress_callback,
211                                        guestfs.EVENT_PROGRESS)
212         self.g.launch()
213         self.guestfs_enabled = True
214         self.g.delete_event_callback(eh)
215         self.progressbar.success('done')
216         self.progressbar = None
217
218         self.out.output('Inspecting Operating System ...', False)
219         roots = self.g.inspect_os()
220         if len(roots) == 0:
221             raise FatalError("No operating system found")
222         if len(roots) > 1:
223             raise FatalError("Multiple operating systems found."
224                              "We only support images with one OS.")
225         self.root = roots[0]
226         self.guestfs_device = self.g.part_to_dev(self.root)
227         self.size = self.g.blockdev_getsize64(self.guestfs_device)
228         self.meta['PARTITION_TABLE'] = \
229             self.g.part_get_parttype(self.guestfs_device)
230
231         self.ostype = self.g.inspect_get_type(self.root)
232         self.distro = self.g.inspect_get_distro(self.root)
233         self.out.success('found a(n) %s system' % self.distro)
234
235     def destroy(self):
236         """Destroy this DiskDevice instance."""
237
238         # In new guestfs versions, there is a handy shutdown method for this
239         try:
240             if self.guestfs_enabled:
241                 self.g.umount_all()
242                 self.g.sync()
243         finally:
244             # Close the guestfs handler if open
245             self.g.close()
246
247     def progress_callback(self, ev, eh, buf, array):
248         position = array[2]
249         total = array[3]
250
251         self.progressbar.goto((position * 100) // total)
252
253     def mount(self, readonly=False):
254         """Mount all disk partitions in a correct order."""
255
256         mount = self.g.mount_ro if readonly else self.g.mount
257         msg = " read-only" if readonly else ""
258         self.out.output("Mounting the media%s..." % msg, False)
259         mps = self.g.inspect_get_mountpoints(self.root)
260
261         # Sort the keys to mount the fs in a correct order.
262         # / should be mounted befor /boot, etc
263         def compare(a, b):
264             if len(a[0]) > len(b[0]):
265                 return 1
266             elif len(a[0]) == len(b[0]):
267                 return 0
268             else:
269                 return -1
270         mps.sort(compare)
271         for mp, dev in mps:
272             try:
273                 mount(dev, mp)
274             except RuntimeError as msg:
275                 self.out.warn("%s (ignored)" % msg)
276         self.out.success("done")
277
278     def umount(self):
279         """Umount all mounted filesystems."""
280         self.g.umount_all()
281
282     def _last_partition(self):
283         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
284             msg = "Unsupported partition table: %s. Only msdos and gpt " \
285                 "partition tables are supported" % self.meta['PARTITION_TABLE']
286             raise FatalError(msg)
287
288         is_extended = lambda p: \
289             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
290         is_logical = lambda p: \
291             self.meta['PARTITION_TABLE'] != 'msdos' and p['part_num'] > 4
292
293         partitions = self.g.part_list(self.guestfs_device)
294         last_partition = partitions[-1]
295
296         if is_logical(last_partition):
297             # The disk contains extended and logical partitions....
298             extended = [p for p in partitions if is_extended(p)][0]
299             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
300
301             # check if extended is the last primary partition
302             if last_primary['part_num'] > extended['part_num']:
303                 last_partition = last_primary
304
305         return last_partition
306
307     def shrink(self):
308         """Shrink the disk.
309
310         This is accomplished by shrinking the last filesystem in the
311         disk and then updating the partition table. The new disk size
312         (in bytes) is returned.
313
314         ATTENTION: make sure unmount is called before shrink
315         """
316         get_fstype = lambda p: \
317             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
318         is_logical = lambda p: \
319             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
320         is_extended = lambda p: \
321             self.meta['PARTITION_TABLE'] == 'msdos' and \
322             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
323
324         part_add = lambda ptype, start, stop: \
325             self.g.part_add(self.guestfs_device, ptype, start, stop)
326         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
327         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
328         part_set_id = lambda p, id: \
329             self.g.part_set_mbr_id(self.guestfs_device, p, id)
330         part_get_bootable = lambda p: \
331             self.g.part_get_bootable(self.guestfs_device, p)
332         part_set_bootable = lambda p, bootable: \
333             self.g.part_set_bootable(self.guestfs_device, p, bootable)
334
335         MB = 2 ** 20
336
337         self.out.output("Shrinking image (this may take a while)...", False)
338
339         sector_size = self.g.blockdev_getss(self.guestfs_device)
340
341         last_part = None
342         fstype = None
343         while True:
344             last_part = self._last_partition()
345             fstype = get_fstype(last_part)
346
347             if fstype == 'swap':
348                 self.meta['SWAP'] = "%d:%s" % \
349                     (last_part['part_num'],
350                      (last_part['part_size'] + MB - 1) // MB)
351                 part_del(last_part['part_num'])
352                 continue
353             elif is_extended(last_part):
354                 part_del(last_part['part_num'])
355                 continue
356
357             # Most disk manipulation programs leave 2048 sectors after the last
358             # partition
359             new_size = last_part['part_end'] + 1 + 2048 * sector_size
360             self.size = min(self.size, new_size)
361             break
362
363         if not re.match("ext[234]", fstype):
364             self.out.warn("Don't know how to resize %s partitions." % fstype)
365             return self.size
366
367         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
368         self.g.e2fsck_f(part_dev)
369         self.g.resize2fs_M(part_dev)
370
371         out = self.g.tune2fs_l(part_dev)
372         block_size = int(
373             filter(lambda x: x[0] == 'Block size', out)[0][1])
374         block_cnt = int(
375             filter(lambda x: x[0] == 'Block count', out)[0][1])
376
377         start = last_part['part_start'] / sector_size
378         end = start + (block_size * block_cnt) / sector_size - 1
379
380         if is_logical(last_part):
381             partitions = self.g.part_list(self.guestfs_device)
382
383             logical = []  # logical partitions
384             for partition in partitions:
385                 if partition['part_num'] < 4:
386                     continue
387                 logical.append({
388                     'num': partition['part_num'],
389                     'start': partition['part_start'] / sector_size,
390                     'end': partition['part_end'] / sector_size,
391                     'id': part_get_(partition['part_num']),
392                     'bootable': part_get_bootable(partition['part_num'])
393                 })
394
395             logical[-1]['end'] = end  # new end after resize
396
397             # Recreate the extended partition
398             extended = [p for p in partitions if self._is_extended(p)][0]
399             part_del(extended['part_num'])
400             part_add('e', extended['part_start'], end)
401
402             # Create all the logical partitions back
403             for l in logical:
404                 part_add('l', l['start'], l['end'])
405                 part_set_id(l['num'], l['id'])
406                 part_set_bootable(l['num'], l['bootable'])
407         else:
408             # Recreate the last partition
409             if self.meta['PARTITION_TABLE'] == 'msdos':
410                 last_part['id'] = part_get_id(last_part['part_num'])
411
412             last_part['bootable'] = part_get_bootable(last_part['part_num'])
413             part_del(last_part['part_num'])
414             part_add('p', start, end)
415             part_set_bootable(last_part['part_num'], last_part['bootable'])
416
417             if self.meta['PARTITION_TABLE'] == 'msdos':
418                 part_set_id(last_part['part_num'], last_part['id'])
419
420         new_size = (end + 1) * sector_size
421
422         assert (new_size <= self.size)
423
424         if self.meta['PARTITION_TABLE'] == 'gpt':
425             ptable = GPTPartitionTable(self.real_device)
426             self.size = ptable.shrink(new_size, self.size)
427         else:
428             self.size = min(new_size + 2048 * sector_size, self.size)
429
430         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
431
432         return self.size
433
434     def dump(self, outfile):
435         """Dumps the content of device into a file.
436
437         This method will only dump the actual payload, found by reading the
438         partition table. Empty space in the end of the device will be ignored.
439         """
440         MB = 2 ** 20
441         blocksize = 4 * MB  # 4MB
442         size = self.size
443         progr_size = (size + MB - 1) // MB  # in MB
444         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
445
446         with open(self.real_device, 'r') as src:
447             with open(outfile, "w") as dst:
448                 left = size
449                 offset = 0
450                 progressbar.next()
451                 while left > 0:
452                     length = min(left, blocksize)
453                     _, sent = sendfile(dst.fileno(), src.fileno(), offset,
454                         length)
455                     offset += sent
456                     left -= sent
457                     progressbar.goto((size - left) // MB)
458         progressbar.success('image file %s was successfully created' % outfile)
459
460 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :