1# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import logging
6
7from autotest_lib.client.common_lib import error
8
9PYSHARK_LOAD_TIMEOUT = 2
10FRAME_FIELD_RADIOTAP_DATARATE = 'radiotap.datarate'
11FRAME_FIELD_RADIOTAP_MCS_INDEX = 'radiotap.mcs_index'
12FRAME_FIELD_WLAN_FRAME_TYPE = 'wlan.fc_type_subtype'
13FRAME_FIELD_WLAN_MGMT_SSID = 'wlan_mgt.ssid'
14RADIOTAP_KNOWN_BAD_FCS_REJECTOR = (
15    'not radiotap.flags.badfcs or radiotap.flags.badfcs==0')
16WLAN_BEACON_FRAME_TYPE = '0x08'
17WLAN_BEACON_ACCEPTOR = 'wlan.fc.type_subtype==0x08'
18WLAN_PROBE_REQ_FRAME_TYPE = '0x04'
19WLAN_PROBE_REQ_ACCEPTOR = 'wlan.fc.type_subtype==0x04'
20PYSHARK_BROADCAST_SSID = 'SSID: '
21BROADCAST_SSID = ''
22
23
24class Frame(object):
25    """A frame from a packet capture."""
26    TIME_FORMAT = "%H:%M:%S.%f"
27
28
29    def __init__(self, frametime, bit_rate, mcs_index, ssid):
30        self._datetime = frametime
31        self._bit_rate = bit_rate
32        self._mcs_index = mcs_index
33        self._ssid = ssid
34
35
36    @property
37    def time_datetime(self):
38        """The time of the frame, as a |datetime| object."""
39        return self._datetime
40
41
42    @property
43    def bit_rate(self):
44        """The bitrate used to transmit the frame, as an int."""
45        return self._bit_rate
46
47
48    @property
49    def mcs_index(self):
50        """
51        The MCS index used to transmit the frame, as an int.
52
53        The value may be None, if the frame was not transmitted
54        using 802.11n modes.
55        """
56        return self._mcs_index
57
58
59    @property
60    def ssid(self):
61        """
62        The SSID of the frame, as a string.
63
64        The value may be None, if the frame does not have an SSID.
65        """
66        return self._ssid
67
68
69    @property
70    def time_string(self):
71        """The time of the frame, in local time, as a string."""
72        return self._datetime.strftime(self.TIME_FORMAT)
73
74
75def _fetch_frame_field_value(frame, field):
76    """
77    Retrieve the value of |field| within the |frame|.
78
79    @param frame: Pyshark packet object corresponding to a captured frame.
80    @param field: Field for which the value needs to be extracted from |frame|.
81
82    @return Value extracted from the frame if the field exists, else None.
83
84    """
85    layer_object = frame
86    for layer in field.split('.'):
87        try:
88            layer_object = getattr(layer_object, layer)
89        except AttributeError:
90            return None
91    return layer_object
92
93
94def _open_capture(pcap_path, display_filter):
95    """
96    Get pyshark packet object parsed contents of a pcap file.
97
98    @param pcap_path: string path to pcap file.
99    @param display_filter: string filter to apply to captured frames.
100
101    @return list of Pyshark packet objects.
102
103    """
104    import pyshark
105    capture = pyshark.FileCapture(
106        input_file=pcap_path, display_filter=display_filter)
107    capture.load_packets(timeout=PYSHARK_LOAD_TIMEOUT)
108    return capture
109
110
111def get_frames(local_pcap_path, display_filter, bad_fcs):
112    """
113    Get a parsed representation of the contents of a pcap file.
114
115    @param local_pcap_path: string path to a local pcap file on the host.
116    @param diplay_filter: string filter to apply to captured frames.
117    @param bad_fcs: string 'include' or 'discard'
118
119    @return list of Frame structs.
120
121    """
122    if bad_fcs == 'include':
123        display_filter = display_filter
124    elif bad_fcs == 'discard':
125        display_filter = '(%s) and (%s)' % (RADIOTAP_KNOWN_BAD_FCS_REJECTOR,
126                                            display_filter)
127    else:
128        raise error.TestError('Invalid value for bad_fcs arg: %s.' % bad_fcs)
129
130    logging.debug('Capture: %s, Filter: %s', local_pcap_path, display_filter)
131    capture_frames = _open_capture(local_pcap_path, display_filter)
132    frames = []
133    logging.info('Parsing frames')
134
135    for frame in capture_frames:
136        rate = _fetch_frame_field_value(frame, FRAME_FIELD_RADIOTAP_DATARATE)
137        if rate:
138            rate = float(rate)
139        else:
140            logging.debug('Found bad capture frame: %s', frame)
141            continue
142
143        frametime = frame.sniff_time
144
145        mcs_index = _fetch_frame_field_value(
146            frame, FRAME_FIELD_RADIOTAP_MCS_INDEX)
147        if mcs_index:
148            mcs_index = int(mcs_index)
149
150        # Get the SSID for any probe requests
151        frame_type = _fetch_frame_field_value(
152            frame, FRAME_FIELD_WLAN_FRAME_TYPE)
153        if (frame_type in [WLAN_BEACON_FRAME_TYPE, WLAN_PROBE_REQ_FRAME_TYPE]):
154            ssid = _fetch_frame_field_value(frame, FRAME_FIELD_WLAN_MGMT_SSID)
155            # Since the SSID name is a variable length field, there seems to be
156            # a bug in the pyshark parsing, it returns 'SSID: ' instead of ''
157            # for broadcast SSID's.
158            if ssid == PYSHARK_BROADCAST_SSID:
159                ssid = BROADCAST_SSID
160        else:
161            ssid = None
162
163        frames.append(Frame(frametime, rate, mcs_index, ssid))
164
165    return frames
166
167
168def get_probe_ssids(local_pcap_path, probe_sender=None):
169    """
170    Get the SSIDs that were named in 802.11 probe requests frames.
171
172    Parse a pcap, returning all the SSIDs named in 802.11 probe
173    request frames. If |probe_sender| is specified, only probes
174    from that MAC address will be considered.
175
176    @param pcap_path: string path to a local pcap file on the host.
177    @param remote_host: Host object (if the file is remote).
178    @param probe_sender: MAC address of the device sending probes.
179
180    @return: A frozenset of the SSIDs that were probed.
181
182    """
183    if probe_sender:
184        diplay_filter = '%s and wlan.addr==%s' % (
185                WLAN_PROBE_REQ_ACCEPTOR, probe_sender)
186    else:
187        diplay_filter = WLAN_PROBE_REQ_ACCEPTOR
188
189    frames = get_frames(local_pcap_path, diplay_filter, bad_fcs='discard')
190
191    return frozenset(
192            [frame.ssid for frame in frames if frame.ssid is not None])
193