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