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