Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 9666a511

History | View | Annotate | Download (16.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
35
from image_creator.util import 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
dd = get_command('dd')
49
dmsetup = get_command('dmsetup')
50
losetup = get_command('losetup')
51
blockdev = get_command('blockdev')
52

    
53

    
54
class Disk(object):
55
    """This class represents a hard disk hosting an Operating System
56

57
    A Disk instance never alters the source media it is created from.
58
    Any change is done on a snapshot created by the device-mapper of
59
    the Linux kernel.
60
    """
61

    
62
    def __init__(self, source, output):
63
        """Create a new Disk instance out of a source media. The source
64
        media can be an image file, a block device or a directory."""
65
        self._cleanup_jobs = []
66
        self._devices = []
67
        self.source = source
68
        self.out = output
69

    
70
    def _add_cleanup(self, job, *args):
71
        self._cleanup_jobs.append((job, args))
72

    
73
    def _losetup(self, fname):
74
        loop = losetup('-f', '--show', fname)
75
        loop = loop.strip()  # remove the new-line char
76
        self._add_cleanup(losetup, '-d', loop)
77
        return loop
78

    
79
    def _dir_to_disk(self):
80
        raise FatalError("Using a directory as media source is not supported "
81
                         "yet!")
82

    
83
    def cleanup(self):
84
        """Cleanup internal data. This needs to be called before the
85
        program ends.
86
        """
87
        while len(self._devices):
88
            device = self._devices.pop()
89
            device.destroy()
90

    
91
        while len(self._cleanup_jobs):
92
            job, args = self._cleanup_jobs.pop()
93
            job(*args)
94

    
95
    def snapshot(self):
96
        """Creates a snapshot of the original source media of the Disk
97
        instance.
98
        """
99

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

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

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

    
137
        finally:
138
            os.unlink(table)
139
        self.out.success('done')
140
        return "/dev/mapper/%s" % snapshot
141

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

    
145
        new_device = DiskDevice(media, self.out)
146
        self._devices.append(new_device)
147
        new_device.enable()
148
        return new_device
149

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

    
157

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

    
163
    def __init__(self, device, output, bootable=True):
164
        """Create a new DiskDevice."""
165

    
166
        self.real_device = device
167
        self.out = output
168
        self.bootable = bootable
169
        self.progress_bar = None
170
        self.guestfs_device = None
171
        self.meta = {}
172

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

    
176
        # Before version 1.17.14 the recovery process, which is a fork of the
177
        # original process that called libguestfs, did not close its inherited
178
        # file descriptors. This can cause problems especially if the parent
179
        # process has opened pipes. Since the recovery process is an optional
180
        # feature of libguestfs, it's better to disable it.
181
        self.g.set_recovery_proc(0)
182
        version = self.g.version()
183
        if version['major'] > 1 or (version['major'] == 1 and
184
            (version['minor'] >= 18 or \
185
            (version['minor'] == 17 and version['release'] >= 14))):
186
            self.g.set_recovery_proc(1)
187
            self.out.output("Enabling recovery proc")
188

    
189
        #self.g.set_trace(1)
190
        #self.g.set_verbose(1)
191

    
192
        self.guestfs_enabled = False
193

    
194
    def enable(self):
195
        """Enable a newly created DiskDevice"""
196
        self.progressbar = self.out.Progress(100, "Launching helper VM",
197
                                             "percent")
198
        eh = self.g.set_event_callback(self.progress_callback,
199
                                                    guestfs.EVENT_PROGRESS)
200
        self.g.launch()
201
        self.guestfs_enabled = True
202
        self.g.delete_event_callback(eh)
203
        self.progressbar.success('done')
204
        self.progressbar = None
205

    
206
        self.out.output('Inspecting Operating System...', False)
207
        roots = self.g.inspect_os()
208
        if len(roots) == 0:
209
            raise FatalError("No operating system found")
210
        if len(roots) > 1:
211
            raise FatalError("Multiple operating systems found."
212
                            "We only support images with one filesystem.")
213
        self.root = roots[0]
214
        self.guestfs_device = self.g.part_to_dev(self.root)
215
        self.meta['SIZE'] = self.g.blockdev_getsize64(self.guestfs_device)
216
        self.meta['PARTITION_TABLE'] = \
217
                                self.g.part_get_parttype(self.guestfs_device)
218

    
219
        self.ostype = self.g.inspect_get_type(self.root)
220
        self.distro = self.g.inspect_get_distro(self.root)
221
        self.out.success('found a(n) %s system' % self.distro)
222

    
223
    def destroy(self):
224
        """Destroy this DiskDevice instance."""
225

    
226
        if self.guestfs_enabled:
227
            self.g.umount_all()
228
            self.g.sync()
229

    
230
        # Close the guestfs handler if open
231
        self.g.close()
232

    
233
    def progress_callback(self, ev, eh, buf, array):
234
        position = array[2]
235
        total = array[3]
236

    
237
        self.progressbar.goto((position * 100) // total)
238

    
239
    def mount(self, readonly=False):
240
        """Mount all disk partitions in a correct order."""
241

    
242
        mount = self.g.mount_ro if readonly else self.g.mount
243
        self.out.output("Mounting image...", False)
244
        mps = self.g.inspect_get_mountpoints(self.root)
245

    
246
        # Sort the keys to mount the fs in a correct order.
247
        # / should be mounted befor /boot, etc
248
        def compare(a, b):
249
            if len(a[0]) > len(b[0]):
250
                return 1
251
            elif len(a[0]) == len(b[0]):
252
                return 0
253
            else:
254
                return -1
255
        mps.sort(compare)
256
        for mp, dev in mps:
257
            try:
258
                mount(dev, mp)
259
            except RuntimeError as msg:
260
                self.out.warn("%s (ignored)" % msg)
261
        self.out.success("done")
262

    
263
    def umount(self):
264
        """Umount all mounted filesystems."""
265
        self.g.umount_all()
266

    
267
    def _last_partition(self):
268
        if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
269
            msg = "Unsupported partition table: %s. Only msdos and gpt " \
270
            "partition tables are supported" % self.meta['PARTITION_TABLE']
271
            raise FatalError(msg)
272

    
273
        is_extended = lambda p: self.g.part_get_mbr_id(
274
                                    self.guestfs_device, p['part_num']) == 5
275
        is_logical = lambda p: self.meta['PARTITION_TABLE'] != 'msdos' and \
276
                                                            p['part_num'] > 4
277

    
278
        partitions = self.g.part_list(self.guestfs_device)
279
        last_partition = partitions[-1]
280

    
281
        if is_logical(last_partition):
282
            # The disk contains extended and logical partitions....
283
            extended = [p for p in partitions if is_extended(p)][0]
284
            last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
285

    
286
            # check if extended is the last primary partition
287
            if last_primary['part_num'] > extended['part_num']:
288
                last_partition = last_primary
289

    
290
        return last_partition
291

    
292
    def shrink(self):
293
        """Shrink the disk.
294

295
        This is accomplished by shrinking the last filesystem in the
296
        disk and then updating the partition table. The new disk size
297
        (in bytes) is returned.
298

299
        ATTENTION: make sure unmount is called before shrink
300
        """
301
        get_fstype = lambda p: self.g.vfs_type("%s%d" % \
302
                                        (self.guestfs_device, p['part_num']))
303
        is_logical = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
304
                                                            p['part_num'] > 4
305
        is_extended = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
306
                self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
307

    
308
        part_add = lambda ptype, start, stop: \
309
                    self.g.part_add(self.guestfs_device, ptype, start, stop)
310
        part_del = lambda p: self.g.part_del(self.guestfs_device, p)
311
        part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
312
        part_set_id = lambda p, id: self.g.part_set_mbr_id(
313
                                                    self.guestfs_device, p, id)
314
        part_get_bootable = lambda p: self.g.part_get_bootable(
315
                                                        self.guestfs_device, p)
316
        part_set_bootable = lambda p, bootable: self.g.part_set_bootable(
317
                                            self.guestfs_device, p, bootable)
318

    
319
        MB = 2 ** 20
320

    
321
        self.out.output("Shrinking image (this may take a while)...", False)
322

    
323
        sector_size = self.g.blockdev_getss(self.guestfs_device)
324

    
325
        last_part = None
326
        fstype = None
327
        while True:
328
            last_part = self._last_partition()
329
            fstype = get_fstype(last_part)
330

    
331
            if fstype == 'swap':
332
                self.meta['SWAP'] = "%d:%s" % \
333
                        (last_part['part_num'],
334
                        (last_part['part_size'] + MB - 1) // MB)
335
                part_del(last_part['part_num'])
336
                continue
337
            elif is_extended(last_part):
338
                part_del(last_part['part_num'])
339
                continue
340

    
341
            # Most disk manipulation programs leave 2048 sectors after the last
342
            # partition
343
            new_size = last_part['part_end'] + 1 + 2048 * sector_size
344
            self.meta['SIZE'] = min(self.meta['SIZE'], new_size)
345
            break
346

    
347
        if not re.match("ext[234]", fstype):
348
            self.out.warn("Don't know how to resize %s partitions." % fstype)
349
            return self.meta['SIZE']
350

    
351
        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
352
        self.g.e2fsck_f(part_dev)
353
        self.g.resize2fs_M(part_dev)
354

    
355
        out = self.g.tune2fs_l(part_dev)
356
        block_size = int(
357
            filter(lambda x: x[0] == 'Block size', out)[0][1])
358
        block_cnt = int(
359
            filter(lambda x: x[0] == 'Block count', out)[0][1])
360

    
361
        start = last_part['part_start'] / sector_size
362
        end = start + (block_size * block_cnt) / sector_size - 1
363

    
364
        if is_logical(last_part):
365
            partitions = self.g.part_list(self.guestfs_device)
366

    
367
            logical = []  # logical partitions
368
            for partition in partitions:
369
                if partition['part_num'] < 4:
370
                    continue
371
                logical.append({
372
                    'num': partition['part_num'],
373
                    'start': partition['part_start'] / sector_size,
374
                    'end': partition['part_end'] / sector_size,
375
                    'id': part_get_(partition['part_num']),
376
                    'bootable': part_get_bootable(partition['part_num'])
377
                })
378

    
379
            logical[-1]['end'] = end  # new end after resize
380

    
381
            # Recreate the extended partition
382
            extended = [p for p in partitions if self._is_extended(p)][0]
383
            part_del(extended['part_num'])
384
            part_add('e', extended['part_start'], end)
385

    
386
            # Create all the logical partitions back
387
            for l in logical:
388
                part_add('l', l['start'], l['end'])
389
                part_set_id(l['num'], l['id'])
390
                part_set_bootable(l['num'], l['bootable'])
391
        else:
392
            # Recreate the last partition
393
            if self.meta['PARTITION_TABLE'] == 'msdos':
394
                last_part['id'] = part_get_id(last_part['part_num'])
395

    
396
            last_part['bootable'] = part_get_bootable(last_part['part_num'])
397
            part_del(last_part['part_num'])
398
            part_add('p', start, end)
399
            part_set_bootable(last_part['part_num'], last_part['bootable'])
400

    
401
            if self.meta['PARTITION_TABLE'] == 'msdos':
402
                part_set_id(last_part['part_num'], last_part['id'])
403

    
404
        new_size = (end + 1) * sector_size
405

    
406
        assert (new_size <= self.meta['SIZE'])
407

    
408
        if self.meta['PARTITION_TABLE'] == 'gpt':
409
            ptable = GPTPartitionTable(self.real_device)
410
            self.meta['SIZE'] = ptable.shrink(new_size, self.meta['SIZE'])
411
        else:
412
            self.meta['SIZE'] = min(new_size + 2048 * sector_size,
413
                                    self.meta['SIZE'])
414

    
415
        self.out.success("new size is %dMB" %
416
                         ((self.meta['SIZE'] + MB - 1) // MB))
417

    
418
        return self.meta['SIZE']
419

    
420
    def dump(self, outfile):
421
        """Dumps the content of device into a file.
422

423
        This method will only dump the actual payload, found by reading the
424
        partition table. Empty space in the end of the device will be ignored.
425
        """
426
        MB = 2 ** 20
427
        blocksize = 4 * MB  # 4MB
428
        size = self.meta['SIZE']
429
        progr_size = (size + MB - 1) // MB  # in MB
430
        progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
431

    
432
        with open(self.real_device, 'r') as src:
433
            with open(outfile, "w") as dst:
434
                left = size
435
                offset = 0
436
                progressbar.next()
437
                while left > 0:
438
                    length = min(left, blocksize)
439
                    sent = sendfile(dst.fileno(), src.fileno(), offset, length)
440
                    offset += sent
441
                    left -= sent
442
                    progressbar.goto((size - left) // MB)
443
        progressbar.success('image file %s was successfully created' % outfile)
444

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