1#!/usr/bin/env python3.4
2#
3#   Copyright 2021 - 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 collections
18import hashlib
19import logging
20import math
21import os
22import re
23import statistics
24import numpy
25import time
26from acts import asserts
27
28SHORT_SLEEP = 1
29MED_SLEEP = 6
30STATION_DUMP = 'iw {} station dump'
31SCAN = 'wpa_cli scan'
32SCAN_RESULTS = 'wpa_cli scan_results'
33SIGNAL_POLL = 'wpa_cli signal_poll'
34WPA_CLI_STATUS = 'wpa_cli status'
35RSSI_ERROR_VAL = float('nan')
36FW_REGEX = re.compile(r'FW:(?P<firmware>\S+) HW:')
37
38
39# Rssi Utilities
40def empty_rssi_result():
41    return collections.OrderedDict([('data', []), ('mean', None),
42                                    ('stdev', None)])
43
44
45def get_connected_rssi(dut,
46                       num_measurements=1,
47                       polling_frequency=SHORT_SLEEP,
48                       first_measurement_delay=0,
49                       disconnect_warning=True,
50                       ignore_samples=0,
51                       interface='wlan0'):
52    # yapf: disable
53    connected_rssi = collections.OrderedDict(
54        [('time_stamp', []),
55         ('bssid', []), ('ssid', []), ('frequency', []),
56         ('signal_poll_rssi', empty_rssi_result()),
57         ('signal_poll_avg_rssi', empty_rssi_result()),
58         ('chain_0_rssi', empty_rssi_result()),
59         ('chain_1_rssi', empty_rssi_result())])
60    # yapf: enable
61    previous_bssid = 'disconnected'
62    t0 = time.time()
63    time.sleep(first_measurement_delay)
64    for idx in range(num_measurements):
65        measurement_start_time = time.time()
66        connected_rssi['time_stamp'].append(measurement_start_time - t0)
67        # Get signal poll RSSI
68        try:
69            status_output = dut.adb.shell(
70                'wpa_cli -i {} status'.format(interface))
71        except:
72            status_output = ''
73        match = re.search('bssid=.*', status_output)
74        if match:
75            current_bssid = match.group(0).split('=')[1]
76            connected_rssi['bssid'].append(current_bssid)
77        else:
78            current_bssid = 'disconnected'
79            connected_rssi['bssid'].append(current_bssid)
80            if disconnect_warning and previous_bssid != 'disconnected':
81                logging.warning('WIFI DISCONNECT DETECTED!')
82        previous_bssid = current_bssid
83        match = re.search('\s+ssid=.*', status_output)
84        if match:
85            ssid = match.group(0).split('=')[1]
86            connected_rssi['ssid'].append(ssid)
87        else:
88            connected_rssi['ssid'].append('disconnected')
89        try:
90            signal_poll_output = dut.adb.shell(
91                'wpa_cli -i {} signal_poll'.format(interface))
92        except:
93            signal_poll_output = ''
94        match = re.search('FREQUENCY=.*', signal_poll_output)
95        if match:
96            frequency = int(match.group(0).split('=')[1])
97            connected_rssi['frequency'].append(frequency)
98        else:
99            connected_rssi['frequency'].append(RSSI_ERROR_VAL)
100        match = re.search('RSSI=.*', signal_poll_output)
101        if match:
102            temp_rssi = int(match.group(0).split('=')[1])
103            if temp_rssi == -9999 or temp_rssi == 0:
104                connected_rssi['signal_poll_rssi']['data'].append(
105                    RSSI_ERROR_VAL)
106            else:
107                connected_rssi['signal_poll_rssi']['data'].append(temp_rssi)
108        else:
109            connected_rssi['signal_poll_rssi']['data'].append(RSSI_ERROR_VAL)
110        match = re.search('AVG_RSSI=.*', signal_poll_output)
111        if match:
112            connected_rssi['signal_poll_avg_rssi']['data'].append(
113                int(match.group(0).split('=')[1]))
114        else:
115            connected_rssi['signal_poll_avg_rssi']['data'].append(
116                RSSI_ERROR_VAL)
117
118        # Get per chain RSSI
119        try:
120            per_chain_rssi = dut.adb.shell(STATION_DUMP.format(interface))
121        except:
122            per_chain_rssi = ''
123        match = re.search('.*signal avg:.*', per_chain_rssi)
124        if match:
125            per_chain_rssi = per_chain_rssi[per_chain_rssi.find('[') +
126                                            1:per_chain_rssi.find(']')]
127            per_chain_rssi = per_chain_rssi.split(', ')
128            connected_rssi['chain_0_rssi']['data'].append(
129                int(per_chain_rssi[0]))
130            connected_rssi['chain_1_rssi']['data'].append(
131                int(per_chain_rssi[1]))
132        else:
133            connected_rssi['chain_0_rssi']['data'].append(RSSI_ERROR_VAL)
134            connected_rssi['chain_1_rssi']['data'].append(RSSI_ERROR_VAL)
135        measurement_elapsed_time = time.time() - measurement_start_time
136        time.sleep(max(0, polling_frequency - measurement_elapsed_time))
137
138    # Compute mean RSSIs. Only average valid readings.
139    # Output RSSI_ERROR_VAL if no valid connected readings found.
140    for key, val in connected_rssi.copy().items():
141        if 'data' not in val:
142            continue
143        filtered_rssi_values = [x for x in val['data'] if not math.isnan(x)]
144        if len(filtered_rssi_values) > ignore_samples:
145            filtered_rssi_values = filtered_rssi_values[ignore_samples:]
146        if filtered_rssi_values:
147            connected_rssi[key]['mean'] = statistics.mean(filtered_rssi_values)
148            if len(filtered_rssi_values) > 1:
149                connected_rssi[key]['stdev'] = statistics.stdev(
150                    filtered_rssi_values)
151            else:
152                connected_rssi[key]['stdev'] = 0
153        else:
154            connected_rssi[key]['mean'] = RSSI_ERROR_VAL
155            connected_rssi[key]['stdev'] = RSSI_ERROR_VAL
156    return connected_rssi
157
158
159def get_scan_rssi(dut, tracked_bssids, num_measurements=1):
160    scan_rssi = collections.OrderedDict()
161    for bssid in tracked_bssids:
162        scan_rssi[bssid] = empty_rssi_result()
163    for idx in range(num_measurements):
164        scan_output = dut.adb.shell(SCAN)
165        time.sleep(MED_SLEEP)
166        scan_output = dut.adb.shell(SCAN_RESULTS)
167        for bssid in tracked_bssids:
168            bssid_result = re.search(bssid + '.*',
169                                     scan_output,
170                                     flags=re.IGNORECASE)
171            if bssid_result:
172                bssid_result = bssid_result.group(0).split('\t')
173                scan_rssi[bssid]['data'].append(int(bssid_result[2]))
174            else:
175                scan_rssi[bssid]['data'].append(RSSI_ERROR_VAL)
176    # Compute mean RSSIs. Only average valid readings.
177    # Output RSSI_ERROR_VAL if no readings found.
178    for key, val in scan_rssi.items():
179        filtered_rssi_values = [x for x in val['data'] if not math.isnan(x)]
180        if filtered_rssi_values:
181            scan_rssi[key]['mean'] = statistics.mean(filtered_rssi_values)
182            if len(filtered_rssi_values) > 1:
183                scan_rssi[key]['stdev'] = statistics.stdev(
184                    filtered_rssi_values)
185            else:
186                scan_rssi[key]['stdev'] = 0
187        else:
188            scan_rssi[key]['mean'] = RSSI_ERROR_VAL
189            scan_rssi[key]['stdev'] = RSSI_ERROR_VAL
190    return scan_rssi
191
192
193def get_sw_signature(dut):
194    bdf_output = dut.adb.shell('cksum /vendor/firmware/bdwlan*')
195    logging.debug('BDF Checksum output: {}'.format(bdf_output))
196    bdf_signature = sum(
197        [int(line.split(' ')[0]) for line in bdf_output.splitlines()]) % 1000
198
199    fw_output = dut.adb.shell('halutil -logger -get fw')
200    logging.debug('Firmware version output: {}'.format(fw_output))
201    fw_version = re.search(FW_REGEX, fw_output).group('firmware')
202    fw_signature = fw_version.split('.')[-3:-1]
203    fw_signature = float('.'.join(fw_signature))
204    serial_hash = int(hashlib.md5(dut.serial.encode()).hexdigest(), 16) % 1000
205    return {
206        'config_signature': bdf_signature,
207        'fw_signature': fw_signature,
208        'serial_hash': serial_hash
209    }
210
211
212def get_country_code(dut):
213    country_code = dut.adb.shell('iw reg get | grep country | head -1')
214    country_code = country_code.split(':')[0].split(' ')[1]
215    if country_code == '00':
216        country_code = 'WW'
217    return country_code
218
219
220def push_config(dut, config_file):
221    config_files_list = dut.adb.shell(
222        'ls /vendor/firmware/bdwlan*').splitlines()
223    for dst_file in config_files_list:
224        dut.push_system_file(config_file, dst_file)
225    dut.reboot()
226
227
228def start_wifi_logging(dut):
229    dut.droid.wifiEnableVerboseLogging(1)
230    msg = "Failed to enable WiFi verbose logging."
231    asserts.assert_equal(dut.droid.wifiGetVerboseLoggingLevel(), 1, msg)
232    logging.info('Starting CNSS logs')
233    dut.adb.shell("find /data/vendor/wifi/wlan_logs/ -type f -delete",
234                  ignore_status=True)
235    dut.adb.shell_nb('cnss_diag -f -s')
236
237
238def stop_wifi_logging(dut):
239    logging.info('Stopping CNSS logs')
240    dut.adb.shell('killall cnss_diag')
241    logs = dut.get_file_names("/data/vendor/wifi/wlan_logs/")
242    if logs:
243        dut.log.info("Pulling cnss_diag logs %s", logs)
244        log_path = os.path.join(dut.device_log_path,
245                                "CNSS_DIAG_%s" % dut.serial)
246        os.makedirs(log_path, exist_ok=True)
247        dut.pull_files(logs, log_path)
248
249
250def push_firmware(dut, firmware_files):
251    """Function to push Wifi firmware files
252
253    Args:
254        dut: dut to push bdf file to
255        firmware_files: path to wlanmdsp.mbn file
256        datamsc_file: path to Data.msc file
257    """
258    for file in firmware_files:
259        dut.push_system_file(file, '/vendor/firmware/')
260    dut.reboot()
261
262
263def _set_ini_fields(ini_file_path, ini_field_dict):
264    template_regex = r'^{}=[0-9,.x-]+'
265    with open(ini_file_path, 'r') as f:
266        ini_lines = f.read().splitlines()
267        for idx, line in enumerate(ini_lines):
268            for field_name, field_value in ini_field_dict.items():
269                line_regex = re.compile(template_regex.format(field_name))
270                if re.match(line_regex, line):
271                    ini_lines[idx] = '{}={}'.format(field_name, field_value)
272                    print(ini_lines[idx])
273    with open(ini_file_path, 'w') as f:
274        f.write('\n'.join(ini_lines) + '\n')
275
276
277def _edit_dut_ini(dut, ini_fields):
278    """Function to edit Wifi ini files."""
279    dut_ini_path = '/vendor/firmware/wlan/qcom_cfg.ini'
280    local_ini_path = os.path.expanduser('~/qcom_cfg.ini')
281    dut.pull_files(dut_ini_path, local_ini_path)
282
283    _set_ini_fields(local_ini_path, ini_fields)
284
285    dut.push_system_file(local_ini_path, dut_ini_path)
286    # For 1x1 mode, we need to wait for sl4a to load (To avoid crashes)
287    dut.reboot(timeout=300, wait_after_reboot_complete=120)
288
289
290def set_chain_mask(dut, chain_mask):
291    curr_mask = getattr(dut, 'chain_mask', '2x2')
292    if curr_mask == chain_mask:
293        return
294    dut.chain_mask = chain_mask
295    if chain_mask == '2x2':
296        ini_fields = {
297            'gEnable2x2': 2,
298            'gSetTxChainmask1x1': 1,
299            'gSetRxChainmask1x1': 1,
300            'gDualMacFeatureDisable': 6,
301            'gDot11Mode': 0
302        }
303    else:
304        ini_fields = {
305            'gEnable2x2': 0,
306            'gSetTxChainmask1x1': chain_mask + 1,
307            'gSetRxChainmask1x1': chain_mask + 1,
308            'gDualMacFeatureDisable': 1,
309            'gDot11Mode': 0
310        }
311    _edit_dut_ini(dut, ini_fields)
312
313
314def set_wifi_mode(dut, mode):
315    TX_MODE_DICT = {
316        'Auto': 0,
317        '11n': 4,
318        '11ac': 9,
319        '11abg': 1,
320        '11b': 2,
321        '11': 3,
322        '11g only': 5,
323        '11n only': 6,
324        '11b only': 7,
325        '11ac only': 8
326    }
327
328    ini_fields = {
329        'gEnable2x2': 2,
330        'gSetTxChainmask1x1': 1,
331        'gSetRxChainmask1x1': 1,
332        'gDualMacFeatureDisable': 6,
333        'gDot11Mode': TX_MODE_DICT[mode]
334    }
335    _edit_dut_ini(dut, ini_fields)
336
337
338class LinkLayerStats():
339
340    LLSTATS_CMD = 'cat /d/wlan0/ll_stats'
341    MOUNT_CMD = 'mount -t debugfs debugfs /sys/kernel/debug'
342    PEER_REGEX = 'LL_STATS_PEER_ALL'
343    MCS_REGEX = re.compile(
344        r'preamble: (?P<mode>\S+), nss: (?P<num_streams>\S+), bw: (?P<bw>\S+), '
345        'mcs: (?P<mcs>\S+), bitrate: (?P<rate>\S+), txmpdu: (?P<txmpdu>\S+), '
346        'rxmpdu: (?P<rxmpdu>\S+), mpdu_lost: (?P<mpdu_lost>\S+), '
347        'retries: (?P<retries>\S+), retries_short: (?P<retries_short>\S+), '
348        'retries_long: (?P<retries_long>\S+)')
349    MCS_ID = collections.namedtuple(
350        'mcs_id', ['mode', 'num_streams', 'bandwidth', 'mcs', 'rate'])
351    MODE_MAP = {'0': '11a/g', '1': '11b', '2': '11n', '3': '11ac', '4': '11ax'}
352    BW_MAP = {'0': 20, '1': 40, '2': 80, '3':160}
353
354    def __init__(self, dut, llstats_enabled=True):
355        self.dut = dut
356        self.llstats_enabled = llstats_enabled
357        self.llstats_cumulative = self._empty_llstats()
358        self.llstats_incremental = self._empty_llstats()
359
360    def update_stats(self):
361        if self.llstats_enabled:
362            # Checking the files to see if the device is mounted to enable
363            # llstats capture
364            mount_check = len(self.dut.get_file_names('/d/wlan0'))
365            if not(mount_check):
366              self.dut.adb.shell(self.MOUNT_CMD, timeout=10)
367
368            try:
369                llstats_output = self.dut.adb.shell(self.LLSTATS_CMD,
370                                                    timeout=0.1)
371            except:
372                llstats_output = ''
373        else:
374            llstats_output = ''
375        self._update_stats(llstats_output)
376
377    def reset_stats(self):
378        self.llstats_cumulative = self._empty_llstats()
379        self.llstats_incremental = self._empty_llstats()
380
381    def _empty_llstats(self):
382        return collections.OrderedDict(mcs_stats=collections.OrderedDict(),
383                                       summary=collections.OrderedDict())
384
385    def _empty_mcs_stat(self):
386        return collections.OrderedDict(txmpdu=0,
387                                       rxmpdu=0,
388                                       mpdu_lost=0,
389                                       retries=0,
390                                       retries_short=0,
391                                       retries_long=0)
392
393    def _mcs_id_to_string(self, mcs_id):
394        mcs_string = '{} {}MHz Nss{} MCS{} {}Mbps'.format(
395            mcs_id.mode, mcs_id.bandwidth, mcs_id.num_streams, mcs_id.mcs,
396            mcs_id.rate)
397        return mcs_string
398
399    def _parse_mcs_stats(self, llstats_output):
400        llstats_dict = {}
401        # Look for per-peer stats
402        match = re.search(self.PEER_REGEX, llstats_output)
403        if not match:
404            self.reset_stats()
405            return collections.OrderedDict()
406        # Find and process all matches for per stream stats
407        match_iter = re.finditer(self.MCS_REGEX, llstats_output)
408        for match in match_iter:
409            current_mcs = self.MCS_ID(self.MODE_MAP[match.group('mode')],
410                                      int(match.group('num_streams')) + 1,
411                                      self.BW_MAP[match.group('bw')],
412                                      int(match.group('mcs'), 16),
413                                      int(match.group('rate'), 16) / 1000)
414            current_stats = collections.OrderedDict(
415                txmpdu=int(match.group('txmpdu')),
416                rxmpdu=int(match.group('rxmpdu')),
417                mpdu_lost=int(match.group('mpdu_lost')),
418                retries=int(match.group('retries')),
419                retries_short=int(match.group('retries_short')),
420                retries_long=int(match.group('retries_long')))
421            llstats_dict[self._mcs_id_to_string(current_mcs)] = current_stats
422        return llstats_dict
423
424    def _diff_mcs_stats(self, new_stats, old_stats):
425        stats_diff = collections.OrderedDict()
426        for stat_key in new_stats.keys():
427            stats_diff[stat_key] = new_stats[stat_key] - old_stats[stat_key]
428        return stats_diff
429
430    def _generate_stats_summary(self, llstats_dict):
431        llstats_summary = collections.OrderedDict(common_tx_mcs=None,
432                                                  common_tx_mcs_count=0,
433                                                  common_tx_mcs_freq=0,
434                                                  common_rx_mcs=None,
435                                                  common_rx_mcs_count=0,
436                                                  common_rx_mcs_freq=0,
437                                                  rx_per=float('nan'))
438
439        phy_rates=[]
440        tx_mpdu=[]
441        rx_mpdu=[]
442        txmpdu_count = 0
443        rxmpdu_count = 0
444        for mcs_id, mcs_stats in llstats_dict['mcs_stats'].items():
445            # Extract the phy-rates
446            mcs_id_split=mcs_id.split();
447            phy_rates.append(float(mcs_id_split[len(mcs_id_split)-1].split('M')[0]))
448            rx_mpdu.append(mcs_stats['rxmpdu'])
449            tx_mpdu.append(mcs_stats['txmpdu'])
450            if mcs_stats['txmpdu'] > llstats_summary['common_tx_mcs_count']:
451                llstats_summary['common_tx_mcs'] = mcs_id
452                llstats_summary['common_tx_mcs_count'] = mcs_stats['txmpdu']
453            if mcs_stats['rxmpdu'] > llstats_summary['common_rx_mcs_count']:
454                llstats_summary['common_rx_mcs'] = mcs_id
455                llstats_summary['common_rx_mcs_count'] = mcs_stats['rxmpdu']
456            txmpdu_count += mcs_stats['txmpdu']
457            rxmpdu_count += mcs_stats['rxmpdu']
458
459        if len(tx_mpdu) == 0 or len(rx_mpdu) == 0:
460            return llstats_summary
461
462        # Calculate the average tx/rx -phy rates
463        if sum(tx_mpdu) and sum(rx_mpdu):
464            llstats_summary['mean_tx_phy_rate'] = numpy.average(phy_rates, weights=tx_mpdu)
465            llstats_summary['mean_rx_phy_rate'] = numpy.average(phy_rates, weights=rx_mpdu)
466
467        if txmpdu_count:
468            llstats_summary['common_tx_mcs_freq'] = (
469                llstats_summary['common_tx_mcs_count'] / txmpdu_count)
470        if rxmpdu_count:
471            llstats_summary['common_rx_mcs_freq'] = (
472                llstats_summary['common_rx_mcs_count'] / rxmpdu_count)
473        return llstats_summary
474
475    def _update_stats(self, llstats_output):
476        # Parse stats
477        new_llstats = self._empty_llstats()
478        new_llstats['mcs_stats'] = self._parse_mcs_stats(llstats_output)
479        # Save old stats and set new cumulative stats
480        old_llstats = self.llstats_cumulative.copy()
481        self.llstats_cumulative = new_llstats.copy()
482        # Compute difference between new and old stats
483        self.llstats_incremental = self._empty_llstats()
484        for mcs_id, new_mcs_stats in new_llstats['mcs_stats'].items():
485            old_mcs_stats = old_llstats['mcs_stats'].get(
486                mcs_id, self._empty_mcs_stat())
487            self.llstats_incremental['mcs_stats'][
488                mcs_id] = self._diff_mcs_stats(new_mcs_stats, old_mcs_stats)
489        # Generate llstats summary
490        self.llstats_incremental['summary'] = self._generate_stats_summary(
491            self.llstats_incremental)
492        self.llstats_cumulative['summary'] = self._generate_stats_summary(
493            self.llstats_cumulative)
494