image.py 15.8 KB
Newer Older
Nikos Skalkotos's avatar
Nikos Skalkotos committed
1 2
# -*- coding: utf-8 -*-
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
3
# Copyright (C) 2011-2014 GRNET S.A.
4
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
5 6 7 8
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
9
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
10 11 12 13
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
14
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
15 16
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
17

18
from image_creator.util import FatalError
19 20 21 22 23
from image_creator.gpt import GPTPartitionTable
from image_creator.os_type import os_cls

import re
import guestfs
24
import hashlib
25 26 27 28 29 30
from sendfile import sendfile


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

31
    def __init__(self, device, output, **kargs):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
32
        """Create a new Image instance"""
33 34 35

        self.device = device
        self.out = output
36 37 38 39 40

        self.meta = kargs['meta'] if 'meta' in kargs else {}
        self.sysprep_params = \
            kargs['sysprep_params'] if 'sysprep_params' in kargs else {}

41 42 43 44 45
        self.progress_bar = None
        self.guestfs_device = None
        self.size = 0

        self.g = guestfs.GuestFS()
46
        self.guestfs_enabled = False
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65
        self.guestfs_version = self.g.version()

    def check_guestfs_version(self, major, minor, release):
        """Checks if the version of the used libguestfs is smaller, equal or
        greater than the one specified by the major, minor and release triplet

        Returns:
            < 0 if the installed version is smaller than the specified one
            = 0 if they are equal
            > 0 if the installed one is greater than the specified one
        """

        for (a, b) in (self.guestfs_version['major'], major), \
                (self.guestfs_version['minor'], minor), \
                (self.guestfs_version['release'], release):
            if a != b:
                return a - b

        return 0
66 67 68 69 70 71 72 73

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

        self.enable_guestfs()

        self.out.output('Inspecting Operating System ...', False)
        roots = self.g.inspect_os()
74 75 76 77 78 79 80

        if len(roots) == 0 or len(roots) > 1:
            self.root = None
            self.ostype = "unsupported"
            self.distro = "unsupported"
            self.guestfs_device = '/dev/sda'
            self.size = self.g.blockdev_getsize64(self.guestfs_device)
81

82
            if len(roots) > 1:
83
                reason = "Multiple operating systems found on the media."
84
            else:
85
                reason = "Unable to detect any operating system on the media."
86

87
            self.set_unsupported(reason)
88 89
            return

90
        self.root = roots[0]
91 92
        self.meta['PARTITION_TABLE'] = self.g.part_get_parttype('/dev/sda')
        self.guestfs_device = '/dev/sda'  # self.g.part_to_dev(self.root)
93 94 95 96 97 98 99 100
        self.size = self.g.blockdev_getsize64(self.guestfs_device)

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

Nikos Skalkotos's avatar
Nikos Skalkotos committed
101 102
        # Inspect the OS
        self.os.inspect()
103 104

    def set_unsupported(self, reason):
105
        """Flag this image as unsupported"""
106 107 108 109 110 111 112 113 114

        self._unsupported = reason
        self.meta['UNSUPPORTED'] = reason
        self.out.warn('Media is not supported. Reason: %s' % reason)

    def is_unsupported(self):
        """Returns if this image is unsupported"""
        return hasattr(self, '_unsupported')

115 116 117 118 119 120 121
    def enable_guestfs(self):
        """Enable the guestfs handler"""

        if self.guestfs_enabled:
            self.out.warn("Guestfs is already enabled")
            return

122
        # Before version 1.18.4 the behavior of kill_subprocess was different
123
        # and you need to reset the guestfs handler to relaunch a previously
124
        # shut down QEMU backend
125
        if self.check_guestfs_version(1, 18, 4) < 0:
126 127
            self.g = guestfs.GuestFS()

128 129 130 131 132 133 134
        self.g.add_drive_opts(self.device, readonly=0, format="raw")

        # Before version 1.17.14 the recovery process, which is a fork of the
        # original process that called libguestfs, did not close its inherited
        # file descriptors. This can cause problems especially if the parent
        # process has opened pipes. Since the recovery process is an optional
        # feature of libguestfs, it's better to disable it.
135
        if self.check_guestfs_version(1, 17, 14) >= 0:
Nikos Skalkotos's avatar
Nikos Skalkotos committed
136
            self.out.output("Enabling recovery process ...", False)
137
            self.g.set_recovery_proc(1)
Nikos Skalkotos's avatar
Nikos Skalkotos committed
138
            self.out.success('done')
139 140
        else:
            self.g.set_recovery_proc(0)
141

Nikos Skalkotos's avatar
Nikos Skalkotos committed
142 143
        # self.g.set_trace(1)
        # self.g.set_verbose(1)
144 145 146 147 148 149

        self.out.output('Launching helper VM (may take a while) ...', False)
        # self.progressbar = self.out.Progress(100, "Launching helper VM",
        #                                     "percent")
        # eh = self.g.set_event_callback(self.progress_callback,
        #                               guestfs.EVENT_PROGRESS)
150 151 152 153 154 155
        try:
            self.g.launch()
        except RuntimeError as e:
            raise FatalError(
                "Launching libguestfs's helper VM failed! Reason: %s" % str(e))

156 157 158 159
        self.guestfs_enabled = True
        # self.g.delete_event_callback(eh)
        # self.progressbar.success('done')
        # self.progressbar = None
160

161
        if self.check_guestfs_version(1, 18, 4) < 0:
162 163
            self.g.inspect_os()  # some calls need this

164 165
        self.out.success('done')

166 167
    def disable_guestfs(self):
        """Disable the guestfs handler"""
168

169 170 171 172 173 174
        if not self.guestfs_enabled:
            self.out.warn("Guestfs is already disabled")
            return

        self.out.output("Shutting down helper VM ...", False)
        self.g.sync()
175
        # guestfs_shutdown which is the preferred way to shutdown the backend
176
        # process was introduced in version 1.19.16
177
        if self.check_guestfs_version(1, 19, 16) >= 0:
178 179 180 181
            self.g.shutdown()
        else:
            self.g.kill_subprocess()

182 183 184 185
        # We will reset the guestfs handler if needed
        if self.check_guestfs_version(1, 18, 4) < 0:
            self.g.close()

186 187
        self.guestfs_enabled = False
        self.out.success('done')
188

189 190
    @property
    def os(self):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
191
        """Return an OS class instance for this image"""
192 193 194 195 196 197
        if hasattr(self, "_os"):
            return self._os

        if not self.guestfs_enabled:
            self.enable()

198
        cls = os_cls(self.distro, self.ostype)
199
        self._os = cls(self, sysprep_params=self.sysprep_params)
200

201
        self._os.collect_metadata()
202 203 204 205

        return self._os

    def destroy(self):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
206
        """Destroy this Image instance."""
207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223

        # In new guestfs versions, there is a handy shutdown method for this
        try:
            if self.guestfs_enabled:
                self.g.umount_all()
                self.g.sync()
        finally:
            # Close the guestfs handler if open
            self.g.close()

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

    def _last_partition(self):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
224
        """Return the last partition of the image disk"""
225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
        if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
            msg = "Unsupported partition table: %s. Only msdos and gpt " \
                "partition tables are supported" % self.meta['PARTITION_TABLE']
            raise FatalError(msg)

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

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

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

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

        return last_partition

Nikos Skalkotos's avatar
Nikos Skalkotos committed
250
    def shrink(self, silent=False):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
251
        """Shrink the image.
252

Nikos Skalkotos's avatar
Nikos Skalkotos committed
253 254
        This is accomplished by shrinking the last file system of the
        image and then updating the partition table. The new disk size
255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280
        (in bytes) is returned.

        ATTENTION: make sure unmount is called before shrink
        """
        get_fstype = lambda p: \
            self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
        is_logical = lambda p: \
            self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
        is_extended = lambda p: \
            self.meta['PARTITION_TABLE'] == 'msdos' and \
            self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
            in (0x5, 0xf)

        part_add = lambda ptype, start, stop: \
            self.g.part_add(self.guestfs_device, ptype, start, stop)
        part_del = lambda p: self.g.part_del(self.guestfs_device, p)
        part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
        part_set_id = lambda p, id: \
            self.g.part_set_mbr_id(self.guestfs_device, p, id)
        part_get_bootable = lambda p: \
            self.g.part_get_bootable(self.guestfs_device, p)
        part_set_bootable = lambda p, bootable: \
            self.g.part_set_bootable(self.guestfs_device, p, bootable)

        MB = 2 ** 20

281
        if self.is_unsupported():
Nikos Skalkotos's avatar
Nikos Skalkotos committed
282 283
            if not silent:
                self.out.warn("Shrinking is disabled for unsupported images")
284 285
            return self.size

286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310
        sector_size = self.g.blockdev_getss(self.guestfs_device)

        last_part = None
        fstype = None
        while True:
            last_part = self._last_partition()
            fstype = get_fstype(last_part)

            if fstype == 'swap':
                self.meta['SWAP'] = "%d:%s" % \
                    (last_part['part_num'],
                     (last_part['part_size'] + MB - 1) // MB)
                part_del(last_part['part_num'])
                continue
            elif is_extended(last_part):
                part_del(last_part['part_num'])
                continue

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

        if not re.match("ext[234]", fstype):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
311 312 313
            if not silent:
                self.out.warn(
                    "Don't know how to shrink %s partitions." % fstype)
314 315 316
            return self.size

        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
317 318 319 320 321 322

        if self.check_guestfs_version(1, 15, 17) >= 0:
            self.g.e2fsck(part_dev, forceall=1)
        else:
            self.g.e2fsck_f(part_dev)

323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381
        self.g.resize2fs_M(part_dev)

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

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

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

            logical = []  # logical partitions
            for partition in partitions:
                if partition['part_num'] < 4:
                    continue
                logical.append({
                    'num': partition['part_num'],
                    'start': partition['part_start'] / sector_size,
                    'end': partition['part_end'] / sector_size,
                    'id': part_get_id(partition['part_num']),
                    'bootable': part_get_bootable(partition['part_num'])
                })

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

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

            # Create all the logical partitions back
            for l in logical:
                part_add('l', l['start'], l['end'])
                part_set_id(l['num'], l['id'])
                part_set_bootable(l['num'], l['bootable'])
        else:
            # Recreate the last partition
            if self.meta['PARTITION_TABLE'] == 'msdos':
                last_part['id'] = part_get_id(last_part['part_num'])

            last_part['bootable'] = part_get_bootable(last_part['part_num'])
            part_del(last_part['part_num'])
            part_add('p', start, end)
            part_set_bootable(last_part['part_num'], last_part['bootable'])

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

        new_size = (end + 1) * sector_size

        assert (new_size <= self.size)

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

Nikos Skalkotos's avatar
Nikos Skalkotos committed
382 383 384
        if not silent:
            self.out.success("Image size is %dMB" %
                             ((self.size + MB - 1) // MB))
385 386 387 388

        return self.size

    def dump(self, outfile):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
389
        """Dumps the content of the image into a file.
390 391 392 393 394

        This method will only dump the actual payload, found by reading the
        partition table. Empty space in the end of the device will be ignored.
        """
        MB = 2 ** 20
395 396
        blocksize = 2 ** 22  # 4MB
        progr_size = (self.size + MB - 1) // MB  # in MB
397 398
        progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')

399 400 401
        with open(self.device, 'rb') as src:
            with open(outfile, "wb") as dst:
                left = self.size
402 403 404 405 406 407 408 409 410
                offset = 0
                progressbar.next()
                while left > 0:
                    length = min(left, blocksize)
                    sent = sendfile(dst.fileno(), src.fileno(), offset, length)

                    # Workaround for python-sendfile API change. In
                    # python-sendfile 1.2.x (py-sendfile) the returning value
                    # of sendfile is a tuple, where in version 2.x (pysendfile)
411
                    # it is just a single integer.
412 413 414 415 416
                    if isinstance(sent, tuple):
                        sent = sent[1]

                    offset += sent
                    left -= sent
417
                    progressbar.goto((self.size - left) // MB)
418 419
        progressbar.success('image file %s was successfully created' % outfile)

420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443
    def md5(self):
        """Computes the MD5 checksum of the image"""

        MB = 2 ** 20
        blocksize = 2 ** 22  # 4MB
        progr_size = ((self.size + MB - 1) // MB)  # in MB
        progressbar = self.out.Progress(progr_size, "Calculating md5sum", 'mb')
        md5 = hashlib.md5()

        with open(self.device, "rb") as src:
            left = self.size
            while left > 0:
                length = min(left, blocksize)
                data = src.read(length)
                md5.update(data)
                left -= length
                progressbar.goto((self.size - left) // MB)

        checksum = md5.hexdigest()
        progressbar.success(checksum)

        return checksum


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