9476dc507c8296d8b5538514553b775c86872536
[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         assert self.progress_bar is not None
178         print 'posisition/total: %s/%s' % (position, total)
179         self.progress_bar.send((position * 100)//total)
180
181         if position == total:
182             self.progress_bar = None
183
184     def mount(self):
185         """Mount all disk partitions in a correct order."""
186         mps = self.g.inspect_get_mountpoints(self.root)
187
188         # Sort the keys to mount the fs in a correct order.
189         # / should be mounted befor /boot, etc
190         def compare(a, b):
191             if len(a[0]) > len(b[0]):
192                 return 1
193             elif len(a[0]) == len(b[0]):
194                 return 0
195             else:
196                 return -1
197         mps.sort(compare)
198         for mp, dev in mps:
199             try:
200                 self.g.mount(dev, mp)
201             except RuntimeError as msg:
202                 print "%s (ignored)" % msg
203
204     def umount(self):
205         """Umount all mounted filesystems."""
206         self.g.umount_all()
207
208     def shrink(self):
209         """Shrink the disk.
210
211         This is accomplished by shrinking the last filesystem in the
212         disk and then updating the partition table. The new disk size
213         (in bytes) is returned.
214         """
215         dev = self.g.part_to_dev(self.root)
216         parttype = self.g.part_get_parttype(dev)
217         if parttype != 'msdos':
218             raise FatalError("You have a %s partition table. "
219                 "Only msdos partitions are supported" % parttype)
220
221         last_partition = self.g.part_list(dev)[-1]
222
223         if last_partition['part_num'] > 4:
224             raise FatalError("This disk contains logical partitions. "
225                 "Only primary partitions are supported.")
226
227         part_dev = "%s%d" % (dev, last_partition['part_num'])
228         fs_type = self.g.vfs_type(part_dev)
229         if not re.match("ext[234]", fs_type):
230             print "Warning: Don't know how to resize %s partitions." % vfs_type
231             return
232
233         self.g.e2fsck_f(part_dev)
234         self.g.resize2fs_M(part_dev)
235         output = self.g.tune2fs_l(part_dev)
236         block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
237         block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
238
239         sector_size = self.g.blockdev_getss(dev)
240
241         start = last_partition['part_start'] / sector_size
242         end = start + (block_size * block_cnt) / sector_size - 1
243
244         self.g.part_del(dev, last_partition['part_num'])
245         self.g.part_add(dev, 'p', start, end)
246
247         return (end + 1) * sector_size
248
249     def size(self):
250         """Returns the "payload" size of the device.
251
252         The size returned by this method is the size of the space occupied by
253         the partitions (including the space before the first partition).
254         """
255         dev = self.g.part_to_dev(self.root)
256         last = self.g.part_list(dev)[-1]
257
258         return last['part_end']
259
260 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :