Statistics
| Branch: | Tag: | Revision:

root / image_creator / image.py @ c9d53ee2

History | View | Annotate | Download (12.5 kB)

1
# -*- coding: utf-8 -*-
2
#
3
# Copyright 2013 GRNET S.A. All rights reserved.
4
#
5
# Redistribution and use in source and binary forms, with or
6
# without modification, are permitted provided that the following
7
# conditions are met:
8
#
9
#   1. Redistributions of source code must retain the above
10
#      copyright notice, this list of conditions and the following
11
#      disclaimer.
12
#
13
#   2. Redistributions in binary form must reproduce the above
14
#      copyright notice, this list of conditions and the following
15
#      disclaimer in the documentation and/or other materials
16
#      provided with the distribution.
17
#
18
# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
19
# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
20
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
21
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
22
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
25
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
26
# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
27
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
28
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29
# POSSIBILITY OF SUCH DAMAGE.
30
#
31
# The views and conclusions contained in the software and
32
# documentation are those of the authors and should not be
33
# interpreted as representing official policies, either expressed
34
# or implied, of GRNET S.A.
35

    
36
from image_creator.util import FatalError, check_guestfs_version
37
from image_creator.gpt import GPTPartitionTable
38
from image_creator.os_type import os_cls
39

    
40
import re
41
import guestfs
42
from sendfile import sendfile
43

    
44

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

    
48
    def __init__(self, device, output, meta={}):
49
        """Create a new Image instance"""
50

    
51
        self.device = device
52
        self.out = output
53
        self.meta = meta
54
        self.progress_bar = None
55
        self.guestfs_device = None
56
        self.size = 0
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
        if check_guestfs_version(self.g, 1, 17, 14) >= 0:
67
            self.out.output("Enabling recovery proc")
68
            self.g.set_recovery_proc(1)
69
        else:
70
            self.g.set_recovery_proc(0)
71

    
72
        #self.g.set_trace(1)
73
        #self.g.set_verbose(1)
74

    
75
        self.guestfs_enabled = False
76

    
77
    def enable(self):
78
        """Enable a newly created Image instance"""
79

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

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

    
105
        self.ostype = self.g.inspect_get_type(self.root)
106
        self.distro = self.g.inspect_get_distro(self.root)
107
        self.out.success(
108
            'found a(n) %s system' %
109
            self.ostype if self.distro == "unknown" else self.distro)
110

    
111
    def _get_os(self):
112
        """Return an OS class instance for this image"""
113
        if hasattr(self, "_os"):
114
            return self._os
115

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

    
119
        cls = os_cls(self.distro, self.ostype)
120
        self._os = cls(self)
121

    
122
        self._os.collect_metadata()
123

    
124
        return self._os
125

    
126
    os = property(_get_os)
127

    
128
    def destroy(self):
129
        """Destroy this Image instance."""
130

    
131
        # In new guestfs versions, there is a handy shutdown method for this
132
        try:
133
            if self.guestfs_enabled:
134
                self.g.umount_all()
135
                self.g.sync()
136
        finally:
137
            # Close the guestfs handler if open
138
            self.g.close()
139

    
140
#    def progress_callback(self, ev, eh, buf, array):
141
#        position = array[2]
142
#        total = array[3]
143
#
144
#        self.progressbar.goto((position * 100) // total)
145

    
146
    def _last_partition(self):
147
        """Return the last partition of the image disk"""
148
        if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
149
            msg = "Unsupported partition table: %s. Only msdos and gpt " \
150
                "partition tables are supported" % self.meta['PARTITION_TABLE']
151
            raise FatalError(msg)
152

    
153
        is_extended = lambda p: \
154
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
155
            in (0x5, 0xf)
156
        is_logical = lambda p: \
157
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
158

    
159
        partitions = self.g.part_list(self.guestfs_device)
160
        last_partition = partitions[-1]
161

    
162
        if is_logical(last_partition):
163
            # The disk contains extended and logical partitions....
164
            extended = filter(is_extended, partitions)[0]
165
            last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
166

    
167
            # check if extended is the last primary partition
168
            if last_primary['part_num'] > extended['part_num']:
169
                last_partition = last_primary
170

    
171
        return last_partition
172

    
173
    def shrink(self):
174
        """Shrink the image.
175

176
        This is accomplished by shrinking the last file system of the
177
        image and then updating the partition table. The new disk size
178
        (in bytes) is returned.
179

180
        ATTENTION: make sure unmount is called before shrink
181
        """
182
        get_fstype = lambda p: \
183
            self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
184
        is_logical = lambda p: \
185
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
186
        is_extended = lambda p: \
187
            self.meta['PARTITION_TABLE'] == 'msdos' and \
188
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
189
            in (0x5, 0xf)
190

    
191
        part_add = lambda ptype, start, stop: \
192
            self.g.part_add(self.guestfs_device, ptype, start, stop)
193
        part_del = lambda p: self.g.part_del(self.guestfs_device, p)
194
        part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
195
        part_set_id = lambda p, id: \
196
            self.g.part_set_mbr_id(self.guestfs_device, p, id)
197
        part_get_bootable = lambda p: \
198
            self.g.part_get_bootable(self.guestfs_device, p)
199
        part_set_bootable = lambda p, bootable: \
200
            self.g.part_set_bootable(self.guestfs_device, p, bootable)
201

    
202
        MB = 2 ** 20
203

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

    
206
        sector_size = self.g.blockdev_getss(self.guestfs_device)
207

    
208
        last_part = None
209
        fstype = None
210
        while True:
211
            last_part = self._last_partition()
212
            fstype = get_fstype(last_part)
213

    
214
            if fstype == 'swap':
215
                self.meta['SWAP'] = "%d:%s" % \
216
                    (last_part['part_num'],
217
                     (last_part['part_size'] + MB - 1) // MB)
218
                part_del(last_part['part_num'])
219
                continue
220
            elif is_extended(last_part):
221
                part_del(last_part['part_num'])
222
                continue
223

    
224
            # Most disk manipulation programs leave 2048 sectors after the last
225
            # partition
226
            new_size = last_part['part_end'] + 1 + 2048 * sector_size
227
            self.size = min(self.size, new_size)
228
            break
229

    
230
        if not re.match("ext[234]", fstype):
231
            self.out.warn("Don't know how to shrink %s partitions." % fstype)
232
            return self.size
233

    
234
        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
235
        self.g.e2fsck_f(part_dev)
236
        self.g.resize2fs_M(part_dev)
237

    
238
        out = self.g.tune2fs_l(part_dev)
239
        block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
240
        block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
241

    
242
        start = last_part['part_start'] / sector_size
243
        end = start + (block_size * block_cnt) / sector_size - 1
244

    
245
        if is_logical(last_part):
246
            partitions = self.g.part_list(self.guestfs_device)
247

    
248
            logical = []  # logical partitions
249
            for partition in partitions:
250
                if partition['part_num'] < 4:
251
                    continue
252
                logical.append({
253
                    'num': partition['part_num'],
254
                    'start': partition['part_start'] / sector_size,
255
                    'end': partition['part_end'] / sector_size,
256
                    'id': part_get_id(partition['part_num']),
257
                    'bootable': part_get_bootable(partition['part_num'])
258
                })
259

    
260
            logical[-1]['end'] = end  # new end after resize
261

    
262
            # Recreate the extended partition
263
            extended = filter(is_extended, partitions)[0]
264
            part_del(extended['part_num'])
265
            part_add('e', extended['part_start'] / sector_size, end)
266

    
267
            # Create all the logical partitions back
268
            for l in logical:
269
                part_add('l', l['start'], l['end'])
270
                part_set_id(l['num'], l['id'])
271
                part_set_bootable(l['num'], l['bootable'])
272
        else:
273
            # Recreate the last partition
274
            if self.meta['PARTITION_TABLE'] == 'msdos':
275
                last_part['id'] = part_get_id(last_part['part_num'])
276

    
277
            last_part['bootable'] = part_get_bootable(last_part['part_num'])
278
            part_del(last_part['part_num'])
279
            part_add('p', start, end)
280
            part_set_bootable(last_part['part_num'], last_part['bootable'])
281

    
282
            if self.meta['PARTITION_TABLE'] == 'msdos':
283
                part_set_id(last_part['part_num'], last_part['id'])
284

    
285
        new_size = (end + 1) * sector_size
286

    
287
        assert (new_size <= self.size)
288

    
289
        if self.meta['PARTITION_TABLE'] == 'gpt':
290
            ptable = GPTPartitionTable(self.device)
291
            self.size = ptable.shrink(new_size, self.size)
292
        else:
293
            self.size = min(new_size + 2048 * sector_size, self.size)
294

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

    
297
        return self.size
298

    
299
    def dump(self, outfile):
300
        """Dumps the content of the image into a file.
301

302
        This method will only dump the actual payload, found by reading the
303
        partition table. Empty space in the end of the device will be ignored.
304
        """
305
        MB = 2 ** 20
306
        blocksize = 4 * MB  # 4MB
307
        size = self.size
308
        progr_size = (size + MB - 1) // MB  # in MB
309
        progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
310

    
311
        with open(self.device, 'r') as src:
312
            with open(outfile, "w") as dst:
313
                left = size
314
                offset = 0
315
                progressbar.next()
316
                while left > 0:
317
                    length = min(left, blocksize)
318
                    sent = sendfile(dst.fileno(), src.fileno(), offset, length)
319

    
320
                    # Workaround for python-sendfile API change. In
321
                    # python-sendfile 1.2.x (py-sendfile) the returning value
322
                    # of sendfile is a tuple, where in version 2.x (pysendfile)
323
                    # it is just a sigle integer.
324
                    if isinstance(sent, tuple):
325
                        sent = sent[1]
326

    
327
                    offset += sent
328
                    left -= sent
329
                    progressbar.goto((size - left) // MB)
330
        progressbar.success('image file %s was successfully created' % outfile)
331

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