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         assert self.progress_bar is not None
 178         print 'posisition/total: %s/%s' % (position, total)
 179         self.progress_bar.send((position * 100)//total)
 180
 181         if position == total:
 182             self.progress_bar = None
 183
 184     def mount(self):
 185         """Mount all disk partitions in a correct order."""
 186         mps = self.g.inspect_get_mountpoints(self.root)
 187
 188         # Sort the keys to mount the fs in a correct order.
 189         # / should be mounted befor /boot, etc
 190         def compare(a, b):
 191             if len(a[0]) > len(b[0]):
 192                 return 1
 193             elif len(a[0]) == len(b[0]):
 194                 return 0
 195             else:
 196                 return -1
 197         mps.sort(compare)
 198         for mp, dev in mps:
 199             try:
 200                 self.g.mount(dev, mp)
 201             except RuntimeError as msg:
 202                 print "%s (ignored)" % msg
 203
 204     def umount(self):
 205         """Umount all mounted filesystems."""
 206         self.g.umount_all()
 207
 208     def shrink(self):
 209         """Shrink the disk.
 210
 211         This is accomplished by shrinking the last filesystem in the
 212         disk and then updating the partition table. The new disk size
 213         (in bytes) is returned.
 214         """
 215         dev = self.g.part_to_dev(self.root)
 216         parttype = self.g.part_get_parttype(dev)
 217         if parttype != 'msdos':
 218             raise FatalError("You have a %s partition table. "
 219                 "Only msdos partitions are supported" % parttype)
 220
 221         last_partition = self.g.part_list(dev)[-1]
 222
 223         if last_partition['part_num'] > 4:
 224             raise FatalError("This disk contains logical partitions. "
 225                 "Only primary partitions are supported.")
 226
 227         part_dev = "%s%d" % (dev, last_partition['part_num'])
 228         fs_type = self.g.vfs_type(part_dev)
 229         if not re.match("ext[234]", fs_type):
 230             print "Warning: Don't know how to resize %s partitions." % vfs_type
 231             return
 232
 233         self.g.e2fsck_f(part_dev)
 234         self.g.resize2fs_M(part_dev)
 235         output = self.g.tune2fs_l(part_dev)
 236         block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
 237         block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
 238
 239         sector_size = self.g.blockdev_getss(dev)
 240
 241         start = last_partition['part_start'] / sector_size
 242         end = start + (block_size * block_cnt) / sector_size - 1
 243
 244         self.g.part_del(dev, last_partition['part_num'])
 245         self.g.part_add(dev, 'p', start, end)
 246
 247         return (end + 1) * sector_size
 248
 249     def size(self):
 250         """Returns the "payload" size of the device.
 251
 252         The size returned by this method is the size of the space occupied by
 253         the partitions (including the space before the first partition).
 254         """
 255         dev = self.g.part_to_dev(self.root)
 256         last = self.g.part_list(dev)[-1]
 257
 258         return last['part_end']
 259
 260 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :