Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ c408053f

History | View | Annotate | Download (6.5 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
from pbs import dmsetup
13
from pbs import blockdev
14
from pbs import dd
15

    
16

    
17
class DiskError(Exception):
18
    pass
19

    
20

    
21
class Disk(object):
22
    """This class represents a hard disk hosting an Operating System
23

24
    A Disk instance never alters the source media it is created from.
25
    Any change is done on a snapshot created by the device-mapper of
26
    the Linux kernel.
27
    """
28

    
29
    def __init__(self, source):
30
        """Create a new Disk instance out of a source media. The source
31
        media can be an image file, a block device or a directory."""
32
        self._cleanup_jobs = []
33
        self._devices = []
34
        self.source = source
35

    
36
    def _add_cleanup(self, job, *args):
37
        self._cleanup_jobs.append((job, args))
38

    
39
    def _losetup(self, fname):
40
        loop = losetup.find_unused_loop_device()
41
        loop.mount(fname)
42
        self._add_cleanup(loop.unmount)
43
        return loop.device
44

    
45
    def _dir_to_disk(self):
46
        raise NotImplementedError
47

    
48
    def cleanup(self):
49
        """Cleanup internal data. This needs to be called before the
50
        program ends.
51
        """
52
        while len(self._devices):
53
            device = self._devices.pop()
54
            device.destroy()
55

    
56
        while len(self._cleanup_jobs):
57
            job, args = self._cleanup_jobs.pop()
58
            job(*args)
59

    
60
    def get_device(self):
61
        """Returns a newly created DiskDevice instance.
62

63
        This instance is a snapshot of the original source media of
64
        the Disk instance.
65
        """
66
        sourcedev = self.source
67
        mode = os.stat(self.source).st_mode
68
        if stat.S_ISDIR(mode):
69
            return self._losetup(self._dir_to_disk())
70
        elif stat.S_ISREG(mode):
71
            sourcedev = self._losetup(self.source)
72
        elif not stat.S_ISBLK(mode):
73
            raise ValueError("Value for self.source is invalid")
74

    
75
        # Take a snapshot and return it to the user
76
        size = blockdev('--getsize', sourcedev)
77
        cowfd, cow = tempfile.mkstemp()
78
        self._add_cleanup(os.unlink, cow)
79
        # Create 1G cow sparse file
80
        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
81
        cowdev = self._losetup(cow)
82

    
83
        snapshot = uuid.uuid4().hex
84
        tablefd, table = tempfile.mkstemp()
85
        try:
86
            os.write(tablefd, "0 %d snapshot %s %s n 8" % \
87
                                        (int(size), sourcedev, cowdev))
88
            dmsetup('create', snapshot, table)
89
            self._add_cleanup(dmsetup, 'remove', snapshot)
90
        finally:
91
            os.unlink(table)
92

    
93
        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
94
        self._devices.append(new_device)
95
        return new_device
96

    
97
    def destroy_device(self, device):
98
        """Destroys a DiskDevice instance previously created by
99
        get_device method.
100
        """
101
        self._devices.remove(device)
102
        device.destroy()
103

    
104

    
105
class DiskDevice(object):
106
    """This class represents a block device hosting an Operating System
107
    as created by the device-mapper.
108
    """
109

    
110
    def __init__(self, device, bootable=True):
111
        """Create a new DiskDevice."""
112
        self.device = device
113
        self.bootable = bootable
114

    
115
        self.g = guestfs.GuestFS()
116

    
117
        self.g.set_trace(1)
118

    
119
        self.g.add_drive_opts(device, readonly=0)
120
        self.g.launch()
121
        roots = self.g.inspect_os()
122
        if len(roots) == 0:
123
            raise DiskError("No operating system found")
124
        if len(roots) > 1:
125
            raise DiskError("Multiple operating systems found")
126

    
127
        self.root = roots[0]
128
        self.ostype = self.g.inspect_get_type(self.root)
129
        self.distro = self.g.inspect_get_distro(self.root)
130

    
131
    def destroy(self):
132
        """Destroy this DiskDevice instance."""
133
        self.g.umount_all()
134
        self.g.sync()
135
        # Close the guestfs handler
136
        self.g.close()
137

    
138
    def mount(self):
139
        """Mount all disk partitions in a correct order."""
140
        mps = self.g.inspect_get_mountpoints(self.root)
141

    
142
        # Sort the keys to mount the fs in a correct order.
143
        # / should be mounted befor /boot, etc
144
        def compare(a, b):
145
            if len(a[0]) > len(b[0]):
146
                return 1
147
            elif len(a[0]) == len(b[0]):
148
                return 0
149
            else:
150
                return -1
151
        mps.sort(compare)
152
        for mp, dev in mps:
153
            try:
154
                self.g.mount(dev, mp)
155
            except RuntimeError as msg:
156
                print "%s (ignored)" % msg
157

    
158
    def umount(self):
159
        """Umount all mounted filesystems."""
160
        self.g.umount_all()
161

    
162
    def shrink(self):
163
        """Shrink the disk.
164

165
        This is accomplished by shrinking the last filesystem in the
166
        disk and then updating the partition table. The new disk size
167
        (in bytes) is returned.
168
        """
169
        dev = self.g.part_to_dev(self.root)
170
        parttype = self.g.part_get_parttype(dev)
171
        if parttype != 'msdos':
172
            raise DiskError("You have a %s partition table. "
173
                "Only msdos partitions are supported" % parttype)
174

    
175
        last_partition = self.g.part_list(dev)[-1]
176

    
177
        if last_partition['part_num'] > 4:
178
            raise DiskError("This disk contains logical partitions. "
179
                "Only primary partitions are supported.")
180

    
181
        part_dev = "%s%d" % (dev, last_partition['part_num'])
182
        fs_type = self.g.vfs_type(part_dev)
183
        if not re.match("ext[234]", fs_type):
184
            print "Warning: Don't know how to resize %s partitions." % vfs_type
185
            return
186

    
187
        self.g.e2fsck_f(part_dev)
188
        self.g.resize2fs_M(part_dev)
189
        output = self.g.tune2fs_l(part_dev)
190
        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
191
        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
192

    
193
        sector_size = self.g.blockdev_getss(dev)
194

    
195
        start = last_partition['part_start'] / sector_size
196
        end = start + (block_size * block_cnt) / sector_size - 1
197

    
198
        self.g.part_del(dev, last_partition['part_num'])
199
        self.g.part_add(dev, 'p', start, end)
200

    
201
        return (end + 1) * sector_size
202

    
203
    def size(self):
204
        """Returns the "payload" size of the device.
205

206
        The size returned by this method is the size of the space occupied by
207
        the partitions (including the space before the first partition).
208
        """
209
        dev = self.g.part_to_dev(self.root)
210
        last = self.g.part_list(dev)[-1]
211

    
212
        return last['part_end']
213

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