code.grnet.gr Git - snf-image-creator/blob - image_creator/disk.py

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