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