Statistics
| Branch: | Tag: | Revision:

root / image_creator / disk.py @ dbf466eb

History | View | Annotate | Download (16.4 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.size = 0
172
        self.meta = {}
173

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

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

    
191
        #self.g.set_trace(1)
192
        #self.g.set_verbose(1)
193

    
194
        self.guestfs_enabled = False
195

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

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

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

    
225
    def destroy(self):
226
        """Destroy this DiskDevice instance."""
227

    
228
        if self.guestfs_enabled:
229
            self.g.umount_all()
230
            self.g.sync()
231

    
232
        # Close the guestfs handler if open
233
        self.g.close()
234

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

    
239
        self.progressbar.goto((position * 100) // total)
240

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

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

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

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

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

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

    
280
        partitions = self.g.part_list(self.guestfs_device)
281
        last_partition = partitions[-1]
282

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

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

    
292
        return last_partition
293

    
294
    def shrink(self):
295
        """Shrink the disk.
296

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

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

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

    
322
        MB = 2 ** 20
323

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

    
326
        sector_size = self.g.blockdev_getss(self.guestfs_device)
327

    
328
        last_part = None
329
        fstype = None
330
        while True:
331
            last_part = self._last_partition()
332
            fstype = get_fstype(last_part)
333

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

    
344
            # Most disk manipulation programs leave 2048 sectors after the last
345
            # partition
346
            new_size = last_part['part_end'] + 1 + 2048 * sector_size
347
            self.size = min(self.size, new_size)
348
            break
349

    
350
        if not re.match("ext[234]", fstype):
351
            self.out.warn("Don't know how to resize %s partitions." % fstype)
352
            return self.size
353

    
354
        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
355
        self.g.e2fsck_f(part_dev)
356
        self.g.resize2fs_M(part_dev)
357

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

    
364
        start = last_part['part_start'] / sector_size
365
        end = start + (block_size * block_cnt) / sector_size - 1
366

    
367
        if is_logical(last_part):
368
            partitions = self.g.part_list(self.guestfs_device)
369

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

    
382
            logical[-1]['end'] = end  # new end after resize
383

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

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

    
399
            last_part['bootable'] = part_get_bootable(last_part['part_num'])
400
            part_del(last_part['part_num'])
401
            part_add('p', start, end)
402
            part_set_bootable(last_part['part_num'], last_part['bootable'])
403

    
404
            if self.meta['PARTITION_TABLE'] == 'msdos':
405
                part_set_id(last_part['part_num'], last_part['id'])
406

    
407
        new_size = (end + 1) * sector_size
408

    
409
        assert (new_size <= self.size)
410

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

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

    
419
        return self.size
420

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

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

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

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