Statistics
| Branch: | Tag: | Revision:

root / image_creator / image.py @ 5e18a927

History | View | Annotate | Download (12.8 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
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, bootable=True, meta={}):
49
        """Create a new Image instance"""
50

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

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

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

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

    
79
        self.guestfs_enabled = False
80

    
81
    def enable(self):
82
        """Enable a newly created Image instance"""
83

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

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

    
109
        self.ostype = self.g.inspect_get_type(self.root)
110
        self.distro = self.g.inspect_get_distro(self.root)
111
        self.out.success(
112
            'found a(n) %s system' %
113
            self.ostype if self.distro == "unknown" else self.distro)
114

    
115
    def _get_os(self):
116
        """Return an OS class instance for this image"""
117
        if hasattr(self, "_os"):
118
            return self._os
119

    
120
        if not self.guestfs_enabled:
121
            self.enable()
122

    
123
        cls = os_cls(self.distro, self.ostype)
124
        self._os = cls(self.root, self.g, self.out)
125

    
126
        self._os.collect_metadata()
127

    
128
        return self._os
129

    
130
    os = property(_get_os)
131

    
132
    def destroy(self):
133
        """Destroy this Image instance."""
134

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

    
144
#    def progress_callback(self, ev, eh, buf, array):
145
#        position = array[2]
146
#        total = array[3]
147
#
148
#        self.progressbar.goto((position * 100) // total)
149

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

    
157
        is_extended = lambda p: \
158
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
159
            in (0x5, 0xf)
160
        is_logical = lambda p: \
161
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
162

    
163
        partitions = self.g.part_list(self.guestfs_device)
164
        last_partition = partitions[-1]
165

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

    
171
            # check if extended is the last primary partition
172
            if last_primary['part_num'] > extended['part_num']:
173
                last_partition = last_primary
174

    
175
        return last_partition
176

    
177
    def shrink(self):
178
        """Shrink the image.
179

180
        This is accomplished by shrinking the last file system of the
181
        image and then updating the partition table. The new disk size
182
        (in bytes) is returned.
183

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

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

    
206
        MB = 2 ** 20
207

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

    
210
        sector_size = self.g.blockdev_getss(self.guestfs_device)
211

    
212
        last_part = None
213
        fstype = None
214
        while True:
215
            last_part = self._last_partition()
216
            fstype = get_fstype(last_part)
217

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

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

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

    
238
        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
239
        self.g.e2fsck_f(part_dev)
240
        self.g.resize2fs_M(part_dev)
241

    
242
        out = self.g.tune2fs_l(part_dev)
243
        block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
244
        block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
245

    
246
        start = last_part['part_start'] / sector_size
247
        end = start + (block_size * block_cnt) / sector_size - 1
248

    
249
        if is_logical(last_part):
250
            partitions = self.g.part_list(self.guestfs_device)
251

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

    
264
            logical[-1]['end'] = end  # new end after resize
265

    
266
            # Recreate the extended partition
267
            extended = filter(is_extended, partitions)[0]
268
            part_del(extended['part_num'])
269
            part_add('e', extended['part_start'] / sector_size, end)
270

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

    
281
            last_part['bootable'] = part_get_bootable(last_part['part_num'])
282
            part_del(last_part['part_num'])
283
            part_add('p', start, end)
284
            part_set_bootable(last_part['part_num'], last_part['bootable'])
285

    
286
            if self.meta['PARTITION_TABLE'] == 'msdos':
287
                part_set_id(last_part['part_num'], last_part['id'])
288

    
289
        new_size = (end + 1) * sector_size
290

    
291
        assert (new_size <= self.size)
292

    
293
        if self.meta['PARTITION_TABLE'] == 'gpt':
294
            ptable = GPTPartitionTable(self.device)
295
            self.size = ptable.shrink(new_size, self.size)
296
        else:
297
            self.size = min(new_size + 2048 * sector_size, self.size)
298

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

    
301
        return self.size
302

    
303
    def dump(self, outfile):
304
        """Dumps the content of the image into a file.
305

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

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

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

    
331
                    offset += sent
332
                    left -= sent
333
                    progressbar.goto((size - left) // MB)
334
        progressbar.success('image file %s was successfully created' % outfile)
335

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