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