1#!/usr/bin/env python3
2#
3#   Copyright 2018 - 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 os
19import subprocess
20import socket
21import threading
22
23from acts import context
24from acts import utils
25from acts.controllers.adb_lib.error import AdbCommandError
26from acts.controllers.android_device import AndroidDevice
27from acts.controllers.iperf_server import _AndroidDeviceBridge
28from acts.controllers.fuchsia_lib.utils_lib import create_ssh_connection
29from acts.controllers.fuchsia_lib.utils_lib import ssh_is_connected
30from acts.controllers.fuchsia_lib.utils_lib import SshResults
31from acts.controllers.utils_lib.ssh import connection
32from acts.controllers.utils_lib.ssh import settings
33from acts.event import event_bus
34from acts.event.decorators import subscribe_static
35from acts.event.event import TestClassBeginEvent
36from acts.event.event import TestClassEndEvent
37from acts.libs.proc import job
38from paramiko.buffered_pipe import PipeTimeout
39MOBLY_CONTROLLER_CONFIG_NAME = 'IPerfClient'
40ACTS_CONTROLLER_REFERENCE_NAME = 'iperf_clients'
41
42
43class IPerfError(Exception):
44    """Raised on execution errors of iPerf."""
45
46
47def create(configs):
48    """Factory method for iperf clients.
49
50    The function creates iperf clients based on at least one config.
51    If configs contain ssh settings or and AndroidDevice, remote iperf clients
52    will be started on those devices, otherwise, a the client will run on the
53    local machine.
54
55    Args:
56        configs: config parameters for the iperf server
57    """
58    results = []
59    for c in configs:
60        if type(c) is dict and 'AndroidDevice' in c:
61            results.append(
62                IPerfClientOverAdb(c['AndroidDevice'],
63                                   test_interface=c.get('test_interface')))
64        elif type(c) is dict and 'ssh_config' in c:
65            results.append(
66                IPerfClientOverSsh(c['ssh_config'],
67                                   use_paramiko=c.get('use_paramiko'),
68                                   test_interface=c.get('test_interface')))
69        else:
70            results.append(IPerfClient())
71    return results
72
73
74def get_info(iperf_clients):
75    """Placeholder for info about iperf clients
76
77    Returns:
78        None
79    """
80    return None
81
82
83def destroy(_):
84    # No cleanup needed.
85    pass
86
87
88class IPerfClientBase(object):
89    """The Base class for all IPerfClients.
90
91    This base class is responsible for synchronizing the logging to prevent
92    multiple IPerfClients from writing results to the same file, as well
93    as providing the interface for IPerfClient objects.
94    """
95    # Keeps track of the number of IPerfClient logs to prevent file name
96    # collisions.
97    __log_file_counter = 0
98
99    __log_file_lock = threading.Lock()
100
101    @staticmethod
102    def _get_full_file_path(tag=''):
103        """Returns the full file path for the IPerfClient log file.
104
105        Note: If the directory for the file path does not exist, it will be
106        created.
107
108        Args:
109            tag: The tag passed in to the server run.
110        """
111        current_context = context.get_current_context()
112        full_out_dir = os.path.join(current_context.get_full_output_path(),
113                                    'iperf_client_files')
114
115        with IPerfClientBase.__log_file_lock:
116            os.makedirs(full_out_dir, exist_ok=True)
117            tags = ['IPerfClient', tag, IPerfClientBase.__log_file_counter]
118            out_file_name = '%s.log' % (','.join(
119                [str(x) for x in tags if x != '' and x is not None]))
120            IPerfClientBase.__log_file_counter += 1
121
122        return os.path.join(full_out_dir, out_file_name)
123
124    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
125        """Starts iperf client, and waits for completion.
126
127        Args:
128            ip: iperf server ip address.
129            iperf_args: A string representing arguments to start iperf
130                client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
131            tag: A string to further identify iperf results file
132            timeout: the maximum amount of time the iperf client can run.
133            iperf_binary: Location of iperf3 binary. If none, it is assumed the
134                the binary is in the path.
135
136        Returns:
137            full_out_path: iperf result path.
138        """
139        raise NotImplementedError('start() must be implemented.')
140
141
142class IPerfClient(IPerfClientBase):
143    """Class that handles iperf3 client operations."""
144    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
145        """Starts iperf client, and waits for completion.
146
147        Args:
148            ip: iperf server ip address.
149            iperf_args: A string representing arguments to start iperf
150            client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
151            tag: tag to further identify iperf results file
152            timeout: unused.
153            iperf_binary: Location of iperf3 binary. If none, it is assumed the
154                the binary is in the path.
155
156        Returns:
157            full_out_path: iperf result path.
158        """
159        if not iperf_binary:
160            logging.debug('No iperf3 binary specified.  '
161                          'Assuming iperf3 is in the path.')
162            iperf_binary = 'iperf3'
163        else:
164            logging.debug('Using iperf3 binary located at %s' % iperf_binary)
165        iperf_cmd = [str(iperf_binary), '-c', ip] + iperf_args.split(' ')
166        full_out_path = self._get_full_file_path(tag)
167
168        with open(full_out_path, 'w') as out_file:
169            subprocess.call(iperf_cmd, stdout=out_file)
170
171        return full_out_path
172
173
174class IPerfClientOverSsh(IPerfClientBase):
175    """Class that handles iperf3 client operations on remote machines."""
176    def __init__(self, ssh_config, use_paramiko=False, test_interface=None):
177        self._ssh_settings = settings.from_config(ssh_config)
178        if not (utils.is_valid_ipv4_address(self._ssh_settings.hostname)
179                or utils.is_valid_ipv6_address(self._ssh_settings.hostname)):
180            mdns_ip = utils.get_fuchsia_mdns_ipv6_address(
181                self._ssh_settings.hostname)
182            if mdns_ip:
183                self._ssh_settings.hostname = mdns_ip
184        # use_paramiko may be passed in as a string (from JSON), so this line
185        # guarantees it is a converted to a bool.
186        self._use_paramiko = str(use_paramiko).lower() == 'true'
187        self._ssh_session = None
188        self.start_ssh()
189
190        self.hostname = self._ssh_settings.hostname
191        self.test_interface = test_interface
192
193    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
194        """Starts iperf client, and waits for completion.
195
196        Args:
197            ip: iperf server ip address.
198            iperf_args: A string representing arguments to start iperf
199            client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
200            tag: tag to further identify iperf results file
201            timeout: the maximum amount of time to allow the iperf client to run
202            iperf_binary: Location of iperf3 binary. If none, it is assumed the
203                the binary is in the path.
204
205        Returns:
206            full_out_path: iperf result path.
207        """
208        if not iperf_binary:
209            logging.debug('No iperf3 binary specified.  '
210                          'Assuming iperf3 is in the path.')
211            iperf_binary = 'iperf3'
212        else:
213            logging.debug('Using iperf3 binary located at %s' % iperf_binary)
214        iperf_cmd = '{} -c {} {}'.format(iperf_binary, ip, iperf_args)
215        full_out_path = self._get_full_file_path(tag)
216
217        try:
218            if not self._ssh_session:
219                self.start_ssh()
220            if self._use_paramiko:
221                if not ssh_is_connected(self._ssh_session):
222                    logging.info('Lost SSH connection to %s. Reconnecting.' %
223                                 self._ssh_settings.hostname)
224                    self._ssh_session.close()
225                    self._ssh_session = create_ssh_connection(
226                        ip_address=self._ssh_settings.hostname,
227                        ssh_username=self._ssh_settings.username,
228                        ssh_config=self._ssh_settings.ssh_config)
229                cmd_result_stdin, cmd_result_stdout, cmd_result_stderr = (
230                    self._ssh_session.exec_command(iperf_cmd, timeout=timeout))
231                iperf_process = SshResults(cmd_result_stdin, cmd_result_stdout,
232                                           cmd_result_stderr,
233                                           cmd_result_stdout.channel)
234            else:
235                iperf_process = self._ssh_session.run(iperf_cmd,
236                                                      timeout=timeout)
237            iperf_output = iperf_process.stdout
238            with open(full_out_path, 'w') as out_file:
239                out_file.write(iperf_output)
240        except PipeTimeout:
241            raise TimeoutError('Paramiko PipeTimeout. Timed out waiting for '
242                               'iperf client to finish.')
243        except socket.timeout:
244            raise TimeoutError('Socket timeout. Timed out waiting for iperf '
245                               'client to finish.')
246        except Exception as e:
247            logging.exception('iperf run failed.')
248
249        return full_out_path
250
251    def start_ssh(self):
252        """Starts an ssh session to the iperf client."""
253        if not self._ssh_session:
254            if self._use_paramiko:
255                self._ssh_session = create_ssh_connection(
256                    ip_address=self._ssh_settings.hostname,
257                    ssh_username=self._ssh_settings.username,
258                    ssh_config=self._ssh_settings.ssh_config)
259            else:
260                self._ssh_session = connection.SshConnection(
261                    self._ssh_settings)
262
263    def close_ssh(self):
264        """Closes the ssh session to the iperf client, if one exists, preventing
265        connection reset errors when rebooting client device.
266        """
267        if self._ssh_session:
268            self._ssh_session.close()
269            self._ssh_session = None
270
271
272class IPerfClientOverAdb(IPerfClientBase):
273    """Class that handles iperf3 operations over ADB devices."""
274    def __init__(self, android_device_or_serial, test_interface=None):
275        """Creates a new IPerfClientOverAdb object.
276
277        Args:
278            android_device_or_serial: Either an AndroidDevice object, or the
279                serial that corresponds to the AndroidDevice. Note that the
280                serial must be present in an AndroidDevice entry in the ACTS
281                config.
282            test_interface: The network interface that will be used to send
283                traffic to the iperf server.
284        """
285        self._android_device_or_serial = android_device_or_serial
286        self.test_interface = test_interface
287
288    @property
289    def _android_device(self):
290        if isinstance(self._android_device_or_serial, AndroidDevice):
291            return self._android_device_or_serial
292        else:
293            return _AndroidDeviceBridge.android_devices()[
294                self._android_device_or_serial]
295
296    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
297        """Starts iperf client, and waits for completion.
298
299        Args:
300            ip: iperf server ip address.
301            iperf_args: A string representing arguments to start iperf
302            client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
303            tag: tag to further identify iperf results file
304            timeout: the maximum amount of time to allow the iperf client to run
305            iperf_binary: Location of iperf3 binary. If none, it is assumed the
306                the binary is in the path.
307
308        Returns:
309            The iperf result file path.
310        """
311        clean_out = ''
312        try:
313            if not iperf_binary:
314                logging.debug('No iperf3 binary specified.  '
315                              'Assuming iperf3 is in the path.')
316                iperf_binary = 'iperf3'
317            else:
318                logging.debug('Using iperf3 binary located at %s' %
319                              iperf_binary)
320            iperf_cmd = '{} -c {} {}'.format(iperf_binary, ip, iperf_args)
321            out = self._android_device.adb.shell(str(iperf_cmd),
322                                                 timeout=timeout)
323            clean_out = out.split('\n')
324            if 'error' in clean_out[0].lower():
325                raise IPerfError(clean_out)
326        except (job.TimeoutError, AdbCommandError):
327            logging.warning('TimeoutError: Iperf measurement failed.')
328
329        full_out_path = self._get_full_file_path(tag)
330        with open(full_out_path, 'w') as out_file:
331            out_file.write('\n'.join(clean_out))
332
333        return full_out_path
334