Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 331aa0ec

History | View | Annotate | Download (12.1 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
from image_creator.gpt import GPTPartitionTable
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.gdev = self.g.part_to_dev(self.root)
204
        self.parttype = self.g.part_get_parttype(self.gdev)
205

    
206
        self.ostype = self.g.inspect_get_type(self.root)
207
        self.distro = self.g.inspect_get_distro(self.root)
208
        success('found a(n) %s system' % self.distro)
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.goto((position * 100) // total)
225

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

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

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

    
262
        if self.parttype not in 'msdos' 'gpt':
263
            raise FatalError("You have a %s partition table. "
264
                "Only msdos and gpt partitions are supported" % self.parttype)
265

    
266
        last_partition = self.g.part_list(self.gdev)[-1]
267

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

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

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

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

    
287
        sector_size = self.g.blockdev_getss(self.gdev)
288

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

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

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

    
299
        if self.parttype == 'gpt':
300
            ptable = GPTPartitionTable(self.device)
301
            return ptable.shrink(new_size)
302

    
303
        return new_size
304

    
305
    def size(self):
306
        """Returns the "payload" size of the device.
307

308
        The size returned by this method is the size of the space occupied by
309
        the partitions (including the space before the first partition).
310
        """
311

    
312
        if self.parttype == 'msdos':
313
            dev = self.g.part_to_dev(self.root)
314
            last = self.g.part_list(dev)[-1]
315
            return last['part_end'] + 1
316
        elif self.parttype == 'gpt':
317
            ptable = GPTPartitionTable(self.device)
318
            return ptable.size()
319
        else:
320
            raise FatalError("Unsupported partition table type: %s" % parttype)
321

    
322
    def dump(self, outfile):
323
        """Dumps the content of device into a file.
324

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

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