#!/usr/bin/python
# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

'''A module containing kernel handler class used by SAFT.'''

import re

TMP_FILE_NAME = 'kernel_header_dump'
MAIN_STORAGE_DEVICE = '/dev/sda'

    # Types of kernel modifications.
KERNEL_BODY_MOD = 1
KERNEL_VERSION_MOD = 2


class KernelHandlerError(Exception):
    pass


class KernelHandler(object):
    '''An object to provide ChromeOS kernel related actions.

    Mostly it allows to corrupt and restore a particular kernel partition
    (designated by the partition name, A or B.
    '''

    # This value is used to alter contents of a byte in the appropriate kernel
    # image. First added to corrupt the image, then subtracted to restore the
    # image.
    DELTA = 1

    def __init__(self):
        self.chros_if = None
        self.dump_file_name = None
        self.partition_map = {}

    def _get_version(self, device):
        '''Get version of the kernel hosted on the passed in partition.'''
        # 16 K should be enough to include headers and keys
        data = self.chros_if.read_partition(device, 0x4000)
        return self.chros_if.retrieve_body_version(data)

    def _get_partition_map(self):
        '''Scan `cgpt show <device> output to find kernel devices.'''
        kernel_partitions = re.compile('KERN-([AB])')
        disk_map = self.chros_if.run_shell_command_get_output(
            'cgpt show %s' % MAIN_STORAGE_DEVICE)

        for line in disk_map:
            matched_line = kernel_partitions.search(line)
            if not matched_line:
                continue
            label = matched_line.group(1)
            part_info = {}
            device = MAIN_STORAGE_DEVICE + line.split()[2]
            part_info['device'] = device
            part_info['version'] = self._get_version(device)
            self.partition_map[label] = part_info

    def _modify_kernel(self, section,
                       delta,
                       modification_type=KERNEL_BODY_MOD):
        '''Modify kernel image on a disk partition.

        This method supports two types of kernel modification. KERNEL_BODY_MOD
        just adds the value of delta to the first byte of the kernel blob.
        This might leave the kernel corrupted (as required by the test).

        The second type, KERNEL_VERSION_MOD - will use 'delta' as the new
        version number, it will put it in the kernel header, and then will
        resign the kernel blob.
        '''
        dev = self.partition_map[section]['device']
        cmd_template = 'dd if=%s of=%s bs=4M count=1'
        self.chros_if.run_shell_command(cmd_template % (
                dev, self.dump_file_name))
        bfile = open(self.dump_file_name, 'r')
        data = list(bfile.read())
        bfile.close()
        if modification_type == KERNEL_BODY_MOD:
            data[0] = '%c' % ((ord(data[0]) + delta) % 0x100)
            dumpf = open(self.dump_file_name, 'w')
            dumpf.write(''.join(data))
            dumpf.close()
            kernel_to_write = self.dump_file_name
        elif modification_type == KERNEL_VERSION_MOD:
            new_version = delta
            kernel_to_write = self.dump_file_name + '.new'
            self.chros_if.run_shell_command(
                'vbutil_kernel --repack %s --version %d '
                '--signprivate %s --oldblob %s' % (
                    kernel_to_write, new_version,
                    'kernel_data_key.vbprivk', self.dump_file_name))
        else:
            return  # Unsupported mode, ignore.
        self.chros_if.run_shell_command(cmd_template % (kernel_to_write, dev))

    def corrupt_kernel(self, section):
        '''Corrupt a kernel section (add DELTA to the first byte).'''
        self._modify_kernel(section.upper(), self.DELTA)

    def restore_kernel(self, section):
        '''Restore the previously corrupted kernel.'''
        self._modify_kernel(section.upper(), -self.DELTA)

    def get_version(self, section):
        '''Return version read from this section blob's header.'''
        return self.partition_map[section.upper()]['version']

    def set_version(self, section, version):
        '''Set version of this kernel blob and re-sign it.'''
        if version < 0:
            raise KernelHandlerError('Bad version value %d' % version)
        self._modify_kernel(section.upper(), version, KERNEL_VERSION_MOD)

    def init (self, chros_if):
        '''Initialize the kernel handler object.

        Input argument is a ChromeOS interface object reference.
        '''
        self.chros_if = chros_if
        self.dump_file_name = chros_if.state_dir_file(TMP_FILE_NAME)
        self._get_partition_map()