Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 586da0a0

History | View | Annotate | Download (7.8 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
from clint.textui import progress
15

    
16

    
17
class DiskError(Exception):
18
    pass
19

    
20

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

    
29

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

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

    
40

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

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

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

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

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

    
65
    def _dir_to_disk(self):
66
        raise NotImplementedError
67

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

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

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

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

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

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

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

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

    
124

    
125
def progress_generator(total):
126
    position = 0;
127
    for i in progress.bar(range(total)):
128
        if i < position:
129
            continue
130
        position = yield
131

    
132

    
133
class DiskDevice(object):
134
    """This class represents a block device hosting an Operating System
135
    as created by the device-mapper.
136
    """
137

    
138
    def __init__(self, device, bootable=True):
139
        """Create a new DiskDevice."""
140
        self.device = device
141
        self.bootable = bootable
142
        self.progress_bar = None
143

    
144
        self.g = guestfs.GuestFS()
145
        self.g.add_drive_opts(device, readonly=0)
146

    
147
        #self.g.set_trace(1)
148
        #self.g.set_verbose(1)
149

    
150
        eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
151
        self.g.launch()
152
        self.g.delete_event_callback(eh)
153
        
154
        roots = self.g.inspect_os()
155
        if len(roots) == 0:
156
            raise DiskError("No operating system found")
157
        if len(roots) > 1:
158
            raise DiskError("Multiple operating systems found")
159

    
160
        self.root = roots[0]
161
        self.ostype = self.g.inspect_get_type(self.root)
162
        self.distro = self.g.inspect_get_distro(self.root)
163

    
164
    def destroy(self):
165
        """Destroy this DiskDevice instance."""
166
        self.g.umount_all()
167
        self.g.sync()
168
        # Close the guestfs handler
169
        self.g.close()
170

    
171
    def progress_callback(self, ev, eh, buf, array):
172
        position = array[2]
173
        total = array[3]
174
        
175
        if self.progress_bar is None:
176
            self.progress_bar = progress_generator(total)
177
            self.progress_bar.next()
178
            if position == 1:
179
                return
180

    
181
        self.progress_bar.send(position)
182

    
183
        if position == total:
184
            self.progress_bar = None
185

    
186
    def mount(self):
187
        """Mount all disk partitions in a correct order."""
188
        mps = self.g.inspect_get_mountpoints(self.root)
189

    
190
        # Sort the keys to mount the fs in a correct order.
191
        # / should be mounted befor /boot, etc
192
        def compare(a, b):
193
            if len(a[0]) > len(b[0]):
194
                return 1
195
            elif len(a[0]) == len(b[0]):
196
                return 0
197
            else:
198
                return -1
199
        mps.sort(compare)
200
        for mp, dev in mps:
201
            try:
202
                self.g.mount(dev, mp)
203
            except RuntimeError as msg:
204
                print "%s (ignored)" % msg
205

    
206
    def umount(self):
207
        """Umount all mounted filesystems."""
208
        self.g.umount_all()
209

    
210
    def shrink(self):
211
        """Shrink the disk.
212

213
        This is accomplished by shrinking the last filesystem in the
214
        disk and then updating the partition table. The new disk size
215
        (in bytes) is returned.
216
        """
217
        dev = self.g.part_to_dev(self.root)
218
        parttype = self.g.part_get_parttype(dev)
219
        if parttype != 'msdos':
220
            raise DiskError("You have a %s partition table. "
221
                "Only msdos partitions are supported" % parttype)
222

    
223
        last_partition = self.g.part_list(dev)[-1]
224

    
225
        if last_partition['part_num'] > 4:
226
            raise DiskError("This disk contains logical partitions. "
227
                "Only primary partitions are supported.")
228

    
229
        part_dev = "%s%d" % (dev, last_partition['part_num'])
230
        fs_type = self.g.vfs_type(part_dev)
231
        if not re.match("ext[234]", fs_type):
232
            print "Warning: Don't know how to resize %s partitions." % vfs_type
233
            return
234

    
235
        self.g.e2fsck_f(part_dev)
236
        self.g.resize2fs_M(part_dev)
237
        output = self.g.tune2fs_l(part_dev)
238
        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
239
        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
240

    
241
        sector_size = self.g.blockdev_getss(dev)
242

    
243
        start = last_partition['part_start'] / sector_size
244
        end = start + (block_size * block_cnt) / sector_size - 1
245

    
246
        self.g.part_del(dev, last_partition['part_num'])
247
        self.g.part_add(dev, 'p', start, end)
248

    
249
        return (end + 1) * sector_size
250

    
251
    def size(self):
252
        """Returns the "payload" size of the device.
253

254
        The size returned by this method is the size of the space occupied by
255
        the partitions (including the space before the first partition).
256
        """
257
        dev = self.g.part_to_dev(self.root)
258
        last = self.g.part_list(dev)[-1]
259

    
260
        return last['part_end']
261

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