be728a2d05ef35ceacd6f697e614b4c5ac501512
[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
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 get_device(self):
97         """Returns a newly created DiskDevice instance.
98
99         This instance is a snapshot of the original source media of
100         the Disk instance.
101         """
102
103         output("Examining source media `%s'..." % self.source, False)
104         sourcedev = self.source
105         mode = os.stat(self.source).st_mode
106         if stat.S_ISDIR(mode):
107             success('looks like a directory')
108             return self._losetup(self._dir_to_disk())
109         elif stat.S_ISREG(mode):
110             success('looks like an image file')
111             sourcedev = self._losetup(self.source)
112         elif not stat.S_ISBLK(mode):
113             raise ValueError("Invalid media source. Only block devices, "
114                             "regular files and directories are supported.")
115         else:
116             success('looks like a block device')
117
118         # Take a snapshot and return it to the user
119         output("Snapshotting media source...", False)
120         size = blockdev('--getsize', sourcedev)
121         cowfd, cow = tempfile.mkstemp()
122         os.close(cowfd)
123         self._add_cleanup(os.unlink, cow)
124         # Create 1G cow sparse file
125         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
126                                         'seek=%d' % (1024 * 1024))
127         cowdev = self._losetup(cow)
128
129         snapshot = uuid.uuid4().hex
130         tablefd, table = tempfile.mkstemp()
131         try:
132             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
133                                         (int(size), sourcedev, cowdev))
134             dmsetup('create', snapshot, table)
135             self._add_cleanup(dmsetup, 'remove', snapshot)
136             # Sometimes dmsetup remove fails with Device or resource busy,
137             # although everything is cleaned up and the snapshot is not
138             # used by anyone. Add a 2 seconds delay to be on the safe side.
139             self._add_cleanup(time.sleep, 2)
140
141         finally:
142             os.unlink(table)
143         success('done')
144         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
145         self._devices.append(new_device)
146         new_device.enable()
147         return new_device
148
149     def destroy_device(self, device):
150         """Destroys a DiskDevice instance previously created by
151         get_device method.
152         """
153         self._devices.remove(device)
154         device.destroy()
155
156
157 class DiskDevice(object):
158     """This class represents a block device hosting an Operating System
159     as created by the device-mapper.
160     """
161
162     def __init__(self, device, bootable=True):
163         """Create a new DiskDevice."""
164
165         self.device = device
166         self.bootable = bootable
167         self.progress_bar = None
168
169         self.g = guestfs.GuestFS()
170         self.g.add_drive_opts(self.device, readonly=0)
171
172         #self.g.set_trace(1)
173         #self.g.set_verbose(1)
174
175         self.guestfs_enabled = False
176
177     def enable(self):
178         """Enable a newly created DiskDevice"""
179         self.progressbar = progress("Launching helper VM: ")
180         self.progressbar.next()
181         eh = self.g.set_event_callback(self.progress_callback,
182                                                     guestfs.EVENT_PROGRESS)
183         self.g.launch()
184         self.guestfs_enabled = True
185         self.g.delete_event_callback(eh)
186         if self.progressbar is not None:
187             self.progressbar.send(100)
188             self.progressbar = None
189
190         output('Inspecting Operating System...', False)
191         roots = self.g.inspect_os()
192         if len(roots) == 0:
193             raise FatalError("No operating system found")
194         if len(roots) > 1:
195             raise FatalError("Multiple operating systems found."
196                             "We only support images with one filesystem.")
197         self.root = roots[0]
198         self.ostype = self.g.inspect_get_type(self.root)
199         self.distro = self.g.inspect_get_distro(self.root)
200         success('found a %s system' % self.distro)
201
202     def destroy(self):
203         """Destroy this DiskDevice instance."""
204
205         if self.guestfs_enabled:
206             self.g.umount_all()
207             self.g.sync()
208
209         # Close the guestfs handler if open
210         self.g.close()
211
212     def progress_callback(self, ev, eh, buf, array):
213         position = array[2]
214         total = array[3]
215
216         self.progressbar.send((position * 100) // total)
217
218         if position == total:
219             self.progressbar = None
220
221     def mount(self):
222         """Mount all disk partitions in a correct order."""
223
224         output("Mounting image...", False)
225         mps = self.g.inspect_get_mountpoints(self.root)
226
227         # Sort the keys to mount the fs in a correct order.
228         # / should be mounted befor /boot, etc
229         def compare(a, b):
230             if len(a[0]) > len(b[0]):
231                 return 1
232             elif len(a[0]) == len(b[0]):
233                 return 0
234             else:
235                 return -1
236         mps.sort(compare)
237         for mp, dev in mps:
238             try:
239                 self.g.mount(dev, mp)
240             except RuntimeError as msg:
241                 warn("%s (ignored)" % msg)
242         success("done")
243
244     def umount(self):
245         """Umount all mounted filesystems."""
246         self.g.umount_all()
247
248     def shrink(self):
249         """Shrink the disk.
250
251         This is accomplished by shrinking the last filesystem in the
252         disk and then updating the partition table. The new disk size
253         (in bytes) is returned.
254         """
255         output("Shrinking image (this may take a while)...", False)
256
257         dev = self.g.part_to_dev(self.root)
258         parttype = self.g.part_get_parttype(dev)
259         if parttype != 'msdos':
260             raise FatalError("You have a %s partition table. "
261                 "Only msdos partitions are supported" % parttype)
262
263         last_partition = self.g.part_list(dev)[-1]
264
265         if last_partition['part_num'] > 4:
266             raise FatalError("This disk contains logical partitions. "
267                 "Only primary partitions are supported.")
268
269         part_dev = "%s%d" % (dev, last_partition['part_num'])
270         fs_type = self.g.vfs_type(part_dev)
271         if not re.match("ext[234]", fs_type):
272             warn("Don't know how to resize %s partitions." % vfs_type)
273             return
274
275         self.g.e2fsck_f(part_dev)
276         self.g.resize2fs_M(part_dev)
277
278         output = self.g.tune2fs_l(part_dev)
279         block_size = int(
280             filter(lambda x: x[0] == 'Block size', output)[0][1])
281         block_cnt = int(
282             filter(lambda x: x[0] == 'Block count', output)[0][1])
283
284         sector_size = self.g.blockdev_getss(dev)
285
286         start = last_partition['part_start'] / sector_size
287         end = start + (block_size * block_cnt) / sector_size - 1
288
289         self.g.part_del(dev, last_partition['part_num'])
290         self.g.part_add(dev, 'p', start, end)
291
292         new_size = (end + 1) * sector_size
293         success("new image size is %dMB" %
294                             ((new_size + 2 ** 20 - 1) // 2 ** 20))
295         return new_size
296
297     def size(self):
298         """Returns the "payload" size of the device.
299
300         The size returned by this method is the size of the space occupied by
301         the partitions (including the space before the first partition).
302         """
303         dev = self.g.part_to_dev(self.root)
304         last = self.g.part_list(dev)[-1]
305
306         return last['part_end'] + 1
307
308     def dump(self, outfile):
309         """Dumps the content of device into a file.
310
311         This method will only dump the actual payload, found by reading the
312         partition table. Empty space in the end of the device will be ignored.
313         """
314         blocksize = 2 ** 22  # 4MB
315         size = self.size()
316         progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
317         progressbar = progress("Dumping image file: ", progress_size)
318
319         source = open(self.device, "r")
320         try:
321             dest = open(outfile, "w")
322             try:
323                 left = size
324                 offset = 0
325                 progressbar.next()
326                 while left > 0:
327                     length = min(left, blocksize)
328                     sent = sendfile(dest.fileno(), source.fileno(), offset,
329                                                                         length)
330                     offset += sent
331                     left -= sent
332                     for i in range((length + 2 ** 20 - 1) // 2 ** 20):
333                         progressbar.next()
334             finally:
335                 dest.close()
336         finally:
337             source.close()
338
339         success('Image file %s was successfully created' % outfile)
340
341 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :