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