Statistics
| Branch: | Tag: | Revision:

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 :