Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ f165adc0

History | View | Annotate | Download (11.5 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 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
        output("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
            success('looks like a directory')
108
            return self._losetup(self._dir_to_disk())
109
        elif stat.S_ISREG(mode):
110
            success('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
            success('looks like a block device')
117

    
118
        # Take a snapshot and return it to the user
119
        output("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
        success('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("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
        output('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
        success('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

    
224
        output("Mounting image...", False)
225
        mps = self.g.inspect_get_mountpoints(self.root)
226

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

    
244
    def umount(self):
245
        """Umount all mounted filesystems."""
246
        self.g.umount_all()
247

    
248
    def shrink(self):
249
        """Shrink the disk.
250

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

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

    
263
        last_partition = self.g.part_list(dev)[-1]
264

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

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

    
275
        self.g.e2fsck_f(part_dev)
276
        self.g.resize2fs_M(part_dev)
277

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

    
284
        sector_size = self.g.blockdev_getss(dev)
285

    
286
        start = last_partition['part_start'] / sector_size
287
        end = start + (block_size * block_cnt) / sector_size - 1
288

    
289
        self.g.part_del(dev, last_partition['part_num'])
290
        self.g.part_add(dev, 'p', start, end)
291

    
292
        new_size = (end + 1) * sector_size
293
        success("new image size is %dMB" %
294
                            ((new_size + 2 ** 20 - 1) // 2 ** 20))
295
        return new_size
296

    
297
    def size(self):
298
        """Returns the "payload" size of the device.
299

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

    
306
        return last['part_end'] + 1
307

    
308
    def dump(self, outfile):
309
        """Dumps the content of device into a file.
310

311
        This method will only dump the actual payload, found by reading the
312
        partition table. Empty space in the end of the device will be ignored.
313
        """
314
        blocksize = 2 ** 22  # 4MB
315
        size = self.size()
316
        progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
317
        progressbar = progress("Dumping image file: ", progress_size)
318

    
319
        source = open(self.device, "r")
320
        try:
321
            dest = open(outfile, "w")
322
            try:
323
                left = size
324
                offset = 0
325
                progressbar.next()
326
                while left > 0:
327
                    length = min(left, blocksize)
328
                    sent = sendfile(dest.fileno(), source.fileno(), offset,
329
                                                                        length)
330
                    offset += sent
331
                    left -= sent
332
                    for i in range((length + 2 ** 20 - 1) // 2 ** 20):
333
                        progressbar.next()
334
            finally:
335
                dest.close()
336
        finally:
337
            source.close()
338

    
339
        success('Image file %s was successfully created' % outfile)
340

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