Leave 2048 sectors space at the end when shrinking
[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 (version['major'] == 1 and
184             (version['minor'] >= 18 or \
185             (version['minor'] == 17 and version['release'] >= 14))):
186             self.g.set_recovery_proc(1)
187             self.out.output("Enabling recovery proc")
188
189         #self.g.set_trace(1)
190         #self.g.set_verbose(1)
191
192         self.guestfs_enabled = False
193
194     def enable(self):
195         """Enable a newly created DiskDevice"""
196         self.progressbar = self.out.Progress(100, "Launching helper VM",
197                                              "percent")
198         eh = self.g.set_event_callback(self.progress_callback,
199                                                     guestfs.EVENT_PROGRESS)
200         self.g.launch()
201         self.guestfs_enabled = True
202         self.g.delete_event_callback(eh)
203         self.progressbar.success('done')
204         self.progressbar = None
205
206         self.out.output('Inspecting Operating System...', False)
207         roots = self.g.inspect_os()
208         if len(roots) == 0:
209             raise FatalError("No operating system found")
210         if len(roots) > 1:
211             raise FatalError("Multiple operating systems found."
212                             "We only support images with one filesystem.")
213         self.root = roots[0]
214         self.guestfs_device = self.g.part_to_dev(self.root)
215         self.meta['SIZE'] = self.g.blockdev_getsize64(self.guestfs_device)
216         self.meta['PARTITION_TABLE'] = \
217                                 self.g.part_get_parttype(self.guestfs_device)
218
219         self.ostype = self.g.inspect_get_type(self.root)
220         self.distro = self.g.inspect_get_distro(self.root)
221         self.out.success('found a(n) %s system' % self.distro)
222
223     def destroy(self):
224         """Destroy this DiskDevice instance."""
225
226         if self.guestfs_enabled:
227             self.g.umount_all()
228             self.g.sync()
229
230         # Close the guestfs handler if open
231         self.g.close()
232
233     def progress_callback(self, ev, eh, buf, array):
234         position = array[2]
235         total = array[3]
236
237         self.progressbar.goto((position * 100) // total)
238
239     def mount(self, readonly=False):
240         """Mount all disk partitions in a correct order."""
241
242         mount = self.g.mount_ro if readonly else self.g.mount
243         self.out.output("Mounting image...", False)
244         mps = self.g.inspect_get_mountpoints(self.root)
245
246         # Sort the keys to mount the fs in a correct order.
247         # / should be mounted befor /boot, etc
248         def compare(a, b):
249             if len(a[0]) > len(b[0]):
250                 return 1
251             elif len(a[0]) == len(b[0]):
252                 return 0
253             else:
254                 return -1
255         mps.sort(compare)
256         for mp, dev in mps:
257             try:
258                 mount(dev, mp)
259             except RuntimeError as msg:
260                 self.out.warn("%s (ignored)" % msg)
261         self.out.success("done")
262
263     def umount(self):
264         """Umount all mounted filesystems."""
265         self.g.umount_all()
266
267     def _last_partition(self):
268         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
269             msg = "Unsupported partition table: %s. Only msdos and gpt " \
270             "partition tables are supported" % self.meta['PARTITION_TABLE']
271             raise FatalError(msg)
272
273         is_extended = lambda p: self.g.part_get_mbr_id(
274                                     self.guestfs_device, p['part_num']) == 5
275         is_logical = lambda p: self.meta['PARTITION_TABLE'] != 'msdos' and \
276                                                             p['part_num'] > 4
277
278         partitions = self.g.part_list(self.guestfs_device)
279         last_partition = partitions[-1]
280
281         if is_logical(last_partition):
282             # The disk contains extended and logical partitions....
283             extended = [p for p in partitions if is_extended(p)][0]
284             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
285
286             # check if extended is the last primary partition
287             if last_primary['part_num'] > extended['part_num']:
288                 last_partition = last_primary
289
290         return last_partition
291
292     def shrink(self):
293         """Shrink the disk.
294
295         This is accomplished by shrinking the last filesystem in the
296         disk and then updating the partition table. The new disk size
297         (in bytes) is returned.
298
299         ATTENTION: make sure unmount is called before shrink
300         """
301         get_fstype = lambda p: self.g.vfs_type("%s%d" % \
302                                         (self.guestfs_device, p['part_num']))
303         is_logical = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
304                                                             p['part_num'] > 4
305         is_extended = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
306                 self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
307
308         part_add = lambda ptype, start, stop: \
309                     self.g.part_add(self.guestfs_device, ptype, start, stop)
310         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
311         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
312         part_set_id = lambda p, id: self.g.part_set_mbr_id(
313                                                     self.guestfs_device, p, id)
314         part_get_bootable = lambda p: self.g.part_get_bootable(
315                                                         self.guestfs_device, p)
316         part_set_bootable = lambda p, bootable: self.g.part_set_bootable(
317                                             self.guestfs_device, p, bootable)
318
319         MB = 2 ** 20
320
321         self.out.output("Shrinking image (this may take a while)...", False)
322
323         sector_size = self.g.blockdev_getss(self.guestfs_device)
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             # Most disk manipulation programs leave 2048 sectors after the last
342             # partition
343             new_size = last_part['part_end'] + 1 + 2048 * sector_size
344             self.meta['SIZE'] = min(self.meta['SIZE'], new_size)
345             break
346
347         if not re.match("ext[234]", fstype):
348             self.out.warn("Don't know how to resize %s partitions." % fstype)
349             return self.meta['SIZE']
350
351         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
352         self.g.e2fsck_f(part_dev)
353         self.g.resize2fs_M(part_dev)
354
355         out = self.g.tune2fs_l(part_dev)
356         block_size = int(
357             filter(lambda x: x[0] == 'Block size', out)[0][1])
358         block_cnt = int(
359             filter(lambda x: x[0] == 'Block count', out)[0][1])
360
361         start = last_part['part_start'] / sector_size
362         end = start + (block_size * block_cnt) / sector_size - 1
363
364         if is_logical(last_part):
365             partitions = self.g.part_list(self.guestfs_device)
366
367             logical = []  # logical partitions
368             for partition in partitions:
369                 if partition['part_num'] < 4:
370                     continue
371                 logical.append({
372                     'num': partition['part_num'],
373                     'start': partition['part_start'] / sector_size,
374                     'end': partition['part_end'] / sector_size,
375                     'id': part_get_(partition['part_num']),
376                     'bootable': part_get_bootable(partition['part_num'])
377                 })
378
379             logical[-1]['end'] = end  # new end after resize
380
381             # Recreate the extended partition
382             extended = [p for p in partitions if self._is_extended(p)][0]
383             part_del(extended['part_num'])
384             part_add('e', extended['part_start'], end)
385
386             # Create all the logical partitions back
387             for l in logical:
388                 part_add('l', l['start'], l['end'])
389                 part_set_id(l['num'], l['id'])
390                 part_set_bootable(l['num'], l['bootable'])
391         else:
392             # Recreate the last partition
393             if self.meta['PARTITION_TABLE'] == 'msdos':
394                 last_part['id'] = part_get_id(last_part['part_num'])
395
396             last_part['bootable'] = part_get_bootable(last_part['part_num'])
397             part_del(last_part['part_num'])
398             part_add('p', start, end)
399             part_set_bootable(last_part['part_num'], last_part['bootable'])
400
401             if self.meta['PARTITION_TABLE'] == 'msdos':
402                 part_set_id(last_part['part_num'], last_part['id'])
403
404         new_size = (end + 1) * sector_size
405
406         assert (new_size <= self.meta['SIZE'])
407
408         if self.meta['PARTITION_TABLE'] == 'gpt':
409             ptable = GPTPartitionTable(self.real_device)
410             self.meta['SIZE'] = ptable.shrink(new_size, self.meta['SIZE'])
411         else:
412             self.meta['SIZE'] = min(new_size + 2048 * sector_size,
413                                     self.meta['SIZE'])
414
415         self.out.success("new size is %dMB" %
416                          ((self.meta['SIZE'] + MB - 1) // MB))
417
418         return self.meta['SIZE']
419
420     def dump(self, outfile):
421         """Dumps the content of device into a file.
422
423         This method will only dump the actual payload, found by reading the
424         partition table. Empty space in the end of the device will be ignored.
425         """
426         MB = 2 ** 20
427         blocksize = 4 * MB  # 4MB
428         size = self.meta['SIZE']
429         progr_size = (size + MB - 1) // MB  # in MB
430         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
431
432         with open(self.real_device, 'r') as src:
433             with open(outfile, "w") as dst:
434                 left = size
435                 offset = 0
436                 progressbar.next()
437                 while left > 0:
438                     length = min(left, blocksize)
439                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
440                     offset += sent
441                     left -= sent
442                     progressbar.goto((size - left) // MB)
443         progressbar.success('image file %s was successfully created' % outfile)
444
445 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :