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