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