firmware_test.py revision 8b377eb0528c632e4cb2d888f7388faf764f9987
1# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import ast
6import ctypes
7import logging
8import os
9import re
10import subprocess
11import time
12import uuid
13
14from autotest_lib.client.bin import utils
15from autotest_lib.client.common_lib import error
16from autotest_lib.server import autotest
17from autotest_lib.server.cros import vboot_constants as vboot
18from autotest_lib.server.cros.faft.config.config import Config as FAFTConfig
19from autotest_lib.server.cros.faft.utils.faft_checkers import FAFTCheckers
20from autotest_lib.server.cros.faft.rpc_proxy import RPCProxy
21from autotest_lib.server.cros.servo import chrome_ec
22from autotest_lib.server.cros.servo_test import ServoTest
23
24
25class ConnectionError(Exception):
26    """Raised on an error of connecting DUT."""
27    pass
28
29
30class FAFTBase(ServoTest):
31    """The base class of FAFT classes.
32
33    It launches the FAFTClient on DUT, such that the test can access its
34    firmware functions and interfaces. It also provides some methods to
35    handle the reboot mechanism, in order to ensure FAFTClient is still
36    connected after reboot.
37    """
38    def initialize(self, host):
39        """Create a FAFTClient object and install the dependency."""
40        super(FAFTBase, self).initialize(host)
41        self._client = host
42        self._autotest_client = autotest.Autotest(self._client)
43        self._autotest_client.install()
44        self.faft_client = RPCProxy(host)
45
46    def wait_for_client(self, install_deps=False, timeout=100):
47        """Wait for the client to come back online.
48
49        New remote processes will be launched if their used flags are enabled.
50
51        @param install_deps: If True, install Autotest dependency when ready.
52        @param timeout: Time in seconds to wait for the client SSH daemon to
53                        come up.
54        @raise ConnectionError: Failed to connect DUT.
55        """
56        if not self._client.wait_up(timeout):
57            raise ConnectionError()
58        if install_deps:
59            self._autotest_client.install()
60        # Check the FAFT client is avaiable.
61        self.faft_client.system.is_available()
62
63    def wait_for_client_offline(self, timeout=60, orig_boot_id=None):
64        """Wait for the client to come offline.
65
66        @param timeout: Time in seconds to wait the client to come offline.
67        @param orig_boot_id: A string containing the original boot id.
68        @raise ConnectionError: Failed to connect DUT.
69        """
70        # When running against panther, we see that sometimes
71        # ping_wait_down() does not work correctly. There needs to
72        # be some investigation to the root cause.
73        # If we sleep for 120s before running get_boot_id(), it
74        # does succeed. But if we change this to ping_wait_down()
75        # there are implications on the wait time when running
76        # commands at the fw screens.
77        if not self._client.ping_wait_down(timeout):
78            if orig_boot_id and self._client.get_boot_id() != orig_boot_id:
79                logging.warn('Reboot done very quickly.')
80                return
81            raise ConnectionError()
82
83
84class FirmwareTest(FAFTBase):
85    """
86    Base class that sets up helper objects/functions for firmware tests.
87
88    TODO: add documentaion as the FAFT rework progresses.
89    """
90    version = 1
91
92    # Mapping of partition number of kernel and rootfs.
93    KERNEL_MAP = {'a':'2', 'b':'4', '2':'2', '4':'4', '3':'2', '5':'4'}
94    ROOTFS_MAP = {'a':'3', 'b':'5', '2':'3', '4':'5', '3':'3', '5':'5'}
95    OTHER_KERNEL_MAP = {'a':'4', 'b':'2', '2':'4', '4':'2', '3':'4', '5':'2'}
96    OTHER_ROOTFS_MAP = {'a':'5', 'b':'3', '2':'5', '4':'3', '3':'5', '5':'3'}
97
98    CHROMEOS_MAGIC = "CHROMEOS"
99    CORRUPTED_MAGIC = "CORRUPTD"
100
101    _SERVOD_LOG = '/var/log/servod.log'
102
103    _ROOTFS_PARTITION_NUMBER = 3
104
105    _backup_firmware_sha = ()
106    _backup_kernel_sha = dict()
107    _backup_cgpt_attr = dict()
108    _backup_gbb_flags = None
109    _backup_dev_mode = None
110
111    # Class level variable, keep track the states of one time setup.
112    # This variable is preserved across tests which inherit this class.
113    _global_setup_done = {
114        'gbb_flags': False,
115        'reimage': False,
116        'usb_check': False,
117    }
118
119    @classmethod
120    def check_setup_done(cls, label):
121        """Check if the given setup is done.
122
123        @param label: The label of the setup.
124        """
125        return cls._global_setup_done[label]
126
127    @classmethod
128    def mark_setup_done(cls, label):
129        """Mark the given setup done.
130
131        @param label: The label of the setup.
132        """
133        cls._global_setup_done[label] = True
134
135    @classmethod
136    def unmark_setup_done(cls, label):
137        """Mark the given setup not done.
138
139        @param label: The label of the setup.
140        """
141        cls._global_setup_done[label] = False
142
143    def initialize(self, host, cmdline_args, ec_wp=None):
144        super(FirmwareTest, self).initialize(host)
145        self.run_id = str(uuid.uuid4())
146        logging.info('FirmwareTest initialize begin (id=%s)', self.run_id)
147        # Parse arguments from command line
148        args = {}
149        self.power_control = host.POWER_CONTROL_RPM
150        for arg in cmdline_args:
151            match = re.search("^(\w+)=(.+)", arg)
152            if match:
153                args[match.group(1)] = match.group(2)
154        if 'power_control' in args:
155            self.power_control = args['power_control']
156            if self.power_control not in host.POWER_CONTROL_VALID_ARGS:
157                raise error.TestError('Valid values for --args=power_control '
158                                      'are %s. But you entered wrong argument '
159                                      'as "%s".'
160                                       % (host.POWER_CONTROL_VALID_ARGS,
161                                       self.power_control))
162
163        self.faft_config = FAFTConfig(
164                self.faft_client.system.get_platform_name())
165        self.checkers = FAFTCheckers(self, self.faft_client)
166
167        if self.faft_config.chrome_ec:
168            self.ec = chrome_ec.ChromeEC(self.servo)
169
170        self._setup_uart_capture()
171        self._setup_servo_log()
172        self._record_system_info()
173        self._setup_gbb_flags()
174        self._stop_service('update-engine')
175        self._setup_ec_write_protect(ec_wp)
176        logging.info('FirmwareTest initialize done (id=%s)', self.run_id)
177
178    def cleanup(self):
179        """Autotest cleanup function."""
180        # Unset state checker in case it's set by subclass
181        logging.info('FirmwareTest cleaning up (id=%s)', self.run_id)
182        try:
183            self.faft_client.system.is_available()
184        except:
185            # Remote is not responding. Revive DUT so that subsequent tests
186            # don't fail.
187            self._restore_routine_from_timeout()
188        self._restore_dev_mode()
189        self._restore_ec_write_protect()
190        self._restore_gbb_flags()
191        self._start_service('update-engine')
192        self._record_servo_log()
193        self._record_faft_client_log()
194        self._cleanup_uart_capture()
195        super(FirmwareTest, self).cleanup()
196        logging.info('FirmwareTest cleanup done (id=%s)', self.run_id)
197
198    def _record_system_info(self):
199        """Record some critical system info to the attr keyval.
200
201        This info is used by generate_test_report and local_dash later.
202        """
203        self.write_attr_keyval({
204            'fw_version': self.faft_client.ec.get_version(),
205            'hwid': self.faft_client.system.get_crossystem_value('hwid'),
206            'fwid': self.faft_client.system.get_crossystem_value('fwid'),
207        })
208
209    def invalidate_firmware_setup(self):
210        """Invalidate all firmware related setup state.
211
212        This method is called when the firmware is re-flashed. It resets all
213        firmware related setup states so that the next test setup properly
214        again.
215        """
216        self.unmark_setup_done('gbb_flags')
217
218    def _retrieve_recovery_reason_from_trap(self):
219        """Try to retrieve the recovery reason from a trapped recovery screen.
220
221        @return: The recovery_reason, 0 if any error.
222        """
223        recovery_reason = 0
224        logging.info('Try to retrieve recovery reason...')
225        if self.servo.get_usbkey_direction() == 'dut':
226            self.wait_fw_screen_and_plug_usb()
227        else:
228            self.servo.switch_usbkey('dut')
229
230        try:
231            self.wait_for_client(install_deps=True)
232            lines = self.faft_client.system.run_shell_command_get_output(
233                        'crossystem recovery_reason')
234            recovery_reason = int(lines[0])
235            logging.info('Got the recovery reason %d.', recovery_reason)
236        except ConnectionError:
237            logging.error('Failed to get the recovery reason due to connection '
238                          'error.')
239        return recovery_reason
240
241    def _reset_client(self):
242        """Reset client to a workable state.
243
244        This method is called when the client is not responsive. It may be
245        caused by the following cases:
246          - halt on a firmware screen without timeout, e.g. REC_INSERT screen;
247          - corrupted firmware;
248          - corrutped OS image.
249        """
250        # DUT may halt on a firmware screen. Try cold reboot.
251        logging.info('Try cold reboot...')
252        self.reboot_cold_trigger()
253        self.wait_for_client_offline()
254        self.wait_dev_screen_and_ctrl_d()
255        try:
256            self.wait_for_client()
257            return
258        except ConnectionError:
259            logging.warn('Cold reboot doesn\'t help, still connection error.')
260
261        # DUT may be broken by a corrupted firmware. Restore firmware.
262        # We assume the recovery boot still works fine. Since the recovery
263        # code is in RO region and all FAFT tests don't change the RO region
264        # except GBB.
265        if self.is_firmware_saved():
266            self._ensure_client_in_recovery()
267            logging.info('Try restore the original firmware...')
268            if self.is_firmware_changed():
269                try:
270                    self.restore_firmware()
271                    return
272                except ConnectionError:
273                    logging.warn('Restoring firmware doesn\'t help, still '
274                                 'connection error.')
275
276        # Perhaps it's kernel that's broken. Let's try restoring it.
277        if self.is_kernel_saved():
278            self._ensure_client_in_recovery()
279            logging.info('Try restore the original kernel...')
280            if self.is_kernel_changed():
281                try:
282                    self.restore_kernel()
283                    return
284                except ConnectionError:
285                    logging.warn('Restoring kernel doesn\'t help, still '
286                                 'connection error.')
287
288        # DUT may be broken by a corrupted OS image. Restore OS image.
289        self._ensure_client_in_recovery()
290        logging.info('Try restore the OS image...')
291        self.faft_client.system.run_shell_command('chromeos-install --yes')
292        self.sync_and_warm_reboot()
293        self.wait_for_client_offline()
294        self.wait_dev_screen_and_ctrl_d()
295        try:
296            self.wait_for_client(install_deps=True)
297            logging.info('Successfully restore OS image.')
298            return
299        except ConnectionError:
300            logging.warn('Restoring OS image doesn\'t help, still connection '
301                         'error.')
302
303    def _ensure_client_in_recovery(self):
304        """Ensure client in recovery boot; reboot into it if necessary.
305
306        @raise TestError: if failed to boot the USB image.
307        """
308        logging.info('Try boot into USB image...')
309        self.servo.switch_usbkey('host')
310        self.enable_rec_mode_and_reboot()
311        self.wait_fw_screen_and_plug_usb()
312        try:
313            self.wait_for_client(install_deps=True)
314        except ConnectionError:
315            raise error.TestError('Failed to boot the USB image.')
316
317    def _restore_routine_from_timeout(self, next_step=None):
318        """A routine to try to restore the system from a timeout error.
319
320        This method is called when FAFT failed to connect DUT after reboot.
321
322        @param next_step: Optional, a FAFT_STEP dict of the next step, which is
323                          used for diagnostic.
324        @raise TestFail: This exception is already raised, with a decription
325                         why it failed.
326        """
327        # DUT is disconnected. Capture the UART output for debug.
328        self._record_uart_capture()
329
330        next_checker_matched = False
331
332        # TODO(waihong@chromium.org): Implement replugging the Ethernet to
333        # identify if it is a network flaky.
334
335        recovery_reason = self._retrieve_recovery_reason_from_trap()
336        if next_step is not None and recovery_reason:
337            if self._call_action(next_test['state_checker']):
338                # Repluging the USB can pass the state_checker of the next step,
339                # meaning that the firmware failed to boot into USB directly.
340                next_checker_matched = True
341
342        # Reset client to a workable state.
343        self._reset_client()
344
345        # Raise the proper TestFail exception.
346        if next_checker_matched:
347            raise error.TestFail('Firmware failed to auto-boot USB in the '
348                                 'recovery boot (reason: %d)' % recovery_reason)
349        elif recovery_reason:
350            raise error.TestFail('Trapped in the recovery screen (reason: %d) '
351                                 'and timed out' % recovery_reason)
352        else:
353            raise error.TestFail('Timed out waiting for DUT reboot')
354
355    def assert_test_image_in_usb_disk(self, usb_dev=None, install_shim=False):
356        """Assert an USB disk plugged-in on servo and a test image inside.
357
358        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
359                        If None, it is detected automatically.
360        @param install_shim: True to verify an install shim instead of a test
361                             image.
362        @raise TestError: if USB disk not detected or not a test (install shim)
363                          image.
364        """
365        if self.check_setup_done('usb_check'):
366            return
367        if usb_dev:
368            assert self.servo.get_usbkey_direction() == 'host'
369        else:
370            self.servo.switch_usbkey('host')
371            usb_dev = self.servo.probe_host_usb_dev()
372            if not usb_dev:
373                raise error.TestError(
374                        'An USB disk should be plugged in the servo board.')
375
376        rootfs = '%s%s' % (usb_dev, self._ROOTFS_PARTITION_NUMBER)
377        logging.info('usb dev is %s', usb_dev)
378        tmpd = self.servo.system_output('mktemp -d -t usbcheck.XXXX')
379        self.servo.system('mount -o ro %s %s' % (rootfs, tmpd))
380
381        if install_shim:
382            dir_list = self.servo.system_output('ls -a %s' %
383                                                os.path.join(tmpd, 'root'))
384            check_passed = '.factory_installer' in dir_list
385        else:
386            check_passed = self.servo.system_output(
387                'grep -i "CHROMEOS_RELEASE_DESCRIPTION=.*test" %s' %
388                os.path.join(tmpd, 'etc/lsb-release'),
389                ignore_status=True)
390        for cmd in ('umount %s' % rootfs, 'sync', 'rm -rf %s' % tmpd):
391            self.servo.system(cmd)
392
393        if not check_passed:
394            raise error.TestError(
395                'No Chrome OS %s found on the USB flash plugged into servo' %
396                'install shim' if install_shim else 'test')
397
398        self.mark_setup_done('usb_check')
399
400    def setup_usbkey(self, usbkey, host=None, install_shim=False):
401        """Setup the USB disk for the test.
402
403        It checks the setup of USB disk and a valid ChromeOS test image inside.
404        It also muxes the USB disk to either the host or DUT by request.
405
406        @param usbkey: True if the USB disk is required for the test, False if
407                       not required.
408        @param host: Optional, True to mux the USB disk to host, False to mux it
409                    to DUT, default to do nothing.
410        @param install_shim: True to verify an install shim instead of a test
411                             image.
412        """
413        if usbkey:
414            self.assert_test_image_in_usb_disk(install_shim=install_shim)
415        elif host is None:
416            # USB disk is not required for the test. Better to mux it to host.
417            host = True
418
419        if host is True:
420            self.servo.switch_usbkey('host')
421        elif host is False:
422            self.servo.switch_usbkey('dut')
423
424    def get_usbdisk_path_on_dut(self):
425        """Get the path of the USB disk device plugged-in the servo on DUT.
426
427        Returns:
428          A string representing USB disk path, like '/dev/sdb', or None if
429          no USB disk is found.
430        """
431        cmd = 'ls -d /dev/s*[a-z]'
432        original_value = self.servo.get_usbkey_direction()
433
434        # Make the dut unable to see the USB disk.
435        self.servo.switch_usbkey('off')
436        no_usb_set = set(
437            self.faft_client.system.run_shell_command_get_output(cmd))
438
439        # Make the dut able to see the USB disk.
440        self.servo.switch_usbkey('dut')
441        time.sleep(self.faft_config.between_usb_plug)
442        has_usb_set = set(
443            self.faft_client.system.run_shell_command_get_output(cmd))
444
445        # Back to its original value.
446        if original_value != self.servo.get_usbkey_direction():
447            self.servo.switch_usbkey(original_value)
448
449        diff_set = has_usb_set - no_usb_set
450        if len(diff_set) == 1:
451            return diff_set.pop()
452        else:
453            return None
454
455    def _stop_service(self, service):
456        """Stops a upstart service on the client.
457
458        @param service: The name of the upstart service.
459        """
460        logging.info('Stopping %s...', service)
461        command = 'status %s | grep stop || stop %s' % (service, service)
462        self.faft_client.system.run_shell_command(command)
463
464    def _start_service(self, service):
465        """Starts a upstart service on the client.
466
467        @param service: The name of the upstart service.
468        """
469        logging.info('Starting %s...', service)
470        command = 'status %s | grep start || start %s' % (service, service)
471        self.faft_client.system.run_shell_command(command)
472
473    def _write_gbb_flags(self, new_flags):
474        """Write the GBB flags to the current firmware.
475
476        @param new_flags: The flags to write.
477        """
478        gbb_flags = self.faft_client.bios.get_gbb_flags()
479        if gbb_flags == new_flags:
480            return
481        logging.info('Changing GBB flags from 0x%x to 0x%x.',
482                     gbb_flags, new_flags)
483        self.faft_client.system.run_shell_command(
484                '/usr/share/vboot/bin/set_gbb_flags.sh 0x%x' % new_flags)
485        self.faft_client.bios.reload()
486        # If changing FORCE_DEV_SWITCH_ON flag, reboot to get a clear state
487        if ((gbb_flags ^ new_flags) & vboot.GBB_FLAG_FORCE_DEV_SWITCH_ON):
488            self.reboot_warm_trigger()
489            self.wait_dev_screen_and_ctrl_d()
490
491    def clear_set_gbb_flags(self, clear_mask, set_mask):
492        """Clear and set the GBB flags in the current flashrom.
493
494        @param clear_mask: A mask of flags to be cleared.
495        @param set_mask: A mask of flags to be set.
496        """
497        gbb_flags = self.faft_client.bios.get_gbb_flags()
498        new_flags = gbb_flags & ctypes.c_uint32(~clear_mask).value | set_mask
499        self._write_gbb_flags(new_flags)
500
501    def check_ec_capability(self, required_cap=None, suppress_warning=False):
502        """Check if current platform has required EC capabilities.
503
504        @param required_cap: A list containing required EC capabilities. Pass in
505                             None to only check for presence of Chrome EC.
506        @param suppress_warning: True to suppress any warning messages.
507        @return: True if requirements are met. Otherwise, False.
508        """
509        if not self.faft_config.chrome_ec:
510            if not suppress_warning:
511                logging.warn('Requires Chrome EC to run this test.')
512            return False
513
514        if not required_cap:
515            return True
516
517        for cap in required_cap:
518            if cap not in self.faft_config.ec_capability:
519                if not suppress_warning:
520                    logging.warn('Requires EC capability "%s" to run this '
521                                 'test.', cap)
522                return False
523
524        return True
525
526    def check_root_part_on_non_recovery(self, part):
527        """Check the partition number of root device and on normal/dev boot.
528
529        @param part: A string of partition number, e.g.'3'.
530        @return: True if the root device matched and on normal/dev boot;
531                 otherwise, False.
532        """
533        return self.checkers.root_part_checker(part) and \
534                self.checkers.crossystem_checker({
535                    'mainfw_type': ('normal', 'developer'),
536                })
537
538    def _join_part(self, dev, part):
539        """Return a concatenated string of device and partition number.
540
541        @param dev: A string of device, e.g.'/dev/sda'.
542        @param part: A string of partition number, e.g.'3'.
543        @return: A concatenated string of device and partition number,
544                 e.g.'/dev/sda3'.
545
546        >>> seq = FirmwareTest()
547        >>> seq._join_part('/dev/sda', '3')
548        '/dev/sda3'
549        >>> seq._join_part('/dev/mmcblk0', '2')
550        '/dev/mmcblk0p2'
551        """
552        if 'mmcblk' in dev:
553            return dev + 'p' + part
554        else:
555            return dev + part
556
557    def copy_kernel_and_rootfs(self, from_part, to_part):
558        """Copy kernel and rootfs from from_part to to_part.
559
560        @param from_part: A string of partition number to be copied from.
561        @param to_part: A string of partition number to be copied to.
562        """
563        root_dev = self.faft_client.system.get_root_dev()
564        logging.info('Copying kernel from %s to %s. Please wait...',
565                     from_part, to_part)
566        self.faft_client.system.run_shell_command('dd if=%s of=%s bs=4M' %
567                (self._join_part(root_dev, self.KERNEL_MAP[from_part]),
568                 self._join_part(root_dev, self.KERNEL_MAP[to_part])))
569        logging.info('Copying rootfs from %s to %s. Please wait...',
570                     from_part, to_part)
571        self.faft_client.system.run_shell_command('dd if=%s of=%s bs=4M' %
572                (self._join_part(root_dev, self.ROOTFS_MAP[from_part]),
573                 self._join_part(root_dev, self.ROOTFS_MAP[to_part])))
574
575    def ensure_kernel_boot(self, part):
576        """Ensure the request kernel boot.
577
578        If not, it duplicates the current kernel to the requested kernel
579        and sets the requested higher priority to ensure it boot.
580
581        @param part: A string of kernel partition number or 'a'/'b'.
582        """
583        if not self.checkers.root_part_checker(part):
584            if self.faft_client.kernel.diff_a_b():
585                self.copy_kernel_and_rootfs(
586                        from_part=self.OTHER_KERNEL_MAP[part],
587                        to_part=part)
588            self.reset_and_prioritize_kernel(part)
589
590    def set_hardware_write_protect(self, enable):
591        """Set hardware write protect pin.
592
593        @param enable: True if asserting write protect pin. Otherwise, False.
594        """
595        self.servo.set('fw_wp_vref', self.faft_config.wp_voltage)
596        self.servo.set('fw_wp_en', 'on')
597        self.servo.set('fw_wp', 'on' if enable else 'off')
598
599    def set_ec_write_protect_and_reboot(self, enable):
600        """Set EC write protect status and reboot to take effect.
601
602        The write protect state is only activated if both hardware write
603        protect pin is asserted and software write protect flag is set.
604        This method asserts/deasserts hardware write protect pin first, and
605        set corresponding EC software write protect flag.
606
607        If the device uses non-Chrome EC, set the software write protect via
608        flashrom.
609
610        If the device uses Chrome EC, a reboot is required for write protect
611        to take effect. Since the software write protect flag cannot be unset
612        if hardware write protect pin is asserted, we need to deasserted the
613        pin first if we are deactivating write protect. Similarly, a reboot
614        is required before we can modify the software flag.
615
616        @param enable: True if activating EC write protect. Otherwise, False.
617        """
618        self.set_hardware_write_protect(enable)
619        if self.faft_config.chrome_ec:
620            self.set_chrome_ec_write_protect_and_reboot(enable)
621        else:
622            self.faft_client.ec.set_write_protect(enable)
623            self.sync_and_warm_reboot()
624
625    def set_chrome_ec_write_protect_and_reboot(self, enable):
626        """Set Chrome EC write protect status and reboot to take effect.
627
628        @param enable: True if activating EC write protect. Otherwise, False.
629        """
630        if enable:
631            # Set write protect flag and reboot to take effect.
632            self.ec.set_flash_write_protect(enable)
633            self.sync_and_ec_reboot()
634        else:
635            # Reboot after deasserting hardware write protect pin to deactivate
636            # write protect. And then remove software write protect flag.
637            self.sync_and_ec_reboot()
638            self.ec.set_flash_write_protect(enable)
639
640    def _setup_ec_write_protect(self, ec_wp):
641        """Setup for EC write-protection.
642
643        It makes sure the EC in the requested write-protection state. If not, it
644        flips the state. Flipping the write-protection requires DUT reboot.
645
646        @param ec_wp: True to request EC write-protected; False to request EC
647                      not write-protected; None to do nothing.
648        """
649        if ec_wp is None:
650            self._old_ec_wp = None
651            return
652        self._old_ec_wp = self.checkers.crossystem_checker({'wpsw_boot': '1'})
653        if ec_wp != self._old_ec_wp:
654            logging.info('The test required EC is %swrite-protected. Reboot '
655                         'and flip the state.', '' if ec_wp else 'not ')
656            self.do_reboot_action(self.set_ec_write_protect_and_reboot, ec_wp)
657            self.wait_dev_screen_and_ctrl_d()
658
659    def _restore_ec_write_protect(self):
660        """Restore the original EC write-protection."""
661        if (not hasattr(self, '_old_ec_wp')) or (self._old_ec_wp is None):
662            return
663        if not self.checkers.crossystem_checker(
664                {'wpsw_boot': '1' if self._old_ec_wp else '0'}):
665            logging.info('Restore original EC write protection and reboot.')
666            self.do_reboot_action(self.set_ec_write_protect_and_reboot,
667                                  self._old_ec_wp)
668            self.wait_dev_screen_and_ctrl_d()
669
670    def press_ctrl_d(self, press_secs=''):
671        """Send Ctrl-D key to DUT.
672
673        @param press_secs : Str. Time to press key.
674        """
675        self.servo.ctrl_d(press_secs)
676
677    def press_ctrl_u(self):
678        """Send Ctrl-U key to DUT.
679
680        @raise TestError: if a non-Chrome EC device or no Ctrl-U command given
681                          on a no-build-in-keyboard device.
682        """
683        if not self.faft_config.has_keyboard:
684            self.servo.ctrl_u()
685        elif self.check_ec_capability(['keyboard'], suppress_warning=True):
686            self.ec.key_down('<ctrl_l>')
687            self.ec.key_down('u')
688            self.ec.key_up('u')
689            self.ec.key_up('<ctrl_l>')
690        elif self.faft_config.has_keyboard:
691            raise error.TestError(
692                    "Can't send Ctrl-U to DUT without using Chrome EC.")
693        else:
694            raise error.TestError(
695                    "Should specify the ctrl_u_cmd argument.")
696
697    def press_enter(self, press_secs=''):
698        """Send Enter key to DUT.
699
700        @param press_secs: Seconds of holding the key.
701        """
702        self.servo.enter_key(press_secs)
703
704    def wait_dev_screen_and_ctrl_d(self):
705        """Wait for firmware warning screen and press Ctrl-D."""
706        time.sleep(self.faft_config.dev_screen)
707        self.press_ctrl_d()
708
709    def wait_fw_screen_and_ctrl_d(self):
710        """Wait for firmware warning screen and press Ctrl-D."""
711        time.sleep(self.faft_config.firmware_screen)
712        self.press_ctrl_d()
713
714    def wait_fw_screen_and_ctrl_u(self):
715        """Wait for firmware warning screen and press Ctrl-U."""
716        time.sleep(self.faft_config.firmware_screen)
717        self.press_ctrl_u()
718
719    def wait_fw_screen_and_trigger_recovery(self, need_dev_transition=False):
720        """Wait for firmware warning screen and trigger recovery boot.
721
722        @param need_dev_transition: True when needs dev mode transition, only
723                                    for Alex/ZGB.
724        """
725        time.sleep(self.faft_config.firmware_screen)
726
727        # Pressing Enter for too long triggers a second key press.
728        # Let's press it without delay
729        self.press_enter(press_secs=0)
730
731        # For Alex/ZGB, there is a dev warning screen in text mode.
732        # Skip it by pressing Ctrl-D.
733        if need_dev_transition:
734            time.sleep(self.faft_config.legacy_text_screen)
735            self.press_ctrl_d()
736
737    def wait_fw_screen_and_unplug_usb(self):
738        """Wait for firmware warning screen and then unplug the servo USB."""
739        time.sleep(self.faft_config.load_usb)
740        self.servo.switch_usbkey('host')
741        time.sleep(self.faft_config.between_usb_plug)
742
743    def wait_fw_screen_and_plug_usb(self):
744        """Wait for firmware warning screen and then unplug and plug the USB."""
745        self.wait_fw_screen_and_unplug_usb()
746        self.servo.switch_usbkey('dut')
747
748    def wait_fw_screen_and_press_power(self):
749        """Wait for firmware warning screen and press power button."""
750        time.sleep(self.faft_config.firmware_screen)
751        # While the firmware screen, the power button probing loop sleeps
752        # 0.25 second on every scan. Use the normal delay (1.2 second) for
753        # power press.
754        self.servo.power_normal_press()
755
756    def wait_longer_fw_screen_and_press_power(self):
757        """Wait for firmware screen without timeout and press power button."""
758        time.sleep(self.faft_config.dev_screen_timeout)
759        self.wait_fw_screen_and_press_power()
760
761    def wait_fw_screen_and_close_lid(self):
762        """Wait for firmware warning screen and close lid."""
763        time.sleep(self.faft_config.firmware_screen)
764        self.servo.lid_close()
765
766    def wait_longer_fw_screen_and_close_lid(self):
767        """Wait for firmware screen without timeout and close lid."""
768        time.sleep(self.faft_config.firmware_screen)
769        self.wait_fw_screen_and_close_lid()
770
771    def _setup_uart_capture(self):
772        """Setup the CPU/EC UART capture."""
773        self.cpu_uart_file = os.path.join(self.resultsdir, 'cpu_uart.txt')
774        self.servo.set('cpu_uart_capture', 'on')
775        self.ec_uart_file = None
776        if self.faft_config.chrome_ec:
777            try:
778                self.servo.set('ec_uart_capture', 'on')
779                self.ec_uart_file = os.path.join(self.resultsdir, 'ec_uart.txt')
780            except error.TestFail as e:
781                if 'No control named' in str(e):
782                    logging.warn('The servod is too old that ec_uart_capture '
783                                 'not supported.')
784        else:
785            logging.info('Not a Google EC, cannot capture ec console output.')
786
787    def _record_uart_capture(self):
788        """Record the CPU/EC UART output stream to files."""
789        if self.cpu_uart_file:
790            with open(self.cpu_uart_file, 'a') as f:
791                f.write(ast.literal_eval(self.servo.get('cpu_uart_stream')))
792        if self.ec_uart_file and self.faft_config.chrome_ec:
793            with open(self.ec_uart_file, 'a') as f:
794                f.write(ast.literal_eval(self.servo.get('ec_uart_stream')))
795
796    def _cleanup_uart_capture(self):
797        """Cleanup the CPU/EC UART capture."""
798        # Flush the remaining UART output.
799        self._record_uart_capture()
800        self.servo.set('cpu_uart_capture', 'off')
801        if self.ec_uart_file and self.faft_config.chrome_ec:
802            self.servo.set('ec_uart_capture', 'off')
803
804    def _fetch_servo_log(self):
805        """Fetch the servo log."""
806        cmd = '[ -e %s ] && cat %s || echo NOTFOUND' % ((self._SERVOD_LOG,) * 2)
807        servo_log = self.servo.system_output(cmd)
808        return None if servo_log == 'NOTFOUND' else servo_log
809
810    def _setup_servo_log(self):
811        """Setup the servo log capturing."""
812        self.servo_log_original_len = -1
813        if self.servo.is_localhost():
814            # No servo log recorded when servod runs locally.
815            return
816
817        servo_log = self._fetch_servo_log()
818        if servo_log:
819            self.servo_log_original_len = len(servo_log)
820        else:
821            logging.warn('Servo log file not found.')
822
823    def _record_servo_log(self):
824        """Record the servo log to the results directory."""
825        if self.servo_log_original_len != -1:
826            servo_log = self._fetch_servo_log()
827            servo_log_file = os.path.join(self.resultsdir, 'servod.log')
828            with open(servo_log_file, 'a') as f:
829                f.write(servo_log[self.servo_log_original_len:])
830
831    def _record_faft_client_log(self):
832        """Record the faft client log to the results directory."""
833        client_log = self.faft_client.system.dump_log(True)
834        client_log_file = os.path.join(self.resultsdir, 'faft_client.log')
835        with open(client_log_file, 'w') as f:
836            f.write(client_log)
837
838    def _setup_gbb_flags(self):
839        """Setup the GBB flags for FAFT test."""
840        if self.faft_config.gbb_version < 1.1:
841            logging.info('Skip modifying GBB on versions older than 1.1.')
842            return
843
844        if self.check_setup_done('gbb_flags'):
845            return
846
847        self._backup_gbb_flags = self.faft_client.bios.get_gbb_flags()
848
849        logging.info('Set proper GBB flags for test.')
850        self.clear_set_gbb_flags(vboot.GBB_FLAG_DEV_SCREEN_SHORT_DELAY |
851                                 vboot.GBB_FLAG_FORCE_DEV_SWITCH_ON |
852                                 vboot.GBB_FLAG_FORCE_DEV_BOOT_USB |
853                                 vboot.GBB_FLAG_DISABLE_FW_ROLLBACK_CHECK,
854                                 vboot.GBB_FLAG_ENTER_TRIGGERS_TONORM |
855                                 vboot.GBB_FLAG_FAFT_KEY_OVERIDE)
856        self.mark_setup_done('gbb_flags')
857
858    def drop_backup_gbb_flags(self):
859        """Drops the backup GBB flags.
860
861        This can be used when a test intends to permanently change GBB flags.
862        """
863        self._backup_gbb_flags = None
864
865    def _restore_gbb_flags(self):
866        """Restore GBB flags to their original state."""
867        if not self._backup_gbb_flags:
868            return
869        self._write_gbb_flags(self._backup_gbb_flags)
870        self.unmark_setup_done('gbb_flags')
871
872    def setup_tried_fwb(self, tried_fwb):
873        """Setup for fw B tried state.
874
875        It makes sure the system in the requested fw B tried state. If not, it
876        tries to do so.
877
878        @param tried_fwb: True if requested in tried_fwb=1;
879                          False if tried_fwb=0.
880        """
881        if tried_fwb:
882            if not self.checkers.crossystem_checker({'tried_fwb': '1'}):
883                logging.info(
884                    'Firmware is not booted with tried_fwb. Reboot into it.')
885                self.faft_client.system.set_try_fw_b()
886        else:
887            if not self.checkers.crossystem_checker({'tried_fwb': '0'}):
888                logging.info(
889                    'Firmware is booted with tried_fwb. Reboot to clear.')
890
891    def power_on(self):
892        """Switch DUT AC power on."""
893        self._client.power_on(self.power_control)
894
895    def power_off(self):
896        """Switch DUT AC power off."""
897        self._client.power_off(self.power_control)
898
899    def power_cycle(self):
900        """Power cycle DUT AC power."""
901        self._client.power_cycle(self.power_control)
902
903    def enable_rec_mode_and_reboot(self):
904        """Switch to rec mode and reboot.
905
906        This method emulates the behavior of the old physical recovery switch,
907        i.e. switch ON + reboot + switch OFF, and the new keyboard controlled
908        recovery mode, i.e. just press Power + Esc + Refresh.
909        """
910        if self.faft_config.chrome_ec:
911            # Reset twice to emulate a long recovery-key-combo hold.
912            cold_reset_num = 2 if self.faft_config.long_rec_combo else 1
913            for i in range(cold_reset_num):
914                if i:
915                    time.sleep(self.faft_config.ec_boot_to_console)
916                # Cold reset to clear EC_IN_RW signal
917                self.servo.set('cold_reset', 'on')
918                time.sleep(self.faft_config.hold_cold_reset)
919                self.servo.set('cold_reset', 'off')
920            time.sleep(self.faft_config.ec_boot_to_console)
921            self.ec.reboot("ap-off")
922            time.sleep(self.faft_config.ec_boot_to_console)
923            self.ec.set_hostevent(chrome_ec.HOSTEVENT_KEYBOARD_RECOVERY)
924            self.servo.power_short_press()
925        elif self.faft_config.broken_rec_mode:
926            self.power_cycle()
927            logging.info('Booting to recovery mode.')
928            self.servo.custom_recovery_mode()
929        else:
930            self.servo.enable_recovery_mode()
931            self.reboot_cold_trigger()
932            time.sleep(self.faft_config.ec_boot_to_console)
933            self.servo.disable_recovery_mode()
934
935    def enable_dev_mode_and_reboot(self):
936        """Switch to developer mode and reboot."""
937        if self.faft_config.keyboard_dev:
938            self.enable_keyboard_dev_mode()
939        else:
940            self.servo.enable_development_mode()
941            self.faft_client.system.run_shell_command(
942                    'chromeos-firmwareupdate --mode todev && reboot')
943
944    def enable_normal_mode_and_reboot(self):
945        """Switch to normal mode and reboot."""
946        if self.faft_config.keyboard_dev:
947            self.disable_keyboard_dev_mode()
948        else:
949            self.servo.disable_development_mode()
950            self.faft_client.system.run_shell_command(
951                    'chromeos-firmwareupdate --mode tonormal && reboot')
952
953    def wait_fw_screen_and_switch_keyboard_dev_mode(self, dev):
954        """Wait for firmware screen and then switch into or out of dev mode.
955
956        @param dev: True if switching into dev mode. Otherwise, False.
957        """
958        time.sleep(self.faft_config.firmware_screen)
959        if dev:
960            self.press_ctrl_d()
961            time.sleep(self.faft_config.confirm_screen)
962            if self.faft_config.rec_button_dev_switch:
963                logging.info('RECOVERY button pressed to switch to dev mode')
964                self.servo.set('rec_mode', 'on')
965                time.sleep(self.faft_config.hold_cold_reset)
966                self.servo.set('rec_mode', 'off')
967            else:
968                logging.info('ENTER pressed to switch to dev mode')
969                self.press_enter()
970        else:
971            self.press_enter()
972            time.sleep(self.faft_config.confirm_screen)
973            self.press_enter()
974
975    def enable_keyboard_dev_mode(self):
976        """Enable keyboard controlled developer mode"""
977        logging.info("Enabling keyboard controlled developer mode")
978        # Plug out USB disk for preventing recovery boot without warning
979        self.servo.switch_usbkey('host')
980        # Rebooting EC with rec mode on. Should power on AP.
981        self.enable_rec_mode_and_reboot()
982        self.wait_for_client_offline()
983        self.wait_fw_screen_and_switch_keyboard_dev_mode(dev=True)
984
985        # TODO (crosbug.com/p/16231) remove this conditional completely if/when
986        # issue is resolved.
987        if self.faft_config.platform == 'Parrot':
988            self.wait_for_client_offline()
989            self.reboot_cold_trigger()
990
991    def disable_keyboard_dev_mode(self):
992        """Disable keyboard controlled developer mode"""
993        logging.info("Disabling keyboard controlled developer mode")
994        if (not self.faft_config.chrome_ec and
995            not self.faft_config.broken_rec_mode):
996            self.servo.disable_recovery_mode()
997        self.reboot_cold_trigger()
998        self.wait_for_client_offline()
999        self.wait_fw_screen_and_switch_keyboard_dev_mode(dev=False)
1000
1001    def setup_dev_mode(self, dev_mode):
1002        """Setup for development mode.
1003
1004        It makes sure the system in the requested normal/dev mode. If not, it
1005        tries to do so.
1006
1007        @param dev_mode: True if requested in dev mode; False if normal mode.
1008        """
1009        if dev_mode:
1010            if (not self.faft_config.keyboard_dev and
1011                not self.checkers.crossystem_checker({'devsw_cur': '1'})):
1012                logging.info('Dev switch is not on. Now switch it on.')
1013                self.servo.enable_development_mode()
1014            if not self.checkers.crossystem_checker({'devsw_boot': '1',
1015                    'mainfw_type': 'developer'}):
1016                logging.info('System is not in dev mode. Reboot into it.')
1017                if self._backup_dev_mode is None:
1018                    self._backup_dev_mode = False
1019                if self.faft_config.keyboard_dev:
1020                    self.faft_client.system.run_shell_command(
1021                             'chromeos-firmwareupdate --mode todev && reboot')
1022                    self.do_reboot_action(self.enable_keyboard_dev_mode)
1023                    # TOOD : Add a ctrl_d to speed through the dev screen.
1024        else:
1025            if (not self.faft_config.keyboard_dev and
1026                not self.checkers.crossystem_checker({'devsw_cur': '0'})):
1027                logging.info('Dev switch is not off. Now switch it off.')
1028                self.servo.disable_development_mode()
1029            if not self.checkers.crossystem_checker({'devsw_boot': '0',
1030                    'mainfw_type': 'normal'}):
1031                logging.info('System is not in normal mode. Reboot into it.')
1032                if self._backup_dev_mode is None:
1033                    self._backup_dev_mode = True
1034                if self.faft_config.keyboard_dev:
1035                    self.faft_client.system.run_shell_command(
1036                        'chromeos-firmwareupdate --mode tonormal && reboot')
1037                    self.do_reboot_action(self.disable_keyboard_dev_mode)
1038
1039    def _restore_dev_mode(self):
1040        """Restores original dev mode status if it has changed."""
1041        if self._backup_dev_mode is not None:
1042            self.setup_dev_mode(self._backup_dev_mode)
1043            self._backup_dev_mode = None
1044
1045    def setup_rw_boot(self, section='a'):
1046        """Make sure firmware is in RW-boot mode.
1047
1048        If the given firmware section is in RO-boot mode, turn off the RO-boot
1049        flag and reboot DUT into RW-boot mode.
1050
1051        @param section: A firmware section, either 'a' or 'b'.
1052        """
1053        flags = self.faft_client.bios.get_preamble_flags(section)
1054        if flags & vboot.PREAMBLE_USE_RO_NORMAL:
1055            flags = flags ^ vboot.PREAMBLE_USE_RO_NORMAL
1056            self.faft_client.bios.set_preamble_flags(section, flags)
1057            self.reboot_warm()
1058
1059    def setup_kernel(self, part):
1060        """Setup for kernel test.
1061
1062        It makes sure both kernel A and B bootable and the current boot is
1063        the requested kernel part.
1064
1065        @param part: A string of kernel partition number or 'a'/'b'.
1066        """
1067        self.ensure_kernel_boot(part)
1068        logging.info('Checking the integrity of kernel B and rootfs B...')
1069        if (self.faft_client.kernel.diff_a_b() or
1070                not self.faft_client.rootfs.verify_rootfs('B')):
1071            logging.info('Copying kernel and rootfs from A to B...')
1072            self.copy_kernel_and_rootfs(from_part=part,
1073                                        to_part=self.OTHER_KERNEL_MAP[part])
1074        self.reset_and_prioritize_kernel(part)
1075
1076    def reset_and_prioritize_kernel(self, part):
1077        """Make the requested partition highest priority.
1078
1079        This function also reset kerenl A and B to bootable.
1080
1081        @param part: A string of partition number to be prioritized.
1082        """
1083        root_dev = self.faft_client.system.get_root_dev()
1084        # Reset kernel A and B to bootable.
1085        self.faft_client.system.run_shell_command(
1086            'cgpt add -i%s -P1 -S1 -T0 %s' % (self.KERNEL_MAP['a'], root_dev))
1087        self.faft_client.system.run_shell_command(
1088            'cgpt add -i%s -P1 -S1 -T0 %s' % (self.KERNEL_MAP['b'], root_dev))
1089        # Set kernel part highest priority.
1090        self.faft_client.system.run_shell_command('cgpt prioritize -i%s %s' %
1091                (self.KERNEL_MAP[part], root_dev))
1092
1093
1094    ################################################
1095    # Reboot APIs
1096
1097    def reboot_warm(self, sync_before_boot=True, wait_for_dut_up=True):
1098        """
1099        Perform a warm reboot.
1100
1101        This is the highest level function that most users will need.
1102        It performs a sync, triggers a reboot and waits for kernel to boot.
1103
1104        @param sync_before_boot: bool, sync to disk before booting.
1105        @param wait_for_dut_up: bool, wait for dut to boot before returning.
1106        """
1107        if sync_before_boot:
1108            self.faft_client.system.run_shell_command('sync')
1109            time.sleep(self.faft_config.sync)
1110        self.reboot_warm_trigger()
1111        if wait_for_dut_up:
1112            self.wait_for_client_offline()
1113            self.wait_for_kernel_up()
1114
1115    def reboot_cold(self, sync_before_boot=True, wait_for_dut_up=True):
1116        """
1117        Perform a cold reboot.
1118
1119        This is the highest level function that most users will need.
1120        It performs a sync, triggers a reboot and waits for kernel to boot.
1121
1122        @param sync_before_boot: bool, sync to disk before booting.
1123        @param wait_for_dut_up: bool, wait for dut to boot before returning.
1124        """
1125        if sync_before_boot:
1126            self.faft_client.system.run_shell_command('sync')
1127            time.sleep(self.faft_config.sync)
1128        self.reboot_cold_trigger()
1129        if wait_for_dut_up:
1130            self.wait_for_client_offline()
1131            self.wait_for_kernel_up()
1132
1133    def do_reboot_action(self, func):
1134        """
1135        Helper function that wraps the reboot function so that we check if the
1136        DUT went down.
1137
1138        @param func: function to trigger the reboot.
1139        """
1140        logging.info("-[FAFT]-[ start do_reboot_action ]----------")
1141        boot_id = self.get_bootid()
1142        self._call_action(func)
1143        self.wait_for_client_offline(orig_boot_id=boot_id)
1144        logging.info("-[FAFT]-[ end do_reboot_action ]------------")
1145
1146    def wait_for_kernel_up(self, install_deps=False):
1147        """
1148        Helper function that waits for the device to boot up to kernel.
1149
1150        @param install_deps: bool, install deps after boot.
1151        """
1152        logging.info("-[FAFT]-[ start wait_for_kernel_up ]---")
1153        try:
1154            logging.info("Installing deps after boot : %s" % install_deps)
1155            self.wait_for_client(install_deps=install_deps)
1156            # Stop update-engine as it may change firmware/kernel.
1157            self._stop_service('update-engine')
1158        except ConnectionError:
1159            logging.error('wait_for_client() timed out.')
1160            self._restore_routine_from_timeout(next_step)
1161        logging.info("-[FAFT]-[ end wait_for_kernel_up ]-----")
1162
1163    def reboot_warm_trigger(self):
1164        """Request a warm reboot.
1165
1166        A wrapper for underlying servo warm reset.
1167        """
1168        # Use cold reset if the warm reset is broken.
1169        if self.faft_config.broken_warm_reset:
1170            logging.info('broken_warm_reset is True. Cold rebooting instead.')
1171            self.reboot_cold_trigger()
1172        else:
1173            self.servo.get_power_state_controller().warm_reset()
1174
1175    def reboot_cold_trigger(self):
1176        """Request a cold reboot.
1177
1178        A wrapper for underlying servo cold reset.
1179        """
1180        if self.faft_config.broken_warm_reset:
1181            self.servo.set('pwr_button', 'press')
1182            self.servo.set('cold_reset', 'on')
1183            self.servo.set('cold_reset', 'off')
1184            time.sleep(self.faft_config.ec_boot_to_pwr_button)
1185            self.servo.set('pwr_button', 'release')
1186        else:
1187            self.servo.get_power_state_controller().cold_reset()
1188
1189    def sync_and_warm_reboot(self):
1190        """Request the client sync and do a warm reboot.
1191
1192        This is the default reboot action on FAFT.
1193        """
1194        self.faft_client.system.run_shell_command('sync')
1195        time.sleep(self.faft_config.sync)
1196        self.reboot_warm_trigger()
1197
1198    def sync_and_cold_reboot(self):
1199        """Request the client sync and do a cold reboot.
1200
1201        This reboot action is used to reset EC for recovery mode.
1202        """
1203        self.faft_client.system.run_shell_command('sync')
1204        time.sleep(self.faft_config.sync)
1205        self.reboot_cold_trigger()
1206
1207    def sync_and_ec_reboot(self, flags=''):
1208        """Request the client sync and do a EC triggered reboot.
1209
1210        @param flags: Optional, a space-separated string of flags passed to EC
1211                      reboot command, including:
1212                          default: EC soft reboot;
1213                          'hard': EC cold/hard reboot.
1214        """
1215        self.faft_client.system.run_shell_command('sync')
1216        time.sleep(self.faft_config.sync)
1217        self.ec.reboot(flags)
1218        time.sleep(self.faft_config.ec_boot_to_console)
1219        self.check_lid_and_power_on()
1220
1221    def reboot_with_factory_install_shim(self):
1222        """Request reboot with factory install shim to reset TPM.
1223
1224        Factory install shim requires dev mode enabled. So this method switches
1225        firmware to dev mode first and reboot. The client uses factory install
1226        shim to reset TPM values.
1227        """
1228        # Unplug USB first to avoid the complicated USB autoboot cases.
1229        self.servo.switch_usbkey('host')
1230        is_dev = self.checkers.crossystem_checker({'devsw_boot': '1'})
1231        if not is_dev:
1232            self.enable_dev_mode_and_reboot()
1233        time.sleep(self.faft_config.sync)
1234        self.enable_rec_mode_and_reboot()
1235        self.wait_fw_screen_and_plug_usb()
1236        time.sleep(self.faft_config.install_shim_done)
1237        self.reboot_warm_trigger()
1238
1239    def full_power_off_and_on(self):
1240        """Shutdown the device by pressing power button and power on again."""
1241        # Press power button to trigger Chrome OS normal shutdown process.
1242        # We use a customized delay since the normal-press 1.2s is not enough.
1243        self.servo.power_key(self.faft_config.hold_pwr_button)
1244        time.sleep(self.faft_config.shutdown)
1245        # Short press power button to boot DUT again.
1246        self.servo.power_short_press()
1247
1248    def check_lid_and_power_on(self):
1249        """
1250        On devices with EC software sync, system powers on after EC reboots if
1251        lid is open. Otherwise, the EC shuts down CPU after about 3 seconds.
1252        This method checks lid switch state and presses power button if
1253        necessary.
1254        """
1255        if self.servo.get("lid_open") == "no":
1256            time.sleep(self.faft_config.software_sync)
1257            self.servo.power_short_press()
1258
1259    def _modify_usb_kernel(self, usb_dev, from_magic, to_magic):
1260        """Modify the kernel header magic in USB stick.
1261
1262        The kernel header magic is the first 8-byte of kernel partition.
1263        We modify it to make it fail on kernel verification check.
1264
1265        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
1266        @param from_magic: A string of magic which we change it from.
1267        @param to_magic: A string of magic which we change it to.
1268        @raise TestError: if failed to change magic.
1269        """
1270        assert len(from_magic) == 8
1271        assert len(to_magic) == 8
1272        # USB image only contains one kernel.
1273        kernel_part = self._join_part(usb_dev, self.KERNEL_MAP['a'])
1274        read_cmd = "sudo dd if=%s bs=8 count=1 2>/dev/null" % kernel_part
1275        current_magic = self.servo.system_output(read_cmd)
1276        if current_magic == to_magic:
1277            logging.info("The kernel magic is already %s.", current_magic)
1278            return
1279        if current_magic != from_magic:
1280            raise error.TestError("Invalid kernel image on USB: wrong magic.")
1281
1282        logging.info('Modify the kernel magic in USB, from %s to %s.',
1283                     from_magic, to_magic)
1284        write_cmd = ("echo -n '%s' | sudo dd of=%s oflag=sync conv=notrunc "
1285                     " 2>/dev/null" % (to_magic, kernel_part))
1286        self.servo.system(write_cmd)
1287
1288        if self.servo.system_output(read_cmd) != to_magic:
1289            raise error.TestError("Failed to write new magic.")
1290
1291    def corrupt_usb_kernel(self, usb_dev):
1292        """Corrupt USB kernel by modifying its magic from CHROMEOS to CORRUPTD.
1293
1294        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
1295        """
1296        self._modify_usb_kernel(usb_dev, self.CHROMEOS_MAGIC,
1297                                self.CORRUPTED_MAGIC)
1298
1299    def restore_usb_kernel(self, usb_dev):
1300        """Restore USB kernel by modifying its magic from CORRUPTD to CHROMEOS.
1301
1302        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
1303        """
1304        self._modify_usb_kernel(usb_dev, self.CORRUPTED_MAGIC,
1305                                self.CHROMEOS_MAGIC)
1306
1307    def _call_action(self, action_tuple, check_status=False):
1308        """Call the action function with/without arguments.
1309
1310        @param action_tuple: A function, or a tuple (function, args, error_msg),
1311                             in which, args and error_msg are optional. args is
1312                             either a value or a tuple if multiple arguments.
1313                             This can also be a list containing multiple
1314                             function or tuple. In this case, these actions are
1315                             called in sequence.
1316        @param check_status: Check the return value of action function. If not
1317                             succeed, raises a TestFail exception.
1318        @return: The result value of the action function.
1319        @raise TestError: An error when the action function is not callable.
1320        @raise TestFail: When check_status=True, action function not succeed.
1321        """
1322        if isinstance(action_tuple, list):
1323            return all([self._call_action(action, check_status=check_status)
1324                        for action in action_tuple])
1325
1326        action = action_tuple
1327        args = ()
1328        error_msg = 'Not succeed'
1329        if isinstance(action_tuple, tuple):
1330            action = action_tuple[0]
1331            if len(action_tuple) >= 2:
1332                args = action_tuple[1]
1333                if not isinstance(args, tuple):
1334                    args = (args,)
1335            if len(action_tuple) >= 3:
1336                error_msg = action_tuple[2]
1337
1338        if action is None:
1339            return
1340
1341        if not callable(action):
1342            raise error.TestError('action is not callable!')
1343
1344        info_msg = 'calling %s' % str(action)
1345        if args:
1346            info_msg += ' with args %s' % str(args)
1347        logging.info(info_msg)
1348        ret = action(*args)
1349
1350        if check_status and not ret:
1351            raise error.TestFail('%s: %s returning %s' %
1352                                 (error_msg, info_msg, str(ret)))
1353        return ret
1354
1355    def run_shutdown_process(self, shutdown_action, pre_power_action=None,
1356            post_power_action=None, shutdown_timeout=None):
1357        """Run shutdown_action(), which makes DUT shutdown, and power it on.
1358
1359        @param shutdown_action: function which makes DUT shutdown, like
1360                                pressing power key.
1361        @param pre_power_action: function which is called before next power on.
1362        @param post_power_action: function which is called after next power on.
1363        @param shutdown_timeout: a timeout to confirm DUT shutdown.
1364        @raise TestFail: if the shutdown_action() failed to turn DUT off.
1365        """
1366        self._call_action(shutdown_action)
1367        logging.info('Wait to ensure DUT shut down...')
1368        try:
1369            if shutdown_timeout is None:
1370                shutdown_timeout = self.faft_config.shutdown_timeout
1371            self.wait_for_client(timeout=shutdown_timeout)
1372            raise error.TestFail(
1373                    'Should shut the device down after calling %s.' %
1374                    str(shutdown_action))
1375        except ConnectionError:
1376            logging.info(
1377                'DUT is surely shutdown. We are going to power it on again...')
1378
1379        if pre_power_action:
1380            self._call_action(pre_power_action)
1381        self.servo.power_short_press()
1382        if post_power_action:
1383            self._call_action(post_power_action)
1384
1385    def get_bootid(self, retry=3):
1386        """
1387        Return the bootid.
1388        """
1389        boot_id = None
1390        while retry:
1391            try:
1392                boot_id = self._client.get_boot_id()
1393                break
1394            except error.AutoservRunError:
1395                retry -= 1
1396                if retry:
1397                    logging.info('Retry to get boot_id...')
1398                else:
1399                    logging.warning('Failed to get boot_id.')
1400        logging.info('boot_id: %s', boot_id)
1401        return boot_id
1402
1403    def check_state(self, func):
1404        """
1405        Wrapper around _call_action with check_status set to True. This is a
1406        helper function to be used by tests and is currently implemented by
1407        calling _call_action with check_status=True.
1408
1409        TODO: This function's arguments need to be made more stringent. And
1410        its functionality should be moved over to check functions directly in
1411        the future.
1412
1413        @param func: A function, or a tuple (function, args, error_msg),
1414                             in which, args and error_msg are optional. args is
1415                             either a value or a tuple if multiple arguments.
1416                             This can also be a list containing multiple
1417                             function or tuple. In this case, these actions are
1418                             called in sequence.
1419        @return: The result value of the action function.
1420        @raise TestFail: If the function does notsucceed.
1421        """
1422        logging.info("-[FAFT]-[ start stepstate_checker ]----------")
1423        self._call_action(func, check_status=True)
1424        logging.info("-[FAFT]-[ end state_checker ]----------------")
1425
1426    def get_current_firmware_sha(self):
1427        """Get current firmware sha of body and vblock.
1428
1429        @return: Current firmware sha follows the order (
1430                 vblock_a_sha, body_a_sha, vblock_b_sha, body_b_sha)
1431        """
1432        current_firmware_sha = (self.faft_client.bios.get_sig_sha('a'),
1433                                self.faft_client.bios.get_body_sha('a'),
1434                                self.faft_client.bios.get_sig_sha('b'),
1435                                self.faft_client.bios.get_body_sha('b'))
1436        if not all(current_firmware_sha):
1437            raise error.TestError('Failed to get firmware sha.')
1438        return current_firmware_sha
1439
1440    def is_firmware_changed(self):
1441        """Check if the current firmware changed, by comparing its SHA.
1442
1443        @return: True if it is changed, otherwise Flase.
1444        """
1445        # Device may not be rebooted after test.
1446        self.faft_client.bios.reload()
1447
1448        current_sha = self.get_current_firmware_sha()
1449
1450        if current_sha == self._backup_firmware_sha:
1451            return False
1452        else:
1453            corrupt_VBOOTA = (current_sha[0] != self._backup_firmware_sha[0])
1454            corrupt_FVMAIN = (current_sha[1] != self._backup_firmware_sha[1])
1455            corrupt_VBOOTB = (current_sha[2] != self._backup_firmware_sha[2])
1456            corrupt_FVMAINB = (current_sha[3] != self._backup_firmware_sha[3])
1457            logging.info("Firmware changed:")
1458            logging.info('VBOOTA is changed: %s', corrupt_VBOOTA)
1459            logging.info('VBOOTB is changed: %s', corrupt_VBOOTB)
1460            logging.info('FVMAIN is changed: %s', corrupt_FVMAIN)
1461            logging.info('FVMAINB is changed: %s', corrupt_FVMAINB)
1462            return True
1463
1464    def backup_firmware(self, suffix='.original'):
1465        """Backup firmware to file, and then send it to host.
1466
1467        @param suffix: a string appended to backup file name
1468        """
1469        remote_temp_dir = self.faft_client.system.create_temp_dir()
1470        self.faft_client.bios.dump_whole(os.path.join(remote_temp_dir, 'bios'))
1471        self._client.get_file(os.path.join(remote_temp_dir, 'bios'),
1472                              os.path.join(self.resultsdir, 'bios' + suffix))
1473
1474        self._backup_firmware_sha = self.get_current_firmware_sha()
1475        logging.info('Backup firmware stored in %s with suffix %s',
1476            self.resultsdir, suffix)
1477
1478    def is_firmware_saved(self):
1479        """Check if a firmware saved (called backup_firmware before).
1480
1481        @return: True if the firmware is backuped; otherwise False.
1482        """
1483        return self._backup_firmware_sha != ()
1484
1485    def clear_saved_firmware(self):
1486        """Clear the firmware saved by the method backup_firmware."""
1487        self._backup_firmware_sha = ()
1488
1489    def restore_firmware(self, suffix='.original'):
1490        """Restore firmware from host in resultsdir.
1491
1492        @param suffix: a string appended to backup file name
1493        """
1494        if not self.is_firmware_changed():
1495            return
1496
1497        # Backup current corrupted firmware.
1498        self.backup_firmware(suffix='.corrupt')
1499
1500        # Restore firmware.
1501        remote_temp_dir = self.faft_client.system.create_temp_dir()
1502        self._client.send_file(os.path.join(self.resultsdir, 'bios' + suffix),
1503                               os.path.join(remote_temp_dir, 'bios'))
1504
1505        self.faft_client.bios.write_whole(
1506            os.path.join(remote_temp_dir, 'bios'))
1507        self.sync_and_warm_reboot()
1508        self.wait_for_client_offline()
1509        self.wait_dev_screen_and_ctrl_d()
1510        self.wait_for_client()
1511
1512        logging.info('Successfully restore firmware.')
1513
1514    def setup_firmwareupdate_shellball(self, shellball=None):
1515        """Deside a shellball to use in firmware update test.
1516
1517        Check if there is a given shellball, and it is a shell script. Then,
1518        send it to the remote host. Otherwise, use
1519        /usr/sbin/chromeos-firmwareupdate.
1520
1521        @param shellball: path of a shellball or default to None.
1522
1523        @return: Path of shellball in remote host. If use default shellball,
1524                 reutrn None.
1525        """
1526        updater_path = None
1527        if shellball:
1528            # Determine the firmware file is a shellball or a raw binary.
1529            is_shellball = (utils.system_output("file %s" % shellball).find(
1530                    "shell script") != -1)
1531            if is_shellball:
1532                logging.info('Device will update firmware with shellball %s',
1533                             shellball)
1534                temp_dir = self.faft_client.system.create_temp_dir(
1535                            'shellball_')
1536                temp_shellball = os.path.join(temp_dir, 'updater.sh')
1537                self._client.send_file(shellball, temp_shellball)
1538                updater_path = temp_shellball
1539            else:
1540                raise error.TestFail(
1541                    'The given shellball is not a shell script.')
1542            return updater_path
1543
1544    def is_kernel_changed(self):
1545        """Check if the current kernel is changed, by comparing its SHA1 hash.
1546
1547        @return: True if it is changed; otherwise, False.
1548        """
1549        changed = False
1550        for p in ('A', 'B'):
1551            backup_sha = self._backup_kernel_sha.get(p, None)
1552            current_sha = self.faft_client.kernel.get_sha(p)
1553            if backup_sha != current_sha:
1554                changed = True
1555                logging.info('Kernel %s is changed', p)
1556        return changed
1557
1558    def backup_kernel(self, suffix='.original'):
1559        """Backup kernel to files, and the send them to host.
1560
1561        @param suffix: a string appended to backup file name.
1562        """
1563        remote_temp_dir = self.faft_client.system.create_temp_dir()
1564        for p in ('A', 'B'):
1565            remote_path = os.path.join(remote_temp_dir, 'kernel_%s' % p)
1566            self.faft_client.kernel.dump(p, remote_path)
1567            self._client.get_file(
1568                    remote_path,
1569                    os.path.join(self.resultsdir, 'kernel_%s%s' % (p, suffix)))
1570            self._backup_kernel_sha[p] = self.faft_client.kernel.get_sha(p)
1571        logging.info('Backup kernel stored in %s with suffix %s',
1572            self.resultsdir, suffix)
1573
1574    def is_kernel_saved(self):
1575        """Check if kernel images are saved (backup_kernel called before).
1576
1577        @return: True if the kernel is saved; otherwise, False.
1578        """
1579        return len(self._backup_kernel_sha) != 0
1580
1581    def clear_saved_kernel(self):
1582        """Clear the kernel saved by backup_kernel()."""
1583        self._backup_kernel_sha = dict()
1584
1585    def restore_kernel(self, suffix='.original'):
1586        """Restore kernel from host in resultsdir.
1587
1588        @param suffix: a string appended to backup file name.
1589        """
1590        if not self.is_kernel_changed():
1591            return
1592
1593        # Backup current corrupted kernel.
1594        self.backup_kernel(suffix='.corrupt')
1595
1596        # Restore kernel.
1597        remote_temp_dir = self.faft_client.system.create_temp_dir()
1598        for p in ('A', 'B'):
1599            remote_path = os.path.join(remote_temp_dir, 'kernel_%s' % p)
1600            self._client.send_file(
1601                    os.path.join(self.resultsdir, 'kernel_%s%s' % (p, suffix)),
1602                    remote_path)
1603            self.faft_client.kernel.write(p, remote_path)
1604
1605        self.sync_and_warm_reboot()
1606        self.wait_for_client_offline()
1607        self.wait_dev_screen_and_ctrl_d()
1608        self.wait_for_client()
1609
1610        logging.info('Successfully restored kernel.')
1611
1612    def backup_cgpt_attributes(self):
1613        """Backup CGPT partition table attributes."""
1614        self._backup_cgpt_attr = self.faft_client.cgpt.get_attributes()
1615
1616    def restore_cgpt_attributes(self):
1617        """Restore CGPT partition table attributes."""
1618        current_table = self.faft_client.cgpt.get_attributes()
1619        if current_table == self._backup_cgpt_attr:
1620            return
1621        logging.info('CGPT table is changed. Original: %r. Current: %r.',
1622                     self._backup_cgpt_attr,
1623                     current_table)
1624        self.faft_client.cgpt.set_attributes(self._backup_cgpt_attr)
1625
1626        self.sync_and_warm_reboot()
1627        self.wait_for_client_offline()
1628        self.wait_dev_screen_and_ctrl_d()
1629        self.wait_for_client()
1630
1631        logging.info('Successfully restored CGPT table.')
1632