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