test_flashrom_handler.py

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

'''Unittest for flashrom_handler module.

This script can be run in a host, chroot or target environment. It requires
the vbutil_firmware utility to be available in the path.

When run as a program from the command line, this script expects one or two
parameters. The first parameter is the file name of the public key which can
be used to verify the firmware image. The second parameter, when specified, is
the file name of a ChromeOS firmware file. If the script is run in the target
environment and the firmaware file name is not specified, the current firmware
is read using 'flashrom -r' command.

See usage() below for command line options.
'''

import os
import shutil
import sys
import tempfile
import unittest

import chromeos_interface
import flashrom_handler
import saft_flashrom_util

progname = os.path.basename(sys.argv[0])
pub_key_file = None
test_fd_file = None


def usage():
    text = '''
usage: %s <key_file> [<fd_image_file>]
    <fd_image_file> can be omitted on the target, in which case
           it is read from the flashrom.
   NOTE: vbutil_firmware utility must be available in PATH.
'''
    print >> sys.stderr, (text % progname).strip()
    sys.exit(1)


class TestFlashromHandler(unittest.TestCase):

    PREAMBLE_USE_RO_NORMAL = 1

    def setUp(self):
        self.fh = flashrom_handler.FlashromHandler()
        self.chros_if = chromeos_interface.ChromeOSInterface(True)
        self.chros_if.init()
        self.tmpd = self.chros_if.init_environment()
        self.fh.init(saft_flashrom_util, self.chros_if,
                     pub_key_file=pub_key_file,
                     dev_key_path=os.getenv('DEV_KEYS_PATH', './'))

        # Suppress error messages generated by the program under test.
        self.stdout = sys.stdout
        sys.stdout = open('/dev/null', 'w')

    def test_pubkey_retireval(self):
        if not test_fd_file:
            # This test requires a file with a firmware image.
            return

        if self.chros_if.target_hosted():
            # This will retrieve the key from the onboard flash image.
            self.fh.new_image()
        self.fh.new_image(test_fd_file)
        self.fh.verify_image()

    def test_image_read(self):
        self.fh.new_image(test_fd_file)
        self.fh.verify_image()

    def test_image_re_sign(self):
        if not test_fd_file:
            if self.chros_if.target_hosted():
                fw_image_file = self.chros_if.state_dir_file('fw_image.fd')
                self.fh.new_image()
                self.fh.dump_whole(fw_image_file)
            else:
                # This test requires a file with a firmware image.
                return
        else:
            fw_image_file = test_fd_file

        backup = self.chros_if.state_dir_file('backup.fd')
        shutil.copyfile(fw_image_file, backup)
        for section in ('a', 'b'):
            self.fh.new_image(backup)
            self.fh.verify_image()
            new_version = self.fh.get_section_version(section) + 1
            flags = self.fh.get_section_flags(section)
            self.fh.set_section_version(section, new_version, flags)
            self.fh.dump_whole(backup)
            self.fh.new_image(backup)
            self.fh.verify_image()
            self.assertEqual(self.fh.get_section_version(section),
                             new_version)

    def test_image_corrupt_restore(self):
        image_name = os.path.join(self.tmpd, 'tmp.image')
        self.fh.new_image(test_fd_file)
        self.fh.verify_image()
        self.fh.dump_whole(image_name)
        self.fh.new_image(image_name)
        self.fh.verify_image()
        for section in ('a', 'b'):
            corrupted_file_name = image_name + section
            corrupted_subsection = self.fh.corrupt_section(section)
            self.fh.dump_whole(corrupted_file_name)
            self.fh.new_image(corrupted_file_name)
            flags = self.fh.get_section_flags(section)
            if flags & self.PREAMBLE_USE_RO_NORMAL:
                self.fh.verify_image  # No assertion when USE_RO_NORMAL
            else:
                self.assertRaises(chromeos_interface.ChromeOSInterfaceError,
                                  self.fh.verify_image)
            self.fh.restore_section(section)
            self.fh.dump_whole(corrupted_file_name)
            self.fh.new_image(corrupted_file_name)
            self.fh.verify_image()
            self.fh.new_image(image_name)

    def tearDown(self):
        sys.stdout = self.stdout
        shutil.rmtree(self.tmpd)


if __name__ == '__main__':
    if len(sys.argv) > 1:
        test_fd_file = sys.argv[1]
        sys.argv = sys.argv[0:1]
    unittest.main()