1# Copyright (c) 2013 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 logging
6
7from autotest_lib.client.common_lib import error
8from autotest_lib.client.common_lib.cros.network import ping_runner
9from autotest_lib.client.common_lib.cros.network import xmlrpc_datatypes
10from autotest_lib.server import hosts
11from autotest_lib.server import site_linux_router
12from autotest_lib.server.cros import dnsname_mangler
13from autotest_lib.server.cros.network import attenuator_controller
14from autotest_lib.server.cros.network import wifi_client
15
16class WiFiTestContextManager(object):
17    """A context manager for state used in WiFi autotests.
18
19    Some of the building blocks we use in WiFi tests need to be cleaned up
20    after use.  For instance, we start an XMLRPC server on the client
21    which should be shut down so that the next test can start its instance.
22    It is convenient to manage this setup and teardown through a context
23    manager rather than building it into the test class logic.
24
25    """
26    CMDLINE_ATTEN_ADDR = 'atten_addr'
27    CMDLINE_CLIENT_PACKET_CAPTURES = 'client_capture'
28    CMDLINE_CONDUCTIVE_RIG = 'conductive_rig'
29    CMDLINE_PACKET_CAPTURE_SNAPLEN = 'capture_snaplen'
30    CMDLINE_ROUTER_ADDR = 'router_addr'
31    CMDLINE_ROUTER_PACKET_CAPTURES = 'router_capture'
32    CMDLINE_USE_WPA_CLI = 'use_wpa_cli'
33
34
35    @property
36    def attenuator(self):
37        """@return attenuator object (e.g. a BeagleBone)."""
38        if self._attenuator is None:
39            raise error.TestNAError('No attenuator available in this setup.')
40
41        return self._attenuator
42
43
44    @property
45    def client(self):
46        """@return WiFiClient object abstracting the DUT."""
47        return self._client_proxy
48
49
50    @property
51    def router(self):
52        """@return router object (e.g. a LinuxCrosRouter)."""
53        return self._router
54
55
56    def __init__(self, test_name, host, cmdline_args, debug_dir):
57        """Construct a WiFiTestContextManager.
58
59        Optionally can pull addresses of the server address, router address,
60        or router port from cmdline_args.
61
62        @param test_name string descriptive name for this test.
63        @param host host object representing the DUT.
64        @param cmdline_args dict of key, value settings from command line.
65
66        """
67        super(WiFiTestContextManager, self).__init__()
68        self._test_name = test_name
69        self._cmdline_args = cmdline_args.copy()
70        self._client_proxy = wifi_client.WiFiClient(
71                host, debug_dir,
72                self._get_bool_cmdline_value(self.CMDLINE_USE_WPA_CLI, True))
73        self._attenuator = None
74        self._router = None
75        self._enable_client_packet_captures = False
76        self._enable_router_packet_captures = False
77        self._packet_capture_snaplen = None
78
79
80    def __enter__(self):
81        self.setup()
82        return self
83
84
85    def __exit__(self, exc_type, exc_value, traceback):
86        self.teardown()
87
88
89    def _get_bool_cmdline_value(self, key, default_value):
90        """Returns a bool value for the given key from the cmdline args.
91
92        @param key string cmdline args key.
93        @param default_value value to return if the key is not specified in the
94               cmdline args.
95
96        @return True/False or default_value if key is not specified in the
97                cmdline args.
98
99        """
100        if key in self._cmdline_args:
101            value = self._cmdline_args[key].lower()
102            if value in ('1', 'true', 'yes', 'y'):
103                return True
104            else:
105                return False
106        else:
107            return default_value
108
109
110    def get_wifi_addr(self, ap_num=0):
111        """Return an IPv4 address pingable by the client on the WiFi subnet.
112
113        @param ap_num int number of AP.  Only used in stumpy cells.
114        @return string IPv4 address.
115
116        """
117        return self.router.local_server_address(ap_num)
118
119
120    def get_wifi_if(self, ap_num=0):
121        """Returns the interface name for the IP address of self.get_wifi_addr.
122
123        @param ap_num int number of AP.  Only used in stumpy cells.
124        @return string interface name "e.g. wlan0".
125
126        """
127        return self.router.get_hostapd_interface(ap_num)
128
129
130    def get_wifi_host(self):
131        """@return host object representing a pingable machine."""
132        return self.router.host
133
134
135    def configure(self, configuration_parameters, multi_interface=None,
136                  is_ibss=None):
137        """Configure a router with the given parameters.
138
139        Configures an AP according to the specified parameters and
140        enables whatever packet captures are appropriate.  Will deconfigure
141        existing APs unless |multi_interface| is specified.
142
143        @param configuration_parameters HostapConfig object.
144        @param multi_interface True iff having multiple configured interfaces
145                is expected for this configure call.
146        @param is_ibss True iff this is an IBSS endpoint.
147
148        """
149        if not self.client.is_frequency_supported(
150                configuration_parameters.frequency):
151            raise error.TestNAError('DUT does not support frequency: %s' %
152                                    configuration_parameters.frequency)
153        configuration_parameters.security_config.install_router_credentials(
154                self.router.host)
155        if is_ibss:
156            if multi_interface:
157                raise error.TestFail('IBSS mode does not support multiple '
158                                     'interfaces.')
159            if not self.client.is_ibss_supported():
160                raise error.TestNAError('DUT does not support IBSS mode')
161            self.router.ibss_configure(configuration_parameters)
162        else:
163            self.router.hostap_configure(configuration_parameters,
164                                         multi_interface=multi_interface)
165        if self._enable_client_packet_captures:
166            self.client.start_capture(configuration_parameters.frequency,
167                                      snaplen=self._packet_capture_snaplen)
168        if self._enable_router_packet_captures:
169            self.router.start_capture(
170                    configuration_parameters.frequency,
171                    ht_type=configuration_parameters.ht_packet_capture_mode,
172                    snaplen=self._packet_capture_snaplen)
173
174
175    def setup(self):
176        """Construct the state used in a WiFi test."""
177        self._router = site_linux_router.build_router_proxy(
178                test_name=self._test_name,
179                client_hostname=self.client.host.hostname,
180                router_addr=self._cmdline_args.get(self.CMDLINE_ROUTER_ADDR,
181                                                   None))
182        # The attenuator host gives us the ability to attenuate particular
183        # antennas on the router.  Most setups don't have this capability
184        # and most tests do not require it.  We use this for RvR
185        # (network_WiFi_AttenuatedPerf) and some roaming tests.
186        attenuator_addr = dnsname_mangler.get_attenuator_addr(
187                self.client.host.hostname,
188                cmdline_override=self._cmdline_args.get(
189                        self.CMDLINE_ATTEN_ADDR, None),
190                allow_failure=True)
191        ping_helper = ping_runner.PingRunner()
192        if attenuator_addr and ping_helper.simple_ping(attenuator_addr):
193            self._attenuator = attenuator_controller.AttenuatorController(
194                    hosts.SSHHost(attenuator_addr, port=22))
195        # Set up a clean context to conduct WiFi tests in.
196        self.client.shill.init_test_network_state()
197        if self.CMDLINE_CLIENT_PACKET_CAPTURES in self._cmdline_args:
198            self._enable_client_packet_captures = True
199        if self.CMDLINE_ROUTER_PACKET_CAPTURES in self._cmdline_args:
200            self._enable_router_packet_captures = True
201        if self.CMDLINE_PACKET_CAPTURE_SNAPLEN in self._cmdline_args:
202            self._packet_capture_snaplen = int(
203                    self._cmdline_args[self.CMDLINE_PACKET_CAPTURE_SNAPLEN])
204        self.client.conductive = self._get_bool_cmdline_value(
205                self.CMDLINE_CONDUCTIVE_RIG, None)
206        for system in (self.client, self.router):
207            system.sync_host_times()
208
209
210    def teardown(self):
211        """Teardown the state used in a WiFi test."""
212        logging.debug('Tearing down the test context.')
213        for system in [self._attenuator, self._client_proxy,
214                       self._router]:
215            if system is not None:
216                system.close()
217
218
219    def assert_connect_wifi(self, wifi_params, description=None):
220        """Connect to a WiFi network and check for success.
221
222        Connect a DUT to a WiFi network and check that we connect successfully.
223
224        @param wifi_params AssociationParameters describing network to connect.
225        @param description string Additional text for logging messages.
226
227        @returns AssociationResult if successful; None if wifi_params
228                 contains expect_failure; asserts otherwise.
229
230        """
231        if description:
232            connect_name = '%s (%s)' % (wifi_params.ssid, description)
233        else:
234            connect_name = '%s' % wifi_params.ssid
235        logging.info('Connecting to %s.', connect_name)
236        assoc_result = xmlrpc_datatypes.deserialize(
237                self.client.shill.connect_wifi(wifi_params))
238        logging.info('Finished connection attempt to %s with times: '
239                     'discovery=%.2f, association=%.2f, configuration=%.2f.',
240                     connect_name,
241                     assoc_result.discovery_time,
242                     assoc_result.association_time,
243                     assoc_result.configuration_time)
244
245        if assoc_result.success and wifi_params.expect_failure:
246            raise error.TestFail(
247                'Expected connection to %s to fail, but it was successful.' %
248                connect_name)
249
250        if not assoc_result.success and not wifi_params.expect_failure:
251            raise error.TestFail(
252                'Expected connection to %s to succeed, '
253                'but it failed with reason: %s.' % (
254                    connect_name, assoc_result.failure_reason))
255
256        if wifi_params.expect_failure:
257            logging.info('Unable to connect to %s, as intended.',
258                         connect_name)
259            return None
260
261        logging.info('Connected successfully to %s.', connect_name)
262        return assoc_result
263
264
265    def assert_ping_from_dut(self, ping_config=None, ap_num=None):
266        """Ping a host on the WiFi network from the DUT.
267
268        Ping a host reachable on the WiFi network from the DUT, and
269        check that the ping is successful.  The host we ping depends
270        on the test setup, sometimes that host may be the server and
271        sometimes it will be the router itself.  Ping-ability may be
272        used to confirm that a WiFi network is operating correctly.
273
274        @param ping_config optional PingConfig object to override defaults.
275        @param ap_num int which AP to ping if more than one is configured.
276
277        """
278        if ap_num is None:
279            ap_num = 0
280        if ping_config is None:
281            ping_ip = self.router.get_wifi_ip(ap_num=ap_num)
282            ping_config = ping_runner.PingConfig(ping_ip)
283        self.client.ping(ping_config)
284
285
286    def assert_ping_from_server(self, ping_config=None):
287        """Ping the DUT across the WiFi network from the server.
288
289        Check that the ping is mostly successful and fail the test if it
290        is not.
291
292        @param ping_config optional PingConfig object to override defaults.
293
294        """
295        logging.info('Pinging from server.')
296        if ping_config is None:
297            ping_ip = self.client.wifi_ip
298            ping_config = ping_runner.PingConfig(ping_ip)
299        self.router.ping(ping_config)
300
301
302    def wait_for_connection(self, ssid, freq=None, ap_num=None,
303                            timeout_seconds=30):
304        """Verifies a connection to network ssid on frequency freq.
305
306        @param ssid string ssid of the network to check.
307        @param freq int frequency of network to check.
308        @param ap_num int AP to which to connect
309        @param timeout_seconds int number of seconds to wait for
310                connection on the given frequency.
311
312        @returns a named tuple of (state, time)
313        """
314        if ap_num is None:
315            ap_num = 0
316        desired_subnet = self.router.get_wifi_ip_subnet(ap_num)
317        wifi_ip = self.router.get_wifi_ip(ap_num)
318        return self.client.wait_for_connection(
319                ssid, timeout_seconds=timeout_seconds, freq=freq,
320                ping_ip=wifi_ip, desired_subnet=desired_subnet)
321