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