Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ c54fc0e8

History | View | Annotate | Download (7.8 kB)

1
#!/usr/bin/env python
2

    
3
from image_creator.util import get_command
4
from image_creator import FatalError
5
from clint.textui import progress
6

    
7
import stat
8
import os
9
import tempfile
10
import uuid
11
import re
12
import sys
13
import guestfs
14
import time
15

    
16

    
17
class DiskError(Exception):
18
    pass
19

    
20
dd = get_command('dd')
21
dmsetup = get_command('dmsetup')
22
losetup = get_command('losetup')
23
blockdev = get_command('blockdev')
24

    
25

    
26
class Disk(object):
27
    """This class represents a hard disk hosting an Operating System
28

29
    A Disk instance never alters the source media it is created from.
30
    Any change is done on a snapshot created by the device-mapper of
31
    the Linux kernel.
32
    """
33

    
34
    def __init__(self, source):
35
        """Create a new Disk instance out of a source media. The source
36
        media can be an image file, a block device or a directory."""
37
        self._cleanup_jobs = []
38
        self._devices = []
39
        self.source = source
40

    
41
    def _add_cleanup(self, job, *args):
42
        self._cleanup_jobs.append((job, args))
43

    
44
    def _losetup(self, fname):
45
        loop = losetup('-f', '--show', fname)
46
        loop = loop.strip() # remove the new-line char
47
        self._add_cleanup(losetup, '-d', loop)
48
        return loop
49

    
50
    def _dir_to_disk(self):
51
        raise NotImplementedError
52

    
53
    def cleanup(self):
54
        """Cleanup internal data. This needs to be called before the
55
        program ends.
56
        """
57
        while len(self._devices):
58
            device = self._devices.pop()
59
            device.destroy()
60

    
61
        while len(self._cleanup_jobs):
62
            job, args = self._cleanup_jobs.pop()
63
            job(*args)
64

    
65
    def get_device(self):
66
        """Returns a newly created DiskDevice instance.
67

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

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

    
89
        snapshot = uuid.uuid4().hex
90
        tablefd, table = tempfile.mkstemp()
91
        try:
92
            os.write(tablefd, "0 %d snapshot %s %s n 8" % \
93
                                        (int(size), sourcedev, cowdev))
94
            dmsetup('create', snapshot, table)
95
            self._add_cleanup(dmsetup, 'remove', snapshot)
96
        finally:
97
            os.unlink(table)
98
        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
99
        self._devices.append(new_device)
100
        new_device.enable()
101
        return new_device
102

    
103
    def destroy_device(self, device):
104
        """Destroys a DiskDevice instance previously created by
105
        get_device method.
106
        """
107
        self._devices.remove(device)
108
        device.destroy()
109

    
110

    
111
def progress_generator(label=''):
112
    position = 0;
113
    for i in progress.bar(range(100),label):
114
        if i < position:
115
            continue
116
        position = yield
117
    yield #suppress the StopIteration exception
118

    
119

    
120
class DiskDevice(object):
121
    """This class represents a block device hosting an Operating System
122
    as created by the device-mapper.
123
    """
124

    
125
    def __init__(self, device, bootable=True):
126
        """Create a new DiskDevice."""
127

    
128
        self.device = device
129
        self.bootable = bootable
130
        self.progress_bar = None
131

    
132
        self.g = guestfs.GuestFS()
133
        self.g.add_drive_opts(self.device, readonly=0)
134

    
135
        #self.g.set_trace(1)
136
        #self.g.set_verbose(1)
137

    
138
        self.guestfs_enabled = False
139
    
140
    def enable(self):
141
        """Enable a newly created DiskDevice"""
142

    
143
        self.progressbar = progress_generator("VM lauch: ")
144
        self.progressbar.next()
145
        eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
146
        self.g.launch()
147
        self.guestfs_enabled = True
148
        self.g.delete_event_callback(eh)
149
        if self.progressbar is not None:
150
            self.progressbar.send(100)
151
            self.progressbar = None
152
        
153
        roots = self.g.inspect_os()
154
        if len(roots) == 0:
155
            raise FatalError("No operating system found")
156
        if len(roots) > 1:
157
            raise FatalError("Multiple operating systems found")
158

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

    
163
    def destroy(self):
164
        """Destroy this DiskDevice instance."""
165

    
166
        if self.guestfs_enabled:
167
            self.g.umount_all()
168
            self.g.sync()
169

    
170
        # Close the guestfs handler if open
171
        self.g.close()
172

    
173
    def progress_callback(self, ev, eh, buf, array):
174
        position = array[2]
175
        total = array[3]
176

    
177
        self.progressbar.send((position * 100)//total)
178

    
179
        if position == total:
180
            self.progressbar = None
181

    
182
    def mount(self):
183
        """Mount all disk partitions in a correct order."""
184
        mps = self.g.inspect_get_mountpoints(self.root)
185

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

    
202
    def umount(self):
203
        """Umount all mounted filesystems."""
204
        self.g.umount_all()
205

    
206
    def shrink(self):
207
        """Shrink the disk.
208

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

    
219
        last_partition = self.g.part_list(dev)[-1]
220

    
221
        if last_partition['part_num'] > 4:
222
            raise FatalError("This disk contains logical partitions. "
223
                "Only primary partitions are supported.")
224

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

    
231
        self.g.e2fsck_f(part_dev)
232
        self.g.resize2fs_M(part_dev)
233
        output = self.g.tune2fs_l(part_dev)
234
        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
235
        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
236

    
237
        sector_size = self.g.blockdev_getss(dev)
238

    
239
        start = last_partition['part_start'] / sector_size
240
        end = start + (block_size * block_cnt) / sector_size - 1
241

    
242
        self.g.part_del(dev, last_partition['part_num'])
243
        self.g.part_add(dev, 'p', start, end)
244

    
245
        return (end + 1) * sector_size
246

    
247
    def size(self):
248
        """Returns the "payload" size of the device.
249

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

    
256
        return last['part_end']
257

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