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
  23     def __init__(self, source):
  24         self._cleanup_jobs = []
  25         self._devices = []
  26         self.source = source
  27
  28     def _add_cleanup(self, job, *args):
  29         self._cleanup_jobs.append((job, args))
  30
  31     def _losetup(self, fname):
  32         loop = losetup.find_unused_loop_device()
  33         loop.mount(fname)
  34         self._add_cleanup(loop.unmount)
  35         return loop.device
  36
  37     def _dir_to_disk(self):
  38         raise NotImplementedError
  39
  40     def cleanup(self):
  41         while len(self._devices):
  42             device = self._devices.pop()
  43             device.destroy()
  44
  45         while len(self._cleanup_jobs):
  46             job, args = self._cleanup_jobs.pop()
  47             job(*args)
  48
  49     def get_device(self):
  50         sourcedev = self.source
  51         mode = os.stat(self.source).st_mode
  52         if stat.S_ISDIR(mode):
  53             return self._losetup(self._dir_to_disk())
  54         elif stat.S_ISREG(mode):
  55             sourcedev = self._losetup(self.source)
  56         elif not stat.S_ISBLK(mode):
  57             raise ValueError("Value for self.source is invalid")
  58
  59         # Take a snapshot and return it to the user
  60         size = blockdev('--getsize', sourcedev)
  61         cowfd, cow = tempfile.mkstemp()
  62         self._add_cleanup(os.unlink, cow)
  63         # Create 1G cow sparse file
  64         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
  65         cowdev = self._losetup(cow)
  66
  67         snapshot = uuid.uuid4().hex
  68         tablefd, table = tempfile.mkstemp()
  69         try:
  70             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
  71                                         (int(size), sourcedev, cowdev))
  72             dmsetup('create', snapshot, table)
  73             self._add_cleanup(dmsetup, 'remove', snapshot)
  74         finally:
  75             os.unlink(table)
  76
  77         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
  78         self._devices.append(new_device)
  79         return new_device
  80
  81     def destroy_device(self, device):
  82         self._devices.remove(device)
  83         device.destroy()
  84
  85
  86 class DiskDevice(object):
  87
  88     def __init__(self, device, bootable=True):
  89         self.device = device
  90         self.bootable = bootable
  91
  92         self.g = guestfs.GuestFS()
  93
  94         self.g.set_trace(1)
  95
  96         self.g.add_drive_opts(device, readonly=0)
  97         self.g.launch()
  98         roots = self.g.inspect_os()
  99         if len(roots) == 0:
 100             raise DiskError("No operating system found")
 101         if len(roots) > 1:
 102             raise DiskError("Multiple operating systems found")
 103
 104         self.root = roots[0]
 105         self.ostype = self.g.inspect_get_type(self.root)
 106         self.distro = self.g.inspect_get_distro(self.root)
 107
 108     def destroy(self):
 109         self.g.umount_all()
 110         self.g.sync()
 111         # Close the guestfs handler
 112         self.g.close()
 113         del self.g
 114
 115     def mount(self):
 116         mps = self.g.inspect_get_mountpoints(self.root)
 117
 118         # Sort the keys to mount the fs in a correct order.
 119         # / should be mounted befor /boot, etc
 120         def compare(a, b):
 121             if len(a[0]) > len(b[0]):
 122                 return 1
 123             elif len(a[0]) == len(b[0]):
 124                 return 0
 125             else:
 126                 return -1
 127         mps.sort(compare)
 128         for mp, dev in mps:
 129             try:
 130                 self.g.mount(dev, mp)
 131             except RuntimeError as msg:
 132                 print "%s (ignored)" % msg
 133
 134     def umount(self):
 135         self.g.umount_all()
 136
 137     def shrink(self):
 138         dev = self.g.part_to_dev(self.root)
 139         parttype = self.g.part_get_parttype(dev)
 140         if parttype != 'msdos':
 141             raise DiskError("You have a %s partition table. "
 142                 "Only msdos partitions are supported" % parttype)
 143
 144         last_partition = self.g.part_list(dev)[-1]
 145
 146         if last_partition['part_num'] > 4:
 147             raise DiskError("This disk contains logical partitions. "
 148                 "Only primary partitions are supported.")
 149
 150         part_dev = "%s%d" % (dev, last_partition['part_num'])
 151         fs_type = self.g.vfs_type(part_dev)
 152         if not re.match("ext[234]", fs_type):
 153             print "Warning, don't know how to resize %s partitions" % vfs_type
 154             return
 155
 156         self.g.e2fsck_f(part_dev)
 157         self.g.resize2fs_M(part_dev)
 158         output = self.g.tune2fs_l(part_dev)
 159         block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
 160         block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
 161
 162         sector_size = self.g.blockdev_getss(dev)
 163
 164         start = last_partition['part_start'] / sector_size
 165         end = start + (block_size * block_cnt) / sector_size - 1
 166
 167
 168
 169 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :