A rapidly-growing script that visualizes WiFi sniffing

Question

Below is one of the versions of my wifimap script:

#!/usr/bin/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

TODO:
    * 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

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."""
    p = subprocess.Popen("SUDO_ASKPASS=/usr/bin/false sudo -A true",
                         shell=True)
    p.communicate()
    return 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."""
        s = subprocess.check_output(["iwconfig", self.iface])
        return re.findall('Frequency:([^ ]+)', s)[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:
                x = subprocess.call(['sudo', '-A',
                                     'iwconfig', self.iface, 'channel',
                                     str(channel)], stderr=subprocess.PIPE)
                if x != 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 Wifibacon(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.
        """
        d = {}
        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')
                d[name] = field.get('showname').split(': ')[1]
                d[name] = d[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 d.get('wlan.ta') != d.get('wlan.sa') \
                and d.get('wlan.ta') is not None \
                and d.get('wlan.sa') is not None:
            sys.stderr.write("WTF: ta=%s != sa=%s\n" %
                             (repr(d.get('wlan.ta')), repr(d.get('wlan.sa'))))
        to_mac = d.get('wlan.ra', '?') \
            if d.get('wlan.ra') != 'ffffffffffff' else '?'
        from_mac = d.get('wlan.ta', '?') or d.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 k1 in self.seen:
            for k2 in self.seen[k1]:
                if skip_broadcast and (k1 == '?' or k2 == '?'):
                    continue
                if k1 in self.announces:
                    k1 += '\\nAnnounces: ' + ',\\n'.join(self.announces[k1])
                if k2 in self.announces:
                    k2 += '\\nAnnounces: ' + ',\\n'.join(self.announces[k2])
                print('"%s" -> "%s";' % (k1, k2))
        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(w, 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.
    """
    def dump_wifimap(*args, **kwargs):
        try:
            sys.stdout.seek(0)
        except IOError:
            pass
        w.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)')
    args = parser.parse_args()

    p = None
    if args.infile:
        args.no_channel_hop = True
    else:
        p = subprocess.Popen(["tshark", "-i", args.iface, "-I",
                              "-y", "IEEE802_11_RADIO", "-T", "pdml"],
                             stdout=subprocess.PIPE)
        args.infile = p.stdout

    w = Wifibacon()
    if not args.no_channel_hop:
        c = ChannelHopper(interval=args.channel_hop_interval, iface=args.iface)
        c.start()
    signal.signal(signal.SIGUSR1, get_dump_wifimap(w, args.skip_broadcast))
    try:
        w.read_from_file(args.infile, args.outfile)
    except KeyboardInterrupt:
        pass
    finally:
        w.print_report(args.skip_broadcast)
        if not args.no_channel_hop:
            c.running = False
            c.join()

if __name__ == '__main__':
    main()

Here's a sample output image:

Are there things that are going to make it harder to develop this script? I am thinking of keeping timestamps of when a particular connection was created and a command-line switch that would enable the displaying of only last few seconds of such connections. Also, I was thinking of making channel hopper more intelligent by staying longer on the channels that carry more packets.

As for design, at first I was thinking of splitting this into two-three programs: wifigrep, datadraw and perhaps c-hopper - the wifigrep would just convert tshark's output to a stream of JSON objects including only data relevant to this project. datadraw is where I would put the graph drawing logic, including USR1 handling - I could probably make it more generic so it would be useful in more cases. c-hopper would be the channel hopper. I quickly realized that this would take a lot of time though and opted for the monolithic approach.

What do you think about my choices? Are there any signs of bad Python programming practice that I should probably change?

The in-development version of the script is being hosted here. And here's a list of changes I made since then.

(of course, I would prefer to get a review of the snippet above)

Note

Mast caught a problem running the script that was related to my non-standard system configuration. You can find the details and a solution in the comments to this answer (mostly the third comment).

Answer 1

Style

You seem to be adhering to PEP8, which is usually a good thing. Many of your functions have descriptive names, but your variables do not. p, s, x, d, w and c are terrible. k1 and k2 are also lousy index names, but since they're used only in a relatively small loop it's less problematic. I have no idea why they're called k1 and k2 while code should be self descriptive. Nick (known on SE as 'nhgrif') wrote a post about this recently.

The following snippet stood out:

def get_dump_wifimap(w, 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.
    """
    def dump_wifimap(*args, **kwargs):
        try:
            sys.stdout.seek(0)
        except IOError:
            pass
        w.print_report(skip_broadcast)
        sys.stdout.flush()
    return dump_wifimap

Are you sure this is working the way you intend it to work? And why is the IOError passed?

Usage

Unfortunately I have not managed to get your script running. This may be my fault, but in my opinion programs and scripts should not be a pain in the behind to get running. We are past those days, even in Linux land.

The example usage: timeout --foreground -s INT 5s wifimap | dot -Tsvg > out.svg seems incomplete. I suppose this should have been timeout --foreground -s INT 5s python wifimap.py | dot -Tsvg > out.svg, assuming wifimap.py as name of the file. However, it will fail on line 227-229:

p = subprocess.Popen(["tshark", "-i", args.iface, "-I",
                      "-y", "IEEE802_11_RADIO", "-T", "pdml"],
                     stdout=subprocess.PIPE)

OSError: [Errno 2] No such file or directory