Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ d57775d4

History | View | Annotate | Download (3 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
        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 :