1#  Copyright (C) 2024 The Android Open Source Project
2#
3#  Licensed under the Apache License, Version 2.0 (the "License");
4#  you may not use this file except in compliance with the License.
5#  You may obtain a copy of the License at
6#
7#       http://www.apache.org/licenses/LICENSE-2.0
8#
9#  Unless required by applicable law or agreed to in writing, software
10#  distributed under the License is distributed on an "AS IS" BASIS,
11#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#  See the License for the specific language governing permissions and
13#  limitations under the License.
14
15"""Utils for Android Wi-Fi operations."""
16
17import dataclasses
18import datetime
19import enum
20import re
21import time
22
23from mobly.controllers import android_device
24from mobly.controllers.android_device_lib import adb
25
26_DELAY_AFTER_CHANGE_WIFI_STATUS = datetime.timedelta(seconds=5)
27_WAIT_FOR_CONNECTION = datetime.timedelta(seconds=30)
28
29_SAVED_WIFI_LIST_PATTERN = re.compile(
30    r'(?P<id>\d+)\s+(?P<ssid>.*)\s+(?P<security>.*)'
31)
32_SSID_PATTERN = re.compile(rb'Wifi is connected to "(?P<ssid>.*?)"')
33
34
35@dataclasses.dataclass(frozen=True)
36class SavedWifiInfo:
37  """Information about a saved Wi-Fi network."""
38
39  id: str
40  ssid: str
41  security: str
42
43
44@enum.unique
45class WiFiSecurity(enum.StrEnum):
46  """Security type of the Wi-Fi network."""
47
48  OPEN = 'open'
49  OWE = 'owe'
50  WPA2 = 'wpa2'
51  WPA3 = 'wpa3'
52  WEP = 'wep'
53
54
55class AndroidWiFiError(Exception):
56  """Error when failed to operate Wi-Fi on Android device."""
57
58  def __init__(self, ad: android_device.AndroidDevice, message: str) -> None:
59    self._ad = ad
60    self._message = message
61
62  def __str__(self) -> str:
63    return self._message if self._ad is None else f'{self._ad} {self._message}'
64
65
66def connect_to_wifi(
67    ad: android_device.AndroidDevice,
68    ssid: str,
69    passphrase: str | None = None,
70    security: WiFiSecurity | None = None,
71) -> None:
72  """Connects to a W-Fi network and adds to saved networks list."""
73  enable_wifi(ad)
74  if get_current_wifi(ad) == ssid:
75    ad.log.info(f'Wi-Fi was already connected to {repr(ssid)}')
76    return
77  cmd = ['cmd', 'wifi', 'connect-network', f'"{ssid}"']
78  if passphrase is None:
79    cmd.append(f'"{security or WiFiSecurity.OPEN}"')
80  else:
81    cmd.extend([f'"{security or WiFiSecurity.WPA2}"', f'"{passphrase}"'])
82  ad.adb.shell(cmd)
83  if not _wait_for_data_connected(ad) or get_current_wifi(ad) != ssid:
84    raise AndroidWiFiError(ad, f'Fail to connect to Wi-Fi {repr(ssid)}')
85  ad.log.info(f'Wi-Fi connected to {repr(ssid)}')
86
87
88def disable_wifi(ad: android_device.AndroidDevice) -> None:
89  """Disables Wi-Fi."""
90  if not is_wifi_enabled(ad):
91    ad.log.info('Wi-Fi was already disabled.')
92    return
93  ad.log.info('Disabling Wi-Fi...')
94  ad.adb.shell(['cmd', 'wifi', 'set-wifi-enabled', 'disabled'])
95  start_time = time.monotonic()
96  timeout = start_time + _DELAY_AFTER_CHANGE_WIFI_STATUS.total_seconds()
97  while time.monotonic() < timeout:
98    if not is_wifi_enabled(ad):
99      ad.log.info('Wi-Fi is disabled.')
100      return
101  raise AndroidWiFiError(
102      ad,
103      'Fail to disable Wi-Fi after waiting for'
104      f' {_DELAY_AFTER_CHANGE_WIFI_STATUS.total_seconds()} seconds',
105  )
106
107
108def enable_wifi(ad: android_device.AndroidDevice) -> None:
109  """Enables Wi-Fi."""
110  if is_wifi_enabled(ad):
111    ad.log.info('Wi-Fi was already enabled.')
112    return
113  ad.log.info('Enabling Wi-Fi...')
114  ad.adb.shell(['cmd', 'wifi', 'set-wifi-enabled', 'enabled'])
115  start_time = time.monotonic()
116  timeout = start_time + _DELAY_AFTER_CHANGE_WIFI_STATUS.total_seconds()
117  while time.monotonic() < timeout:
118    if is_wifi_enabled(ad):
119      ad.log.info('Wi-Fi is enabled.')
120      return
121  raise AndroidWiFiError(
122      ad,
123      'Fail to enable Wi-Fi after waiting for'
124      f' {_DELAY_AFTER_CHANGE_WIFI_STATUS.total_seconds()} seconds',
125  )
126
127
128def forget_all_wifi(ad: android_device.AndroidDevice) -> None:
129  """Forgets all Wi-Fi from saved networks list."""
130  for saved_wifi in list_saved_wifi(ad):
131    ad.adb.shell(['cmd', 'wifi', 'forget-network', saved_wifi.id])
132  saved_wifis = list_saved_wifi(ad)
133  if saved_wifis:
134    raise AndroidWiFiError(
135        ad,
136        'Fail to forget all Wi-Fi networks, remaining in the list:'
137        f' {saved_wifis}',
138    )
139
140
141def forget_wifi(ad: android_device.AndroidDevice, ssid: str) -> None:
142  """Forgets Wi-Fi from saved networks list."""
143  saved_wifis = list_saved_wifi(ad)
144  for saved_wifi in saved_wifis:
145    if saved_wifi.ssid == ssid:
146      stdout = ad.adb.shell(['cmd', 'wifi', 'forget-network', saved_wifi.id])
147      if stdout == b'Forget successful\n':
148        ad.log.info(f'Wi-Fi {repr(ssid)} was forgotten from saved networks.')
149        return
150      raise AndroidWiFiError(ad, f'Fail to forget Wi-Fi {repr(ssid)}')
151  ad.log.info(f'Nothing was deleted since Wi-Fi {repr(ssid)} was not saved')
152  return
153
154
155def get_current_wifi(ad: android_device.AndroidDevice) -> str:
156  """Returns current Wi-Fi network."""
157  if match := _SSID_PATTERN.search(ad.adb.shell(['cmd', 'wifi', 'status'])):
158    return match.group('ssid').decode()
159  return ''
160
161
162def is_wifi_enabled(ad: android_device.AndroidDevice) -> bool:
163  """Returns True if Wi-Fi is enabled, False otherwise."""
164  return ad.adb.shell(['cmd', 'wifi', 'status']).startswith(b'Wifi is enabled')
165
166
167def list_saved_wifi(ad: android_device.AndroidDevice) -> list[SavedWifiInfo]:
168  """Returns list of saved Wi-Fi networks."""
169  saved_wifis = []
170  stdout = ad.adb.shell(['cmd', 'wifi', 'list-networks']).decode()
171  if stdout == 'No networks\n':
172    return saved_wifis
173  for line in stdout.splitlines()[1:]:
174    if match := _SAVED_WIFI_LIST_PATTERN.search(line):
175      saved_wifis.append(
176          SavedWifiInfo(
177              id=match.group('id'),
178              ssid=match.group('ssid').strip(),
179              security=match.group('security'),
180          )
181      )
182  return saved_wifis
183
184
185def _is_data_connected(ad: android_device.AndroidDevice) -> bool:
186  """Returns True if data is connected, False otherwise."""
187  try:
188    return b'5 received' in ad.adb.shell(['ping', '-c', '5', '8.8.8.8'])
189  except adb.AdbError:
190    return False
191
192
193def _wait_for_data_connected(
194    ad: android_device.AndroidDevice,
195    timeout: datetime.timedelta = _WAIT_FOR_CONNECTION,
196) -> bool:
197  """Returns True if data is connected before timeout, False otherwise."""
198  start_time = time.monotonic()
199  timeout = start_time + timeout.total_seconds()
200  while time.monotonic() < timeout:
201    if _is_data_connected(ad):
202      return True
203  return False
204