Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ 61d14323

History | View | Annotate | Download (18.2 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.util import try_fail_repeat
37
from image_creator.util import free_space
38
from image_creator.gpt import GPTPartitionTable
39
from image_creator.bundle_volume import BundleVolume
40

    
41
import stat
42
import os
43
import tempfile
44
import uuid
45
import re
46
import guestfs
47
import shutil
48
from sendfile import sendfile
49

    
50

    
51
dd = get_command('dd')
52
dmsetup = get_command('dmsetup')
53
losetup = get_command('losetup')
54
blockdev = get_command('blockdev')
55

    
56

    
57
TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']
58

    
59

    
60
class Disk(object):
61
    """This class represents a hard disk hosting an Operating System
62

63
    A Disk instance never alters the source media it is created from.
64
    Any change is done on a snapshot created by the device-mapper of
65
    the Linux kernel.
66
    """
67

    
68
    def __init__(self, source, output, tmp=None):
69
        """Create a new Disk instance out of a source media. The source
70
        media can be an image file, a block device or a directory.
71
        """
72
        self._cleanup_jobs = []
73
        self._devices = []
74
        self.source = source
75
        self.out = output
76
        self.meta = {}
77
        self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.',
78
                                    dir=self._get_tmp_dir(tmp))
79

    
80
        self._add_cleanup(shutil.rmtree, self.tmp)
81

    
82
    def _get_tmp_dir(self, default=None):
83
        if default is not None:
84
            return default
85

    
86
        space = map(free_space, TMP_CANDIDATES)
87

    
88
        max_idx = 0
89
        max_val = space[0]
90
        for i, val in zip(range(len(space)), space):
91
            if val > max_val:
92
                max_val = val
93
                max_idx = i
94

    
95
        # Return the candidate path with more available space
96
        return TMP_CANDIDATES[max_idx]
97

    
98
    def _add_cleanup(self, job, *args):
99
        self._cleanup_jobs.append((job, args))
100

    
101
    def _losetup(self, fname):
102
        loop = losetup('-f', '--show', fname)
103
        loop = loop.strip()  # remove the new-line char
104
        self._add_cleanup(try_fail_repeat, losetup, '-d', loop)
105
        return loop
106

    
107
    def _dir_to_disk(self):
108
        if self.source == '/':
109
            bundle = BundleVolume(self.out, self.meta)
110
            image = '%s/%s.diskdump' % (self.tmp, uuid.uuid4().hex)
111

    
112
            def check_unlink(path):
113
                if os.path.exists(path):
114
                    os.unlink(path)
115

    
116
            self._add_cleanup(check_unlink, image)
117
            bundle.create_image(image)
118
            return self._losetup(image)
119
        raise FatalError("Using a directory as media source is supported")
120

    
121
    def cleanup(self):
122
        """Cleanup internal data. This needs to be called before the
123
        program ends.
124
        """
125
        try:
126
            while len(self._devices):
127
                device = self._devices.pop()
128
                device.destroy()
129
        finally:
130
            # Make sure those are executed even if one of the device.destroy
131
            # methods throws exeptions.
132
            while len(self._cleanup_jobs):
133
                job, args = self._cleanup_jobs.pop()
134
                job(*args)
135

    
136
    def snapshot(self):
137
        """Creates a snapshot of the original source media of the Disk
138
        instance.
139
        """
140

    
141
        self.out.output("Examining source media `%s' ..." % self.source, False)
142
        sourcedev = self.source
143
        mode = os.stat(self.source).st_mode
144
        if stat.S_ISDIR(mode):
145
            self.out.success('looks like a directory')
146
            return self._dir_to_disk()
147
        elif stat.S_ISREG(mode):
148
            self.out.success('looks like an image file')
149
            sourcedev = self._losetup(self.source)
150
        elif not stat.S_ISBLK(mode):
151
            raise ValueError("Invalid media source. Only block devices, "
152
                             "regular files and directories are supported.")
153
        else:
154
            self.out.success('looks like a block device')
155

    
156
        # Take a snapshot and return it to the user
157
        self.out.output("Snapshotting media source...", False)
158
        size = blockdev('--getsz', sourcedev)
159
        cowfd, cow = tempfile.mkstemp(dir=self.tmp)
160
        os.close(cowfd)
161
        self._add_cleanup(os.unlink, cow)
162
        # Create cow sparse file
163
        dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
164
        cowdev = self._losetup(cow)
165

    
166
        snapshot = uuid.uuid4().hex
167
        tablefd, table = tempfile.mkstemp()
168
        try:
169
            os.write(tablefd, "0 %d snapshot %s %s n 8" %
170
                              (int(size), sourcedev, cowdev))
171
            dmsetup('create', snapshot, table)
172
            self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
173

    
174
        finally:
175
            os.unlink(table)
176
        self.out.success('done')
177
        return "/dev/mapper/%s" % snapshot
178

    
179
    def get_device(self, media):
180
        """Returns a newly created DiskDevice instance."""
181

    
182
        new_device = DiskDevice(media, self.out)
183
        self._devices.append(new_device)
184
        new_device.enable()
185
        return new_device
186

    
187
    def destroy_device(self, device):
188
        """Destroys a DiskDevice instance previously created by
189
        get_device method.
190
        """
191
        self._devices.remove(device)
192
        device.destroy()
193

    
194

    
195
class DiskDevice(object):
196
    """This class represents a block device hosting an Operating System
197
    as created by the device-mapper.
198
    """
199

    
200
    def __init__(self, device, output, bootable=True, meta={}):
201
        """Create a new DiskDevice."""
202

    
203
        self.real_device = device
204
        self.out = output
205
        self.bootable = bootable
206
        self.meta = meta
207
        self.progress_bar = None
208
        self.guestfs_device = None
209
        self.size = 0
210

    
211
        self.g = guestfs.GuestFS()
212
        self.g.add_drive_opts(self.real_device, readonly=0, format="raw")
213

    
214
        # Before version 1.17.14 the recovery process, which is a fork of the
215
        # original process that called libguestfs, did not close its inherited
216
        # file descriptors. This can cause problems especially if the parent
217
        # process has opened pipes. Since the recovery process is an optional
218
        # feature of libguestfs, it's better to disable it.
219
        self.g.set_recovery_proc(0)
220
        version = self.g.version()
221
        if version['major'] > 1 or \
222
            (version['major'] == 1 and (version['minor'] >= 18 or
223
                                        (version['minor'] == 17 and
224
                                         version['release'] >= 14))):
225
            self.g.set_recovery_proc(1)
226
            self.out.output("Enabling recovery proc")
227

    
228
        #self.g.set_trace(1)
229
        #self.g.set_verbose(1)
230

    
231
        self.guestfs_enabled = False
232

    
233
    def enable(self):
234
        """Enable a newly created DiskDevice"""
235

    
236
        self.out.output('Launching helper VM (may take a while) ...', False)
237
        # self.progressbar = self.out.Progress(100, "Launching helper VM",
238
        #                                     "percent")
239
        # eh = self.g.set_event_callback(self.progress_callback,
240
        #                               guestfs.EVENT_PROGRESS)
241
        self.g.launch()
242
        self.guestfs_enabled = True
243
        # self.g.delete_event_callback(eh)
244
        # self.progressbar.success('done')
245
        # self.progressbar = None
246
        self.out.success('done')
247

    
248
        self.out.output('Inspecting Operating System ...', False)
249
        roots = self.g.inspect_os()
250
        if len(roots) == 0:
251
            raise FatalError("No operating system found")
252
        if len(roots) > 1:
253
            raise FatalError("Multiple operating systems found."
254
                             "We only support images with one OS.")
255
        self.root = roots[0]
256
        self.guestfs_device = self.g.part_to_dev(self.root)
257
        self.size = self.g.blockdev_getsize64(self.guestfs_device)
258
        self.meta['PARTITION_TABLE'] = \
259
            self.g.part_get_parttype(self.guestfs_device)
260

    
261
        self.ostype = self.g.inspect_get_type(self.root)
262
        self.distro = self.g.inspect_get_distro(self.root)
263
        self.out.success('found a(n) %s system' % self.distro)
264

    
265
    def destroy(self):
266
        """Destroy this DiskDevice instance."""
267

    
268
        # In new guestfs versions, there is a handy shutdown method for this
269
        try:
270
            if self.guestfs_enabled:
271
                self.g.umount_all()
272
                self.g.sync()
273
        finally:
274
            # Close the guestfs handler if open
275
            self.g.close()
276

    
277
#    def progress_callback(self, ev, eh, buf, array):
278
#        position = array[2]
279
#        total = array[3]
280
#
281
#        self.progressbar.goto((position * 100) // total)
282

    
283
    def mount(self, readonly=False):
284
        """Mount all disk partitions in a correct order."""
285

    
286
        mount = self.g.mount_ro if readonly else self.g.mount
287
        msg = " read-only" if readonly else ""
288
        self.out.output("Mounting the media%s ..." % msg, False)
289
        mps = self.g.inspect_get_mountpoints(self.root)
290

    
291
        # Sort the keys to mount the fs in a correct order.
292
        # / should be mounted befor /boot, etc
293
        def compare(a, b):
294
            if len(a[0]) > len(b[0]):
295
                return 1
296
            elif len(a[0]) == len(b[0]):
297
                return 0
298
            else:
299
                return -1
300
        mps.sort(compare)
301
        for mp, dev in mps:
302
            try:
303
                mount(dev, mp)
304
            except RuntimeError as msg:
305
                self.out.warn("%s (ignored)" % msg)
306
        self.out.success("done")
307

    
308
    def umount(self):
309
        """Umount all mounted filesystems."""
310
        self.g.umount_all()
311

    
312
    def _last_partition(self):
313
        if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
314
            msg = "Unsupported partition table: %s. Only msdos and gpt " \
315
                "partition tables are supported" % self.meta['PARTITION_TABLE']
316
            raise FatalError(msg)
317

    
318
        is_extended = lambda p: \
319
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
320
            in (0x5, 0xf)
321
        is_logical = lambda p: \
322
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
323

    
324
        partitions = self.g.part_list(self.guestfs_device)
325
        last_partition = partitions[-1]
326

    
327
        if is_logical(last_partition):
328
            # The disk contains extended and logical partitions....
329
            extended = filter(is_extended, partitions)[0]
330
            last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
331

    
332
            # check if extended is the last primary partition
333
            if last_primary['part_num'] > extended['part_num']:
334
                last_partition = last_primary
335

    
336
        return last_partition
337

    
338
    def shrink(self):
339
        """Shrink the disk.
340

341
        This is accomplished by shrinking the last filesystem in the
342
        disk and then updating the partition table. The new disk size
343
        (in bytes) is returned.
344

345
        ATTENTION: make sure unmount is called before shrink
346
        """
347
        get_fstype = lambda p: \
348
            self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
349
        is_logical = lambda p: \
350
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
351
        is_extended = lambda p: \
352
            self.meta['PARTITION_TABLE'] == 'msdos' and \
353
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
354
            in (0x5, 0xf)
355

    
356
        part_add = lambda ptype, start, stop: \
357
            self.g.part_add(self.guestfs_device, ptype, start, stop)
358
        part_del = lambda p: self.g.part_del(self.guestfs_device, p)
359
        part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
360
        part_set_id = lambda p, id: \
361
            self.g.part_set_mbr_id(self.guestfs_device, p, id)
362
        part_get_bootable = lambda p: \
363
            self.g.part_get_bootable(self.guestfs_device, p)
364
        part_set_bootable = lambda p, bootable: \
365
            self.g.part_set_bootable(self.guestfs_device, p, bootable)
366

    
367
        MB = 2 ** 20
368

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

    
371
        sector_size = self.g.blockdev_getss(self.guestfs_device)
372

    
373
        last_part = None
374
        fstype = None
375
        while True:
376
            last_part = self._last_partition()
377
            fstype = get_fstype(last_part)
378

    
379
            if fstype == 'swap':
380
                self.meta['SWAP'] = "%d:%s" % \
381
                    (last_part['part_num'],
382
                     (last_part['part_size'] + MB - 1) // MB)
383
                part_del(last_part['part_num'])
384
                continue
385
            elif is_extended(last_part):
386
                part_del(last_part['part_num'])
387
                continue
388

    
389
            # Most disk manipulation programs leave 2048 sectors after the last
390
            # partition
391
            new_size = last_part['part_end'] + 1 + 2048 * sector_size
392
            self.size = min(self.size, new_size)
393
            break
394

    
395
        if not re.match("ext[234]", fstype):
396
            self.out.warn("Don't know how to resize %s partitions." % fstype)
397
            return self.size
398

    
399
        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
400
        self.g.e2fsck_f(part_dev)
401
        self.g.resize2fs_M(part_dev)
402

    
403
        out = self.g.tune2fs_l(part_dev)
404
        block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
405
        block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
406

    
407
        start = last_part['part_start'] / sector_size
408
        end = start + (block_size * block_cnt) / sector_size - 1
409

    
410
        if is_logical(last_part):
411
            partitions = self.g.part_list(self.guestfs_device)
412

    
413
            logical = []  # logical partitions
414
            for partition in partitions:
415
                if partition['part_num'] < 4:
416
                    continue
417
                logical.append({
418
                    'num': partition['part_num'],
419
                    'start': partition['part_start'] / sector_size,
420
                    'end': partition['part_end'] / sector_size,
421
                    'id': part_get_id(partition['part_num']),
422
                    'bootable': part_get_bootable(partition['part_num'])
423
                })
424

    
425
            logical[-1]['end'] = end  # new end after resize
426

    
427
            # Recreate the extended partition
428
            extended = filter(is_extended, partitions)[0]
429
            part_del(extended['part_num'])
430
            part_add('e', extended['part_start'] / sector_size, end)
431

    
432
            # Create all the logical partitions back
433
            for l in logical:
434
                part_add('l', l['start'], l['end'])
435
                part_set_id(l['num'], l['id'])
436
                part_set_bootable(l['num'], l['bootable'])
437
        else:
438
            # Recreate the last partition
439
            if self.meta['PARTITION_TABLE'] == 'msdos':
440
                last_part['id'] = part_get_id(last_part['part_num'])
441

    
442
            last_part['bootable'] = part_get_bootable(last_part['part_num'])
443
            part_del(last_part['part_num'])
444
            part_add('p', start, end)
445
            part_set_bootable(last_part['part_num'], last_part['bootable'])
446

    
447
            if self.meta['PARTITION_TABLE'] == 'msdos':
448
                part_set_id(last_part['part_num'], last_part['id'])
449

    
450
        new_size = (end + 1) * sector_size
451

    
452
        assert (new_size <= self.size)
453

    
454
        if self.meta['PARTITION_TABLE'] == 'gpt':
455
            ptable = GPTPartitionTable(self.real_device)
456
            self.size = ptable.shrink(new_size, self.size)
457
        else:
458
            self.size = min(new_size + 2048 * sector_size, self.size)
459

    
460
        self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
461

    
462
        return self.size
463

    
464
    def dump(self, outfile):
465
        """Dumps the content of device into a file.
466

467
        This method will only dump the actual payload, found by reading the
468
        partition table. Empty space in the end of the device will be ignored.
469
        """
470
        MB = 2 ** 20
471
        blocksize = 4 * MB  # 4MB
472
        size = self.size
473
        progr_size = (size + MB - 1) // MB  # in MB
474
        progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
475

    
476
        with open(self.real_device, 'r') as src:
477
            with open(outfile, "w") as dst:
478
                left = size
479
                offset = 0
480
                progressbar.next()
481
                while left > 0:
482
                    length = min(left, blocksize)
483
                    sent = sendfile(dst.fileno(), src.fileno(), offset, length)
484

    
485
                    # Workaround for python-sendfile API change. In
486
                    # python-sendfile 1.2.x (py-sendfile) the returning value
487
                    # of sendfile is a tuple, where in version 2.x (pysendfile)
488
                    # it is just a sigle integer.
489
                    if isinstance(sent, tuple):
490
                        sent = sent[1]
491

    
492
                    offset += sent
493
                    left -= sent
494
                    progressbar.goto((size - left) // MB)
495
        progressbar.success('image file %s was successfully created' % outfile)
496

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