Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ d144e954

History | View | Annotate | Download (10.3 kB)

1
# Copyright 2012 GRNET S.A. All rights reserved.
2
#
3
# Redistribution and use in source and binary forms, with or
4
# without modification, are permitted provided that the following
5
# conditions are met:
6
#
7
#   1. Redistributions of source code must retain the above
8
#      copyright notice, this list of conditions and the following
9
#      disclaimer.
10
#
11
#   2. Redistributions in binary form must reproduce the above
12
#      copyright notice, this list of conditions and the following
13
#      disclaimer in the documentation and/or other materials
14
#      provided with the distribution.
15
#
16
# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
17
# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
18
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
19
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
20
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
21
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
22
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
23
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
24
# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
25
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
26
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
27
# POSSIBILITY OF SUCH DAMAGE.
28
#
29
# The views and conclusions contained in the software and
30
# documentation are those of the authors and should not be
31
# interpreted as representing official policies, either expressed
32
# or implied, of GRNET S.A.
33

    
34
from image_creator.util import get_command, warn, progress_generator
35
from image_creator import FatalError
36
from clint.textui import indent, puts, colored
37

    
38
import stat
39
import os
40
import tempfile
41
import uuid
42
import re
43
import sys
44
import guestfs
45
import time
46

    
47

    
48
class DiskError(Exception):
49
    pass
50

    
51
dd = get_command('dd')
52
dmsetup = get_command('dmsetup')
53
losetup = get_command('losetup')
54
blockdev = get_command('blockdev')
55

    
56

    
57
class Disk(object):
58
    """This class represents a hard disk hosting an Operating System
59

60
    A Disk instance never alters the source media it is created from.
61
    Any change is done on a snapshot created by the device-mapper of
62
    the Linux kernel.
63
    """
64

    
65
    def __init__(self, source):
66
        """Create a new Disk instance out of a source media. The source
67
        media can be an image file, a block device or a directory."""
68
        self._cleanup_jobs = []
69
        self._devices = []
70
        self.source = source
71

    
72
    def _add_cleanup(self, job, *args):
73
        self._cleanup_jobs.append((job, args))
74

    
75
    def _losetup(self, fname):
76
        loop = losetup('-f', '--show', fname)
77
        loop = loop.strip()  # remove the new-line char
78
        self._add_cleanup(losetup, '-d', loop)
79
        return loop
80

    
81
    def _dir_to_disk(self):
82
        raise NotImplementedError
83

    
84
    def cleanup(self):
85
        """Cleanup internal data. This needs to be called before the
86
        program ends.
87
        """
88
        while len(self._devices):
89
            device = self._devices.pop()
90
            device.destroy()
91

    
92
        while len(self._cleanup_jobs):
93
            job, args = self._cleanup_jobs.pop()
94
            job(*args)
95

    
96
    def get_device(self):
97
        """Returns a newly created DiskDevice instance.
98

99
        This instance is a snapshot of the original source media of
100
        the Disk instance.
101
        """
102

    
103
        puts("Examining source media `%s'..." % self.source, False)
104
        sourcedev = self.source
105
        mode = os.stat(self.source).st_mode
106
        if stat.S_ISDIR(mode):
107
            puts(colored.green('looks like a directory'))
108
            return self._losetup(self._dir_to_disk())
109
        elif stat.S_ISREG(mode):
110
            puts(colored.green('looks like an image file'))
111
            sourcedev = self._losetup(self.source)
112
        elif not stat.S_ISBLK(mode):
113
            raise ValueError("Invalid media source. Only block devices, "
114
                            "regular files and directories are supported.")
115
        else:
116
            puts(colored.green('looks like a block device'))
117

    
118
        # Take a snapshot and return it to the user
119
        puts("Snapshotting media source...", False)
120
        size = blockdev('--getsize', sourcedev)
121
        cowfd, cow = tempfile.mkstemp()
122
        os.close(cowfd)
123
        self._add_cleanup(os.unlink, cow)
124
        # Create 1G cow sparse file
125
        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
126
                                        'seek=%d' % (1024 * 1024))
127
        cowdev = self._losetup(cow)
128

    
129
        snapshot = uuid.uuid4().hex
130
        tablefd, table = tempfile.mkstemp()
131
        try:
132
            os.write(tablefd, "0 %d snapshot %s %s n 8" % \
133
                                        (int(size), sourcedev, cowdev))
134
            dmsetup('create', snapshot, table)
135
            self._add_cleanup(dmsetup, 'remove', snapshot)
136
            # Sometimes dmsetup remove fails with Device or resource busy,
137
            # although everything is cleaned up and the snapshot is not
138
            # used by anyone. Add a 2 seconds delay to be on the safe side.
139
            self._add_cleanup(time.sleep, 2)
140

    
141
        finally:
142
            os.unlink(table)
143
        puts(colored.green('done'))
144
        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
145
        self._devices.append(new_device)
146
        new_device.enable()
147
        return new_device
148

    
149
    def destroy_device(self, device):
150
        """Destroys a DiskDevice instance previously created by
151
        get_device method.
152
        """
153
        self._devices.remove(device)
154
        device.destroy()
155

    
156

    
157
class DiskDevice(object):
158
    """This class represents a block device hosting an Operating System
159
    as created by the device-mapper.
160
    """
161

    
162
    def __init__(self, device, bootable=True):
163
        """Create a new DiskDevice."""
164

    
165
        self.device = device
166
        self.bootable = bootable
167
        self.progress_bar = None
168

    
169
        self.g = guestfs.GuestFS()
170
        self.g.add_drive_opts(self.device, readonly=0)
171

    
172
        #self.g.set_trace(1)
173
        #self.g.set_verbose(1)
174

    
175
        self.guestfs_enabled = False
176

    
177
    def enable(self):
178
        """Enable a newly created DiskDevice"""
179
        self.progressbar = progress_generator("Launching helper VM: ")
180
        self.progressbar.next()
181
        eh = self.g.set_event_callback(self.progress_callback,
182
                                                    guestfs.EVENT_PROGRESS)
183
        self.g.launch()
184
        self.guestfs_enabled = True
185
        self.g.delete_event_callback(eh)
186
        if self.progressbar is not None:
187
            self.progressbar.send(100)
188
            self.progressbar = None
189

    
190
        puts('Inspecting Operating System...', False)
191
        roots = self.g.inspect_os()
192
        if len(roots) == 0:
193
            raise FatalError("No operating system found")
194
        if len(roots) > 1:
195
            raise FatalError("Multiple operating systems found."
196
                            "We only support images with one filesystem.")
197
        self.root = roots[0]
198
        self.ostype = self.g.inspect_get_type(self.root)
199
        self.distro = self.g.inspect_get_distro(self.root)
200
        puts(colored.green('found a %s system' % self.distro))
201

    
202
    def destroy(self):
203
        """Destroy this DiskDevice instance."""
204

    
205
        if self.guestfs_enabled:
206
            self.g.umount_all()
207
            self.g.sync()
208

    
209
        # Close the guestfs handler if open
210
        self.g.close()
211

    
212
    def progress_callback(self, ev, eh, buf, array):
213
        position = array[2]
214
        total = array[3]
215

    
216
        self.progressbar.send((position * 100) // total)
217

    
218
        if position == total:
219
            self.progressbar = None
220

    
221
    def mount(self):
222
        """Mount all disk partitions in a correct order."""
223
        mps = self.g.inspect_get_mountpoints(self.root)
224

    
225
        # Sort the keys to mount the fs in a correct order.
226
        # / should be mounted befor /boot, etc
227
        def compare(a, b):
228
            if len(a[0]) > len(b[0]):
229
                return 1
230
            elif len(a[0]) == len(b[0]):
231
                return 0
232
            else:
233
                return -1
234
        mps.sort(compare)
235
        for mp, dev in mps:
236
            try:
237
                self.g.mount(dev, mp)
238
            except RuntimeError as msg:
239
                print "%s (ignored)" % msg
240

    
241
    def umount(self):
242
        """Umount all mounted filesystems."""
243
        self.g.umount_all()
244

    
245
    def shrink(self):
246
        """Shrink the disk.
247

248
        This is accomplished by shrinking the last filesystem in the
249
        disk and then updating the partition table. The new disk size
250
        (in bytes) is returned.
251
        """
252
        puts("Shrinking image (this may take a while)...", False)
253

    
254
        dev = self.g.part_to_dev(self.root)
255
        parttype = self.g.part_get_parttype(dev)
256
        if parttype != 'msdos':
257
            raise FatalError("You have a %s partition table. "
258
                "Only msdos partitions are supported" % parttype)
259

    
260
        last_partition = self.g.part_list(dev)[-1]
261

    
262
        if last_partition['part_num'] > 4:
263
            raise FatalError("This disk contains logical partitions. "
264
                "Only primary partitions are supported.")
265

    
266
        part_dev = "%s%d" % (dev, last_partition['part_num'])
267
        fs_type = self.g.vfs_type(part_dev)
268
        if not re.match("ext[234]", fs_type):
269
            warn("Don't know how to resize %s partitions." % vfs_type)
270
            return
271

    
272
        self.g.e2fsck_f(part_dev)
273
        self.g.resize2fs_M(part_dev)
274

    
275
        output = self.g.tune2fs_l(part_dev)
276
        block_size = int(
277
            filter(lambda x: x[0] == 'Block size', output)[0][1])
278
        block_cnt = int(
279
            filter(lambda x: x[0] == 'Block count', output)[0][1])
280

    
281
        sector_size = self.g.blockdev_getss(dev)
282

    
283
        start = last_partition['part_start'] / sector_size
284
        end = start + (block_size * block_cnt) / sector_size - 1
285

    
286
        self.g.part_del(dev, last_partition['part_num'])
287
        self.g.part_add(dev, 'p', start, end)
288

    
289
        new_size = (end + 1) * sector_size
290
        puts(colored.green("new image size is %dMB\n" % (new_size // 2 ** 20)))
291

    
292
        return new_size
293

    
294
    def size(self):
295
        """Returns the "payload" size of the device.
296

297
        The size returned by this method is the size of the space occupied by
298
        the partitions (including the space before the first partition).
299
        """
300
        dev = self.g.part_to_dev(self.root)
301
        last = self.g.part_list(dev)[-1]
302

    
303
        return last['part_end'] + 1
304

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