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
 132
 133 class DiskDevice(object):
 134     """This class represents a block device hosting an Operating System
 135     as created by the device-mapper.
 136     """
 137
 138     def __init__(self, device, bootable=True):
 139         """Create a new DiskDevice."""
 140         self.device = device
 141         self.bootable = bootable
 142         self.progress_bar = None
 143
 144         self.g = guestfs.GuestFS()
 145         self.g.add_drive_opts(device, readonly=0)
 146
 147         #self.g.set_trace(1)
 148         #self.g.set_verbose(1)
 149
 150         eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
 151         self.g.launch()
 152         self.g.delete_event_callback(eh)
 153
 154         roots = self.g.inspect_os()
 155         if len(roots) == 0:
 156             raise DiskError("No operating system found")
 157         if len(roots) > 1:
 158             raise DiskError("Multiple operating systems found")
 159
 160         self.root = roots[0]
 161         self.ostype = self.g.inspect_get_type(self.root)
 162         self.distro = self.g.inspect_get_distro(self.root)
 163
 164     def destroy(self):
 165         """Destroy this DiskDevice instance."""
 166         self.g.umount_all()
 167         self.g.sync()
 168         # Close the guestfs handler
 169         self.g.close()
 170
 171     def progress_callback(self, ev, eh, buf, array):
 172         position = array[2]
 173         total = array[3]
 174
 175         if self.progress_bar is None:
 176             self.progress_bar = progress_generator(total)
 177             self.progress_bar.next()
 178             if position == 1:
 179                 return
 180
 181         self.progress_bar.send(position)
 182
 183         if position == total:
 184             self.progress_bar = None
 185
 186     def mount(self):
 187         """Mount all disk partitions in a correct order."""
 188         mps = self.g.inspect_get_mountpoints(self.root)
 189
 190         # Sort the keys to mount the fs in a correct order.
 191         # / should be mounted befor /boot, etc
 192         def compare(a, b):
 193             if len(a[0]) > len(b[0]):
 194                 return 1
 195             elif len(a[0]) == len(b[0]):
 196                 return 0
 197             else:
 198                 return -1
 199         mps.sort(compare)
 200         for mp, dev in mps:
 201             try:
 202                 self.g.mount(dev, mp)
 203             except RuntimeError as msg:
 204                 print "%s (ignored)" % msg
 205
 206     def umount(self):
 207         """Umount all mounted filesystems."""
 208         self.g.umount_all()
 209
 210     def shrink(self):
 211         """Shrink the disk.
 212
 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.
 216         """
 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)
 222
 223         last_partition = self.g.part_list(dev)[-1]
 224
 225         if last_partition['part_num'] > 4:
 226             raise DiskError("This disk contains logical partitions. "
 227                 "Only primary partitions are supported.")
 228
 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
 233             return
 234
 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])
 240
 241         sector_size = self.g.blockdev_getss(dev)
 242
 243         start = last_partition['part_start'] / sector_size
 244         end = start + (block_size * block_cnt) / sector_size - 1
 245
 246         self.g.part_del(dev, last_partition['part_num'])
 247         self.g.part_add(dev, 'p', start, end)
 248
 249         return (end + 1) * sector_size
 250
 251     def size(self):
 252         """Returns the "payload" size of the device.
 253
 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).
 256         """
 257         dev = self.g.part_to_dev(self.root)
 258         last = self.g.part_list(dev)[-1]
 259
 260         return last['part_end']
 261
 262 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :