wifimap

#!/usr/bin/env python

"""
Sniffs on the WiFi network and generates a graph showing the communication
between the devices, including information about the announced SSIDs. Will
loop infinitely - kill the program with SIGINT in order to stop it
gracefully. You can also use SIGUSR1 to force the generation of the graph
without killing the program - useful for the periodic generation of the
visuals.

Example usage (the "timeout" kills wifimap after 5s):

    timeout --foreground -s INT 5s wifimap | dot -Tsvg > out.svg

The example assumes that you saved this script as "wifimap" with rights to
being executed somewhere to your $PATH.

TODO:
    * make channel hopping more intelligent, staying longer on better channels
    * split wifimap into sniffing and graphing programs, maybe also ch-hopping
    * verify if there's a bug related to 'Announces' being printed with wrong
      source
    * maybe keep stats about the connection recency and allow to only display
      connections from last N seconds?
    * detect whether wlan0 or wlp3s0 should be used by default

Requires "tshark" application in PATH, which is provided by Wireshark. Also,
a Python library "lxml" needs to be installed.

tshark also needs to grant sniffing privileges without running as root for
the user running it - you can add these by calling the following command:

sudo setcap cap_net_raw,cap_net_admin,cap_net_bind_service+eip /usr/sbin/tshark

NOTE: This will only work if your Wi-Fi card supports sniffing.

Author: Jacek Wielemborek, licensed under WTFPL
"""

import subprocess
import sys
import time
import re
import threading
from StringIO import StringIO
from lxml import etree
from collections import defaultdict
import signal


def can_do_sudo():
    """Tells whether we can run sudo without being asked for password."""
    sudo_true_p = subprocess.Popen("SUDO_ASKPASS=/usr/bin/false sudo -A true",
                                   shell=True)
    sudo_true_p.communicate()
    return sudo_true_p.returncode == 0


class ChannelHopper(threading.Thread):
    """A channel hopper thread. Switches the Wireless channel periodically."""

    def __init__(self, interval, iface):
        """
        Constructs a ChannelHopper instance.

        Args:

          interval - the time between channel switches
          iface - the interface to perform switching on
        """
        threading.Thread.__init__(self)
        self.running = True
        self.interval = interval
        self.iface = iface

    def get_frequency(self):
        """Returns the current frequency of the interface."""
        iwconfig_output = subprocess.check_output(["iwconfig", self.iface])
        return re.findall('Frequency:([^ ]+)', iwconfig_output)[0]

    def run(self):
        if not can_do_sudo():
            sys.stderr.write("Could not run passwordless sudo - channel hopper"
                             " skipped.\n")
            return
        else:
            sys.stderr.write("Channel hopper started.\n")
        start = self.get_frequency()
        channel = 1
        try:
            while self.running:
                retcode = subprocess.call(['sudo', '-A',
                                           'iwconfig', self.iface, 'channel',
                                           str(channel)],
                                          stderr=subprocess.PIPE)
                if retcode != 0:
                    channel = 0
                time.sleep(self.interval)
                channel += 1
        finally:
            subprocess.call(['sudo', '-A',
                             'iwconfig', self.iface, 'freq', start])
            end = self.get_frequency()
            if start != end:
                sys.stderr.write("WTF: Could not restore the "
                                 "frequency (%s vs %s)\n" % (start, end))
            else:
                sys.stderr.write("Channel hopper stopped.\n")


class Wifimap(object):

    def __init__(self):
        self.seen = defaultdict(lambda: defaultdict(int))
        self.announces = defaultdict(list)

    def parse_packet(self, packet_str):
        """
        Parses a <packet></packet> XML string, returning information about the
        sender, receiver and the announced networks. If sender or receiver is
        not known, returns '?' in its place. If SSID is not announced, None
        is returned.
        """
        packet_dict = {}
        ssid = None
        for line in packet_str.split("\n"):
            if '"wlan.ra"' in line or '"wlan.ta"' in line \
                    or '"wlan.sa"' in line:
                field = etree.fromstring(line)
                if 'ff:ff:ff:ff:ff:ff' in field.get('showname'):
                    continue
                name = field.get('name')
                packet_dict[name] = field.get('showname').split(': ')[1]
                packet_dict[name] = packet_dict[name].replace(' ', '\\n')
            if '"wlan_mgt.ssid"' in line:
                field = etree.fromstring(line)
                if ssid is not None and field.get('show') != ssid:
                    sys.stderr.write("WTF: SSID: %s vs %s" % (repr(ssid),
                                     repr(field.get('show'))))
                ssid = field.get('show')
        if packet_dict.get('wlan.ta') != packet_dict.get('wlan.sa') \
                and packet_dict.get('wlan.ta') is not None \
                and packet_dict.get('wlan.sa') is not None:
            sys.stderr.write("WTF: ta=%s != sa=%s\n" %
                             (repr(packet_dict.get('wlan.ta')),
                              repr(packet_dict.get('wlan.sa'))))
        to_mac = packet_dict.get('wlan.ra', '?') \
            if packet_dict.get('wlan.ra') != 'ffffffffffff' else '?'
        from_mac = packet_dict.get('wlan.ta', '?') or \
            packet_dict.get('wlan.sa', '?')
        return from_mac, to_mac, ssid

    def handle_packet(self, from_mac, to_mac, ssid):
        """
        Handles information about noticing a given packet in order to prepare
        it for reporting.
        """
        self.seen[from_mac][to_mac] += 1
        if ssid:
            found = self.announces[from_mac]
            if len(found) != 0 and ssid not in found:
                sys.stderr.write('WTF: two ssids: %s, %s, %s\n' %
                                 (from_mac, ssid, found))
            if ssid not in found:
                self.announces[from_mac] += [ssid]

    def print_report(self, skip_broadcast=False):
        """
        Prints out a DOT file based on the gathered information.
        """
        print("strict digraph {")
        for from_mac in self.seen:
            for to_mac in self.seen[from_mac]:
                if skip_broadcast and (from_mac == '?' or to_mac == '?'):
                    continue
                from_mac_display = from_mac
                if from_mac in self.announces:
                    from_mac_display += '\\nAnnounces: '
                    from_mac_display += ',\\n'.join(self.announces[from_mac])
                to_mac_display = to_mac
                if to_mac in self.announces:
                    to_mac_display += '\\nAnnounces: '
                    to_mac_display += ',\\n'.join(self.announces[to_mac])
                print('"%s" -> "%s";' % (from_mac_display, to_mac_display))
        print("}")

    def read_from_file(self, infile, outfile=None):
        """
        Reads the output of tshark -T pdml. If outfile is specified,
        the information is also saved to the outfile.
        """
        packet = StringIO()
        while True:
            line = infile.readline()
            if line == '':
                break
            packet.write(line)
            if outfile:
                outfile.write(line)
            if '</packet>' in line:
                packet_str = packet.getvalue()
                packet_info = self.parse_packet(packet_str)
                self.handle_packet(*packet_info)
                packet = StringIO()


def get_dump_wifimap(wmap, skip_broadcast):
    """
    Returns a closure that is supposed to work as a signal handler. It can be
    used when SIGUSR1 is received to force the generation of the report at the
    given time.

    Args:

      wmap - a Wifimap instance that will be used for printing the report
      skip_broadcast - a boolean value telling whether broadcasts should not
        be reported
    """
    def dump_wifimap(*args, **kwargs):
        # If sys.stdout is redirected to a file (as opposed to the terminal),
        # this will make wifimap rewrite it.
        try:
            sys.stdout.seek(0)
        except IOError:
            pass
        wmap.print_report(skip_broadcast)
        sys.stdout.flush()
    return dump_wifimap


def main():

    from argparse import ArgumentParser, RawTextHelpFormatter, FileType
    parser = ArgumentParser(description=__doc__,
                            formatter_class=RawTextHelpFormatter)
    parser.add_argument('--infile', help='file to read the PDML data'
                        ' from instead of sniffing (implies'
                        ' --no-channel-hop)', type=FileType('r'))
    parser.add_argument('--outfile', help='file to save the a copy of PDML'
                        ' data to while sniffing', type=FileType('w'))
    parser.add_argument('--skip-broadcast', action='store_true', help='do not'
                        ' draw broadcast connections - this will remove some'
                        ' results')
    parser.add_argument('--no-channel-hop', action='store_true', help='do not'
                        ' attempt channel hopping even if possible - might'
                        ' give more results')
    parser.add_argument('--channel-hop-interval', type=float, default=5,
                        help='channel hopping interval in seconds'
                        ' (default: 5.0)')
    parser.add_argument('--iface', default='wlp3s0', help='name of the WLAN'
                        ' interface to perform sniffing and hopping on'
                        ' (default: wlp3s0)')
    parser.add_argument('--justhop', action='store_true', help="don't do any"
                        "sniffing, just run the channel hopper")
    args = parser.parse_args()

    if not args.no_channel_hop:
        ch_hopper = ChannelHopper(interval=args.channel_hop_interval,
                                  iface=args.iface)
        ch_hopper.start()

    if args.justhop:
        try:
            while True:
                time.sleep(1)
        finally:
            ch_hopper.running = False
            ch_hopper.join()

    tshark_p = None
    if args.infile:
        args.no_channel_hop = True
    else:
        try:
            tshark_p = subprocess.Popen(["tshark", "-i", args.iface,
                                         "-I", "-y", "IEEE802_11_RADIO",
                                         "-T", "pdml"],
                                        stdout=subprocess.PIPE)
        except OSError:
            sys.exit("ERROR: Attempt to call tshark failed. "
                     "Have you installed Wireshark? Is tshark in your $PATH?")
        args.infile = tshark_p.stdout

    wmap = Wifimap()

    signal.signal(signal.SIGUSR1, get_dump_wifimap(wmap, args.skip_broadcast))
    try:
        wmap.read_from_file(args.infile, args.outfile)
    except KeyboardInterrupt:
        pass
    finally:
        wmap.print_report(args.skip_broadcast)
        if not args.no_channel_hop:
            ch_hopper.running = False
            ch_hopper.join()

if __name__ == '__main__':
    main()