root / image_creator / disk.py @ 333ff548
History | View | Annotate | Download (3.4 kB)
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 :
|