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

   1 #!/usr/bin/env python
   2
   3 from image_creator.util import get_command
   4 from clint.textui import progress
   5
   6 import stat
   7 import os
   8 import tempfile
   9 import uuid
  10 import re
  11 import sys
  12 import guestfs
  13
  14
  15 class DiskError(Exception):
  16     pass
  17
  18 dd = get_command('dd')
  19 dmsetup = get_command('dmsetup')
  20 losetup = get_command('losetup')
  21 blockdev = get_command('blockdev')
  22
  23
  24 class Disk(object):
  25     """This class represents a hard disk hosting an Operating System
  26
  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
  29     the Linux kernel.
  30     """
  31
  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 = []
  36         self._devices = []
  37         self.source = source
  38
  39     def _add_cleanup(self, job, *args):
  40         self._cleanup_jobs.append((job, args))
  41
  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)
  46         return loop
  47
  48     def _dir_to_disk(self):
  49         raise NotImplementedError
  50
  51     def cleanup(self):
  52         """Cleanup internal data. This needs to be called before the
  53         program ends.
  54         """
  55         while len(self._devices):
  56             device = self._devices.pop()
  57             device.destroy()
  58
  59         while len(self._cleanup_jobs):
  60             job, args = self._cleanup_jobs.pop()
  61             job(*args)
  62
  63     def get_device(self):
  64         """Returns a newly created DiskDevice instance.
  65
  66         This instance is a snapshot of the original source media of
  67         the Disk instance.
  68         """
  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")
  77
  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)
  85
  86         snapshot = uuid.uuid4().hex
  87         tablefd, table = tempfile.mkstemp()
  88         try:
  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)
  93         finally:
  94             os.unlink(table)
  95
  96         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
  97         self._devices.append(new_device)
  98         return new_device
  99
 100     def destroy_device(self, device):
 101         """Destroys a DiskDevice instance previously created by
 102         get_device method.
 103         """
 104         self._devices.remove(device)
 105         device.destroy()
 106
 107
 108 def progress_generator(total):
 109     position = 0;
 110     for i in progress.bar(range(total)):
 111         if i < position:
 112             continue
 113         position = yield
 114     yield #suppress the StopIteration exception
 115
 116
 117 class DiskDevice(object):
 118     """This class represents a block device hosting an Operating System
 119     as created by the device-mapper.
 120     """
 121
 122     def __init__(self, device, bootable=True):
 123         """Create a new DiskDevice."""
 124         self.device = device
 125         self.bootable = bootable
 126         self.progress_bar = None
 127
 128         self.g = guestfs.GuestFS()
 129         self.g.add_drive_opts(device, readonly=0)
 130
 131         #self.g.set_trace(1)
 132         #self.g.set_verbose(1)
 133
 134         eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
 135         self.g.launch()
 136         self.g.delete_event_callback(eh)
 137
 138         roots = self.g.inspect_os()
 139         if len(roots) == 0:
 140             raise DiskError("No operating system found")
 141         if len(roots) > 1:
 142             raise DiskError("Multiple operating systems found")
 143
 144         self.root = roots[0]
 145         self.ostype = self.g.inspect_get_type(self.root)
 146         self.distro = self.g.inspect_get_distro(self.root)
 147
 148     def destroy(self):
 149         """Destroy this DiskDevice instance."""
 150         self.g.umount_all()
 151         self.g.sync()
 152         # Close the guestfs handler
 153         self.g.close()
 154
 155     def progress_callback(self, ev, eh, buf, array):
 156         position = array[2]
 157         total = array[3]
 158
 159         if self.progress_bar is None:
 160             self.progress_bar = progress_generator(total)
 161             self.progress_bar.next()
 162
 163         self.progress_bar.send(position)
 164
 165         if position == total:
 166             self.progress_bar = None
 167
 168     def mount(self):
 169         """Mount all disk partitions in a correct order."""
 170         mps = self.g.inspect_get_mountpoints(self.root)
 171
 172         # Sort the keys to mount the fs in a correct order.
 173         # / should be mounted befor /boot, etc
 174         def compare(a, b):
 175             if len(a[0]) > len(b[0]):
 176                 return 1
 177             elif len(a[0]) == len(b[0]):
 178                 return 0
 179             else:
 180                 return -1
 181         mps.sort(compare)
 182         for mp, dev in mps:
 183             try:
 184                 self.g.mount(dev, mp)
 185             except RuntimeError as msg:
 186                 print "%s (ignored)" % msg
 187
 188     def umount(self):
 189         """Umount all mounted filesystems."""
 190         self.g.umount_all()
 191
 192     def shrink(self):
 193         """Shrink the disk.
 194
 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.
 198         """
 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)
 204
 205         last_partition = self.g.part_list(dev)[-1]
 206
 207         if last_partition['part_num'] > 4:
 208             raise DiskError("This disk contains logical partitions. "
 209                 "Only primary partitions are supported.")
 210
 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
 215             return
 216
 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])
 222
 223         sector_size = self.g.blockdev_getss(dev)
 224
 225         start = last_partition['part_start'] / sector_size
 226         end = start + (block_size * block_cnt) / sector_size - 1
 227
 228         self.g.part_del(dev, last_partition['part_num'])
 229         self.g.part_add(dev, 'p', start, end)
 230
 231         return (end + 1) * sector_size
 232
 233     def size(self):
 234         """Returns the "payload" size of the device.
 235
 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).
 238         """
 239         dev = self.g.part_to_dev(self.root)
 240         last = self.g.part_list(dev)[-1]
 241
 242         return last['part_end']
 243
 244 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :