1#!/usr/bin/env python3
2#
3#   Copyright 2023 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import logging
18import paramiko
19import time
20from typing import List
21
22from acts.test_decorators import test_tracker_info
23import acts_contrib.test_utils.wifi.wifi_test_utils as wutils
24from acts_contrib.test_utils.wifi.WifiBaseTest import WifiBaseTest
25from acts import signals
26from acts.controllers.utils_lib.ssh import connection
27
28
29_POLL_AP_RETRY_INTERVAL_SEC = 1
30_WAIT_OPENWRT_AP_BOOT_SEC = 30
31
32class WifiPreTest(WifiBaseTest):
33  """ Wi-Fi PreTest."""
34  def __init__(self, configs):
35    super().__init__(configs)
36    self.enable_packet_log = True
37
38  def setup_class(self):
39    super().setup_class()
40
41    req_params = ["OpenWrtAP"]
42    self.unpack_userparams(req_param_names=req_params)
43
44    self.dut = self.android_devices[0]
45
46    # Reboot OpenWrt APs.
47    if "OpenWrtAP" in self.user_params:
48      for i, openwrt in enumerate(self.access_points):
49        logging.info(f"Rebooting OpenWrt AP: {openwrt.ssh_settings.hostname}")
50        openwrt.reboot()
51    time.sleep(_WAIT_OPENWRT_AP_BOOT_SEC)
52
53    # Polling OpenWrt APs until they are ready.
54    for i, openwrt in enumerate(self.access_points):
55      if self.poll_openwrt_over_ssh(openwrt):
56        continue
57      else:
58        raise signals.TestFailure(
59          f"Unable to connect to OpenWrt AP: {openwrt.ssh_settings.hostname}")
60
61    self.start_openwrt()
62
63    wutils.list_scan_results(self.dut, wait_time=30)
64
65  def setup_test(self):
66    super().setup_test()
67    self.dut.droid.wakeLockAcquireBright()
68    self.dut.droid.wakeUpNow()
69    wutils.wifi_toggle_state(self.dut, True)
70    wutils.reset_wifi(self.dut)
71
72  def teardown_test(self):
73    super().teardown_test()
74    self.dut.droid.wakeLockRelease()
75    self.dut.droid.goToSleepNow()
76    wutils.reset_wifi(self.dut)
77
78  def poll_openwrt_over_ssh(self,openwrt,
79                            retry_duration: int=60):
80    """
81    Attempt to establish an SSH connection with the device at the given IP address and port.
82
83    Args:
84      ip: The IP address of the device to connect to.
85      port: The port number for SSH connection.
86      username: The username for SSH authentication.
87      password: The password for SSH authentication.
88      retry_duration: The maximum duration in seconds to attempt reconnection
89                                      before giving up.
90
91    Returns:
92      bool: True if the connection was successful, False otherwise.
93    """
94    ip = openwrt.ssh_settings.hostname
95    start_time = time.time()
96    while time.time() - start_time < retry_duration:
97      try:
98        logging.info(f"Attempt to connect to {ip}")
99        openwrt.close()
100        openwrt.ssh = connection.SshConnection(openwrt.ssh_settings)
101        openwrt.ssh.setup_master_ssh()
102        return True
103      except (
104          paramiko.ssh_exception.NoValidConnectionsError,
105          paramiko.ssh_exception.AuthenticationException,
106          paramiko.ssh_exception.SSHException,
107          connection.Error,
108          TimeoutError,
109      ) as e:
110        logging.info(f"Connection error: {e}, reconnecting {ip} "
111                      f"in {retry_duration} seconds.")
112        time.sleep(_POLL_AP_RETRY_INTERVAL_SEC)
113    logging.info(f"Connection attempts exhausted. Unable to connect to {ip}.")
114    return False
115
116  def start_openwrt(self):
117    """Enable OpenWrts to generate Open Wi-Fi networks."""
118
119    if "OpenWrtAP" in self.user_params:
120      logging.info("Launching OpenWrt AP...")
121      self.configure_openwrt_ap_and_start(open_network=True,
122                                          ap_count=len(self.access_points))
123      self.open_networks = []
124      for i in range(len(self.access_points)):
125        self.open_networks.append(self.open_network[i]["2g"])
126        self.open_networks.append(self.open_network[i]["5g"])
127
128      # stdout APs' information.
129      for i, openwrt in enumerate(self.access_points):
130        openwrt.log.info(f"AP_{i} Info: ")
131        openwrt.log.info(f"IP address: {openwrt.ssh_settings.hostname}")
132
133        radios = ["radio0", "radio1"]
134        for radio in radios:
135          ssid_radio_map = openwrt.get_ifnames_for_ssids(radio)
136          for ssid, radio_ifname in ssid_radio_map.items():
137            openwrt.log.info(f"{radio_ifname}:  {ssid}")
138
139        band_bssid_map = openwrt.get_bssids_for_wifi_networks()
140        openwrt.log.info(band_bssid_map)
141
142  @test_tracker_info(uuid="913605ea-38bf-492c-b634-d1823caed4b3")
143  def test_connects_all_testbed_wifi_networks(self):
144    """Test whether the DUT can successfully connect to restarted APs."""
145    for network in self.open_networks:
146      wutils.connect_to_wifi_network(self.dut, network)
147