14 from clint.textui import progress
17 class DiskError(Exception):
21 def find_sbin_command(command, exception):
22 search_paths = ['/usr/local/sbin', '/usr/sbin', '/sbin']
23 for fullpath in map(lambda x: "%s/%s" % (x, command), search_paths):
24 if os.path.exists(fullpath) and os.access(fullpath, os.X_OK):
25 return pbs.Command(fullpath)
31 from pbs import dmsetup
32 except pbs.CommandNotFound as e:
33 dmsetup = find_sbin_command('dmsetup', e)
36 from pbs import blockdev
37 except pbs.CommandNotFound as e:
38 blockdev = find_sbin_command('blockdev', e)
42 """This class represents a hard disk hosting an Operating System
44 A Disk instance never alters the source media it is created from.
45 Any change is done on a snapshot created by the device-mapper of
49 def __init__(self, source):
50 """Create a new Disk instance out of a source media. The source
51 media can be an image file, a block device or a directory."""
52 self._cleanup_jobs = []
56 def _add_cleanup(self, job, *args):
57 self._cleanup_jobs.append((job, args))
59 def _losetup(self, fname):
60 loop = losetup.find_unused_loop_device()
62 self._add_cleanup(loop.unmount)
65 def _dir_to_disk(self):
66 raise NotImplementedError
69 """Cleanup internal data. This needs to be called before the
72 while len(self._devices):
73 device = self._devices.pop()
76 while len(self._cleanup_jobs):
77 job, args = self._cleanup_jobs.pop()
81 """Returns a newly created DiskDevice instance.
83 This instance is a snapshot of the original source media of
86 sourcedev = self.source
87 mode = os.stat(self.source).st_mode
88 if stat.S_ISDIR(mode):
89 return self._losetup(self._dir_to_disk())
90 elif stat.S_ISREG(mode):
91 sourcedev = self._losetup(self.source)
92 elif not stat.S_ISBLK(mode):
93 raise ValueError("Value for self.source is invalid")
95 # Take a snapshot and return it to the user
96 size = blockdev('--getsize', sourcedev)
97 cowfd, cow = tempfile.mkstemp()
98 self._add_cleanup(os.unlink, cow)
99 # Create 1G cow sparse file
100 dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
101 cowdev = self._losetup(cow)
103 snapshot = uuid.uuid4().hex
104 tablefd, table = tempfile.mkstemp()
106 os.write(tablefd, "0 %d snapshot %s %s n 8" % \
107 (int(size), sourcedev, cowdev))
108 dmsetup('create', snapshot, table)
109 self._add_cleanup(dmsetup, 'remove', snapshot)
113 new_device = DiskDevice("/dev/mapper/%s" % snapshot)
114 self._devices.append(new_device)
117 def destroy_device(self, device):
118 """Destroys a DiskDevice instance previously created by
121 self._devices.remove(device)
125 def progress_generator(total):
127 for i in progress.bar(range(total)):
133 class DiskDevice(object):
134 """This class represents a block device hosting an Operating System
135 as created by the device-mapper.
138 def __init__(self, device, bootable=True):
139 """Create a new DiskDevice."""
141 self.bootable = bootable
142 self.progress_bar = None
144 self.g = guestfs.GuestFS()
145 self.g.add_drive_opts(device, readonly=0)
148 #self.g.set_verbose(1)
150 eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
152 self.g.delete_event_callback(eh)
154 roots = self.g.inspect_os()
156 raise DiskError("No operating system found")
158 raise DiskError("Multiple operating systems found")
161 self.ostype = self.g.inspect_get_type(self.root)
162 self.distro = self.g.inspect_get_distro(self.root)
165 """Destroy this DiskDevice instance."""
168 # Close the guestfs handler
171 def progress_callback(self, ev, eh, buf, array):
175 if self.progress_bar is None:
176 self.progress_bar = progress_generator(total)
177 self.progress_bar.next()
181 self.progress_bar.send(position)
183 if position == total:
184 self.progress_bar = None
187 """Mount all disk partitions in a correct order."""
188 mps = self.g.inspect_get_mountpoints(self.root)
190 # Sort the keys to mount the fs in a correct order.
191 # / should be mounted befor /boot, etc
193 if len(a[0]) > len(b[0]):
195 elif len(a[0]) == len(b[0]):
202 self.g.mount(dev, mp)
203 except RuntimeError as msg:
204 print "%s (ignored)" % msg
207 """Umount all mounted filesystems."""
213 This is accomplished by shrinking the last filesystem in the
214 disk and then updating the partition table. The new disk size
215 (in bytes) is returned.
217 dev = self.g.part_to_dev(self.root)
218 parttype = self.g.part_get_parttype(dev)
219 if parttype != 'msdos':
220 raise DiskError("You have a %s partition table. "
221 "Only msdos partitions are supported" % parttype)
223 last_partition = self.g.part_list(dev)[-1]
225 if last_partition['part_num'] > 4:
226 raise DiskError("This disk contains logical partitions. "
227 "Only primary partitions are supported.")
229 part_dev = "%s%d" % (dev, last_partition['part_num'])
230 fs_type = self.g.vfs_type(part_dev)
231 if not re.match("ext[234]", fs_type):
232 print "Warning: Don't know how to resize %s partitions." % vfs_type
235 self.g.e2fsck_f(part_dev)
236 self.g.resize2fs_M(part_dev)
237 output = self.g.tune2fs_l(part_dev)
238 block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
239 block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
241 sector_size = self.g.blockdev_getss(dev)
243 start = last_partition['part_start'] / sector_size
244 end = start + (block_size * block_cnt) / sector_size - 1
246 self.g.part_del(dev, last_partition['part_num'])
247 self.g.part_add(dev, 'p', start, end)
249 return (end + 1) * sector_size
252 """Returns the "payload" size of the device.
254 The size returned by this method is the size of the space occupied by
255 the partitions (including the space before the first partition).
257 dev = self.g.part_to_dev(self.root)
258 last = self.g.part_list(dev)[-1]
260 return last['part_end']
262 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :