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