Add support for GUID Partition Tables
[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 warn, progress, success, output, 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):
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
72     def _add_cleanup(self, job, *args):
73         self._cleanup_jobs.append((job, args))
74
75     def _losetup(self, fname):
76         loop = losetup('-f', '--show', fname)
77         loop = loop.strip()  # remove the new-line char
78         self._add_cleanup(losetup, '-d', loop)
79         return loop
80
81     def _dir_to_disk(self):
82         raise NotImplementedError
83
84     def cleanup(self):
85         """Cleanup internal data. This needs to be called before the
86         program ends.
87         """
88         while len(self._devices):
89             device = self._devices.pop()
90             device.destroy()
91
92         while len(self._cleanup_jobs):
93             job, args = self._cleanup_jobs.pop()
94             job(*args)
95
96     def snapshot(self):
97         """Creates a snapshot of the original source media of the Disk
98         instance.
99         """
100
101         output("Examining source media `%s'..." % self.source, False)
102         sourcedev = self.source
103         mode = os.stat(self.source).st_mode
104         if stat.S_ISDIR(mode):
105             success('looks like a directory')
106             return self._losetup(self._dir_to_disk())
107         elif stat.S_ISREG(mode):
108             success('looks like an image file')
109             sourcedev = self._losetup(self.source)
110         elif not stat.S_ISBLK(mode):
111             raise ValueError("Invalid media source. Only block devices, "
112                             "regular files and directories are supported.")
113         else:
114             success('looks like a block device')
115
116         # Take a snapshot and return it to the user
117         output("Snapshotting media source...", False)
118         size = blockdev('--getsize', sourcedev)
119         cowfd, cow = tempfile.mkstemp()
120         os.close(cowfd)
121         self._add_cleanup(os.unlink, cow)
122         # Create 1G cow sparse file
123         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
124                                         '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         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)
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, bootable=True):
166         """Create a new DiskDevice."""
167
168         self.device = device
169         self.bootable = bootable
170         self.progress_bar = None
171
172         self.g = guestfs.GuestFS()
173         self.g.add_drive_opts(self.device, readonly=0)
174
175         #self.g.set_trace(1)
176         #self.g.set_verbose(1)
177
178         self.guestfs_enabled = False
179
180     def enable(self):
181         """Enable a newly created DiskDevice"""
182         self.progressbar = progress("Launching helper VM: ", "percent")
183         self.progressbar.max = 100
184         self.progressbar.goto(1)
185         eh = self.g.set_event_callback(self.progress_callback,
186                                                     guestfs.EVENT_PROGRESS)
187         self.g.launch()
188         self.guestfs_enabled = True
189         self.g.delete_event_callback(eh)
190         if self.progressbar is not None:
191             output("\rLaunching helper VM...\033[K", False)
192             success("done")
193             self.progressbar = None
194
195         output('Inspecting Operating System...', False)
196         roots = self.g.inspect_os()
197         if len(roots) == 0:
198             raise FatalError("No operating system found")
199         if len(roots) > 1:
200             raise FatalError("Multiple operating systems found."
201                             "We only support images with one filesystem.")
202         self.root = roots[0]
203         self.gdev = self.g.part_to_dev(self.root)
204         self.parttype = self.g.part_get_parttype(self.gdev)
205
206         self.ostype = self.g.inspect_get_type(self.root)
207         self.distro = self.g.inspect_get_distro(self.root)
208         success('found a(n) %s system' % self.distro)
209
210     def destroy(self):
211         """Destroy this DiskDevice instance."""
212
213         if self.guestfs_enabled:
214             self.g.umount_all()
215             self.g.sync()
216
217         # Close the guestfs handler if open
218         self.g.close()
219
220     def progress_callback(self, ev, eh, buf, array):
221         position = array[2]
222         total = array[3]
223
224         self.progressbar.goto((position * 100) // total)
225
226     def mount(self):
227         """Mount all disk partitions in a correct order."""
228
229         output("Mounting image...", False)
230         mps = self.g.inspect_get_mountpoints(self.root)
231
232         # Sort the keys to mount the fs in a correct order.
233         # / should be mounted befor /boot, etc
234         def compare(a, b):
235             if len(a[0]) > len(b[0]):
236                 return 1
237             elif len(a[0]) == len(b[0]):
238                 return 0
239             else:
240                 return -1
241         mps.sort(compare)
242         for mp, dev in mps:
243             try:
244                 self.g.mount(dev, mp)
245             except RuntimeError as msg:
246                 warn("%s (ignored)" % msg)
247         success("done")
248
249     def umount(self):
250         """Umount all mounted filesystems."""
251         self.g.umount_all()
252
253     def shrink(self):
254         """Shrink the disk.
255
256         This is accomplished by shrinking the last filesystem in the
257         disk and then updating the partition table. The new disk size
258         (in bytes) is returned.
259         """
260         output("Shrinking image (this may take a while)...", False)
261
262         if self.parttype not in 'msdos' 'gpt':
263             raise FatalError("You have a %s partition table. "
264                 "Only msdos and gpt partitions are supported" % self.parttype)
265
266         last_partition = self.g.part_list(self.gdev)[-1]
267
268         if last_partition['part_num'] > 4:
269             raise FatalError("This disk contains logical partitions. "
270                 "Only primary partitions are supported.")
271
272         part_dev = "%s%d" % (self.gdev, last_partition['part_num'])
273         fs_type = self.g.vfs_type(part_dev)
274         if not re.match("ext[234]", fs_type):
275             warn("Don't know how to resize %s partitions." % vfs_type)
276             return self.size()
277
278         self.g.e2fsck_f(part_dev)
279         self.g.resize2fs_M(part_dev)
280
281         out = self.g.tune2fs_l(part_dev)
282         block_size = int(
283             filter(lambda x: x[0] == 'Block size', out)[0][1])
284         block_cnt = int(
285             filter(lambda x: x[0] == 'Block count', out)[0][1])
286
287         sector_size = self.g.blockdev_getss(self.gdev)
288
289         start = last_partition['part_start'] / sector_size
290         end = start + (block_size * block_cnt) / sector_size - 1
291
292         self.g.part_del(self.gdev, last_partition['part_num'])
293         self.g.part_add(self.gdev, 'p', start, end)
294
295         new_size = (end + 1) * sector_size
296         success("new image size is %dMB" %
297                             ((new_size + 2 ** 20 - 1) // 2 ** 20))
298
299         if self.parttype == 'gpt':
300             ptable = GPTPartitionTable(self.device)
301             return ptable.shrink(new_size)
302
303         return new_size
304
305     def size(self):
306         """Returns the "payload" size of the device.
307
308         The size returned by this method is the size of the space occupied by
309         the partitions (including the space before the first partition).
310         """
311
312         if self.parttype == 'msdos':
313             dev = self.g.part_to_dev(self.root)
314             last = self.g.part_list(dev)[-1]
315             return last['part_end'] + 1
316         elif self.parttype == 'gpt':
317             ptable = GPTPartitionTable(self.device)
318             return ptable.size()
319         else:
320             raise FatalError("Unsupported partition table type: %s" % parttype)
321
322     def dump(self, outfile):
323         """Dumps the content of device into a file.
324
325         This method will only dump the actual payload, found by reading the
326         partition table. Empty space in the end of the device will be ignored.
327         """
328         blocksize = 2 ** 22  # 4MB
329         size = self.size()
330         progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
331         progressbar = progress("Dumping image file: ", 'mb')
332         progressbar.max = progress_size
333         source = open(self.device, "r")
334         try:
335             dest = open(outfile, "w")
336             try:
337                 left = size
338                 offset = 0
339                 progressbar.next()
340                 while left > 0:
341                     length = min(left, blocksize)
342                     sent = sendfile(dest.fileno(), source.fileno(), offset,
343                                                                         length)
344                     offset += sent
345                     left -= sent
346                     progressbar.goto((size - left) // 2 ** 20)
347             finally:
348                 dest.close()
349         finally:
350             source.close()
351         output("\rDumping image file...\033[K", False)
352         success('image file %s was successfully created' % outfile)
353
354 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :