Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 22a6d232

History | View | Annotate | Download (10.7 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)
104
        with indent(4):
105
            sourcedev = self.source
106
            mode = os.stat(self.source).st_mode
107
            if stat.S_ISDIR(mode):
108
                puts(colored.green('Looks like a directory'))
109
                return self._losetup(self._dir_to_disk())
110
            elif stat.S_ISREG(mode):
111
                puts(colored.green('Looks like an image file'))
112
                sourcedev = self._losetup(self.source)
113
            elif not stat.S_ISBLK(mode):
114
                raise ValueError("Invalid media source. Only block devices, "
115
                                "regular files and directories are supported.")
116
            else:
117
                puts(colored.green('Looks like a block device'))
118
            #puts()
119

    
120
        # Take a snapshot and return it to the user
121
        puts("Snapshotting media source")
122
        with indent(4):
123
            size = blockdev('--getsize', sourcedev)
124
            cowfd, cow = tempfile.mkstemp()
125
            os.close(cowfd)
126
            self._add_cleanup(os.unlink, cow)
127
            # Create 1G cow sparse file
128
            dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
129
                                            'seek=%d' % (1024 * 1024))
130
            cowdev = self._losetup(cow)
131
    
132
            snapshot = uuid.uuid4().hex
133
            tablefd, table = tempfile.mkstemp()
134
            try:
135
                os.write(tablefd, "0 %d snapshot %s %s n 8" % \
136
                                            (int(size), sourcedev, cowdev))
137
                dmsetup('create', snapshot, table)
138
                self._add_cleanup(dmsetup, 'remove', snapshot)
139
                # Sometimes dmsetup remove fails with Device or resource busy,
140
                # although everything is cleaned up and the snapshot is not
141
                # used by anyone. Add a 2 seconds delay to be on the safe side.
142
                self._add_cleanup(time.sleep, 2)
143

    
144
            finally:
145
                os.unlink(table)
146
            puts(colored.green('Done'))
147
        # puts()
148
        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
149
        self._devices.append(new_device)
150
        new_device.enable()
151
        return new_device
152

    
153
    def destroy_device(self, device):
154
        """Destroys a DiskDevice instance previously created by
155
        get_device method.
156
        """
157
        self._devices.remove(device)
158
        device.destroy()
159

    
160

    
161
class DiskDevice(object):
162
    """This class represents a block device hosting an Operating System
163
    as created by the device-mapper.
164
    """
165

    
166
    def __init__(self, device, bootable=True):
167
        """Create a new DiskDevice."""
168

    
169
        self.device = device
170
        self.bootable = bootable
171
        self.progress_bar = None
172

    
173
        self.g = guestfs.GuestFS()
174
        self.g.add_drive_opts(self.device, readonly=0)
175

    
176
        #self.g.set_trace(1)
177
        #self.g.set_verbose(1)
178

    
179
        self.guestfs_enabled = False
180

    
181
    def enable(self):
182
        """Enable a newly created DiskDevice"""
183
        self.progressbar = progress_generator("Launching helper VM: ")
184
        with indent(4):
185
            self.progressbar.next()
186
            eh = self.g.set_event_callback(self.progress_callback,
187
                                                        guestfs.EVENT_PROGRESS)
188
            self.g.launch()
189
            self.guestfs_enabled = True
190
            self.g.delete_event_callback(eh)
191
            if self.progressbar is not None:
192
                self.progressbar.send(100)
193
                self.progressbar = None
194
            puts(colored.green('Done'))
195

    
196
        puts('Inspecting Operating System')
197
        with indent(4):
198
            roots = self.g.inspect_os()
199
            if len(roots) == 0:
200
                raise FatalError("No operating system found")
201
            if len(roots) > 1:
202
                raise FatalError("Multiple operating systems found."
203
                                "We only support images with one filesystem.")
204
            self.root = roots[0]
205
            self.ostype = self.g.inspect_get_type(self.root)
206
            self.distro = self.g.inspect_get_distro(self.root)
207
            puts(colored.green('Found a %s system' % self.distro))
208
        puts()
209

    
210
    def destroy(self):
211
        """Destroy this DiskDevice instance."""
212

    
213
        if self.guestfs_enabled:
214
            self.g.umount_all()
215
            self.g.sync()
216

    
217
        # Close the guestfs handler if open
218
        self.g.close()
219

    
220
    def progress_callback(self, ev, eh, buf, array):
221
        position = array[2]
222
        total = array[3]
223

    
224
        self.progressbar.send((position * 100) // total)
225

    
226
        if position == total:
227
            self.progressbar = None
228

    
229
    def mount(self):
230
        """Mount all disk partitions in a correct order."""
231
        mps = self.g.inspect_get_mountpoints(self.root)
232

    
233
        # Sort the keys to mount the fs in a correct order.
234
        # / should be mounted befor /boot, etc
235
        def compare(a, b):
236
            if len(a[0]) > len(b[0]):
237
                return 1
238
            elif len(a[0]) == len(b[0]):
239
                return 0
240
            else:
241
                return -1
242
        mps.sort(compare)
243
        for mp, dev in mps:
244
            try:
245
                self.g.mount(dev, mp)
246
            except RuntimeError as msg:
247
                print "%s (ignored)" % msg
248

    
249
    def umount(self):
250
        """Umount all mounted filesystems."""
251
        self.g.umount_all()
252

    
253
    def shrink(self):
254
        """Shrink the disk.
255

256
        This is accomplished by shrinking the last filesystem in the
257
        disk and then updating the partition table. The new disk size
258
        (in bytes) is returned.
259
        """
260
        puts("Shrinking image (this may take a while)")
261
        
262
        dev = self.g.part_to_dev(self.root)
263
        parttype = self.g.part_get_parttype(dev)
264
        if parttype != 'msdos':
265
            raise FatalError("You have a %s partition table. "
266
                "Only msdos partitions are supported" % parttype)
267

    
268
        last_partition = self.g.part_list(dev)[-1]
269

    
270
        if last_partition['part_num'] > 4:
271
            raise FatalError("This disk contains logical partitions. "
272
                "Only primary partitions are supported.")
273

    
274
        part_dev = "%s%d" % (dev, last_partition['part_num'])
275
        fs_type = self.g.vfs_type(part_dev)
276
        if not re.match("ext[234]", fs_type):
277
            warn("Don't know how to resize %s partitions." % vfs_type)
278
            return
279

    
280
        with indent(4):
281
            self.g.e2fsck_f(part_dev)
282
            self.g.resize2fs_M(part_dev)
283

    
284
            output = self.g.tune2fs_l(part_dev)
285
            block_size = int(
286
                filter(lambda x: x[0] == 'Block size', output)[0][1])
287
            block_cnt = int(
288
                filter(lambda x: x[0] == 'Block count', output)[0][1])
289

    
290
            sector_size = self.g.blockdev_getss(dev)
291

    
292
            start = last_partition['part_start'] / sector_size
293
            end = start + (block_size * block_cnt) / sector_size - 1
294

    
295
            self.g.part_del(dev, last_partition['part_num'])
296
            self.g.part_add(dev, 'p', start, end)
297

    
298
            new_size = (end + 1) * sector_size
299
            puts("  New image size is %dMB\n" % (new_size // 2 ** 20))
300
        return new_size
301

    
302
    def size(self):
303
        """Returns the "payload" size of the device.
304

305
        The size returned by this method is the size of the space occupied by
306
        the partitions (including the space before the first partition).
307
        """
308
        dev = self.g.part_to_dev(self.root)
309
        last = self.g.part_list(dev)[-1]
310

    
311
        return last['part_end']
312

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