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         # Create 1G cow file
  55         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k' ,'seek=%d' % (1024*1024))
  56         cowdev = self._losetup(cow)
  57
  58         snapshot = uuid.uuid4().hex
  59         tablefd, table = tempfile.mkstemp()
  60         try:
  61             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
  62                                         (int(size), sourcedev, cowdev))
  63             dmsetup('create', snapshot, table)
  64             self._add_cleanup(dmsetup, 'remove', snapshot)
  65         finally:
  66             os.unlink(table)
  67
  68         new_device = DiskDevice(self, "/dev/mapper/%s" % snapshot)
  69         self._devices.append(new_device)
  70         return new_device
  71
  72 class DiskDevice(object):
  73
  74     def __init__(self, disk, device, bootable = True):
  75         self.disk = disk
  76         self.device = device
  77         self.is_bootable = bootable
  78         self.partitions_mapped = False
  79         self.magic_number = uuid.uuid4().hex
  80
  81     def list_partitions(self):
  82         if not self.partitions_mapped:
  83             kpartx("-a", "-p", self.magic_number, self.dev)
  84             self.disk._cleanup_jobs.append(kpartx, "-d", "-p",
  85                         self.magic_number, self.dev)
  86             self.partitions_mapped = True
  87
  88         output = kpartx("-l", "-p", self.magic_number, self.dev)
  89         return [ "/dev/mapper/%s" % x for x in
  90                 re.findall('^\S+', str(output), flags=re.MULTILINE)]
  91
  92     def mount(self, partition):
  93         if not self.partitions_mapped:
  94             self.list_partitions()
  95             kpartx("-a", "-p", self.magic_number, self.dev)
  96             self.disk._cleanup_jobs.append(kpartx, "-d", "-p",
  97                         self.magic_number, self.dev)
  98             self.partitions_mapped = True
  99
 100         targetfd, target = tempfile.mkdtemp()
 101         try:
 102             mount(dev, partition)
 103         except:
 104             os.rmdir(table)
 105             raise
 106         return target
 107
 108     def unmount(self, partition):
 109         umount(target)
 110
 111         mode = os.stat(self.source).st_mode
 112         if stat.S_ISDIR(mode):
 113             os.rmdir(target)
 114
 115 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :