Statistics
| Branch: | Tag: | Revision:

root / image_creator / image.py @ f5174d2c

History | View | Annotate | Download (13.8 kB)

1
# Copyright 2013 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 FatalError
35
from image_creator.gpt import GPTPartitionTable
36
from image_creator.os_type import os_cls
37

    
38
import re
39
import guestfs
40
from sendfile import sendfile
41

    
42

    
43
class Image(object):
44
    """The instances of this class can create images out of block devices."""
45

    
46
    def __init__(self, device, output, bootable=True, meta={}):
47
        """Create a new ImageCreator."""
48

    
49
        self.device = device
50
        self.out = output
51
        self.bootable = bootable
52
        self.meta = meta
53
        self.progress_bar = None
54
        self.guestfs_device = None
55
        self.size = 0
56
        self.mounted = False
57

    
58
        self.g = guestfs.GuestFS()
59
        self.g.add_drive_opts(self.device, readonly=0, format="raw")
60

    
61
        # Before version 1.17.14 the recovery process, which is a fork of the
62
        # original process that called libguestfs, did not close its inherited
63
        # file descriptors. This can cause problems especially if the parent
64
        # process has opened pipes. Since the recovery process is an optional
65
        # feature of libguestfs, it's better to disable it.
66
        self.g.set_recovery_proc(0)
67
        version = self.g.version()
68
        if version['major'] > 1 or \
69
            (version['major'] == 1 and (version['minor'] >= 18 or
70
                                        (version['minor'] == 17 and
71
                                         version['release'] >= 14))):
72
            self.g.set_recovery_proc(1)
73
            self.out.output("Enabling recovery proc")
74

    
75
        #self.g.set_trace(1)
76
        #self.g.set_verbose(1)
77

    
78
        self.guestfs_enabled = False
79

    
80
    def enable(self):
81
        """Enable a newly created ImageCreator"""
82

    
83
        self.out.output('Launching helper VM (may take a while) ...', False)
84
        # self.progressbar = self.out.Progress(100, "Launching helper VM",
85
        #                                     "percent")
86
        # eh = self.g.set_event_callback(self.progress_callback,
87
        #                               guestfs.EVENT_PROGRESS)
88
        self.g.launch()
89
        self.guestfs_enabled = True
90
        # self.g.delete_event_callback(eh)
91
        # self.progressbar.success('done')
92
        # self.progressbar = None
93
        self.out.success('done')
94

    
95
        self.out.output('Inspecting Operating System ...', False)
96
        roots = self.g.inspect_os()
97
        if len(roots) == 0:
98
            raise FatalError("No operating system found")
99
        if len(roots) > 1:
100
            raise FatalError("Multiple operating systems found."
101
                             "We only support images with one OS.")
102
        self.root = roots[0]
103
        self.guestfs_device = self.g.part_to_dev(self.root)
104
        self.size = self.g.blockdev_getsize64(self.guestfs_device)
105
        self.meta['PARTITION_TABLE'] = \
106
            self.g.part_get_parttype(self.guestfs_device)
107

    
108
        self.ostype = self.g.inspect_get_type(self.root)
109
        self.distro = self.g.inspect_get_distro(self.root)
110
        self.out.success('found a(n) %s system' % self.distro)
111

    
112
    def _get_os(self):
113
        if hasattr(self, "_os"):
114
            return self._os
115

    
116
        if not self.guestfs_enabled:
117
            self.enable()
118

    
119
        if not self.mounted:
120
            do_unmount = True
121
            self.mount(readonly=True)
122
        else:
123
            do_unmount = False
124

    
125
        try:
126
            cls = os_cls(self.distro, self.ostype)
127
            self._os = cls(self.root, self.g, self.out)
128

    
129
        finally:
130
            if do_unmount:
131
                self.umount()
132

    
133
        return self._os
134

    
135
    os = property(_get_os)
136

    
137
    def destroy(self):
138
        """Destroy this ImageCreator instance."""
139

    
140
        # In new guestfs versions, there is a handy shutdown method for this
141
        try:
142
            if self.guestfs_enabled:
143
                self.g.umount_all()
144
                self.g.sync()
145
        finally:
146
            # Close the guestfs handler if open
147
            self.g.close()
148

    
149
#    def progress_callback(self, ev, eh, buf, array):
150
#        position = array[2]
151
#        total = array[3]
152
#
153
#        self.progressbar.goto((position * 100) // total)
154

    
155
    def mount(self, readonly=False):
156
        """Mount all disk partitions in a correct order."""
157

    
158
        mount = self.g.mount_ro if readonly else self.g.mount
159
        msg = " read-only" if readonly else ""
160
        self.out.output("Mounting the media%s ..." % msg, False)
161
        mps = self.g.inspect_get_mountpoints(self.root)
162

    
163
        # Sort the keys to mount the fs in a correct order.
164
        # / should be mounted befor /boot, etc
165
        def compare(a, b):
166
            if len(a[0]) > len(b[0]):
167
                return 1
168
            elif len(a[0]) == len(b[0]):
169
                return 0
170
            else:
171
                return -1
172
        mps.sort(compare)
173
        for mp, dev in mps:
174
            try:
175
                mount(dev, mp)
176
            except RuntimeError as msg:
177
                self.out.warn("%s (ignored)" % msg)
178

    
179
        self.mounted = True
180
        self.out.success("done")
181

    
182
    def umount(self):
183
        """Umount all mounted filesystems."""
184
        self.g.umount_all()
185
        self.mounted = False
186

    
187
    def _last_partition(self):
188
        if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
189
            msg = "Unsupported partition table: %s. Only msdos and gpt " \
190
                "partition tables are supported" % self.meta['PARTITION_TABLE']
191
            raise FatalError(msg)
192

    
193
        is_extended = lambda p: \
194
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
195
            in (0x5, 0xf)
196
        is_logical = lambda p: \
197
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
198

    
199
        partitions = self.g.part_list(self.guestfs_device)
200
        last_partition = partitions[-1]
201

    
202
        if is_logical(last_partition):
203
            # The disk contains extended and logical partitions....
204
            extended = filter(is_extended, partitions)[0]
205
            last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
206

    
207
            # check if extended is the last primary partition
208
            if last_primary['part_num'] > extended['part_num']:
209
                last_partition = last_primary
210

    
211
        return last_partition
212

    
213
    def shrink(self):
214
        """Shrink the disk.
215

216
        This is accomplished by shrinking the last filesystem in the
217
        disk and then updating the partition table. The new disk size
218
        (in bytes) is returned.
219

220
        ATTENTION: make sure unmount is called before shrink
221
        """
222
        get_fstype = lambda p: \
223
            self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
224
        is_logical = lambda p: \
225
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
226
        is_extended = lambda p: \
227
            self.meta['PARTITION_TABLE'] == 'msdos' and \
228
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
229
            in (0x5, 0xf)
230

    
231
        part_add = lambda ptype, start, stop: \
232
            self.g.part_add(self.guestfs_device, ptype, start, stop)
233
        part_del = lambda p: self.g.part_del(self.guestfs_device, p)
234
        part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
235
        part_set_id = lambda p, id: \
236
            self.g.part_set_mbr_id(self.guestfs_device, p, id)
237
        part_get_bootable = lambda p: \
238
            self.g.part_get_bootable(self.guestfs_device, p)
239
        part_set_bootable = lambda p, bootable: \
240
            self.g.part_set_bootable(self.guestfs_device, p, bootable)
241

    
242
        MB = 2 ** 20
243

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

    
246
        sector_size = self.g.blockdev_getss(self.guestfs_device)
247

    
248
        last_part = None
249
        fstype = None
250
        while True:
251
            last_part = self._last_partition()
252
            fstype = get_fstype(last_part)
253

    
254
            if fstype == 'swap':
255
                self.meta['SWAP'] = "%d:%s" % \
256
                    (last_part['part_num'],
257
                     (last_part['part_size'] + MB - 1) // MB)
258
                part_del(last_part['part_num'])
259
                continue
260
            elif is_extended(last_part):
261
                part_del(last_part['part_num'])
262
                continue
263

    
264
            # Most disk manipulation programs leave 2048 sectors after the last
265
            # partition
266
            new_size = last_part['part_end'] + 1 + 2048 * sector_size
267
            self.size = min(self.size, new_size)
268
            break
269

    
270
        if not re.match("ext[234]", fstype):
271
            self.out.warn("Don't know how to resize %s partitions." % fstype)
272
            return self.size
273

    
274
        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
275
        self.g.e2fsck_f(part_dev)
276
        self.g.resize2fs_M(part_dev)
277

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

    
282
        start = last_part['part_start'] / sector_size
283
        end = start + (block_size * block_cnt) / sector_size - 1
284

    
285
        if is_logical(last_part):
286
            partitions = self.g.part_list(self.guestfs_device)
287

    
288
            logical = []  # logical partitions
289
            for partition in partitions:
290
                if partition['part_num'] < 4:
291
                    continue
292
                logical.append({
293
                    'num': partition['part_num'],
294
                    'start': partition['part_start'] / sector_size,
295
                    'end': partition['part_end'] / sector_size,
296
                    'id': part_get_id(partition['part_num']),
297
                    'bootable': part_get_bootable(partition['part_num'])
298
                })
299

    
300
            logical[-1]['end'] = end  # new end after resize
301

    
302
            # Recreate the extended partition
303
            extended = filter(is_extended, partitions)[0]
304
            part_del(extended['part_num'])
305
            part_add('e', extended['part_start'] / sector_size, end)
306

    
307
            # Create all the logical partitions back
308
            for l in logical:
309
                part_add('l', l['start'], l['end'])
310
                part_set_id(l['num'], l['id'])
311
                part_set_bootable(l['num'], l['bootable'])
312
        else:
313
            # Recreate the last partition
314
            if self.meta['PARTITION_TABLE'] == 'msdos':
315
                last_part['id'] = part_get_id(last_part['part_num'])
316

    
317
            last_part['bootable'] = part_get_bootable(last_part['part_num'])
318
            part_del(last_part['part_num'])
319
            part_add('p', start, end)
320
            part_set_bootable(last_part['part_num'], last_part['bootable'])
321

    
322
            if self.meta['PARTITION_TABLE'] == 'msdos':
323
                part_set_id(last_part['part_num'], last_part['id'])
324

    
325
        new_size = (end + 1) * sector_size
326

    
327
        assert (new_size <= self.size)
328

    
329
        if self.meta['PARTITION_TABLE'] == 'gpt':
330
            ptable = GPTPartitionTable(self.device)
331
            self.size = ptable.shrink(new_size, self.size)
332
        else:
333
            self.size = min(new_size + 2048 * sector_size, self.size)
334

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

    
337
        return self.size
338

    
339
    def dump(self, outfile):
340
        """Dumps the content of device into a file.
341

342
        This method will only dump the actual payload, found by reading the
343
        partition table. Empty space in the end of the device will be ignored.
344
        """
345
        MB = 2 ** 20
346
        blocksize = 4 * MB  # 4MB
347
        size = self.size
348
        progr_size = (size + MB - 1) // MB  # in MB
349
        progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
350

    
351
        with open(self.device, 'r') as src:
352
            with open(outfile, "w") as dst:
353
                left = size
354
                offset = 0
355
                progressbar.next()
356
                while left > 0:
357
                    length = min(left, blocksize)
358
                    sent = sendfile(dst.fileno(), src.fileno(), offset, length)
359

    
360
                    # Workaround for python-sendfile API change. In
361
                    # python-sendfile 1.2.x (py-sendfile) the returning value
362
                    # of sendfile is a tuple, where in version 2.x (pysendfile)
363
                    # it is just a sigle integer.
364
                    if isinstance(sent, tuple):
365
                        sent = sent[1]
366

    
367
                    offset += sent
368
                    left -= sent
369
                    progressbar.goto((size - left) // MB)
370
        progressbar.success('image file %s was successfully created' % outfile)
371

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