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