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