• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 pprint
10import re
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 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 import mode_switcher
21from autotest_lib.server.cros.faft.utils.faft_checkers import FAFTCheckers
22from autotest_lib.server.cros.servo import chrome_cr50, chrome_ec
23
24ConnectionError = mode_switcher.ConnectionError
25
26
27class FAFTBase(test.test):
28    """The base class of FAFT classes.
29
30    It launches the FAFTClient on DUT, such that the test can access its
31    firmware functions and interfaces. It also provides some methods to
32    handle the reboot mechanism, in order to ensure FAFTClient is still
33    connected after reboot.
34    """
35    def initialize(self, host):
36        """Create a FAFTClient object and install the dependency."""
37        self.servo = host.servo
38        self.servo.initialize_dut()
39        self._client = host
40        self.faft_client = RPCProxy(host)
41        self.lockfile = '/var/tmp/faft/lock'
42
43
44class FirmwareTest(FAFTBase):
45    """
46    Base class that sets up helper objects/functions for firmware tests.
47
48    TODO: add documentaion as the FAFT rework progresses.
49    """
50    version = 1
51
52    # Mapping of partition number of kernel and rootfs.
53    KERNEL_MAP = {'a':'2', 'b':'4', '2':'2', '4':'4', '3':'2', '5':'4'}
54    ROOTFS_MAP = {'a':'3', 'b':'5', '2':'3', '4':'5', '3':'3', '5':'5'}
55    OTHER_KERNEL_MAP = {'a':'4', 'b':'2', '2':'4', '4':'2', '3':'4', '5':'2'}
56    OTHER_ROOTFS_MAP = {'a':'5', 'b':'3', '2':'5', '4':'3', '3':'5', '5':'3'}
57
58    CHROMEOS_MAGIC = "CHROMEOS"
59    CORRUPTED_MAGIC = "CORRUPTD"
60
61    # Delay for waiting client to return before EC suspend
62    EC_SUSPEND_DELAY = 5
63
64    # Delay between EC suspend and wake
65    WAKE_DELAY = 10
66
67    # Delay between closing and opening lid
68    LID_DELAY = 1
69
70    _SERVOD_LOG = '/var/log/servod.log'
71
72    _ROOTFS_PARTITION_NUMBER = 3
73
74    _backup_firmware_sha = ()
75    _backup_kernel_sha = dict()
76    _backup_cgpt_attr = dict()
77    _backup_gbb_flags = None
78    _backup_dev_mode = None
79
80    # Class level variable, keep track the states of one time setup.
81    # This variable is preserved across tests which inherit this class.
82    _global_setup_done = {
83        'gbb_flags': False,
84        'reimage': False,
85        'usb_check': False,
86    }
87
88    @classmethod
89    def check_setup_done(cls, label):
90        """Check if the given setup is done.
91
92        @param label: The label of the setup.
93        """
94        return cls._global_setup_done[label]
95
96    @classmethod
97    def mark_setup_done(cls, label):
98        """Mark the given setup done.
99
100        @param label: The label of the setup.
101        """
102        cls._global_setup_done[label] = True
103
104    @classmethod
105    def unmark_setup_done(cls, label):
106        """Mark the given setup not done.
107
108        @param label: The label of the setup.
109        """
110        cls._global_setup_done[label] = False
111
112    def initialize(self, host, cmdline_args, ec_wp=None):
113        super(FirmwareTest, self).initialize(host)
114        self.run_id = str(uuid.uuid4())
115        logging.info('FirmwareTest initialize begin (id=%s)', self.run_id)
116        # Parse arguments from command line
117        args = {}
118        self.power_control = host.POWER_CONTROL_RPM
119        for arg in cmdline_args:
120            match = re.search("^(\w+)=(.+)", arg)
121            if match:
122                args[match.group(1)] = match.group(2)
123        if 'power_control' in args:
124            self.power_control = args['power_control']
125            if self.power_control not in host.POWER_CONTROL_VALID_ARGS:
126                raise error.TestError('Valid values for --args=power_control '
127                                      'are %s. But you entered wrong argument '
128                                      'as "%s".'
129                                       % (host.POWER_CONTROL_VALID_ARGS,
130                                       self.power_control))
131
132        self.faft_config = FAFTConfig(
133                self.faft_client.system.get_platform_name())
134        self.checkers = FAFTCheckers(self)
135        self.switcher = mode_switcher.create_mode_switcher(self)
136
137        if self.faft_config.chrome_ec:
138            self.ec = chrome_ec.ChromeEC(self.servo)
139        # Check for presence of a USBPD console
140        if self.faft_config.chrome_usbpd:
141            self.usbpd = chrome_ec.ChromeUSBPD(self.servo)
142        elif self.faft_config.chrome_ec:
143            # If no separate USBPD console, then PD exists on EC console
144            self.usbpd = self.ec
145        # Get plankton console
146        self.plankton = host.plankton
147        self.plankton_host = host._plankton_host
148
149        self._setup_uart_capture()
150        self._setup_servo_log()
151        self._record_system_info()
152        self.fw_vboot2 = self.faft_client.system.get_fw_vboot2()
153        logging.info('vboot version: %d', 2 if self.fw_vboot2 else 1)
154        if self.fw_vboot2:
155            self.faft_client.system.set_fw_try_next('A')
156            if self.faft_client.system.get_crossystem_value('mainfw_act') == 'B':
157                logging.info('mainfw_act is B. rebooting to set it A')
158                self.switcher.mode_aware_reboot()
159        self._setup_gbb_flags()
160        self._stop_service('update-engine')
161        self._create_faft_lockfile()
162        self._setup_ec_write_protect(ec_wp)
163        # See chromium:239034 regarding needing this sync.
164        self.blocking_sync()
165        logging.info('FirmwareTest initialize done (id=%s)', self.run_id)
166
167    def cleanup(self):
168        """Autotest cleanup function."""
169        # Unset state checker in case it's set by subclass
170        logging.info('FirmwareTest cleaning up (id=%s)', self.run_id)
171        try:
172            self.faft_client.system.is_available()
173        except:
174            # Remote is not responding. Revive DUT so that subsequent tests
175            # don't fail.
176            self._restore_routine_from_timeout()
177        self.switcher.restore_mode()
178        self._restore_ec_write_protect()
179        self._restore_gbb_flags()
180        self._start_service('update-engine')
181        self._remove_faft_lockfile()
182        self._record_servo_log()
183        self._record_faft_client_log()
184        self._cleanup_uart_capture()
185        super(FirmwareTest, self).cleanup()
186        logging.info('FirmwareTest cleanup done (id=%s)', self.run_id)
187
188    def _record_system_info(self):
189        """Record some critical system info to the attr keyval.
190
191        This info is used by generate_test_report later.
192        """
193        system_info = {
194            'hwid': self.faft_client.system.get_crossystem_value('hwid'),
195            'ec_version': self.faft_client.ec.get_version(),
196            'ro_fwid': self.faft_client.system.get_crossystem_value('ro_fwid'),
197            'rw_fwid': self.faft_client.system.get_crossystem_value('fwid'),
198            'servod_version': self._client._servo_host.run(
199                'servod --version').stdout.strip(),
200        }
201
202        if hasattr(self, 'cr50'):
203            system_info['cr50_version'] = self.servo.get('cr50_version')
204
205        logging.info('System info:\n' + pprint.pformat(system_info))
206        self.write_attr_keyval(system_info)
207
208    def invalidate_firmware_setup(self):
209        """Invalidate all firmware related setup state.
210
211        This method is called when the firmware is re-flashed. It resets all
212        firmware related setup states so that the next test setup properly
213        again.
214        """
215        self.unmark_setup_done('gbb_flags')
216
217    def _retrieve_recovery_reason_from_trap(self):
218        """Try to retrieve the recovery reason from a trapped recovery screen.
219
220        @return: The recovery_reason, 0 if any error.
221        """
222        recovery_reason = 0
223        logging.info('Try to retrieve recovery reason...')
224        if self.servo.get_usbkey_direction() == 'dut':
225            self.switcher.bypass_rec_mode()
226        else:
227            self.servo.switch_usbkey('dut')
228
229        try:
230            self.switcher.wait_for_client()
231            lines = self.faft_client.system.run_shell_command_get_output(
232                        'crossystem recovery_reason')
233            recovery_reason = int(lines[0])
234            logging.info('Got the recovery reason %d.', recovery_reason)
235        except ConnectionError:
236            logging.error('Failed to get the recovery reason due to connection '
237                          'error.')
238        return recovery_reason
239
240    def _reset_client(self):
241        """Reset client to a workable state.
242
243        This method is called when the client is not responsive. It may be
244        caused by the following cases:
245          - halt on a firmware screen without timeout, e.g. REC_INSERT screen;
246          - corrupted firmware;
247          - corrutped OS image.
248        """
249        # DUT may halt on a firmware screen. Try cold reboot.
250        logging.info('Try cold reboot...')
251        self.switcher.mode_aware_reboot(reboot_type='cold',
252                                        sync_before_boot=False,
253                                        wait_for_dut_up=False)
254        self.switcher.wait_for_client_offline()
255        self.switcher.bypass_dev_mode()
256        try:
257            self.switcher.wait_for_client()
258            return
259        except ConnectionError:
260            logging.warn('Cold reboot doesn\'t help, still connection error.')
261
262        # DUT may be broken by a corrupted firmware. Restore firmware.
263        # We assume the recovery boot still works fine. Since the recovery
264        # code is in RO region and all FAFT tests don't change the RO region
265        # except GBB.
266        if self.is_firmware_saved():
267            self._ensure_client_in_recovery()
268            logging.info('Try restore the original firmware...')
269            if self.is_firmware_changed():
270                try:
271                    self.restore_firmware()
272                    return
273                except ConnectionError:
274                    logging.warn('Restoring firmware doesn\'t help, still '
275                                 'connection error.')
276
277        # Perhaps it's kernel that's broken. Let's try restoring it.
278        if self.is_kernel_saved():
279            self._ensure_client_in_recovery()
280            logging.info('Try restore the original kernel...')
281            if self.is_kernel_changed():
282                try:
283                    self.restore_kernel()
284                    return
285                except ConnectionError:
286                    logging.warn('Restoring kernel doesn\'t help, still '
287                                 'connection error.')
288
289        # DUT may be broken by a corrupted OS image. Restore OS image.
290        self._ensure_client_in_recovery()
291        logging.info('Try restore the OS image...')
292        self.faft_client.system.run_shell_command('chromeos-install --yes')
293        self.switcher.mode_aware_reboot(wait_for_dut_up=False)
294        self.switcher.wait_for_client_offline()
295        self.switcher.bypass_dev_mode()
296        try:
297            self.switcher.wait_for_client()
298            logging.info('Successfully restore OS image.')
299            return
300        except ConnectionError:
301            logging.warn('Restoring OS image doesn\'t help, still connection '
302                         'error.')
303
304    def _ensure_client_in_recovery(self):
305        """Ensure client in recovery boot; reboot into it if necessary.
306
307        @raise TestError: if failed to boot the USB image.
308        """
309        logging.info('Try boot into USB image...')
310        self.switcher.reboot_to_mode(to_mode='rec', sync_before_boot=False,
311                                     wait_for_dut_up=False)
312        self.servo.switch_usbkey('host')
313        self.switcher.bypass_rec_mode()
314        try:
315            self.switcher.wait_for_client()
316        except ConnectionError:
317            raise error.TestError('Failed to boot the USB image.')
318
319    def _restore_routine_from_timeout(self):
320        """A routine to try to restore the system from a timeout error.
321
322        This method is called when FAFT failed to connect DUT after reboot.
323
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        # TODO(waihong@chromium.org): Implement replugging the Ethernet to
331        # identify if it is a network flaky.
332
333        recovery_reason = self._retrieve_recovery_reason_from_trap()
334
335        # Reset client to a workable state.
336        self._reset_client()
337
338        # Raise the proper TestFail exception.
339        if recovery_reason:
340            raise error.TestFail('Trapped in the recovery screen (reason: %d) '
341                                 'and timed out' % recovery_reason)
342        else:
343            raise error.TestFail('Timed out waiting for DUT reboot')
344
345    def assert_test_image_in_usb_disk(self, usb_dev=None):
346        """Assert an USB disk plugged-in on servo and a test image inside.
347
348        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
349                        If None, it is detected automatically.
350        @raise TestError: if USB disk not detected or not a test image.
351        """
352        if self.check_setup_done('usb_check'):
353            return
354        if usb_dev:
355            assert self.servo.get_usbkey_direction() == 'host'
356        else:
357            self.servo.switch_usbkey('host')
358            usb_dev = self.servo.probe_host_usb_dev()
359            if not usb_dev:
360                raise error.TestError(
361                        'An USB disk should be plugged in the servo board.')
362
363        rootfs = '%s%s' % (usb_dev, self._ROOTFS_PARTITION_NUMBER)
364        logging.info('usb dev is %s', usb_dev)
365        tmpd = self.servo.system_output('mktemp -d -t usbcheck.XXXX')
366        self.servo.system('mount -o ro %s %s' % (rootfs, tmpd))
367
368        try:
369            usb_lsb = self.servo.system_output('cat %s' %
370                os.path.join(tmpd, 'etc/lsb-release'))
371            logging.debug('Dumping lsb-release on USB stick:\n%s', usb_lsb)
372            dut_lsb = '\n'.join(self.faft_client.system.
373                run_shell_command_get_output('cat /etc/lsb-release'))
374            logging.debug('Dumping lsb-release on DUT:\n%s', dut_lsb)
375            if not re.search(r'RELEASE_TRACK=.*test', usb_lsb):
376                raise error.TestError('USB stick in servo is no test image')
377            usb_board = re.search(r'BOARD=(.*)', usb_lsb).group(1)
378            dut_board = re.search(r'BOARD=(.*)', dut_lsb).group(1)
379            if usb_board != dut_board:
380                raise error.TestError('USB stick in servo contains a %s '
381                    'image, but DUT is a %s' % (usb_board, dut_board))
382        finally:
383            for cmd in ('umount -l %s' % rootfs, 'sync', 'rm -rf %s' % tmpd):
384                self.servo.system(cmd)
385
386        self.mark_setup_done('usb_check')
387
388    def setup_usbkey(self, usbkey, host=None):
389        """Setup the USB disk for the test.
390
391        It checks the setup of USB disk and a valid ChromeOS test image inside.
392        It also muxes the USB disk to either the host or DUT by request.
393
394        @param usbkey: True if the USB disk is required for the test, False if
395                       not required.
396        @param host: Optional, True to mux the USB disk to host, False to mux it
397                    to DUT, default to do nothing.
398        """
399        if usbkey:
400            self.assert_test_image_in_usb_disk()
401        elif host is None:
402            # USB disk is not required for the test. Better to mux it to host.
403            host = True
404
405        if host is True:
406            self.servo.switch_usbkey('host')
407        elif host is False:
408            self.servo.switch_usbkey('dut')
409
410    def get_usbdisk_path_on_dut(self):
411        """Get the path of the USB disk device plugged-in the servo on DUT.
412
413        Returns:
414          A string representing USB disk path, like '/dev/sdb', or None if
415          no USB disk is found.
416        """
417        cmd = 'ls -d /dev/s*[a-z]'
418        original_value = self.servo.get_usbkey_direction()
419
420        # Make the dut unable to see the USB disk.
421        self.servo.switch_usbkey('off')
422        no_usb_set = set(
423            self.faft_client.system.run_shell_command_get_output(cmd))
424
425        # Make the dut able to see the USB disk.
426        self.servo.switch_usbkey('dut')
427        time.sleep(self.faft_config.usb_plug)
428        has_usb_set = set(
429            self.faft_client.system.run_shell_command_get_output(cmd))
430
431        # Back to its original value.
432        if original_value != self.servo.get_usbkey_direction():
433            self.servo.switch_usbkey(original_value)
434
435        diff_set = has_usb_set - no_usb_set
436        if len(diff_set) == 1:
437            return diff_set.pop()
438        else:
439            return None
440
441    def _create_faft_lockfile(self):
442        """Creates the FAFT lockfile."""
443        logging.info('Creating FAFT lockfile...')
444        command = 'touch %s' % (self.lockfile)
445        self.faft_client.system.run_shell_command(command)
446
447    def _remove_faft_lockfile(self):
448        """Removes the FAFT lockfile."""
449        logging.info('Removing FAFT lockfile...')
450        command = 'rm -f %s' % (self.lockfile)
451        self.faft_client.system.run_shell_command(command)
452
453    def _stop_service(self, service):
454        """Stops a upstart service on the client.
455
456        @param service: The name of the upstart service.
457        """
458        logging.info('Stopping %s...', service)
459        command = 'status %s | grep stop || stop %s' % (service, service)
460        self.faft_client.system.run_shell_command(command)
461
462    def _start_service(self, service):
463        """Starts a upstart service on the client.
464
465        @param service: The name of the upstart service.
466        """
467        logging.info('Starting %s...', service)
468        command = 'status %s | grep start || start %s' % (service, service)
469        self.faft_client.system.run_shell_command(command)
470
471    def clear_set_gbb_flags(self, clear_mask, set_mask):
472        """Clear and set the GBB flags in the current flashrom.
473
474        @param clear_mask: A mask of flags to be cleared.
475        @param set_mask: A mask of flags to be set.
476        """
477        gbb_flags = self.faft_client.bios.get_gbb_flags()
478        new_flags = gbb_flags & ctypes.c_uint32(~clear_mask).value | set_mask
479        if new_flags != gbb_flags:
480            self._backup_gbb_flags = gbb_flags
481            logging.info('Changing GBB flags from 0x%x to 0x%x.',
482                         gbb_flags, new_flags)
483            self.faft_client.bios.set_gbb_flags(new_flags)
484            # If changing FORCE_DEV_SWITCH_ON flag, reboot to get a clear state
485            if ((gbb_flags ^ new_flags) & vboot.GBB_FLAG_FORCE_DEV_SWITCH_ON):
486                self.switcher.mode_aware_reboot()
487        else:
488            logging.info('Current GBB flags look good for test: 0x%x.',
489                         gbb_flags)
490
491    def check_ec_capability(self, required_cap=None, suppress_warning=False):
492        """Check if current platform has required EC capabilities.
493
494        @param required_cap: A list containing required EC capabilities. Pass in
495                             None to only check for presence of Chrome EC.
496        @param suppress_warning: True to suppress any warning messages.
497        @return: True if requirements are met. Otherwise, False.
498        """
499        if not self.faft_config.chrome_ec:
500            if not suppress_warning:
501                logging.warn('Requires Chrome EC to run this test.')
502            return False
503
504        if not required_cap:
505            return True
506
507        for cap in required_cap:
508            if cap not in self.faft_config.ec_capability:
509                if not suppress_warning:
510                    logging.warn('Requires EC capability "%s" to run this '
511                                 'test.', cap)
512                return False
513
514        return True
515
516    def check_root_part_on_non_recovery(self, part):
517        """Check the partition number of root device and on normal/dev boot.
518
519        @param part: A string of partition number, e.g.'3'.
520        @return: True if the root device matched and on normal/dev boot;
521                 otherwise, False.
522        """
523        return self.checkers.root_part_checker(part) and \
524                self.checkers.crossystem_checker({
525                    'mainfw_type': ('normal', 'developer'),
526                })
527
528    def _join_part(self, dev, part):
529        """Return a concatenated string of device and partition number.
530
531        @param dev: A string of device, e.g.'/dev/sda'.
532        @param part: A string of partition number, e.g.'3'.
533        @return: A concatenated string of device and partition number,
534                 e.g.'/dev/sda3'.
535
536        >>> seq = FirmwareTest()
537        >>> seq._join_part('/dev/sda', '3')
538        '/dev/sda3'
539        >>> seq._join_part('/dev/mmcblk0', '2')
540        '/dev/mmcblk0p2'
541        """
542        if 'mmcblk' in dev:
543            return dev + 'p' + part
544        else:
545            return dev + part
546
547    def copy_kernel_and_rootfs(self, from_part, to_part):
548        """Copy kernel and rootfs from from_part to to_part.
549
550        @param from_part: A string of partition number to be copied from.
551        @param to_part: A string of partition number to be copied to.
552        """
553        root_dev = self.faft_client.system.get_root_dev()
554        logging.info('Copying kernel from %s to %s. Please wait...',
555                     from_part, to_part)
556        self.faft_client.system.run_shell_command('dd if=%s of=%s bs=4M' %
557                (self._join_part(root_dev, self.KERNEL_MAP[from_part]),
558                 self._join_part(root_dev, self.KERNEL_MAP[to_part])))
559        logging.info('Copying rootfs from %s to %s. Please wait...',
560                     from_part, to_part)
561        self.faft_client.system.run_shell_command('dd if=%s of=%s bs=4M' %
562                (self._join_part(root_dev, self.ROOTFS_MAP[from_part]),
563                 self._join_part(root_dev, self.ROOTFS_MAP[to_part])))
564
565    def ensure_kernel_boot(self, part):
566        """Ensure the request kernel boot.
567
568        If not, it duplicates the current kernel to the requested kernel
569        and sets the requested higher priority to ensure it boot.
570
571        @param part: A string of kernel partition number or 'a'/'b'.
572        """
573        if not self.checkers.root_part_checker(part):
574            if self.faft_client.kernel.diff_a_b():
575                self.copy_kernel_and_rootfs(
576                        from_part=self.OTHER_KERNEL_MAP[part],
577                        to_part=part)
578            self.reset_and_prioritize_kernel(part)
579            self.switcher.mode_aware_reboot()
580
581    def set_hardware_write_protect(self, enable):
582        """Set hardware write protect pin.
583
584        @param enable: True if asserting write protect pin. Otherwise, False.
585        """
586        try:
587            self.servo.set('fw_wp_state', 'force_on' if enable else 'force_off')
588        except:
589            # TODO(waihong): Remove this fallback when all servos have the
590            # above new fw_wp_state control.
591            self.servo.set('fw_wp_vref', self.faft_config.wp_voltage)
592            self.servo.set('fw_wp_en', 'on')
593            self.servo.set('fw_wp', 'on' if enable else 'off')
594
595    def set_ec_write_protect_and_reboot(self, enable):
596        """Set EC write protect status and reboot to take effect.
597
598        The write protect state is only activated if both hardware write
599        protect pin is asserted and software write protect flag is set.
600        This method asserts/deasserts hardware write protect pin first, and
601        set corresponding EC software write protect flag.
602
603        If the device uses non-Chrome EC, set the software write protect via
604        flashrom.
605
606        If the device uses Chrome EC, a reboot is required for write protect
607        to take effect. Since the software write protect flag cannot be unset
608        if hardware write protect pin is asserted, we need to deasserted the
609        pin first if we are deactivating write protect. Similarly, a reboot
610        is required before we can modify the software flag.
611
612        @param enable: True if activating EC write protect. Otherwise, False.
613        """
614        self.set_hardware_write_protect(enable)
615        if self.faft_config.chrome_ec:
616            self.set_chrome_ec_write_protect_and_reboot(enable)
617        else:
618            self.faft_client.ec.set_write_protect(enable)
619            self.switcher.mode_aware_reboot()
620
621    def set_chrome_ec_write_protect_and_reboot(self, enable):
622        """Set Chrome EC write protect status and reboot to take effect.
623
624        @param enable: True if activating EC write protect. Otherwise, False.
625        """
626        if enable:
627            # Set write protect flag and reboot to take effect.
628            self.ec.set_flash_write_protect(enable)
629            self.sync_and_ec_reboot()
630        else:
631            # Reboot after deasserting hardware write protect pin to deactivate
632            # write protect. And then remove software write protect flag.
633            self.sync_and_ec_reboot()
634            self.ec.set_flash_write_protect(enable)
635
636    def _setup_ec_write_protect(self, ec_wp):
637        """Setup for EC write-protection.
638
639        It makes sure the EC in the requested write-protection state. If not, it
640        flips the state. Flipping the write-protection requires DUT reboot.
641
642        @param ec_wp: True to request EC write-protected; False to request EC
643                      not write-protected; None to do nothing.
644        """
645        if ec_wp is None:
646            self._old_ec_wp = None
647            return
648        self._old_ec_wp = self.checkers.crossystem_checker({'wpsw_boot': '1'})
649        if ec_wp != self._old_ec_wp:
650            logging.info('The test required EC is %swrite-protected. Reboot '
651                         'and flip the state.', '' if ec_wp else 'not ')
652            self.switcher.mode_aware_reboot(
653                    'custom',
654                     lambda:self.set_ec_write_protect_and_reboot(ec_wp))
655
656    def _restore_ec_write_protect(self):
657        """Restore the original EC write-protection."""
658        if (not hasattr(self, '_old_ec_wp')) or (self._old_ec_wp is None):
659            return
660        if not self.checkers.crossystem_checker(
661                {'wpsw_boot': '1' if self._old_ec_wp else '0'}):
662            logging.info('Restore original EC write protection and reboot.')
663            self.switcher.mode_aware_reboot(
664                    'custom',
665                    lambda:self.set_ec_write_protect_and_reboot(
666                            self._old_ec_wp))
667
668    def _setup_uart_capture(self):
669        """Setup the CPU/EC/PD UART capture."""
670        self.cpu_uart_file = os.path.join(self.resultsdir, 'cpu_uart.txt')
671        self.servo.set('cpu_uart_capture', 'on')
672        self.cr50_console_file = None
673        self.ec_uart_file = None
674        self.usbpd_uart_file = None
675        try:
676            self.servo.set('cr50_console_capture', 'on')
677            self.cr50_console_file = os.path.join(self.resultsdir,
678                                                  'cr50_console.txt')
679            # Check that the console works before declaring the cr50 console
680            # connection exists.
681            self.servo.get('ccd_lock')
682            self.cr50 = chrome_cr50.ChromeCr50(self.servo)
683        except error.TestFail as e:
684            if 'No control named' in str(e):
685                logging.warn('cr50 console not supported.')
686        if self.faft_config.chrome_ec:
687            try:
688                self.servo.set('ec_uart_capture', 'on')
689                self.ec_uart_file = os.path.join(self.resultsdir, 'ec_uart.txt')
690            except error.TestFail as e:
691                if 'No control named' in str(e):
692                    logging.warn('The servod is too old that ec_uart_capture '
693                                 'not supported.')
694            # Log separate PD console if supported
695            if self.check_ec_capability(['usbpd_uart'], suppress_warning=True):
696                try:
697                    self.servo.set('usbpd_uart_capture', 'on')
698                    self.usbpd_uart_file = os.path.join(self.resultsdir,
699                                                        'usbpd_uart.txt')
700                except error.TestFail as e:
701                    if 'No control named' in str(e):
702                        logging.warn('The servod is too old that '
703                                     'usbpd_uart_capture is not supported.')
704        else:
705            logging.info('Not a Google EC, cannot capture ec console output.')
706
707    def _record_uart_capture(self):
708        """Record the CPU/EC/PD UART output stream to files."""
709        if self.cpu_uart_file:
710            with open(self.cpu_uart_file, 'a') as f:
711                f.write(ast.literal_eval(self.servo.get('cpu_uart_stream')))
712        if self.cr50_console_file:
713            with open(self.cr50_console_file, 'a') as f:
714                f.write(ast.literal_eval(self.servo.get('cr50_console_stream')))
715        if self.ec_uart_file and self.faft_config.chrome_ec:
716            with open(self.ec_uart_file, 'a') as f:
717                f.write(ast.literal_eval(self.servo.get('ec_uart_stream')))
718        if (self.usbpd_uart_file and self.faft_config.chrome_ec and
719            self.check_ec_capability(['usbpd_uart'], suppress_warning=True)):
720            with open(self.usbpd_uart_file, 'a') as f:
721                f.write(ast.literal_eval(self.servo.get('usbpd_uart_stream')))
722
723    def _cleanup_uart_capture(self):
724        """Cleanup the CPU/EC/PD UART capture."""
725        # Flush the remaining UART output.
726        self._record_uart_capture()
727        self.servo.set('cpu_uart_capture', 'off')
728        if self.cr50_console_file:
729            self.servo.set('cr50_console_capture', 'off')
730        if self.ec_uart_file and self.faft_config.chrome_ec:
731            self.servo.set('ec_uart_capture', 'off')
732        if (self.usbpd_uart_file and self.faft_config.chrome_ec and
733            self.check_ec_capability(['usbpd_uart'], suppress_warning=True)):
734            self.servo.set('usbpd_uart_capture', 'off')
735
736    def _get_power_state(self, power_state):
737        """
738        Return the current power state of the AP
739        """
740        return self.ec.send_command_get_output("powerinfo", [power_state])
741
742    def wait_power_state(self, power_state, retries):
743        """
744        Wait for certain power state.
745
746        @param power_state: power state you are expecting
747        @param retries: retries.  This is necessary if AP is powering down
748        and transitioning through different states.
749        """
750        logging.info('Checking power state "%s" maximum %d times.',
751                     power_state, retries)
752        while retries > 0:
753            logging.info("try count: %d", retries)
754            try:
755                retries = retries - 1
756                ret = self._get_power_state(power_state)
757                return True
758            except error.TestFail:
759                pass
760        return False
761
762    def suspend(self):
763        """Suspends the DUT."""
764        cmd = '(sleep %d; powerd_dbus_suspend) &' % self.EC_SUSPEND_DELAY
765        self.faft_client.system.run_shell_command(cmd)
766        time.sleep(self.EC_SUSPEND_DELAY)
767
768    def _fetch_servo_log(self):
769        """Fetch the servo log."""
770        cmd = '[ -e %s ] && cat %s || echo NOTFOUND' % ((self._SERVOD_LOG,) * 2)
771        servo_log = self.servo.system_output(cmd)
772        return None if servo_log == 'NOTFOUND' else servo_log
773
774    def _setup_servo_log(self):
775        """Setup the servo log capturing."""
776        self.servo_log_original_len = -1
777        if self.servo.is_localhost():
778            # No servo log recorded when servod runs locally.
779            return
780
781        servo_log = self._fetch_servo_log()
782        if servo_log:
783            self.servo_log_original_len = len(servo_log)
784        else:
785            logging.warn('Servo log file not found.')
786
787    def _record_servo_log(self):
788        """Record the servo log to the results directory."""
789        if self.servo_log_original_len != -1:
790            servo_log = self._fetch_servo_log()
791            servo_log_file = os.path.join(self.resultsdir, 'servod.log')
792            with open(servo_log_file, 'a') as f:
793                f.write(servo_log[self.servo_log_original_len:])
794
795    def _record_faft_client_log(self):
796        """Record the faft client log to the results directory."""
797        client_log = self.faft_client.system.dump_log(True)
798        client_log_file = os.path.join(self.resultsdir, 'faft_client.log')
799        with open(client_log_file, 'w') as f:
800            f.write(client_log)
801
802    def _setup_gbb_flags(self):
803        """Setup the GBB flags for FAFT test."""
804        if self.faft_config.gbb_version < 1.1:
805            logging.info('Skip modifying GBB on versions older than 1.1.')
806            return
807
808        if self.check_setup_done('gbb_flags'):
809            return
810
811        logging.info('Set proper GBB flags for test.')
812        self.clear_set_gbb_flags(vboot.GBB_FLAG_DEV_SCREEN_SHORT_DELAY |
813                                 vboot.GBB_FLAG_FORCE_DEV_SWITCH_ON |
814                                 vboot.GBB_FLAG_FORCE_DEV_BOOT_USB |
815                                 vboot.GBB_FLAG_DISABLE_FW_ROLLBACK_CHECK |
816                                 vboot.GBB_FLAG_FORCE_DEV_BOOT_FASTBOOT_FULL_CAP,
817                                 vboot.GBB_FLAG_ENTER_TRIGGERS_TONORM |
818                                 vboot.GBB_FLAG_FAFT_KEY_OVERIDE)
819        self.mark_setup_done('gbb_flags')
820
821    def drop_backup_gbb_flags(self):
822        """Drops the backup GBB flags.
823
824        This can be used when a test intends to permanently change GBB flags.
825        """
826        self._backup_gbb_flags = None
827
828    def _restore_gbb_flags(self):
829        """Restore GBB flags to their original state."""
830        if self._backup_gbb_flags is None:
831            return
832        # Setting up and restoring the GBB flags take a lot of time. For
833        # speed-up purpose, don't restore it.
834        logging.info('***')
835        logging.info('*** Please manually restore the original GBB flags to: '
836                     '0x%x ***', self._backup_gbb_flags)
837        logging.info('***')
838        self.unmark_setup_done('gbb_flags')
839
840    def setup_tried_fwb(self, tried_fwb):
841        """Setup for fw B tried state.
842
843        It makes sure the system in the requested fw B tried state. If not, it
844        tries to do so.
845
846        @param tried_fwb: True if requested in tried_fwb=1;
847                          False if tried_fwb=0.
848        """
849        if tried_fwb:
850            if not self.checkers.crossystem_checker({'tried_fwb': '1'}):
851                logging.info(
852                    'Firmware is not booted with tried_fwb. Reboot into it.')
853                self.faft_client.system.set_try_fw_b()
854        else:
855            if not self.checkers.crossystem_checker({'tried_fwb': '0'}):
856                logging.info(
857                    'Firmware is booted with tried_fwb. Reboot to clear.')
858
859    def power_on(self):
860        """Switch DUT AC power on."""
861        self._client.power_on(self.power_control)
862
863    def power_off(self):
864        """Switch DUT AC power off."""
865        self._client.power_off(self.power_control)
866
867    def power_cycle(self):
868        """Power cycle DUT AC power."""
869        self._client.power_cycle(self.power_control)
870
871    def setup_rw_boot(self, section='a'):
872        """Make sure firmware is in RW-boot mode.
873
874        If the given firmware section is in RO-boot mode, turn off the RO-boot
875        flag and reboot DUT into RW-boot mode.
876
877        @param section: A firmware section, either 'a' or 'b'.
878        """
879        flags = self.faft_client.bios.get_preamble_flags(section)
880        if flags & vboot.PREAMBLE_USE_RO_NORMAL:
881            flags = flags ^ vboot.PREAMBLE_USE_RO_NORMAL
882            self.faft_client.bios.set_preamble_flags(section, flags)
883            self.switcher.mode_aware_reboot()
884
885    def setup_kernel(self, part):
886        """Setup for kernel test.
887
888        It makes sure both kernel A and B bootable and the current boot is
889        the requested kernel part.
890
891        @param part: A string of kernel partition number or 'a'/'b'.
892        """
893        self.ensure_kernel_boot(part)
894        logging.info('Checking the integrity of kernel B and rootfs B...')
895        if (self.faft_client.kernel.diff_a_b() or
896                not self.faft_client.rootfs.verify_rootfs('B')):
897            logging.info('Copying kernel and rootfs from A to B...')
898            self.copy_kernel_and_rootfs(from_part=part,
899                                        to_part=self.OTHER_KERNEL_MAP[part])
900        self.reset_and_prioritize_kernel(part)
901
902    def reset_and_prioritize_kernel(self, part):
903        """Make the requested partition highest priority.
904
905        This function also reset kerenl A and B to bootable.
906
907        @param part: A string of partition number to be prioritized.
908        """
909        root_dev = self.faft_client.system.get_root_dev()
910        # Reset kernel A and B to bootable.
911        self.faft_client.system.run_shell_command(
912            'cgpt add -i%s -P1 -S1 -T0 %s' % (self.KERNEL_MAP['a'], root_dev))
913        self.faft_client.system.run_shell_command(
914            'cgpt add -i%s -P1 -S1 -T0 %s' % (self.KERNEL_MAP['b'], root_dev))
915        # Set kernel part highest priority.
916        self.faft_client.system.run_shell_command('cgpt prioritize -i%s %s' %
917                (self.KERNEL_MAP[part], root_dev))
918
919    def blocking_sync(self):
920        """Run a blocking sync command."""
921        # The double calls to sync fakes a blocking call
922        # since the first call returns before the flush
923        # is complete, but the second will wait for the
924        # first to finish.
925        self.faft_client.system.run_shell_command('sync')
926        self.faft_client.system.run_shell_command('sync')
927
928        # sync only sends SYNCHRONIZE_CACHE but doesn't
929        # check the status. For mmc devices, use `mmc
930        # status get` command to send an empty command to
931        # wait for the disk to be available again.  For
932        # other devices, hdparm sends TUR to check if
933        # a device is ready for transfer operation.
934        root_dev = self.faft_client.system.get_root_dev()
935        if 'mmcblk' in root_dev:
936            self.faft_client.system.run_shell_command('mmc status get %s' %
937                                                      root_dev)
938        else:
939            self.faft_client.system.run_shell_command('hdparm -f %s' % root_dev)
940
941    def sync_and_ec_reboot(self, flags=''):
942        """Request the client sync and do a EC triggered reboot.
943
944        @param flags: Optional, a space-separated string of flags passed to EC
945                      reboot command, including:
946                          default: EC soft reboot;
947                          'hard': EC cold/hard reboot.
948        """
949        self.blocking_sync()
950        self.ec.reboot(flags)
951        time.sleep(self.faft_config.ec_boot_to_console)
952        self.check_lid_and_power_on()
953
954    def reboot_and_reset_tpm(self):
955        """Reboot into recovery mode, reset TPM, then reboot back to disk."""
956        self.switcher.reboot_to_mode(to_mode='rec')
957        self.faft_client.system.run_shell_command('chromeos-tpm-recovery')
958        self.switcher.mode_aware_reboot()
959
960    def full_power_off_and_on(self):
961        """Shutdown the device by pressing power button and power on again."""
962        boot_id = self.get_bootid()
963        # Press power button to trigger Chrome OS normal shutdown process.
964        # We use a customized delay since the normal-press 1.2s is not enough.
965        self.servo.power_key(self.faft_config.hold_pwr_button_poweroff)
966        # device can take 44-51 seconds to restart,
967        # add buffer from the default timeout of 60 seconds.
968        self.switcher.wait_for_client_offline(timeout=100, orig_boot_id=boot_id)
969        time.sleep(self.faft_config.shutdown)
970        # Short press power button to boot DUT again.
971        self.servo.power_key(self.faft_config.hold_pwr_button_poweron)
972
973    def check_lid_and_power_on(self):
974        """
975        On devices with EC software sync, system powers on after EC reboots if
976        lid is open. Otherwise, the EC shuts down CPU after about 3 seconds.
977        This method checks lid switch state and presses power button if
978        necessary.
979        """
980        if self.servo.get("lid_open") == "no":
981            time.sleep(self.faft_config.software_sync)
982            self.servo.power_short_press()
983
984    def _modify_usb_kernel(self, usb_dev, from_magic, to_magic):
985        """Modify the kernel header magic in USB stick.
986
987        The kernel header magic is the first 8-byte of kernel partition.
988        We modify it to make it fail on kernel verification check.
989
990        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
991        @param from_magic: A string of magic which we change it from.
992        @param to_magic: A string of magic which we change it to.
993        @raise TestError: if failed to change magic.
994        """
995        assert len(from_magic) == 8
996        assert len(to_magic) == 8
997        # USB image only contains one kernel.
998        kernel_part = self._join_part(usb_dev, self.KERNEL_MAP['a'])
999        read_cmd = "sudo dd if=%s bs=8 count=1 2>/dev/null" % kernel_part
1000        current_magic = self.servo.system_output(read_cmd)
1001        if current_magic == to_magic:
1002            logging.info("The kernel magic is already %s.", current_magic)
1003            return
1004        if current_magic != from_magic:
1005            raise error.TestError("Invalid kernel image on USB: wrong magic.")
1006
1007        logging.info('Modify the kernel magic in USB, from %s to %s.',
1008                     from_magic, to_magic)
1009        write_cmd = ("echo -n '%s' | sudo dd of=%s oflag=sync conv=notrunc "
1010                     " 2>/dev/null" % (to_magic, kernel_part))
1011        self.servo.system(write_cmd)
1012
1013        if self.servo.system_output(read_cmd) != to_magic:
1014            raise error.TestError("Failed to write new magic.")
1015
1016    def corrupt_usb_kernel(self, usb_dev):
1017        """Corrupt USB kernel by modifying its magic from CHROMEOS to CORRUPTD.
1018
1019        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
1020        """
1021        self._modify_usb_kernel(usb_dev, self.CHROMEOS_MAGIC,
1022                                self.CORRUPTED_MAGIC)
1023
1024    def restore_usb_kernel(self, usb_dev):
1025        """Restore USB kernel by modifying its magic from CORRUPTD to CHROMEOS.
1026
1027        @param usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
1028        """
1029        self._modify_usb_kernel(usb_dev, self.CORRUPTED_MAGIC,
1030                                self.CHROMEOS_MAGIC)
1031
1032    def _call_action(self, action_tuple, check_status=False):
1033        """Call the action function with/without arguments.
1034
1035        @param action_tuple: A function, or a tuple (function, args, error_msg),
1036                             in which, args and error_msg are optional. args is
1037                             either a value or a tuple if multiple arguments.
1038                             This can also be a list containing multiple
1039                             function or tuple. In this case, these actions are
1040                             called in sequence.
1041        @param check_status: Check the return value of action function. If not
1042                             succeed, raises a TestFail exception.
1043        @return: The result value of the action function.
1044        @raise TestError: An error when the action function is not callable.
1045        @raise TestFail: When check_status=True, action function not succeed.
1046        """
1047        if isinstance(action_tuple, list):
1048            return all([self._call_action(action, check_status=check_status)
1049                        for action in action_tuple])
1050
1051        action = action_tuple
1052        args = ()
1053        error_msg = 'Not succeed'
1054        if isinstance(action_tuple, tuple):
1055            action = action_tuple[0]
1056            if len(action_tuple) >= 2:
1057                args = action_tuple[1]
1058                if not isinstance(args, tuple):
1059                    args = (args,)
1060            if len(action_tuple) >= 3:
1061                error_msg = action_tuple[2]
1062
1063        if action is None:
1064            return
1065
1066        if not callable(action):
1067            raise error.TestError('action is not callable!')
1068
1069        info_msg = 'calling %s' % action.__name__
1070        if args:
1071            info_msg += ' with args %s' % str(args)
1072        logging.info(info_msg)
1073        ret = action(*args)
1074
1075        if check_status and not ret:
1076            raise error.TestFail('%s: %s returning %s' %
1077                                 (error_msg, info_msg, str(ret)))
1078        return ret
1079
1080    def run_shutdown_process(self, shutdown_action, pre_power_action=None,
1081                             run_power_action=True, post_power_action=None,
1082                             shutdown_timeout=None):
1083        """Run shutdown_action(), which makes DUT shutdown, and power it on.
1084
1085        @param shutdown_action: function which makes DUT shutdown, like
1086                                pressing power key.
1087        @param pre_power_action: function which is called before next power on.
1088        @param run_power_action: power_key press by default, set to None to skip.
1089        @param post_power_action: function which is called after next power on.
1090        @param shutdown_timeout: a timeout to confirm DUT shutdown.
1091        @raise TestFail: if the shutdown_action() failed to turn DUT off.
1092        """
1093        self._call_action(shutdown_action)
1094        logging.info('Wait to ensure DUT shut down...')
1095        try:
1096            if shutdown_timeout is None:
1097                shutdown_timeout = self.faft_config.shutdown_timeout
1098            self.switcher.wait_for_client(timeout=shutdown_timeout)
1099            raise error.TestFail(
1100                    'Should shut the device down after calling %s.' %
1101                    shutdown_action.__name__)
1102        except ConnectionError:
1103            logging.info(
1104                'DUT is surely shutdown. We are going to power it on again...')
1105
1106        if pre_power_action:
1107            self._call_action(pre_power_action)
1108        if run_power_action:
1109            self.servo.power_key(self.faft_config.hold_pwr_button_poweron)
1110        if post_power_action:
1111            self._call_action(post_power_action)
1112
1113    def get_bootid(self, retry=3):
1114        """
1115        Return the bootid.
1116        """
1117        boot_id = None
1118        while retry:
1119            try:
1120                boot_id = self._client.get_boot_id()
1121                break
1122            except error.AutoservRunError:
1123                retry -= 1
1124                if retry:
1125                    logging.info('Retry to get boot_id...')
1126                else:
1127                    logging.warning('Failed to get boot_id.')
1128        logging.info('boot_id: %s', boot_id)
1129        return boot_id
1130
1131    def check_state(self, func):
1132        """
1133        Wrapper around _call_action with check_status set to True. This is a
1134        helper function to be used by tests and is currently implemented by
1135        calling _call_action with check_status=True.
1136
1137        TODO: This function's arguments need to be made more stringent. And
1138        its functionality should be moved over to check functions directly in
1139        the future.
1140
1141        @param func: A function, or a tuple (function, args, error_msg),
1142                             in which, args and error_msg are optional. args is
1143                             either a value or a tuple if multiple arguments.
1144                             This can also be a list containing multiple
1145                             function or tuple. In this case, these actions are
1146                             called in sequence.
1147        @return: The result value of the action function.
1148        @raise TestFail: If the function does notsucceed.
1149        """
1150        logging.info("-[FAFT]-[ start stepstate_checker ]----------")
1151        self._call_action(func, check_status=True)
1152        logging.info("-[FAFT]-[ end state_checker ]----------------")
1153
1154    def get_current_firmware_sha(self):
1155        """Get current firmware sha of body and vblock.
1156
1157        @return: Current firmware sha follows the order (
1158                 vblock_a_sha, body_a_sha, vblock_b_sha, body_b_sha)
1159        """
1160        current_firmware_sha = (self.faft_client.bios.get_sig_sha('a'),
1161                                self.faft_client.bios.get_body_sha('a'),
1162                                self.faft_client.bios.get_sig_sha('b'),
1163                                self.faft_client.bios.get_body_sha('b'))
1164        if not all(current_firmware_sha):
1165            raise error.TestError('Failed to get firmware sha.')
1166        return current_firmware_sha
1167
1168    def is_firmware_changed(self):
1169        """Check if the current firmware changed, by comparing its SHA.
1170
1171        @return: True if it is changed, otherwise Flase.
1172        """
1173        # Device may not be rebooted after test.
1174        self.faft_client.bios.reload()
1175
1176        current_sha = self.get_current_firmware_sha()
1177
1178        if current_sha == self._backup_firmware_sha:
1179            return False
1180        else:
1181            corrupt_VBOOTA = (current_sha[0] != self._backup_firmware_sha[0])
1182            corrupt_FVMAIN = (current_sha[1] != self._backup_firmware_sha[1])
1183            corrupt_VBOOTB = (current_sha[2] != self._backup_firmware_sha[2])
1184            corrupt_FVMAINB = (current_sha[3] != self._backup_firmware_sha[3])
1185            logging.info('Firmware changed:')
1186            logging.info('VBOOTA is changed: %s', corrupt_VBOOTA)
1187            logging.info('VBOOTB is changed: %s', corrupt_VBOOTB)
1188            logging.info('FVMAIN is changed: %s', corrupt_FVMAIN)
1189            logging.info('FVMAINB is changed: %s', corrupt_FVMAINB)
1190            return True
1191
1192    def backup_firmware(self, suffix='.original'):
1193        """Backup firmware to file, and then send it to host.
1194
1195        @param suffix: a string appended to backup file name
1196        """
1197        remote_temp_dir = self.faft_client.system.create_temp_dir()
1198        remote_bios_path = os.path.join(remote_temp_dir, 'bios')
1199        self.faft_client.bios.dump_whole(remote_bios_path)
1200        self._client.get_file(remote_bios_path,
1201                              os.path.join(self.resultsdir, 'bios' + suffix))
1202
1203        if self.faft_config.chrome_ec:
1204            remote_ec_path = os.path.join(remote_temp_dir, 'ec')
1205            self.faft_client.ec.dump_whole(remote_ec_path)
1206            self._client.get_file(remote_ec_path,
1207                              os.path.join(self.resultsdir, 'ec' + suffix))
1208
1209        self._client.run('rm -rf %s' % remote_temp_dir)
1210        logging.info('Backup firmware stored in %s with suffix %s',
1211            self.resultsdir, suffix)
1212
1213        self._backup_firmware_sha = self.get_current_firmware_sha()
1214
1215    def is_firmware_saved(self):
1216        """Check if a firmware saved (called backup_firmware before).
1217
1218        @return: True if the firmware is backuped; otherwise False.
1219        """
1220        return self._backup_firmware_sha != ()
1221
1222    def clear_saved_firmware(self):
1223        """Clear the firmware saved by the method backup_firmware."""
1224        self._backup_firmware_sha = ()
1225
1226    def restore_firmware(self, suffix='.original'):
1227        """Restore firmware from host in resultsdir.
1228
1229        @param suffix: a string appended to backup file name
1230        """
1231        if not self.is_firmware_changed():
1232            return
1233
1234        # Backup current corrupted firmware.
1235        self.backup_firmware(suffix='.corrupt')
1236
1237        # Restore firmware.
1238        remote_temp_dir = self.faft_client.system.create_temp_dir()
1239        self._client.send_file(os.path.join(self.resultsdir, 'bios' + suffix),
1240                               os.path.join(remote_temp_dir, 'bios'))
1241
1242        self.faft_client.bios.write_whole(
1243            os.path.join(remote_temp_dir, 'bios'))
1244
1245        if self.faft_config.chrome_ec:
1246            self._client.send_file(os.path.join(self.resultsdir, 'ec' + suffix),
1247                os.path.join(remote_temp_dir, 'ec'))
1248            self.faft_client.ec.write_whole(
1249                os.path.join(remote_temp_dir, 'ec'))
1250
1251        self.switcher.mode_aware_reboot()
1252        logging.info('Successfully restore firmware.')
1253
1254    def setup_firmwareupdate_shellball(self, shellball=None):
1255        """Setup a shellball to use in firmware update test.
1256
1257        Check if there is a given shellball, and it is a shell script. Then,
1258        send it to the remote host. Otherwise, use the
1259        /usr/sbin/chromeos-firmwareupdate in the image and replace its inside
1260        BIOS and EC images with the active firmware images.
1261
1262        @param shellball: path of a shellball or default to None.
1263        """
1264        if shellball:
1265            # Determine the firmware file is a shellball or a raw binary.
1266            is_shellball = (utils.system_output("file %s" % shellball).find(
1267                    "shell script") != -1)
1268            if is_shellball:
1269                logging.info('Device will update firmware with shellball %s',
1270                             shellball)
1271                temp_path = self.faft_client.updater.get_temp_path()
1272                working_shellball = os.path.join(temp_path,
1273                                                 'chromeos-firmwareupdate')
1274                self._client.send_file(shellball, working_shellball)
1275                self.faft_client.updater.extract_shellball()
1276            else:
1277                raise error.TestFail(
1278                    'The given shellball is not a shell script.')
1279        else:
1280            logging.info('No shellball given, use the original shellball and '
1281                         'replace its BIOS and EC images.')
1282            work_path = self.faft_client.updater.get_work_path()
1283            bios_in_work_path = os.path.join(work_path, 'bios.bin')
1284            ec_in_work_path = os.path.join(work_path, 'ec.bin')
1285            self.faft_client.bios.dump_whole(bios_in_work_path)
1286            if self.faft_config.chrome_ec:
1287                self.faft_client.ec.dump_firmware(ec_in_work_path)
1288            self.faft_client.updater.repack_shellball()
1289
1290    def is_kernel_changed(self):
1291        """Check if the current kernel is changed, by comparing its SHA1 hash.
1292
1293        @return: True if it is changed; otherwise, False.
1294        """
1295        changed = False
1296        for p in ('A', 'B'):
1297            backup_sha = self._backup_kernel_sha.get(p, None)
1298            current_sha = self.faft_client.kernel.get_sha(p)
1299            if backup_sha != current_sha:
1300                changed = True
1301                logging.info('Kernel %s is changed', p)
1302        return changed
1303
1304    def backup_kernel(self, suffix='.original'):
1305        """Backup kernel to files, and the send them to host.
1306
1307        @param suffix: a string appended to backup file name.
1308        """
1309        remote_temp_dir = self.faft_client.system.create_temp_dir()
1310        for p in ('A', 'B'):
1311            remote_path = os.path.join(remote_temp_dir, 'kernel_%s' % p)
1312            self.faft_client.kernel.dump(p, remote_path)
1313            self._client.get_file(
1314                    remote_path,
1315                    os.path.join(self.resultsdir, 'kernel_%s%s' % (p, suffix)))
1316            self._backup_kernel_sha[p] = self.faft_client.kernel.get_sha(p)
1317        logging.info('Backup kernel stored in %s with suffix %s',
1318            self.resultsdir, suffix)
1319
1320    def is_kernel_saved(self):
1321        """Check if kernel images are saved (backup_kernel called before).
1322
1323        @return: True if the kernel is saved; otherwise, False.
1324        """
1325        return len(self._backup_kernel_sha) != 0
1326
1327    def clear_saved_kernel(self):
1328        """Clear the kernel saved by backup_kernel()."""
1329        self._backup_kernel_sha = dict()
1330
1331    def restore_kernel(self, suffix='.original'):
1332        """Restore kernel from host in resultsdir.
1333
1334        @param suffix: a string appended to backup file name.
1335        """
1336        if not self.is_kernel_changed():
1337            return
1338
1339        # Backup current corrupted kernel.
1340        self.backup_kernel(suffix='.corrupt')
1341
1342        # Restore kernel.
1343        remote_temp_dir = self.faft_client.system.create_temp_dir()
1344        for p in ('A', 'B'):
1345            remote_path = os.path.join(remote_temp_dir, 'kernel_%s' % p)
1346            self._client.send_file(
1347                    os.path.join(self.resultsdir, 'kernel_%s%s' % (p, suffix)),
1348                    remote_path)
1349            self.faft_client.kernel.write(p, remote_path)
1350
1351        self.switcher.mode_aware_reboot()
1352        logging.info('Successfully restored kernel.')
1353
1354    def backup_cgpt_attributes(self):
1355        """Backup CGPT partition table attributes."""
1356        self._backup_cgpt_attr = self.faft_client.cgpt.get_attributes()
1357
1358    def restore_cgpt_attributes(self):
1359        """Restore CGPT partition table attributes."""
1360        current_table = self.faft_client.cgpt.get_attributes()
1361        if current_table == self._backup_cgpt_attr:
1362            return
1363        logging.info('CGPT table is changed. Original: %r. Current: %r.',
1364                     self._backup_cgpt_attr,
1365                     current_table)
1366        self.faft_client.cgpt.set_attributes(self._backup_cgpt_attr)
1367
1368        self.switcher.mode_aware_reboot()
1369        logging.info('Successfully restored CGPT table.')
1370
1371    def try_fwb(self, count=0):
1372        """set to try booting FWB count # times
1373
1374        Wrapper to set fwb_tries for vboot1 and fw_try_count,fw_try_next for
1375        vboot2
1376
1377        @param count: an integer specifying value to program into
1378                      fwb_tries(vb1)/fw_try_next(vb2)
1379        """
1380        if self.fw_vboot2:
1381            self.faft_client.system.set_fw_try_next('B', count)
1382        else:
1383            # vboot1: we need to boot into fwb at least once
1384            if not count:
1385                count = count + 1
1386            self.faft_client.system.set_try_fw_b(count)
1387
1388