Fix wrong variable name in progress_callback
[snf-image-creator] / image_creator / disk.py
1 #!/usr/bin/env python
2
3 from image_creator.util import get_command
4 from image_creator import FatalError
5 from clint.textui import progress
6
7 import stat
8 import os
9 import tempfile
10 import uuid
11 import re
12 import sys
13 import guestfs
14 import time
15
16
17 class DiskError(Exception):
18     pass
19
20 dd = get_command('dd')
21 dmsetup = get_command('dmsetup')
22 losetup = get_command('losetup')
23 blockdev = get_command('blockdev')
24
25
26 class Disk(object):
27     """This class represents a hard disk hosting an Operating System
28
29     A Disk instance never alters the source media it is created from.
30     Any change is done on a snapshot created by the device-mapper of
31     the Linux kernel.
32     """
33
34     def __init__(self, source):
35         """Create a new Disk instance out of a source media. The source
36         media can be an image file, a block device or a directory."""
37         self._cleanup_jobs = []
38         self._devices = []
39         self.source = source
40
41     def _add_cleanup(self, job, *args):
42         self._cleanup_jobs.append((job, args))
43
44     def _losetup(self, fname):
45         loop = losetup('-f', '--show', fname)
46         loop = loop.strip() # remove the new-line char
47         self._add_cleanup(losetup, '-d', loop)
48         return loop
49
50     def _dir_to_disk(self):
51         raise NotImplementedError
52
53     def cleanup(self):
54         """Cleanup internal data. This needs to be called before the
55         program ends.
56         """
57         while len(self._devices):
58             device = self._devices.pop()
59             device.destroy()
60
61         while len(self._cleanup_jobs):
62             job, args = self._cleanup_jobs.pop()
63             job(*args)
64
65     def get_device(self):
66         """Returns a newly created DiskDevice instance.
67
68         This instance is a snapshot of the original source media of
69         the Disk instance.
70         """
71         sourcedev = self.source
72         mode = os.stat(self.source).st_mode
73         if stat.S_ISDIR(mode):
74             return self._losetup(self._dir_to_disk())
75         elif stat.S_ISREG(mode):
76             sourcedev = self._losetup(self.source)
77         elif not stat.S_ISBLK(mode):
78             raise ValueError("Value for self.source is invalid")
79
80         # Take a snapshot and return it to the user
81         size = blockdev('--getsize', sourcedev)
82         cowfd, cow = tempfile.mkstemp()
83         os.close(cowfd)
84         self._add_cleanup(os.unlink, cow)
85         # Create 1G cow sparse file
86         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
87         cowdev = self._losetup(cow)
88
89         snapshot = uuid.uuid4().hex
90         tablefd, table = tempfile.mkstemp()
91         try:
92             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
93                                         (int(size), sourcedev, cowdev))
94             dmsetup('create', snapshot, table)
95             self._add_cleanup(dmsetup, 'remove', snapshot)
96         finally:
97             os.unlink(table)
98         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
99         self._devices.append(new_device)
100         new_device.enable()
101         return new_device
102
103     def destroy_device(self, device):
104         """Destroys a DiskDevice instance previously created by
105         get_device method.
106         """
107         self._devices.remove(device)
108         device.destroy()
109
110
111 def progress_generator(label=''):
112     position = 0;
113     for i in progress.bar(range(100),label):
114         if i < position:
115             continue
116         position = yield
117     yield #suppress the StopIteration exception
118
119
120 class DiskDevice(object):
121     """This class represents a block device hosting an Operating System
122     as created by the device-mapper.
123     """
124
125     def __init__(self, device, bootable=True):
126         """Create a new DiskDevice."""
127
128         self.device = device
129         self.bootable = bootable
130         self.progress_bar = None
131
132         self.g = guestfs.GuestFS()
133         self.g.add_drive_opts(self.device, readonly=0)
134
135         #self.g.set_trace(1)
136         #self.g.set_verbose(1)
137
138         self.guestfs_enabled = False
139     
140     def enable(self):
141         """Enable a newly created DiskDevice"""
142
143         self.progressbar = progress_generator("VM lauch: ")
144         self.progressbar.next()
145         eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
146         self.g.launch()
147         self.guestfs_enabled = True
148         self.g.delete_event_callback(eh)
149         if self.progressbar is not None:
150             self.progressbar.send(100)
151             self.progressbar = None
152         
153         roots = self.g.inspect_os()
154         if len(roots) == 0:
155             raise FatalError("No operating system found")
156         if len(roots) > 1:
157             raise FatalError("Multiple operating systems found")
158
159         self.root = roots[0]
160         self.ostype = self.g.inspect_get_type(self.root)
161         self.distro = self.g.inspect_get_distro(self.root)
162
163     def destroy(self):
164         """Destroy this DiskDevice instance."""
165
166         if self.guestfs_enabled:
167             self.g.umount_all()
168             self.g.sync()
169
170         # Close the guestfs handler if open
171         self.g.close()
172
173     def progress_callback(self, ev, eh, buf, array):
174         position = array[2]
175         total = array[3]
176
177         self.progressbar.send((position * 100)//total)
178
179         if position == total:
180             self.progressbar = None
181
182     def mount(self):
183         """Mount all disk partitions in a correct order."""
184         mps = self.g.inspect_get_mountpoints(self.root)
185
186         # Sort the keys to mount the fs in a correct order.
187         # / should be mounted befor /boot, etc
188         def compare(a, b):
189             if len(a[0]) > len(b[0]):
190                 return 1
191             elif len(a[0]) == len(b[0]):
192                 return 0
193             else:
194                 return -1
195         mps.sort(compare)
196         for mp, dev in mps:
197             try:
198                 self.g.mount(dev, mp)
199             except RuntimeError as msg:
200                 print "%s (ignored)" % msg
201
202     def umount(self):
203         """Umount all mounted filesystems."""
204         self.g.umount_all()
205
206     def shrink(self):
207         """Shrink the disk.
208
209         This is accomplished by shrinking the last filesystem in the
210         disk and then updating the partition table. The new disk size
211         (in bytes) is returned.
212         """
213         dev = self.g.part_to_dev(self.root)
214         parttype = self.g.part_get_parttype(dev)
215         if parttype != 'msdos':
216             raise FatalError("You have a %s partition table. "
217                 "Only msdos partitions are supported" % parttype)
218
219         last_partition = self.g.part_list(dev)[-1]
220
221         if last_partition['part_num'] > 4:
222             raise FatalError("This disk contains logical partitions. "
223                 "Only primary partitions are supported.")
224
225         part_dev = "%s%d" % (dev, last_partition['part_num'])
226         fs_type = self.g.vfs_type(part_dev)
227         if not re.match("ext[234]", fs_type):
228             print "Warning: Don't know how to resize %s partitions." % vfs_type
229             return
230
231         self.g.e2fsck_f(part_dev)
232         self.g.resize2fs_M(part_dev)
233         output = self.g.tune2fs_l(part_dev)
234         block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
235         block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
236
237         sector_size = self.g.blockdev_getss(dev)
238
239         start = last_partition['part_start'] / sector_size
240         end = start + (block_size * block_cnt) / sector_size - 1
241
242         self.g.part_del(dev, last_partition['part_num'])
243         self.g.part_add(dev, 'p', start, end)
244
245         return (end + 1) * sector_size
246
247     def size(self):
248         """Returns the "payload" size of the device.
249
250         The size returned by this method is the size of the space occupied by
251         the partitions (including the space before the first partition).
252         """
253         dev = self.g.part_to_dev(self.root)
254         last = self.g.part_list(dev)[-1]
255
256         return last['part_end']
257
258 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :