Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ b1395967

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
        self.progressbar = progress("Launching helper VM: ", "percent")
183
        self.progressbar.max = 100
184
        self.progressbar.goto(1)
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
            output("\rLaunching helper VM...\033[K", False)
192
            success("done")
193
            self.progressbar = None
194

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

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

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

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

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

    
221
        self.progressbar.goto((position * 100) // total)
222

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

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

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

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

    
250
    def shrink(self):
251
        """Shrink the disk.
252

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

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

    
265
        last_partition = self.g.part_list(dev)[-1]
266

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

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

    
277
        self.g.e2fsck_f(part_dev)
278
        self.g.resize2fs_M(part_dev)
279

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

    
286
        sector_size = self.g.blockdev_getss(dev)
287

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

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

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

    
299
    def size(self):
300
        """Returns the "payload" size of the device.
301

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

    
308
        return last['part_end'] + 1
309

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

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

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