Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 8e3065a0

History | View | Annotate | Download (11.6 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
35
from image_creator.util import warn, progress, success, output, FatalError
36

    
37
import stat
38
import os
39
import tempfile
40
import uuid
41
import re
42
import sys
43
import guestfs
44
import time
45
from sendfile import sendfile
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 snapshot(self):
97
        """Creates a snapshot of the original source media of the Disk
98
        instance.
99
        """
100

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

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

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

    
139
        finally:
140
            os.unlink(table)
141
        success('done')
142
        return "/dev/mapper/%s" % snapshot
143

    
144
    def get_device(self, media):
145
        """Returns a newly created DiskDevice instance."""
146

    
147
        new_device = DiskDevice(media)
148
        self._devices.append(new_device)
149
        new_device.enable()
150
        return new_device
151

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

    
159

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

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

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

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

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

    
178
        self.guestfs_enabled = False
179

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

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

    
206
    def destroy(self):
207
        """Destroy this DiskDevice instance."""
208

    
209
        if self.guestfs_enabled:
210
            self.g.umount_all()
211
            self.g.sync()
212

    
213
        # Close the guestfs handler if open
214
        self.g.close()
215

    
216
    def progress_callback(self, ev, eh, buf, array):
217
        position = array[2]
218
        total = array[3]
219

    
220
        self.progressbar.send((position * 100) // total)
221

    
222
        if position == total:
223
            self.progressbar = None
224

    
225
    def mount(self):
226
        """Mount all disk partitions in a correct order."""
227

    
228
        output("Mounting image...", False)
229
        mps = self.g.inspect_get_mountpoints(self.root)
230

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

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

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

255
        This is accomplished by shrinking the last filesystem in the
256
        disk and then updating the partition table. The new disk size
257
        (in bytes) is returned.
258
        """
259
        output("Shrinking image (this may take a while)...", False)
260

    
261
        dev = self.g.part_to_dev(self.root)
262
        parttype = self.g.part_get_parttype(dev)
263
        if parttype != 'msdos':
264
            raise FatalError("You have a %s partition table. "
265
                "Only msdos partitions are supported" % parttype)
266

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

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

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

    
279
        self.g.e2fsck_f(part_dev)
280
        self.g.resize2fs_M(part_dev)
281

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

    
288
        sector_size = self.g.blockdev_getss(dev)
289

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

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

    
296
        new_size = (end + 1) * sector_size
297
        success("new image size is %dMB" %
298
                            ((new_size + 2 ** 20 - 1) // 2 ** 20))
299
        return new_size
300

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

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

    
310
        return last['part_end'] + 1
311

    
312
    def dump(self, outfile):
313
        """Dumps the content of device into a file.
314

315
        This method will only dump the actual payload, found by reading the
316
        partition table. Empty space in the end of the device will be ignored.
317
        """
318
        blocksize = 2 ** 22  # 4MB
319
        size = self.size()
320
        progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
321
        new_progress = progress("Dumping image file: ")
322
        progressbar = new_progress(progress_size)
323
        source = open(self.device, "r")
324
        try:
325
            dest = open(outfile, "w")
326
            try:
327
                left = size
328
                offset = 0
329
                progressbar.next()
330
                while left > 0:
331
                    length = min(left, blocksize)
332
                    sent = sendfile(dest.fileno(), source.fileno(), offset,
333
                                                                        length)
334
                    offset += sent
335
                    left -= sent
336
                    for i in range((length + 2 ** 20 - 1) // 2 ** 20):
337
                        progressbar.next()
338
            finally:
339
                dest.close()
340
        finally:
341
            source.close()
342

    
343
        success('Image file %s was successfully created' % outfile)
344

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