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