#!/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.

'''Test utility for verifying ChromeOS firmware.'''

import datetime
import functools
import getopt
import os
import re
import shutil
import sys
import tempfile
import traceback

import cgpt_handler
import cgpt_state
import chromeos_interface
import flashrom_handler
import kernel_handler
import saft_flashrom_util
import tpm_handler

#
# We need to know the names of two files:
#
# - the SAFT shell script used by the upstart script (/etc/init/saft.conf) to
#   determine if SAFT is in progress
#
# - the name of the log file where the upstart script redirects its output.
#
# These file names are defined in the upstart script as <NAME>=<VALUE>, one
# per line, using names below:

# SAFT shell script
SAFT_SCRIPT_VAR_NAME = 'SAFT_SCRIPT'

# Upstart log file
CONF_LOG_VAR_NAME = 'SAFT_LOG'

# This is a template used to populate the SAFT shell script which
# /etc/init/saft.conf is looking for.
SAFT_SCRIPT_TEMPLATE = '''
#!/bin/sh
# 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.

# Do not edit, this file was autogenerated.

# The only parameter is the name of the Flash block device. The root main root
# file system is presumed to be /dev/<flash_block_dev>3.
#
# TODO(vbendeb): instead of using the hardcoded device name look for block
#                labeled  C-KEYFOB

set -e

saft_dev="/dev/${1}3"
root_dev=$(rootdev -s)

if [ "${saft_dev}" = "${root_dev}" ]; then
    mount_point=""
else
    mount_point='/tmp/saft_device'
    mkdir "${mount_point}"
    mount "${saft_dev}" "${mount_point}"
fi

cd "${mount_point}%s"

# Start our own X server to make it easier to control the screen
X_PORT=1
X :"${X_PORT}" -config saft.xorg.conf vt02 > /dev/null 2>&1 &
xpid="$!"
initctl stop update-engine  # Prevent messing with cgpt attributes
export DISPLAY=":${X_PORT}.0"
./%s --next_step
kill "${xpid}"
'''

# Subdirectory to keep the state of this test over reboots. Created in the
# /var on the USB flash drive the test is running from.
STATE_SUBDIR = '.fw_test'

# Files storing SAFT state over reboots, located in state_dir defined below.
LOG_FILE = 'fw_test_log.txt'  # Test log.
STEP_FILE = 'fw_test_step'  # The next test step number.
FW_BACKUP_FILE = 'flashrom.bak'  # Preserved original flashrom contents.
FW_COPY_FILE = 'flashrom.new'  # A copy of the flashrom contents being tested.
FWID_BACKUP_FILE = 'fwid.bak'  # FWID reported by the original firmware.
FWID_NEW_FILE = 'fwid.new'  # FWID reported by the firmware being tested.

BASE_STORAGE_DEVICE = '/dev/sda'

# The list of shell executables necessary for this program to work.
REQUIRED_PROGRAMS = '''
cgpt blkid flashrom reboot_mode readlink rootdev vbutil_firmware vbutil_kernel
mosys
'''


FLASHROM_HANDLER = flashrom_handler.FlashromHandler()
CHROS_IF = chromeos_interface.ChromeOSInterface(__name__ != '__main__')
cgpt_st = cgpt_state.CgptState('COMPLETE', CHROS_IF, BASE_STORAGE_DEVICE)

def allow_multiple_section_input(image_operator):
    @functools.wraps(image_operator)
    def wrapper(self, section):
        if type(section) is not tuple:
            image_operator(self, section)
        else:
            for sec in section:
                image_operator(self, sec)
    return wrapper

class FwError(Exception):
    '''A class to encapsulate this module specific exceptions.'''
    pass


def RetriveSaftConfDefinion(name):
    '''Find a name definition in the SAFT upstart script.

    This function scans the /etc/init/saft.conf for <name>=<value> lines. and
    if found - returns the <value> string.

    It will throw an IOError exception in case the upstart script is not
    present or FwError() in case <name> definition is not in the file.
    '''
    pattern = name + '='
    conf_file = '/etc/init/saft.conf'

    for line in open(conf_file, 'r').readlines():
        line = line.strip()
        if not line.startswith(pattern):
            continue
        pieces = line.split('=')
        if len(pieces) != 2 or pieces[0] != name:
            continue
        return pieces[1].strip('\'"')

    # If we did not find it - something is very wrong.
    raise FwError('name %s not found in %s' % (name, conf_file))


class FirmwareTest(object):
    '''An object to represent a Firmware Semi Automated test (SAFT).

    It follows the steps as enumerated in test_state_sequence below,
    restarting the target after each step. The first time around the
    init_fw_test() method must be called to set up the persistent environment.
    All following invocations should initialize the object (using .init()) and
    the invoke the next_step() method.
    '''

    def __init__(self):
        '''Object initializer, does nothing to make mocking easier.'''
        self.mydir = None
        self.base_partition = None
        self.chros_if = None
        self.progname = None
        self.test_state_sequence = None
        self.kern_handler = None
        self.step_failed = False
        self.window = None
        self.tpm_handler = None

    def _verify_fw_id(self, compare_to_file):
        '''Verify if the current firmware ID matches the contents a file.

        compare_to_file - a string, name of the file in the state directory.
        '''
        old_fwid = open(
            self.chros_if.state_dir_file(compare_to_file), 'r').read()
        now_fwid = open(self.chros_if.acpi_file('FWID'), 'r').read()
        return old_fwid == now_fwid

    def _get_step(self):
        '''Set the current value of SAFT step number.'''
        step_file = self.chros_if.state_dir_file(STEP_FILE)
        step = open(step_file, 'r').read().strip()
        return int(step)

    def _set_step(self, step):
        '''Set the SAFT step number to control the next test pass.'''
        step_file = self.chros_if.state_dir_file(STEP_FILE)
        open(step_file, 'w').write('%d' % step)

    def _handle_saft_script(self, install):
        '''Install or remove the SAFT shell script.

      On SAFT initialization, this function is invoked to install the shell
      script which the upstart script /etc/init/saft.conf will be passing
      control to after each restart. The presence of the shell script is an
      indication that SAFT is in progress.

      The shell script is expected to be installed in the stateful partition
      of the flash drive. When invoked, the shell script completes setting up
      the environment and in turn invokes this program passing it --next_step
      as the command line parameter.

      When SAFT finishes all steps, this function is called to remove the
      shell script so that next time the traget restarts it does not try to
      execute SAFT anymore.
      '''

        if not self.chros_if.target_hosted():
            print 'bypassing upstart management'
            return

        label_pattern = re.compile('/dev/(.+): LABEL="C-STATE"')
        tmp_dir = tempfile.mkdtemp()

        # Find out our file name off the root of the device
        mount_point = (self.chros_if.run_shell_command_get_output(
                'df %s' % self.mydir)[-1]).split()[-1]
        this_dir = self.mydir[len(mount_point):]

        for line in self.chros_if.run_shell_command_get_output('blkid'):
            match = label_pattern.search(line)
            if not match:
                continue
            dev = match.groups(0)[0]
            mount_point = self.chros_if.get_writeable_mount_point(
                '/dev/' + dev, tmp_dir)
            script_name = RetriveSaftConfDefinion(SAFT_SCRIPT_VAR_NAME)

            # Find location of the script in the currently mounted file
            # system.
            file_name = os.path.join(mount_point, script_name.lstrip('/'))
            self.chros_if.log('Handling ' + file_name)
            if install:
                saft_script = open(file_name, 'w')
                saft_script.write(SAFT_SCRIPT_TEMPLATE % (
                        this_dir, self.progname))
                saft_script.close()
                os.chmod(file_name, 0700)
            else:
                os.unlink(file_name)
            if mount_point.startswith(tmp_dir):
                self.chros_if.run_shell_command('umount %s' % tmp_dir)

        os.rmdir(tmp_dir)

    def _check_runtime_env(self):
        '''Ensure that the script is running in proper environment.

      This involves checking that the script is running off a removable
      device, configuring proper file names for logging, etc.
      '''
        line = self.chros_if.run_shell_command_get_output(
            'df %s' % self.mydir)[-1]

        self.base_partition = line.split()[0]
        if self.base_partition == '/dev/root':
            self.base_partition = self.chros_if.get_root_dev()

        if not self.chros_if.is_removable_device(self.base_partition):
            raise FwError(
                'This test must run off a removable device, not %s'
                % self.base_partition)

        env_root = '/var'
        state_fs = '%s1' % self.base_partition[:-1]

      # is state file system mounted?
        for line in self.chros_if.run_shell_command_get_output('mount'):
            if line.startswith('%s on ' % state_fs):
                mount_point = line.split()[2]
                state_root = mount_point + env_root
                break
        else:
            tmp_dir = tempfile.mkdtemp()
            self.chros_if.run_shell_command('mount %s %s' % (state_fs, tmp_dir))
            state_root = '%s%s' % (tmp_dir, env_root)

        self.chros_if.init(os.path.join(state_root, STATE_SUBDIR), LOG_FILE)

    def _cmp_boot_vector(self, vector1, vector2):
        '''Compare if the two boot vectors are the same

        Note: a wildcard (*) will match any value.

        '''
        list1 = vector1.split(cgpt_st.DELIMIT)
        list2 = vector2.split(cgpt_st.DELIMIT)
        if len(list1) != len(list2):
            raise FwError('Boot vectors (%s %s) should be of the same length'
                          % (vecotr1, vector2))
        for i in range(len(list1)):
            if list1[i] != list2[i] and list1[i] != '*' and list2[i] != '*':
                return False
        return True

    def init(self, progname, chros_if,
             kern_handler=None, state_sequence=None, tpm_handler=None):
        '''Initialize the Firmware self test instance.

        progname - a string, name of this program as it was invoked.

        chros_if - an object of ChromeOSInterface type to be initialized and
                   used by this instance of FirmwareTest.

        kern_handler - an object providing SAFT with services manipulating
                   kernel images. Unittest does not provide this object.

        test_state_sequence - a tuple of three-tuples driving test execution,
                   see description below. Unittest does not provide this
                   sequence.
        '''
        real_name = os.path.realpath(progname)
        self.mydir = os.path.dirname(real_name)
        self.progname = os.path.basename(real_name)
        self.chros_if = chros_if
        self._check_runtime_env()
        self.test_state_sequence = state_sequence
        if kern_handler:
            self.kern_handler = kern_handler
            self.kern_handler.init(self.chros_if)
        if tpm_handler:
            self.tpm_handler = tpm_handler
            self.tpm_handler.init(self.chros_if)

        if self.tpm_handler:
            if not self.chros_if.is_removable_device(self.base_partition):
                # On each non USB flash rooted start confirm that kernel and
                # BIOS versions and TPM setting match.
                vers_a = FLASHROM_HANDLER.get_version('a')
                vers_b = FLASHROM_HANDLER.get_version('b')
                if not self.tpm_handler.fw_version_good(vers_a, vers_b):
                    raise FwError('TPM firmware version mismatch (%d %d %s)' % (
                            vers_a, vers_b, self.base_partition))
                if self.kern_handler:
                    vers_a = self.kern_handler.get_version('a')
                    vers_b = self.kern_handler.get_version('b')
                    if not self.tpm_handler.kernel_version_good(vers_a, vers_b):
                        raise FwError('TPM kernel version mismatch (%d %d %s)' % (
                                vers_a, vers_b, self.base_partition))

    def set_try_fw_b(self):
        '''Request running firmware B on the next restart.'''
        self.chros_if.log('Requesting restart with FW B')
        self.chros_if.run_shell_command('reboot_mode --try_firmware_b=1')

    def request_recovery_boot(self):
        '''Request running in recovery mode on the restart.'''
        self.chros_if.log('Requesting restart in recovery mode')
        self.chros_if.run_shell_command('reboot_mode --recovery=1')

    @allow_multiple_section_input
    def restore_firmware(self, section):
        '''Restore the requested firmware section (previously corrupted).'''
        self.chros_if.log('Restoring firmware %s' % section)
        FLASHROM_HANDLER.restore_firmware(section)

    @allow_multiple_section_input
    def corrupt_firmware(self, section):
        '''Corrupt the requested firmware section.'''
        self.chros_if.log('Corrupting firmware %s' % section)
        FLASHROM_HANDLER.corrupt_firmware(section)

    @allow_multiple_section_input
    def corrupt_kernel(self, section):
        '''Corrupt the requested kernel section.'''
        self.chros_if.log('Corrupting kernel %s' % section)
        self.kern_handler.corrupt_kernel(section)

    @allow_multiple_section_input
    def restore_kernel(self, section):
        '''restore the requested kernel section.'''
        self.chros_if.log('restoring kernel %s' % section)
        self.kern_handler.restore_kernel(section)

    def cgpt_test_start(self):
        '''Set up CGPT test sequence.'''
        cgpth = cgpt_handler.CgptHandler(self.chros_if)
        cgpth.read_device_info(BASE_STORAGE_DEVICE)

        # Set gpt attributes for kernels A and B such that after reset it comes
        # back up using kernel A, with tries A set to 14 (decremented) and
        # other A and B attributes unchanged.
        properties = {'successful': 0,
                      'tries': 15,
                      'priority': 10,
                      }
        cgpth.set_partition(BASE_STORAGE_DEVICE, 'KERN-A', properties)
        properties['priority'] = 9
        cgpth.set_partition(BASE_STORAGE_DEVICE, 'KERN-B', properties)

    def cgpt_test_1(self):
        '''CGPT test checking the attributes settings by the firmware.'''
        expected_cgpt = {
            'KERN-A' : {'priority': 10, 'tries': 14, 'successful': 0},
            'KERN-B' : {'priority': 9, 'tries': 15, 'successful': 0}
            }
        cgpth = cgpt_handler.CgptHandler(self.chros_if)
        cgpth.read_device_info(BASE_STORAGE_DEVICE)
        for part, state in expected_cgpt.iteritems():
            props = cgpth.get_partition(BASE_STORAGE_DEVICE, part)
            for prop, value in state.iteritems():
                if value != props[prop]:
                    self.chros_if.log('wrong partition %s value' % part)
                    self.chros_if.log(cgpth.dump_partition(BASE_STORAGE_DEVICE,
                                                           part))
                    self.step_failed = True
                    break

    def move_kernel_backward(self, section):
        '''Decrement kernel version for the requested section.'''
        new_version = self.kern_handler.get_version(section) - 1
        self.chros_if.log(
            'setting kernel section %s version to %d' % (section, new_version))
        self.kern_handler.set_version(section, new_version)

    def move_firmware_backward(self, section):
        '''Decrement firmware version for the requested section.'''
        new_version = FLASHROM_HANDLER.get_section_version(section) - 1
        self.chros_if.log(
            'setting firmware section %s version to %d' % (
                section, new_version))
        FLASHROM_HANDLER.set_section_version(section, new_version, True)

    def jump_kernels_forward(self):
        '''Add two to both kernels' versions.

        This compensates for the previous decrement and increases the version
        number by one as compared to the original state (at the start of the
        test).
        '''
        for section in ('a', 'b'):
            new_version = self.kern_handler.get_version(section) + 2
            self.chros_if.log(
                'setting section %s version to %d' % (section, new_version))
            self.kern_handler.set_version(section, new_version)

    def jump_firmwares_forward(self):
        '''Add two to both firmwares' versions.

        This compensates for the previous decrement and increases the version
        number by one as compared to the original state (at the start of the
        test).
        '''
        for section in ('a', 'b'):
            new_version = FLASHROM_HANDLER.get_section_version(section) + 2
            self.chros_if.log(
                'setting firmware section %s version to %d' % (
                    section, new_version))
            FLASHROM_HANDLER.set_section_version(section, new_version, True)

    def terminate_tpm_tests(self):
        '''Restore TPM and kernel states.

        Move kernel and firmware versions to where they were before SAFT
        started. Set the TPM's values of kernel and firmware versions to the
        lowest of the two kernels/firmwares available on the system. Prevent
        future TPM write accesses by restoring the upstart file.
        '''
        new_tpm_kern_version  = 0xffff
        new_tpm_fw_version  = 0xffff
        for section in ('a', 'b'):
            new_version = self.kern_handler.get_version(section) - 1
            self.chros_if.log(
                'setting kernel section %s version to %d' % (
                    section, new_version))
            self.kern_handler.set_version(section, new_version)
            new_tpm_kern_version = min(new_version, new_tpm_kern_version)

            new_version = FLASHROM_HANDLER.get_section_version(section)  - 1
            self.chros_if.log(
                'setting firmware section %s version to %d' % (
                    section, new_version))
            FLASHROM_HANDLER.set_section_version(section, new_version, True)
            new_tpm_fw_version = min(new_version, new_tpm_fw_version)

        self.tpm_handler.set_kernel_version(new_tpm_kern_version)
        self.tpm_handler.set_fw_version(new_tpm_fw_version)
        self.tpm_handler.disable_write_access()

    def revert_firmware(self):
        '''Restore firmware to the image backed up when SAFT started.'''
        self.chros_if.log('restoring original firmware image')
        self.chros_if.run_shell_command(
            'flashrom -w %s' % self.chros_if.state_dir_file(FW_BACKUP_FILE))

    def new_fw_image(self, image_file = None):
        FLASHROM_HANDLER.new_image(image_file)
        FLASHROM_HANDLER.verify_image()
        if self.tpm_handler:
            vers_a = FLASHROM_HANDLER.get_section_version('a')
            vers_b = FLASHROM_HANDLER.get_section_version('b')
            if not self.tpm_handler.fw_version_good(vers_a, vers_b):
                raise FwError('TPM firmware version mismatch')

    def prepare_tpm_tests(self):
        '''Prepare TPM for testing.

        Enable write access on the next reboot, and roll back the kernel we
        are running now (the other kernel is expected to be used on the next
        restart).
        '''
        mount_point = self.chros_if.run_shell_command_get_output(
            'df %s' % self.mydir)[-1].split()[-1]
        cfg_file = os.path.join(mount_point, 'etc/init/tcsd.conf')
        self.tpm_handler.enable_write_access(cfg_file)
        self.move_kernel_backward('a')

    def init_fw_test(self, opt_dictionary, chros_if):
        '''Prepare firmware test context.

      This function tries creating the state directory for the fw test and
      initializes the test state machine.

      Return
        True on success
        False on any failure or if the directory already exists
      '''
        chros_if.init_environment()
        chros_if.log('Automated firmware test log generated on %s' % (
                datetime.datetime.strftime(
                    datetime.datetime.now(), '%b %d %Y')))
        chros_if.log('Original boot state %s' % chros_if.boot_state_vector())
        self.chros_if = chros_if
        fw_image = opt_dictionary['image_file']
        self.new_fw_image()
        FLASHROM_HANDLER.dump_whole(
            self.chros_if.state_dir_file(FW_BACKUP_FILE))
        self.new_fw_image(fw_image)
        self._handle_saft_script(True)
        shutil.copyfile(self.chros_if.acpi_file('FWID'),
                        self.chros_if.state_dir_file(FWID_BACKUP_FILE))
        shutil.copyfile(fw_image, self.chros_if.state_dir_file(FW_COPY_FILE))

        self._set_step(0)
        cgpt_st.set_step(0)

    def next_step(self):
        '''Function to execute a single SAFT step.

      This function is running after each reboot. It determines the current
      step the SAFT is on, executes the appropriate action, increments the
      step value and then restats the machine.
      '''

        this_step = self._get_step()

        # Import when needed, because otherwise running test generates a
        # warning exception about the unavailable display.
        window = __import__('window')
        self.window = window.GraphThread()

        if this_step == 0:
            shutil.copyfile(self.chros_if.acpi_file('FWID'),
                            self.chros_if.state_dir_file(FWID_NEW_FILE))

            if self._verify_fw_id(FWID_BACKUP_FILE):
          # we expected FWID to change, but it did not - have the firmware
          # been even replaced?
                self.chros_if.log('New firmware - old FWID')
                self.finish_saft(False)
        test_state_tuple = self.test_state_sequence[this_step]
        expected_vector = test_state_tuple[0]
        action = test_state_tuple[1]

        # Display cgpt sub step if the action is a cgpt test
        step_display = "SAFT's performing step #%d"
        step_tuple = (this_step,)
        if action == cgpt_st.test_loop:
            step_display += " (cgpt test step #%d)"
            step_tuple += (cgpt_st.get_step(),)
        self.window.display_text(step_display % step_tuple)

        boot_vector = self.chros_if.boot_state_vector()
        self.chros_if.log('\nRebooted into state %s on step %d' % (
                boot_vector, this_step))

        conf_log_file = RetriveSaftConfDefinion(CONF_LOG_VAR_NAME)

        if os.path.exists(conf_log_file):
            contents = open(conf_log_file).read().rstrip()
            if len(contents):
                self.chros_if.log('startup log contents:')
                self.chros_if.log(contents)

        if not self._cmp_boot_vector(boot_vector, expected_vector):
            self.chros_if.log('Error: Wrong boot vector, %s was expected'
                % expected_vector)
            self.finish_saft(False)
        if this_step == (len(self.test_state_sequence) - 1):
            self.finish_saft(True)
        if action and not self._verify_fw_id(FWID_NEW_FILE):
            self.chros_if.log('Error: Wrong FWID value')
            self.finish_saft(False)

        inc_step = 1
        if action:
            FLASHROM_HANDLER.new_image()
            if action == cgpt_st.test_loop:
                inc_step = action()
            elif len(test_state_tuple) > 2:
                self.chros_if.log('calling %s with parameter %s' % (
                        str(action), str(test_state_tuple[2])))
                action(test_state_tuple[2])
            else:
                self.chros_if.log('calling %s' % str(action))
                action()
            if self.step_failed:
                self.finish_saft(False)

        self._set_step(this_step + inc_step)
        self.chros_if.run_shell_command('reboot')

    def finish_saft(self, success):
        '''SAFT completed, restore the environment and expose the results.

        Delete the startup script from the removable device, copy the log to
        '/var' directory on the current drive and exit.

        Input:
          success - a Boolean set to True if the test has passed through all
            steps. Affects placing of the last line 'we are done' in the log.
        '''

        if success and not self._verify_fw_id(FWID_BACKUP_FILE):
            self.chros_if.log(
                'Error: Failed to restore to original firmware')
            success = False
        self.chros_if.log('Removing upstart scripts')
        self._handle_saft_script(False)
        if success:
            self.chros_if.log('we are done!')
            # Have chros_if.shutdown move the log into the '/var' directory to
            # make it easier to see SAFT results.
            self.chros_if.shut_down(os.path.join('/var', LOG_FILE))

        self.window.stop()
        sys.exit(0)

# Firmware self test instance controlling this module.
FST = FirmwareTest()

# This is a tuple of tuples controlling the SAFT state machine. The states
# are expected to be passed strictly in order. The states are identified
# by the contents of BINF.[012] files in the sys fs ACPI directory. The
# BINF files store information about the reason for reboot, what
# firmware/kernel partitions were used, etc.
#
# The first element of each component tuple is the expected state of the
# machine (a ':' concatenation of the BINF files' contents).
#
# The second element of the component tuples is the action to take to
# advance the test. The action is a function to call. The last line has
# action set to None, which indicates to the state machine that the test
# is over.
#
# The third component, if present, is the parameter to pass to the action
# function.

TEST_STATE_SEQUENCE = (
    ('1:1:1:0:3', FST.set_try_fw_b),                   # Step 0
    ('1:2:1:0:3', None),
    ('1:1:1:0:3', FST.corrupt_firmware, 'a'),
    ('1:2:1:0:3', FST.restore_firmware, 'a'),
    ('1:1:1:0:3', FST.corrupt_firmware, ('a', 'b')),
    ('5:0:1:1:3', FST.restore_firmware, ('a', 'b')),   # Step 5
    ('1:1:1:0:3', FST.corrupt_kernel, 'a'),
    ('1:1:1:0:5', FST.corrupt_kernel, 'b'),
    ('6:0:1:1:3', FST.restore_kernel, ('a', 'b')),
    ('1:1:1:0:3', FST.request_recovery_boot),
    ('8:0:1:1:3', FST.prepare_tpm_tests),              # Step 10
    ('1:1:1:0:5', FST.move_kernel_backward, 'b'),
    ('6:0:1:1:3', FST.jump_kernels_forward),
    ('1:1:1:0:3', FST.move_firmware_backward, 'a'),
    ('1:2:1:0:3', FST.move_firmware_backward, 'b'),
    ('5:0:1:1:3', FST.jump_firmwares_forward),         # Step 15
    ('1:1:1:0:3', FST.request_recovery_boot),
    ('8:0:1:1:3', FST.terminate_tpm_tests),
    ('*:*:*:*:*', cgpt_st.test_loop),
    ('1:1:1:0:3', FST.revert_firmware),
    ('1:1:1:0:3', None),                               # Step 20
    )


# The string below serves two purposes:
#
# - spell out the usage string for this program
#
# - provide text for parsing to derive command line flags. The text is split
#   to words, then the words which have -- in them are stripped off the
#   leading/traling brackets and dashes and used to prepare inut parameters
#   for getopt. This is the only way to pass the parameters to getopt, this
#   enforces that each accepted parameter is mentioned in the usage() output.
USAGE_STRING = '''
 [--image_file=<firmware_image_file>] [--pub_key=<file>] [--next_step]

  The program can be invoked in two modes.

  When invoked for the first time, most of the parameters are required to set
  up the test context and start it.

  Specifying --next_step means that the program is being invoked by the
  restarted system, all the context is expected to be available. No other
  parameters are required or expected in that case.
'''


def usage(msg='', retv=0):
    '''Print error message (if any), usage string and exit.

    Depending on the passed in return value use stdout (if retv is 0) or
    stderr (if otherwise).
    '''
    progname = os.path.basename(sys.argv[0])
    if retv:
        ofile = sys.stderr
    else:
        ofile = sys.stdout
    if msg:
        print >> ofile, '%s: %s' % (progname, msg)
        CHROS_IF.log(msg)
    print >> ofile, 'usage: %s %s' % (progname, USAGE_STRING.strip())
    sys.exit(retv)


def get_options_set():
    '''Generate the list of command line options accepted by this program.

  This function derives the set of accepted command line options from
  usage_string. The string is split into words, all words starting with -- are
  considered accepted command line parameters. The presence of an '=' sign in
  the word,means that the command line parameter requires a value.

  Returns a list of strings suitable for use by the getopt module.
  '''

    drop_tail = re.compile('=.*$')
    option_set = []
    items = (' '.join(USAGE_STRING.split('\n'))).split(' ')
    for item in items:
        if '--' in item:
            if item.startswith("'"):
                continue
            option = drop_tail.sub('=', item.strip('-[]'))
            if option not in option_set:
                option_set.append(option)
    return option_set


def main(argv):
    '''Process command line options and invoke the proper test entry point.'''
    (opts, params) = getopt.gnu_getopt(argv[1:], '', get_options_set())
    if params:
        raise FwError('unrecognized parameters: %s' % ' '.join(params))

    opt_dictionary = {}
    for (name, value) in opts:
        opt_dictionary[name.lstrip('-')] = value

    FST.init(argv[0], CHROS_IF,
             kernel_handler.KernelHandler(), TEST_STATE_SEQUENCE,
             tpm_handler.TpmHandler())

    FLASHROM_HANDLER.init(saft_flashrom_util, CHROS_IF,
                         opt_dictionary.get('pub_key'))
    if 'next_step' in opt_dictionary:
        if len(opt_dictionary) != 1:
            usage('--next_step (when specified) must be the only parameter', 1)
        try:
            FST.next_step()
        except SystemExit:
            pass  # No real exception, this is just an exit (for whatever
                  # reason), gtk window already closed.
        except:
            # Whatever error that might be, gtk window must be shut.
            FST.window.stop()

            # Make sure exception information is saved in the log.
            exc_type, exc_info, exc_trace = sys.exc_info()
            print 'exception type:', str(exc_type)
            print 'exception info:', str(exc_info)
            traceback.print_tb(exc_trace)
        sys.exit(0)

    # check if all executables are available
    missing_execs = []
    for prog in REQUIRED_PROGRAMS.split():
        if not prog:
            continue # ignore line breaks
        if not CHROS_IF.exec_exists(prog):
            missing_execs.append(prog)
    if missing_execs:
        usage('Program(s) %s not found in PATH' % ' '.join(missing_execs), 1)

    if 'image_file' not in opt_dictionary:
        usage(retv=1)

    FST.init_fw_test(opt_dictionary, CHROS_IF)

    if not FLASHROM_HANDLER.firmware_sections_equal():
        # This is a temporary measure needed to address the fact that the
        # development BIOS image is built with 'normal' firmware in section B
        # and 'development' firmware in section A. SAFT operation requires
        # 'normal' flavor in both sections. Note that this fix does not affect
        # the file of the new image, only the flashrom contents.
        CHROS_IF.log('modify firmware A to match B')
        FLASHROM_HANDLER.copy_from_to('b', 'a')
    CHROS_IF.log('program the new image')
    FLASHROM_HANDLER.write_whole()
    CHROS_IF.log('restart')
    CHROS_IF.run_shell_command('reboot')
    return 0


if __name__ == '__main__':
    try:
        main(sys.argv)
    except (getopt.GetoptError, ImportError):
        usage(sys.exc_info()[1], 1)
    except (FwError, flashrom_handler.FlashromHandlerError):
        MSG = 'Error: %s' % str(sys.exc_info()[1])
        print MSG
        CHROS_IF.log(MSG)
        sys.exit(1)

    sys.exit(0)