Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 01a7cff3

History | View | Annotate | Download (7 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
import sys
10
import guestfs
11

    
12
import pbs
13
from pbs import dd
14

    
15

    
16
class DiskError(Exception):
17
    pass
18

    
19

    
20
def find_sbin_command(command, exception):
21
    search_paths = ['/usr/local/sbin', '/usr/sbin', '/sbin']
22
    for fullpath in map(lambda x: "%s/%s" % (x, command), search_paths):
23
        if os.path.exists(fullpath) and os.access(fullpath, os.X_OK):
24
            return pbs.Command(fullpath)
25
        continue
26
    raise exception
27

    
28

    
29
try:
30
    from pbs import dmsetup
31
except pbs.CommandNotFound as e:
32
    dmsetup = find_sbin_command('dmsetup', e)
33

    
34
try:
35
    from pbs import blockdev
36
except pbs.CommandNotFound as e:
37
    blockdev = find_sbin_command('blockdev', e)
38

    
39

    
40
class Disk(object):
41
    """This class represents a hard disk hosting an Operating System
42

43
    A Disk instance never alters the source media it is created from.
44
    Any change is done on a snapshot created by the device-mapper of
45
    the Linux kernel.
46
    """
47

    
48
    def __init__(self, source):
49
        """Create a new Disk instance out of a source media. The source
50
        media can be an image file, a block device or a directory."""
51
        self._cleanup_jobs = []
52
        self._devices = []
53
        self.source = source
54

    
55
    def _add_cleanup(self, job, *args):
56
        self._cleanup_jobs.append((job, args))
57

    
58
    def _losetup(self, fname):
59
        loop = losetup.find_unused_loop_device()
60
        loop.mount(fname)
61
        self._add_cleanup(loop.unmount)
62
        return loop.device
63

    
64
    def _dir_to_disk(self):
65
        raise NotImplementedError
66

    
67
    def cleanup(self):
68
        """Cleanup internal data. This needs to be called before the
69
        program ends.
70
        """
71
        while len(self._devices):
72
            device = self._devices.pop()
73
            device.destroy()
74

    
75
        while len(self._cleanup_jobs):
76
            job, args = self._cleanup_jobs.pop()
77
            job(*args)
78

    
79
    def get_device(self):
80
        """Returns a newly created DiskDevice instance.
81

82
        This instance is a snapshot of the original source media of
83
        the Disk instance.
84
        """
85
        sourcedev = self.source
86
        mode = os.stat(self.source).st_mode
87
        if stat.S_ISDIR(mode):
88
            return self._losetup(self._dir_to_disk())
89
        elif stat.S_ISREG(mode):
90
            sourcedev = self._losetup(self.source)
91
        elif not stat.S_ISBLK(mode):
92
            raise ValueError("Value for self.source is invalid")
93

    
94
        # Take a snapshot and return it to the user
95
        size = blockdev('--getsize', sourcedev)
96
        cowfd, cow = tempfile.mkstemp()
97
        self._add_cleanup(os.unlink, cow)
98
        # Create 1G cow sparse file
99
        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
100
        cowdev = self._losetup(cow)
101

    
102
        snapshot = uuid.uuid4().hex
103
        tablefd, table = tempfile.mkstemp()
104
        try:
105
            os.write(tablefd, "0 %d snapshot %s %s n 8" % \
106
                                        (int(size), sourcedev, cowdev))
107
            dmsetup('create', snapshot, table)
108
            self._add_cleanup(dmsetup, 'remove', snapshot)
109
        finally:
110
            os.unlink(table)
111

    
112
        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
113
        self._devices.append(new_device)
114
        return new_device
115

    
116
    def destroy_device(self, device):
117
        """Destroys a DiskDevice instance previously created by
118
        get_device method.
119
        """
120
        self._devices.remove(device)
121
        device.destroy()
122

    
123

    
124
class DiskDevice(object):
125
    """This class represents a block device hosting an Operating System
126
    as created by the device-mapper.
127
    """
128

    
129
    def __init__(self, device, bootable=True):
130
        """Create a new DiskDevice."""
131
        self.device = device
132
        self.bootable = bootable
133

    
134
        self.g = guestfs.GuestFS()
135

    
136
        self.g.set_trace(1)
137

    
138
        self.g.add_drive_opts(device, readonly=0)
139
        self.g.launch()
140
        roots = self.g.inspect_os()
141
        if len(roots) == 0:
142
            raise DiskError("No operating system found")
143
        if len(roots) > 1:
144
            raise DiskError("Multiple operating systems found")
145

    
146
        self.root = roots[0]
147
        self.ostype = self.g.inspect_get_type(self.root)
148
        self.distro = self.g.inspect_get_distro(self.root)
149

    
150
    def destroy(self):
151
        """Destroy this DiskDevice instance."""
152
        self.g.umount_all()
153
        self.g.sync()
154
        # Close the guestfs handler
155
        self.g.close()
156

    
157
    def mount(self):
158
        """Mount all disk partitions in a correct order."""
159
        mps = self.g.inspect_get_mountpoints(self.root)
160

    
161
        # Sort the keys to mount the fs in a correct order.
162
        # / should be mounted befor /boot, etc
163
        def compare(a, b):
164
            if len(a[0]) > len(b[0]):
165
                return 1
166
            elif len(a[0]) == len(b[0]):
167
                return 0
168
            else:
169
                return -1
170
        mps.sort(compare)
171
        for mp, dev in mps:
172
            try:
173
                self.g.mount(dev, mp)
174
            except RuntimeError as msg:
175
                print "%s (ignored)" % msg
176

    
177
    def umount(self):
178
        """Umount all mounted filesystems."""
179
        self.g.umount_all()
180

    
181
    def shrink(self):
182
        """Shrink the disk.
183

184
        This is accomplished by shrinking the last filesystem in the
185
        disk and then updating the partition table. The new disk size
186
        (in bytes) is returned.
187
        """
188
        dev = self.g.part_to_dev(self.root)
189
        parttype = self.g.part_get_parttype(dev)
190
        if parttype != 'msdos':
191
            raise DiskError("You have a %s partition table. "
192
                "Only msdos partitions are supported" % parttype)
193

    
194
        last_partition = self.g.part_list(dev)[-1]
195

    
196
        if last_partition['part_num'] > 4:
197
            raise DiskError("This disk contains logical partitions. "
198
                "Only primary partitions are supported.")
199

    
200
        part_dev = "%s%d" % (dev, last_partition['part_num'])
201
        fs_type = self.g.vfs_type(part_dev)
202
        if not re.match("ext[234]", fs_type):
203
            print "Warning: Don't know how to resize %s partitions." % vfs_type
204
            return
205

    
206
        self.g.e2fsck_f(part_dev)
207
        self.g.resize2fs_M(part_dev)
208
        output = self.g.tune2fs_l(part_dev)
209
        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
210
        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
211

    
212
        sector_size = self.g.blockdev_getss(dev)
213

    
214
        start = last_partition['part_start'] / sector_size
215
        end = start + (block_size * block_cnt) / sector_size - 1
216

    
217
        self.g.part_del(dev, last_partition['part_num'])
218
        self.g.part_add(dev, 'p', start, end)
219

    
220
        return (end + 1) * sector_size
221

    
222
    def size(self):
223
        """Returns the "payload" size of the device.
224

225
        The size returned by this method is the size of the space occupied by
226
        the partitions (including the space before the first partition).
227
        """
228
        dev = self.g.part_to_dev(self.root)
229
        last = self.g.part_list(dev)[-1]
230

    
231
        return last['part_end']
232

    
233
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :