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