1# Copyright 2015 The Chromium 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
5"""Provides a variety of device interactions based on fastboot."""
6# pylint: disable=unused-argument
7
8import collections
9import contextlib
10import fnmatch
11import logging
12import os
13import re
14
15from devil.android import decorators
16from devil.android import device_errors
17from devil.android.sdk import fastboot
18from devil.utils import timeout_retry
19
20_DEFAULT_TIMEOUT = 30
21_DEFAULT_RETRIES = 3
22_FASTBOOT_REBOOT_TIMEOUT = 10 * _DEFAULT_TIMEOUT
23_KNOWN_PARTITIONS = collections.OrderedDict([
24      ('bootloader', {'image': 'bootloader*.img', 'restart': True}),
25      ('radio', {'image': 'radio*.img', 'restart': True}),
26      ('boot', {'image': 'boot.img'}),
27      ('recovery', {'image': 'recovery.img'}),
28      ('system', {'image': 'system.img'}),
29      ('userdata', {'image': 'userdata.img', 'wipe_only': True}),
30      ('cache', {'image': 'cache.img', 'wipe_only': True}),
31      ('vendor', {'image': 'vendor*.img', 'optional': True}),
32  ])
33ALL_PARTITIONS = _KNOWN_PARTITIONS.keys()
34
35
36def _FindAndVerifyPartitionsAndImages(partitions, directory):
37  """Validate partitions and images.
38
39  Validate all partition names and partition directories. Cannot stop mid
40  flash so its important to validate everything first.
41
42  Args:
43    Partitions: partitions to be tested.
44    directory: directory containing the images.
45
46  Returns:
47    Dictionary with exact partition, image name mapping.
48  """
49
50  files = os.listdir(directory)
51  return_dict = collections.OrderedDict()
52
53  def find_file(pattern):
54    for filename in files:
55      if fnmatch.fnmatch(filename, pattern):
56        return os.path.join(directory, filename)
57    return None
58  for partition in partitions:
59    partition_info = _KNOWN_PARTITIONS[partition]
60    image_file = find_file(partition_info['image'])
61    if image_file:
62      return_dict[partition] = image_file
63    elif not partition_info.get('optional'):
64      raise device_errors.FastbootCommandFailedError(
65          'Failed to flash device. Could not find image for %s.',
66          partition_info['image'])
67  return return_dict
68
69
70class FastbootUtils(object):
71
72  _FASTBOOT_WAIT_TIME = 1
73  _BOARD_VERIFICATION_FILE = 'android-info.txt'
74
75  def __init__(self, device, fastbooter=None, default_timeout=_DEFAULT_TIMEOUT,
76               default_retries=_DEFAULT_RETRIES):
77    """FastbootUtils constructor.
78
79    Example Usage to flash a device:
80      fastboot = fastboot_utils.FastbootUtils(device)
81      fastboot.FlashDevice('/path/to/build/directory')
82
83    Args:
84      device: A DeviceUtils instance.
85      fastbooter: Optional fastboot object. If none is passed, one will
86        be created.
87      default_timeout: An integer containing the default number of seconds to
88        wait for an operation to complete if no explicit value is provided.
89      default_retries: An integer containing the default number or times an
90        operation should be retried on failure if no explicit value is provided.
91    """
92    self._device = device
93    self._board = device.product_board
94    self._serial = str(device)
95    self._default_timeout = default_timeout
96    self._default_retries = default_retries
97    if fastbooter:
98      self.fastboot = fastbooter
99    else:
100      self.fastboot = fastboot.Fastboot(self._serial)
101
102  @decorators.WithTimeoutAndRetriesFromInstance()
103  def WaitForFastbootMode(self, timeout=None, retries=None):
104    """Wait for device to boot into fastboot mode.
105
106    This waits for the device serial to show up in fastboot devices output.
107    """
108    def fastboot_mode():
109      return self._serial in self.fastboot.Devices()
110
111    timeout_retry.WaitFor(fastboot_mode, wait_period=self._FASTBOOT_WAIT_TIME)
112
113  @decorators.WithTimeoutAndRetriesFromInstance(
114      min_default_timeout=_FASTBOOT_REBOOT_TIMEOUT)
115  def EnableFastbootMode(self, timeout=None, retries=None):
116    """Reboots phone into fastboot mode.
117
118    Roots phone if needed, then reboots phone into fastboot mode and waits.
119    """
120    self._device.EnableRoot()
121    self._device.adb.Reboot(to_bootloader=True)
122    self.WaitForFastbootMode()
123
124  @decorators.WithTimeoutAndRetriesFromInstance(
125      min_default_timeout=_FASTBOOT_REBOOT_TIMEOUT)
126  def Reboot(
127      self, bootloader=False, wait_for_reboot=True, timeout=None, retries=None):
128    """Reboots out of fastboot mode.
129
130    It reboots the phone either back into fastboot, or to a regular boot. It
131    then blocks until the device is ready.
132
133    Args:
134      bootloader: If set to True, reboots back into bootloader.
135    """
136    if bootloader:
137      self.fastboot.RebootBootloader()
138      self.WaitForFastbootMode()
139    else:
140      self.fastboot.Reboot()
141      if wait_for_reboot:
142        self._device.WaitUntilFullyBooted(timeout=_FASTBOOT_REBOOT_TIMEOUT)
143
144  def _VerifyBoard(self, directory):
145    """Validate as best as possible that the android build matches the device.
146
147    Goes through build files and checks if the board name is mentioned in the
148    |self._BOARD_VERIFICATION_FILE| or in the build archive.
149
150    Args:
151      directory: directory where build files are located.
152    """
153    files = os.listdir(directory)
154    board_regex = re.compile(r'require board=(\w+)')
155    if self._BOARD_VERIFICATION_FILE in files:
156      with open(os.path.join(directory, self._BOARD_VERIFICATION_FILE)) as f:
157        for line in f:
158          m = board_regex.match(line)
159          if m:
160            board_name = m.group(1)
161            if board_name == self._board:
162              return True
163            elif board_name:
164              return False
165            else:
166              logging.warning('No board type found in %s.',
167                              self._BOARD_VERIFICATION_FILE)
168    else:
169      logging.warning('%s not found. Unable to use it to verify device.',
170                      self._BOARD_VERIFICATION_FILE)
171
172    zip_regex = re.compile(r'.*%s.*\.zip' % re.escape(self._board))
173    for f in files:
174      if zip_regex.match(f):
175        return True
176
177    return False
178
179  def _FlashPartitions(self, partitions, directory, wipe=False, force=False):
180    """Flashes all given partiitons with all given images.
181
182    Args:
183      partitions: List of partitions to flash.
184      directory: Directory where all partitions can be found.
185      wipe: If set to true, will automatically detect if cache and userdata
186          partitions are sent, and if so ignore them.
187      force: boolean to decide to ignore board name safety checks.
188
189    Raises:
190      device_errors.CommandFailedError(): If image cannot be found or if bad
191          partition name is give.
192    """
193    if not self._VerifyBoard(directory):
194      if force:
195        logging.warning('Could not verify build is meant to be installed on '
196                        'the current device type, but force flag is set. '
197                        'Flashing device. Possibly dangerous operation.')
198      else:
199        raise device_errors.CommandFailedError(
200            'Could not verify build is meant to be installed on the current '
201            'device type. Run again with force=True to force flashing with an '
202            'unverified board.')
203
204    flash_image_files = _FindAndVerifyPartitionsAndImages(partitions, directory)
205    partitions = flash_image_files.keys()
206    for partition in partitions:
207      if _KNOWN_PARTITIONS[partition].get('wipe_only') and not wipe:
208        logging.info(
209            'Not flashing in wipe mode. Skipping partition %s.', partition)
210      else:
211        logging.info(
212            'Flashing %s with %s', partition, flash_image_files[partition])
213        self.fastboot.Flash(partition, flash_image_files[partition])
214        if _KNOWN_PARTITIONS[partition].get('restart', False):
215          self.Reboot(bootloader=True)
216
217  @contextlib.contextmanager
218  def FastbootMode(self, wait_for_reboot=True, timeout=None, retries=None):
219    """Context manager that enables fastboot mode, and reboots after.
220
221    Example usage:
222      with FastbootMode():
223        Flash Device
224      # Anything that runs after flashing.
225    """
226    self.EnableFastbootMode()
227    self.fastboot.SetOemOffModeCharge(False)
228    try:
229      yield self
230    finally:
231      self.fastboot.SetOemOffModeCharge(True)
232      self.Reboot(wait_for_reboot=wait_for_reboot)
233
234  def FlashDevice(self, directory, partitions=None, wipe=False):
235    """Flash device with build in |directory|.
236
237    Directory must contain bootloader, radio, boot, recovery, system, userdata,
238    and cache .img files from an android build. This is a dangerous operation so
239    use with care.
240
241    Args:
242      fastboot: A FastbootUtils instance.
243      directory: Directory with build files.
244      wipe: Wipes cache and userdata if set to true.
245      partitions: List of partitions to flash. Defaults to all.
246    """
247    if partitions is None:
248      partitions = ALL_PARTITIONS
249    # If a device is wiped, then it will no longer have adb keys so it cannot be
250    # communicated with to verify that it is rebooted. It is up to the user of
251    # this script to ensure that the adb keys are set on the device after using
252    # this to wipe a device.
253    with self.FastbootMode(wait_for_reboot=not wipe):
254      self._FlashPartitions(partitions, directory, wipe=wipe)
255