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