3 from image_creator.util import get_command
4 from clint.textui import progress
15 class DiskError(Exception):
18 dd = get_command('dd')
19 dmsetup = get_command('dmsetup')
20 losetup = get_command('losetup')
21 blockdev = get_command('blockdev')
25 """This class represents a hard disk hosting an Operating System
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
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 = []
39 def _add_cleanup(self, job, *args):
40 self._cleanup_jobs.append((job, args))
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)
48 def _dir_to_disk(self):
49 raise NotImplementedError
52 """Cleanup internal data. This needs to be called before the
55 while len(self._devices):
56 device = self._devices.pop()
59 while len(self._cleanup_jobs):
60 job, args = self._cleanup_jobs.pop()
64 """Returns a newly created DiskDevice instance.
66 This instance is a snapshot of the original source media of
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")
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)
86 snapshot = uuid.uuid4().hex
87 tablefd, table = tempfile.mkstemp()
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)
96 new_device = DiskDevice("/dev/mapper/%s" % snapshot)
97 self._devices.append(new_device)
100 def destroy_device(self, device):
101 """Destroys a DiskDevice instance previously created by
104 self._devices.remove(device)
108 def progress_generator(total):
110 for i in progress.bar(range(total)):
114 yield #suppress the StopIteration exception
117 class DiskDevice(object):
118 """This class represents a block device hosting an Operating System
119 as created by the device-mapper.
122 def __init__(self, device, bootable=True):
123 """Create a new DiskDevice."""
125 self.bootable = bootable
126 self.progress_bar = None
128 self.g = guestfs.GuestFS()
129 self.g.add_drive_opts(device, readonly=0)
132 #self.g.set_verbose(1)
134 eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
136 self.g.delete_event_callback(eh)
138 roots = self.g.inspect_os()
140 raise DiskError("No operating system found")
142 raise DiskError("Multiple operating systems found")
145 self.ostype = self.g.inspect_get_type(self.root)
146 self.distro = self.g.inspect_get_distro(self.root)
149 """Destroy this DiskDevice instance."""
152 # Close the guestfs handler
155 def progress_callback(self, ev, eh, buf, array):
159 if self.progress_bar is None:
160 self.progress_bar = progress_generator(total)
161 self.progress_bar.next()
163 self.progress_bar.send(position)
165 if position == total:
166 self.progress_bar = None
169 """Mount all disk partitions in a correct order."""
170 mps = self.g.inspect_get_mountpoints(self.root)
172 # Sort the keys to mount the fs in a correct order.
173 # / should be mounted befor /boot, etc
175 if len(a[0]) > len(b[0]):
177 elif len(a[0]) == len(b[0]):
184 self.g.mount(dev, mp)
185 except RuntimeError as msg:
186 print "%s (ignored)" % msg
189 """Umount all mounted filesystems."""
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.
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)
205 last_partition = self.g.part_list(dev)[-1]
207 if last_partition['part_num'] > 4:
208 raise DiskError("This disk contains logical partitions. "
209 "Only primary partitions are supported.")
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
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])
223 sector_size = self.g.blockdev_getss(dev)
225 start = last_partition['part_start'] / sector_size
226 end = start + (block_size * block_cnt) / sector_size - 1
228 self.g.part_del(dev, last_partition['part_num'])
229 self.g.part_add(dev, 'p', start, end)
231 return (end + 1) * sector_size
234 """Returns the "payload" size of the device.
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).
239 dev = self.g.part_to_dev(self.root)
240 last = self.g.part_list(dev)[-1]
242 return last['part_end']
244 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :