b20ca748d213fd3d4e27f224a3e702b0f7e66830
[snf-image-creator] / 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 :