1#!/usr/bin/env python3
2#
3#   Copyright 2018 - Google, Inc.
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
17from acts import logger
18from acts.controllers.ap_lib.hostapd_constants import AP_DEFAULT_CHANNEL_2G
19from acts.controllers.ap_lib.hostapd_constants import AP_DEFAULT_CHANNEL_5G
20from acts.controllers.ap_lib.hostapd_constants import CHANNEL_MAP
21from acts.controllers.ap_lib.hostapd_constants import FREQUENCY_MAP
22from acts.controllers.ap_lib.hostapd_constants import CENTER_CHANNEL_MAP
23from acts.controllers.ap_lib.hostapd_constants import VHT_CHANNEL
24from acts.controllers.utils_lib.ssh import connection
25from acts.controllers.utils_lib.ssh import formatter
26from acts.controllers.utils_lib.ssh import settings
27from acts.libs.logging import log_stream
28from acts.libs.proc.process import Process
29from acts import asserts
30
31import os
32import threading
33import time
34
35MOBLY_CONTROLLER_CONFIG_NAME = 'PacketCapture'
36ACTS_CONTROLLER_REFERENCE_NAME = 'packet_capture'
37BSS = 'BSS'
38BSSID = 'BSSID'
39FREQ = 'freq'
40FREQUENCY = 'frequency'
41LEVEL = 'level'
42MON_2G = 'mon0'
43MON_5G = 'mon1'
44BAND_IFACE = {'2G': MON_2G, '5G': MON_5G}
45SCAN_IFACE = 'wlan2'
46SCAN_TIMEOUT = 60
47SEP = ':'
48SIGNAL = 'signal'
49SSID = 'SSID'
50
51
52def create(configs):
53    return [PacketCapture(c) for c in configs]
54
55
56def destroy(pcaps):
57    for pcap in pcaps:
58        pcap.close()
59
60
61def get_info(pcaps):
62    return [pcap.ssh_settings.hostname for pcap in pcaps]
63
64
65class PcapProperties(object):
66    """Class to maintain packet capture properties after starting tcpdump.
67
68    Attributes:
69        proc: Process object of tcpdump
70        pcap_fname: File name of the tcpdump output file
71        pcap_file: File object for the tcpdump output file
72    """
73
74    def __init__(self, proc, pcap_fname, pcap_file):
75        """Initialize object."""
76        self.proc = proc
77        self.pcap_fname = pcap_fname
78        self.pcap_file = pcap_file
79
80
81class PacketCaptureError(Exception):
82    """Error related to Packet capture."""
83
84
85class PacketCapture(object):
86    """Class representing packet capturer.
87
88    An instance of this class creates and configures two interfaces for monitor
89    mode; 'mon0' for 2G and 'mon1' for 5G and one interface for scanning for
90    wifi networks; 'wlan2' which is a dual band interface.
91
92    Attributes:
93        pcap_properties: dict that specifies packet capture properties for a
94            band.
95    """
96
97    def __init__(self, configs):
98        """Initialize objects.
99
100        Args:
101            configs: config for the packet capture.
102        """
103        self.ssh_settings = settings.from_config(configs['ssh_config'])
104        self.ssh = connection.SshConnection(self.ssh_settings)
105        self.log = logger.create_logger(lambda msg: '[%s|%s] %s' % (
106            MOBLY_CONTROLLER_CONFIG_NAME, self.ssh_settings.hostname, msg))
107
108        self._create_interface(MON_2G, 'monitor')
109        self._create_interface(MON_5G, 'monitor')
110        self.managed_mode = True
111        result = self.ssh.run('ifconfig -a', ignore_status=True)
112        if result.stderr or SCAN_IFACE not in result.stdout:
113            self.managed_mode = False
114        if self.managed_mode:
115            self._create_interface(SCAN_IFACE, 'managed')
116
117        self.pcap_properties = dict()
118        self._pcap_stop_lock = threading.Lock()
119
120    def _create_interface(self, iface, mode):
121        """Create interface of monitor/managed mode.
122
123        Create mon0/mon1 for 2G/5G monitor mode and wlan2 for managed mode.
124        """
125        if mode == 'monitor':
126            self.ssh.run('ifconfig wlan%s down' % iface[-1],
127                         ignore_status=True)
128        self.ssh.run('iw dev %s del' % iface, ignore_status=True)
129        self.ssh.run('iw phy%s interface add %s type %s' %
130                     (iface[-1], iface, mode),
131                     ignore_status=True)
132        self.ssh.run('ip link set %s up' % iface, ignore_status=True)
133        result = self.ssh.run('iw dev %s info' % iface, ignore_status=True)
134        if result.stderr or iface not in result.stdout:
135            raise PacketCaptureError('Failed to configure interface %s' %
136                                     iface)
137
138    def _cleanup_interface(self, iface):
139        """Clean up monitor mode interfaces."""
140        self.ssh.run('iw dev %s del' % iface, ignore_status=True)
141        result = self.ssh.run('iw dev %s info' % iface, ignore_status=True)
142        if not result.stderr or 'No such device' not in result.stderr:
143            raise PacketCaptureError('Failed to cleanup monitor mode for %s' %
144                                     iface)
145
146    def _parse_scan_results(self, scan_result):
147        """Parses the scan dump output and returns list of dictionaries.
148
149        Args:
150            scan_result: scan dump output from scan on mon interface.
151
152        Returns:
153            Dictionary of found network in the scan.
154            The attributes returned are
155                a.) SSID - SSID of the network.
156                b.) LEVEL - signal level.
157                c.) FREQUENCY - WiFi band the network is on.
158                d.) BSSID - BSSID of the network.
159        """
160        scan_networks = []
161        network = {}
162        for line in scan_result.splitlines():
163            if SEP not in line:
164                continue
165            if BSS in line:
166                network[BSSID] = line.split('(')[0].split()[-1]
167            field, value = line.lstrip().rstrip().split(SEP)[0:2]
168            value = value.lstrip()
169            if SIGNAL in line:
170                network[LEVEL] = int(float(value.split()[0]))
171            elif FREQ in line:
172                network[FREQUENCY] = int(value)
173            elif SSID in line:
174                network[SSID] = value
175                scan_networks.append(network)
176                network = {}
177        return scan_networks
178
179    def get_wifi_scan_results(self):
180        """Starts a wifi scan on wlan2 interface.
181
182        Returns:
183            List of dictionaries each representing a found network.
184        """
185        if not self.managed_mode:
186            raise PacketCaptureError('Managed mode not setup')
187        result = self.ssh.run('iw dev %s scan' % SCAN_IFACE)
188        if result.stderr:
189            raise PacketCaptureError('Failed to get scan dump')
190        if not result.stdout:
191            return []
192        return self._parse_scan_results(result.stdout)
193
194    def start_scan_and_find_network(self, ssid):
195        """Start a wifi scan on wlan2 interface and find network.
196
197        Args:
198            ssid: SSID of the network.
199
200        Returns:
201            True/False if the network if found or not.
202        """
203        curr_time = time.time()
204        while time.time() < curr_time + SCAN_TIMEOUT:
205            found_networks = self.get_wifi_scan_results()
206            for network in found_networks:
207                if network[SSID] == ssid:
208                    return True
209            time.sleep(3)  # sleep before next scan
210        return False
211
212    def configure_monitor_mode(self, band, channel, bandwidth=20):
213        """Configure monitor mode.
214
215        Args:
216            band: band to configure monitor mode for.
217            channel: channel to set for the interface.
218            bandwidth : bandwidth for VHT channel as 40,80,160
219
220        Returns:
221            True if configure successful.
222            False if not successful.
223        """
224
225        band = band.upper()
226        if band not in BAND_IFACE:
227            self.log.error('Invalid band. Must be 2g/2G or 5g/5G')
228            return False
229
230        iface = BAND_IFACE[band]
231        if bandwidth == 20:
232            self.ssh.run('iw dev %s set channel %s' % (iface, channel),
233                         ignore_status=True)
234        else:
235            center_freq = None
236            for i, j in CENTER_CHANNEL_MAP[VHT_CHANNEL[bandwidth]]["channels"]:
237                if channel in range(i, j + 1):
238                    center_freq = (FREQUENCY_MAP[i] + FREQUENCY_MAP[j]) / 2
239                    break
240            asserts.assert_true(center_freq,
241                                "No match channel in VHT channel list.")
242            self.ssh.run(
243                'iw dev %s set freq %s %s %s' %
244                (iface, FREQUENCY_MAP[channel], bandwidth, center_freq),
245                ignore_status=True)
246
247        result = self.ssh.run('iw dev %s info' % iface, ignore_status=True)
248        if result.stderr or 'channel %s' % channel not in result.stdout:
249            self.log.error("Failed to configure monitor mode for %s" % band)
250            return False
251        return True
252
253    def start_packet_capture(self, band, log_path, pcap_fname):
254        """Start packet capture for band.
255
256        band = 2G starts tcpdump on 'mon0' interface.
257        band = 5G starts tcpdump on 'mon1' interface.
258
259        Args:
260            band: '2g' or '2G' and '5g' or '5G'.
261            log_path: test log path to save the pcap file.
262            pcap_fname: name of the pcap file.
263
264        Returns:
265            pcap_proc: Process object of the tcpdump.
266        """
267        band = band.upper()
268        if band not in BAND_IFACE.keys() or band in self.pcap_properties:
269            self.log.error("Invalid band or packet capture already running")
270            return None
271
272        pcap_name = '%s_%s.pcap' % (pcap_fname, band)
273        pcap_fname = os.path.join(log_path, pcap_name)
274        pcap_file = open(pcap_fname, 'w+b')
275
276        tcpdump_cmd = 'tcpdump -i %s -w - -U 2>/dev/null' % (BAND_IFACE[band])
277        cmd = formatter.SshFormatter().format_command(tcpdump_cmd,
278                                                      None,
279                                                      self.ssh_settings,
280                                                      extra_flags={'-q': None})
281        pcap_proc = Process(cmd)
282        pcap_proc.set_on_output_callback(lambda msg: pcap_file.write(msg),
283                                         binary=True)
284        pcap_proc.start()
285
286        self.pcap_properties[band] = PcapProperties(pcap_proc, pcap_fname,
287                                                    pcap_file)
288        return pcap_proc
289
290    def stop_packet_capture(self, proc):
291        """Stop the packet capture.
292
293        Args:
294            proc: Process object of tcpdump to kill.
295        """
296        for key, val in self.pcap_properties.items():
297            if val.proc is proc:
298                break
299        else:
300            self.log.error("Failed to stop tcpdump. Invalid process.")
301            return
302
303        proc.stop()
304        with self._pcap_stop_lock:
305            self.pcap_properties[key].pcap_file.close()
306            del self.pcap_properties[key]
307
308    def close(self):
309        """Cleanup.
310
311        Cleans up all the monitor mode interfaces and closes ssh connections.
312        """
313        self._cleanup_interface(MON_2G)
314        self._cleanup_interface(MON_5G)
315        self.ssh.close()
316