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 from pbs import dmsetup
  10 from pbs import blockdev
  11 from pbs import dd
  12 from pbs import kpartx
  13 from pbs import mount
  14 from pbs import umount
  15
  16 class Disk(object):
  17
  18     def __init__(self, source):
  19         self._cleanup_jobs = []
  20         self._devices = []
  21         self.source = source
  22
  23     def _add_cleanup(self, job, *args):
  24         self._cleanup_jobs.append((job, args))
  25
  26     def _losetup(self, fname):
  27         loop = losetup.find_unused_loop_device()
  28         loop.mount(fname)
  29         self._add_cleanup(loop.unmount)
  30         return loop.device
  31
  32     def _dir_to_disk(self):
  33         raise NotImplementedError
  34
  35     def cleanup(self):
  36         while len(self._cleanup_jobs):
  37             job, args = self._cleanup_jobs.pop()
  38             job(*args)
  39
  40     def get_device(self):
  41         sourcedev = self.source
  42         mode = os.stat(self.source).st_mode
  43         if stat.S_ISDIR(mode):
  44             return self._losetup(self._dir_to_disk())
  45         elif stat.S_ISREG(mode):
  46             sourcedev = self._losetup(self.source)
  47         elif not stat.S_ISBLK(mode):
  48             raise ValueError("Value for self.source is invalid")
  49
  50         # Take a snapshot and return it to the user
  51         size = blockdev('--getsize', sourcedev)
  52         cowfd, cow = tempfile.mkstemp()
  53         self._add_cleanup(os.unlink, cow)
  54         dd('if=/dev/zero', 'of=%s' % cow, 'count=%d' % (1024*1024))#(int(size)/4))
  55         cowdev = self._losetup(cow)
  56
  57         snapshot = uuid.uuid4().hex
  58         tablefd, table = tempfile.mkstemp()
  59         try:
  60             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
  61                                         (int(size), sourcedev, cowdev))
  62             dmsetup('create', snapshot, table)
  63             self._add_cleanup(dmsetup, 'remove', snapshot)
  64         finally:
  65             os.unlink(table)
  66
  67         new_device = DiskDevice(self, "/dev/mapper/%s" % snapshot)
  68         self._devices.append(new_device)
  69         return new_device
  70
  71 class DiskDevice(object):
  72
  73     def __init__(self, disk, device):
  74         self.disk = disk
  75         self.dev = device
  76         self.partitions_mapped = False
  77         self.magic_number = uuid.uuid4().hex
  78
  79     def list_partitions(self):
  80         output = kpartx("-l", "-p", self.magic_number, self.dev)
  81         return [ "/dev/mapper/%s" % x for x in
  82                 re.findall('^\S+', str(output), flags=re.MULTILINE)]
  83
  84     def mount(self, partition):
  85         if not self.partitions_mapped:
  86             kpartx("-a", "-p", self.magic_number, self.dev)
  87             self.disk._cleanup_jobs.append(kpartx, "-d", "-p",
  88                         self.magic_number, self.dev)
  89             self.partitions_mapped = True
  90
  91         targetfd, target = tempfile.mkdtemp()
  92         try:
  93             mount(dev, partition)
  94         except:
  95             os.rmdir(table)
  96             raise
  97         return target
  98
  99     def unmount(self, partition):
 100         umount(target)
 101
 102         mode = os.stat(self.source).st_mode
 103         if stat.S_ISDIR(mode):
 104             os.rmdir(target)
 105
 106 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :