Remove the image size from the image metadata
[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.size = 0
172         self.meta = {}
173
174         self.g = guestfs.GuestFS()
175         self.g.add_drive_opts(self.real_device, readonly=0)
176
177         # Before version 1.17.14 the recovery process, which is a fork of the
178         # original process that called libguestfs, did not close its inherited
179         # file descriptors. This can cause problems especially if the parent
180         # process has opened pipes. Since the recovery process is an optional
181         # feature of libguestfs, it's better to disable it.
182         self.g.set_recovery_proc(0)
183         version = self.g.version()
184         if version['major'] > 1 or \
185             (version['major'] == 1 and (version['minor'] >= 18 or
186                                         (version['minor'] == 17 and
187                                          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 OS.")
215         self.root = roots[0]
216         self.guestfs_device = self.g.part_to_dev(self.root)
217         self.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: \
276             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
277         is_logical = lambda p: \
278             self.meta['PARTITION_TABLE'] != 'msdos' and 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: \
304             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
305         is_logical = lambda p: \
306             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
307         is_extended = lambda p: \
308             self.meta['PARTITION_TABLE'] == 'msdos' and \
309             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
310
311         part_add = lambda ptype, start, stop: \
312             self.g.part_add(self.guestfs_device, ptype, start, stop)
313         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
314         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
315         part_set_id = lambda p, id: \
316             self.g.part_set_mbr_id(self.guestfs_device, p, id)
317         part_get_bootable = lambda p: \
318             self.g.part_get_bootable(self.guestfs_device, p)
319         part_set_bootable = lambda p, bootable: \
320             self.g.part_set_bootable(self.guestfs_device, p, bootable)
321
322         MB = 2 ** 20
323
324         self.out.output("Shrinking image (this may take a while)...", False)
325
326         sector_size = self.g.blockdev_getss(self.guestfs_device)
327
328         last_part = None
329         fstype = None
330         while True:
331             last_part = self._last_partition()
332             fstype = get_fstype(last_part)
333
334             if fstype == 'swap':
335                 self.meta['SWAP'] = "%d:%s" % \
336                     (last_part['part_num'],
337                      (last_part['part_size'] + MB - 1) // MB)
338                 part_del(last_part['part_num'])
339                 continue
340             elif is_extended(last_part):
341                 part_del(last_part['part_num'])
342                 continue
343
344             # Most disk manipulation programs leave 2048 sectors after the last
345             # partition
346             new_size = last_part['part_end'] + 1 + 2048 * sector_size
347             self.size = min(self.size, new_size)
348             break
349
350         if not re.match("ext[234]", fstype):
351             self.out.warn("Don't know how to resize %s partitions." % fstype)
352             return self.size
353
354         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
355         self.g.e2fsck_f(part_dev)
356         self.g.resize2fs_M(part_dev)
357
358         out = self.g.tune2fs_l(part_dev)
359         block_size = int(
360             filter(lambda x: x[0] == 'Block size', out)[0][1])
361         block_cnt = int(
362             filter(lambda x: x[0] == 'Block count', out)[0][1])
363
364         start = last_part['part_start'] / sector_size
365         end = start + (block_size * block_cnt) / sector_size - 1
366
367         if is_logical(last_part):
368             partitions = self.g.part_list(self.guestfs_device)
369
370             logical = []  # logical partitions
371             for partition in partitions:
372                 if partition['part_num'] < 4:
373                     continue
374                 logical.append({
375                     'num': partition['part_num'],
376                     'start': partition['part_start'] / sector_size,
377                     'end': partition['part_end'] / sector_size,
378                     'id': part_get_(partition['part_num']),
379                     'bootable': part_get_bootable(partition['part_num'])
380                 })
381
382             logical[-1]['end'] = end  # new end after resize
383
384             # Recreate the extended partition
385             extended = [p for p in partitions if self._is_extended(p)][0]
386             part_del(extended['part_num'])
387             part_add('e', extended['part_start'], end)
388
389             # Create all the logical partitions back
390             for l in logical:
391                 part_add('l', l['start'], l['end'])
392                 part_set_id(l['num'], l['id'])
393                 part_set_bootable(l['num'], l['bootable'])
394         else:
395             # Recreate the last partition
396             if self.meta['PARTITION_TABLE'] == 'msdos':
397                 last_part['id'] = part_get_id(last_part['part_num'])
398
399             last_part['bootable'] = part_get_bootable(last_part['part_num'])
400             part_del(last_part['part_num'])
401             part_add('p', start, end)
402             part_set_bootable(last_part['part_num'], last_part['bootable'])
403
404             if self.meta['PARTITION_TABLE'] == 'msdos':
405                 part_set_id(last_part['part_num'], last_part['id'])
406
407         new_size = (end + 1) * sector_size
408
409         assert (new_size <= self.size)
410
411         if self.meta['PARTITION_TABLE'] == 'gpt':
412             ptable = GPTPartitionTable(self.real_device)
413             self.size = ptable.shrink(new_size, self.size)
414         else:
415             self.size = min(new_size + 2048 * sector_size, self.size)
416
417         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
418
419         return self.size
420
421     def dump(self, outfile):
422         """Dumps the content of device into a file.
423
424         This method will only dump the actual payload, found by reading the
425         partition table. Empty space in the end of the device will be ignored.
426         """
427         MB = 2 ** 20
428         blocksize = 4 * MB  # 4MB
429         size = self.size
430         progr_size = (size + MB - 1) // MB  # in MB
431         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
432
433         with open(self.real_device, 'r') as src:
434             with open(outfile, "w") as dst:
435                 left = size
436                 offset = 0
437                 progressbar.next()
438                 while left > 0:
439                     length = min(left, blocksize)
440                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
441                     offset += sent
442                     left -= sent
443                     progressbar.goto((size - left) // MB)
444         progressbar.success('image file %s was successfully created' % outfile)
445
446 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :