Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 3ccb2618

History | View | Annotate | Download (7.4 kB)

1
#!/usr/bin/env python
2

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

    
6
import stat
7
import os
8
import tempfile
9
import uuid
10
import re
11
import sys
12
import guestfs
13

    
14

    
15
class DiskError(Exception):
16
    pass
17

    
18
dd = get_command('dd')
19
dmsetup = get_command('dmsetup')
20
losetup = get_command('losetup')
21
blockdev = get_command('blockdev')
22

    
23

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

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

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

    
39
    def _add_cleanup(self, job, *args):
40
        self._cleanup_jobs.append((job, args))
41

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

    
48
    def _dir_to_disk(self):
49
        raise NotImplementedError
50

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

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

    
63
    def get_device(self):
64
        """Returns a newly created DiskDevice instance.
65

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

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

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

    
96
        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
97
        self._devices.append(new_device)
98
        return new_device
99

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

    
107

    
108
def progress_generator(total):
109
    position = 0;
110
    for i in progress.bar(range(total)):
111
        if i < position:
112
            continue
113
        position = yield
114
    yield #suppress the StopIteration exception
115

    
116

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

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

    
128
        self.g = guestfs.GuestFS()
129
        self.g.add_drive_opts(device, readonly=0)
130

    
131
        #self.g.set_trace(1)
132
        #self.g.set_verbose(1)
133

    
134
        eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
135
        self.g.launch()
136
        self.g.delete_event_callback(eh)
137
        
138
        roots = self.g.inspect_os()
139
        if len(roots) == 0:
140
            raise DiskError("No operating system found")
141
        if len(roots) > 1:
142
            raise DiskError("Multiple operating systems found")
143

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

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

    
155
    def progress_callback(self, ev, eh, buf, array):
156
        position = array[2]
157
        total = array[3]
158
        
159
        if self.progress_bar is None:
160
            self.progress_bar = progress_generator(total)
161
            self.progress_bar.next()
162

    
163
        self.progress_bar.send(position)
164

    
165
        if position == total:
166
            self.progress_bar = None
167

    
168
    def mount(self):
169
        """Mount all disk partitions in a correct order."""
170
        mps = self.g.inspect_get_mountpoints(self.root)
171

    
172
        # Sort the keys to mount the fs in a correct order.
173
        # / should be mounted befor /boot, etc
174
        def compare(a, b):
175
            if len(a[0]) > len(b[0]):
176
                return 1
177
            elif len(a[0]) == len(b[0]):
178
                return 0
179
            else:
180
                return -1
181
        mps.sort(compare)
182
        for mp, dev in mps:
183
            try:
184
                self.g.mount(dev, mp)
185
            except RuntimeError as msg:
186
                print "%s (ignored)" % msg
187

    
188
    def umount(self):
189
        """Umount all mounted filesystems."""
190
        self.g.umount_all()
191

    
192
    def shrink(self):
193
        """Shrink the disk.
194

195
        This is accomplished by shrinking the last filesystem in the
196
        disk and then updating the partition table. The new disk size
197
        (in bytes) is returned.
198
        """
199
        dev = self.g.part_to_dev(self.root)
200
        parttype = self.g.part_get_parttype(dev)
201
        if parttype != 'msdos':
202
            raise DiskError("You have a %s partition table. "
203
                "Only msdos partitions are supported" % parttype)
204

    
205
        last_partition = self.g.part_list(dev)[-1]
206

    
207
        if last_partition['part_num'] > 4:
208
            raise DiskError("This disk contains logical partitions. "
209
                "Only primary partitions are supported.")
210

    
211
        part_dev = "%s%d" % (dev, last_partition['part_num'])
212
        fs_type = self.g.vfs_type(part_dev)
213
        if not re.match("ext[234]", fs_type):
214
            print "Warning: Don't know how to resize %s partitions." % vfs_type
215
            return
216

    
217
        self.g.e2fsck_f(part_dev)
218
        self.g.resize2fs_M(part_dev)
219
        output = self.g.tune2fs_l(part_dev)
220
        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
221
        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
222

    
223
        sector_size = self.g.blockdev_getss(dev)
224

    
225
        start = last_partition['part_start'] / sector_size
226
        end = start + (block_size * block_cnt) / sector_size - 1
227

    
228
        self.g.part_del(dev, last_partition['part_num'])
229
        self.g.part_add(dev, 'p', start, end)
230

    
231
        return (end + 1) * sector_size
232

    
233
    def size(self):
234
        """Returns the "payload" size of the device.
235

236
        The size returned by this method is the size of the space occupied by
237
        the partitions (including the space before the first partition).
238
        """
239
        dev = self.g.part_to_dev(self.root)
240
        last = self.g.part_list(dev)[-1]
241

    
242
        return last['part_end']
243

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