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 class DiskError(Exception): pass
  17
  18 class Disk(object):
  19
  20     def __init__(self, source):
  21         self._cleanup_jobs = []
  22         self._devices = []
  23         self.source = source
  24
  25     def _add_cleanup(self, job, *args):
  26         self._cleanup_jobs.append((job, args))
  27
  28     def _losetup(self, fname):
  29         loop = losetup.find_unused_loop_device()
  30         loop.mount(fname)
  31         self._add_cleanup(loop.unmount)
  32         return loop.device
  33
  34     def _dir_to_disk(self):
  35         raise NotImplementedError
  36
  37     def cleanup(self):
  38         while len(self._devices):
  39             device = self._devices.pop()
  40             device.destroy()
  41
  42         while len(self._cleanup_jobs):
  43             job, args = self._cleanup_jobs.pop()
  44             job(*args)
  45
  46     def get_device(self):
  47         sourcedev = self.source
  48         mode = os.stat(self.source).st_mode
  49         if stat.S_ISDIR(mode):
  50             return self._losetup(self._dir_to_disk())
  51         elif stat.S_ISREG(mode):
  52             sourcedev = self._losetup(self.source)
  53         elif not stat.S_ISBLK(mode):
  54             raise ValueError("Value for self.source is invalid")
  55
  56         # Take a snapshot and return it to the user
  57         size = blockdev('--getsize', sourcedev)
  58         cowfd, cow = tempfile.mkstemp()
  59         self._add_cleanup(os.unlink, cow)
  60         # Create 1G cow file
  61         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k' ,'seek=%d' % (1024*1024))
  62         cowdev = self._losetup(cow)
  63
  64         snapshot = uuid.uuid4().hex
  65         tablefd, table = tempfile.mkstemp()
  66         try:
  67             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
  68                                         (int(size), sourcedev, cowdev))
  69             dmsetup('create', snapshot, table)
  70             self._add_cleanup(dmsetup, 'remove', snapshot)
  71         finally:
  72             os.unlink(table)
  73
  74         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
  75         self._devices.append(new_device)
  76         return new_device
  77
  78     def destroy_device(self, device):
  79         self._devices.remove(device)
  80         device.destroy()
  81
  82 class DiskDevice(object):
  83
  84     def __init__(self, device, bootable = True):
  85         self.device = device
  86         self.bootable = bootable
  87
  88         self.g = guestfs.GuestFS()
  89
  90         self.g.set_trace(1)
  91
  92         self.g.add_drive_opts(device, readonly = 0)
  93         self.g.launch()
  94         roots = self.g.inspect_os()
  95         if len(roots) == 0:
  96             raise DiskError("No operating system found")
  97         if len(roots) > 1:
  98             raise DiskError("Multiple operating systems found")
  99
 100         self.root = roots[0]
 101         self.ostype = self.g.inspect_get_type(self.root)
 102         self.distro = self.g.inspect_get_distro(self.root)
 103
 104     def destroy(self):
 105         self.g.umount_all()
 106         self.g.sync()
 107         # Close the guestfs handler
 108         self.g.close()
 109         del self.g
 110
 111     def mount(self):
 112         mps = self.g.inspect_get_mountpoints(self.root)
 113         # Sort the keys to mount the fs in a correct order.
 114         # / should be mounted befor /boot, etc
 115         def compare (a, b):
 116             if len(a[0]) > len(b[0]): return 1
 117             elif len(a[0]) == len(b[0]): return 0
 118             else: return -1
 119         mps.sort(compare)
 120         for mp, dev in mps:
 121             try:
 122                 self.g.mount(dev, mp)
 123             except RuntimeError as msg:
 124                 print "%s (ignored)" % msg
 125
 126 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :