7
\$\begingroup\$

I developed a program that collects probe requests from WiFi devices. It already works but I think that a programmer can improve it.

#!/usr/bin/env python

import platform
import threading
import signal
import sys
import time
import subprocess
import re
import os
import os.path
import argparse
import distutils.spawn 

if (platform.system() == 'Windows'):
    print "Windows does not seem to be supported, try it at your own risk!"
    sys.exit(1)

formatString = "{0: <18} {1: <32} {2: <18}\n"
dontstop = False

def flush_data():
    with open(output,'w+') as f:
        if (onlyprobes == False):
            f.write(formatString.format("MAC", "SSID", "Last seen") )
        for key, value in entries.iteritems():
            if (onlyprobes == False):
                f.write(formatString.format(value.mac, value.ssid, time.strftime("%Y%m%d-%H:%M:%S", value.timeLastSeen)))
            else:
                f.write(value.ssid + "\n")
            f.flush()

def signal_handler(signal, frame):
    exit_text = ''
    if (flush == True): exit_text=" flushing all to file and"
    print os.linesep + 'You pressed CTRL+C,' + exit_text + ' exiting...'
    if (dontstop == True):
        switchThread.running = False
        switchThread.join()
    if ((flush == True) and (is_stdout == False)):
        flush_data()
    sys.exit(0)


class switchChannelThread (threading.Thread):
    def __init__(self, threadID, name, delayInSeconds, channels):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = name
        self.delayInSeconds = delayInSeconds
        self.channels = channels
        self.running = True


    def run(self):
        #print 'Starting switch channel thread using a delay of %d seconds' % self.delayInSeconds
        while self.running:
            for channel in self.channels:
                if verbose: 
                    print 'Switching to channel %d' % (channel)
                if osname != "Darwin":
                    if subprocess.call([iwconfigPath, interface, "channel", str(channel)]) != 0:
                        self.running = False
                        sys.exit(4)
                else:
                    if subprocess.call([airportPath, interface, "-c%d" % channel]) != 0:
                        self.running = False
                        sys.exit(4)
                time.sleep(float(self.delayInSeconds))
                if not self.running:
                    return        

class Entry (object):
    def __init__(self, mac, ssid, time):
        self.mac = mac
        self.ssid = ssid
        self.timeLastSeen = time


osname = os.uname()[0]
if osname != "Darwin":
    defaultInterface = ""
else:
    defaultInterface = "en1"


# command line parsing:
parser = argparse.ArgumentParser(description='Show and collect wlan request probes')

parser.add_argument('-i', '--interface', default=defaultInterface, 
    help='the interface used for monitoring')

parser.add_argument('--tshark-path', default=distutils.spawn.find_executable("tshark"), 
    help='path to tshark binary')

parser.add_argument('--ifconfig-path', default=distutils.spawn.find_executable("ifconfig"), 
    help='path to ifconfig')

parser.add_argument('--iwconfig-path', default=distutils.spawn.find_executable("iwconfig"), 
    help='path to iwconfig')

parser.add_argument('-o', '--output', default='-', 
    help='output file (path or - for stdout)')

parser.add_argument('-c', '--channel', default='all', 
    help='channel/s to hop (i.e. 3 or 3,5,7 or 3-9 or all or 0 for current')

parser.add_argument('--verbose', action='store_true', 
    help='verbose information')

parser.add_argument('-p', '--only-probes', action='store_true',
    help='only saves probe data spit by newline')

parser.add_argument('--flush', action='store_true',
    help='stores the data on the file only when interrupted')

parser.add_argument('--delay', default=5,
    help='delay between channel change')

args = parser.parse_args()

tsharkPath = args.tshark_path
ifconfigPath = args.ifconfig_path
iwconfigPath = args.iwconfig_path
interface = args.interface
verbose = args.verbose
onlyprobes = args.only_probes
output = args.output
flush = args.flush
channel = args.channel
delay = args.delay

is_stdout = not ( (output != '') and (output != '-') )
if (interface == ""):
    print "Please specify interface"
    sys.exit(0)

# only on osx:
airportPath = "/System/Library/PrivateFrameworks/Apple80211.framework/Versions/Current/Resources/airport";

# check all params
if not os.path.isfile(tsharkPath):
    print "tshark not found at path {0}".format(tsharkPath)
    sys.exit(1)
if not os.path.isfile(ifconfigPath):
    print "ifconfig not found at path {0}".format(ifconfigPath)
    sys.exit(1)
if osname != "Darwin":
    if not os.path.isfile(iwconfigPath):
        print "iwconfig not found at path {0}".format(iwconfigPath)
        sys.exit(1)

# start interface
if subprocess.call([ifconfigPath, interface, 'up']) != 0:
    print "cannot start interface: {0}".format(interface)
    sys.exit(2)

# Set interface in monitor mode
retVal = 0
if osname != 'Darwin':
    retVal = subprocess.call([iwconfigPath, interface, "mode", "monitor"])
else:
    retVal = subprocess.call([airportPath, interface, "-z"])

if retVal != 0:
    print "cannot set interface to monitor mode: {0}".format(interface)
    sys.exit(3)

# start thread that switches channels

#Regex made with regex101.com
c_list = re.compile('^(([1-9]|1[0-4]),){2,14}$')
c_range = re.compile('^([1-9]|1[0-4])-([1-9]|1[0-4])$')
c_single = re.compile('^([1-9]|1[0-4])$')
d_valid = re.compile('^(\d*\.*\d*)$')

if (channel != '0'):
    dontstop = True
    try:
        float(delay)
    except:
        print "Wrong delay specified!"
        sys.exit(1)

    if (d_valid.match(str(delay)) and (float(delay) > 0)):
        if (channel == 'all'):
            channel = range(1, 13 if (osname == "Darwin") else 15)
            switchThread = switchChannelThread(1, 'SwitchChannel', delay, channel)
            switchThread.start()
        elif c_single.match(channel):
            channel = [int(channel)]
            switchThread = switchChannelThread(1, 'SwitchChannel', delay, channel)
            switchThread.start()
        elif c_range.match(channel):
            rchannel = channel.split('-')
            schannel = int(rchannel[0])
            echannel = int(rchannel[1])
            if (schannel > echannel):
                channel = range(echannel, schannel + 1)
            else:
                channel = range(schannel, echannel + 1) 
            switchThread = switchChannelThread(1, 'SwitchChannel', delay, channel)
            switchThread.start()
        elif c_list.match(channel + ','):
            channel = channel.split(',')
            channel = [int(i) for i in channel]
            switchThread = switchChannelThread(1, 'SwitchChannel', delay, channel)
            switchThread.start()
        else:
            print "Wrong channel/s specified!"
            sys.exit(1)
    else:
        print "Wrong delay specified!"
        sys.exit(1)

signal.signal(signal.SIGINT, signal_handler)
print "Running..."

# start tshark and read the results
displayFilter = "wlan.fcs_good==1 and not wlan_mgt.ssid==\\\"\\\"";
fieldParams = "-T fields -e wlan.sa -e wlan_mgt.ssid -Eseparator=,";
tsharkCommandLine = "{0} -i {1} -n -l {2}"

if (osname != 'Darwin'):
    tsharkCommandLine += " subtype probereq -2 -R \"{3}\""
else:
    tsharkCommandLine += " -y PPI -2 -R \"wlan.fc.type_subtype==4 and {3}\""

tsharkCommandLine = tsharkCommandLine.format(tsharkPath, interface, fieldParams, displayFilter)

if verbose: 
    print 'tshark command: %s\n' % tsharkCommandLine, 

DEVNULL = open(os.devnull, 'w')
popen = subprocess.Popen(tsharkCommandLine, shell=True, stdout=subprocess.PIPE, stderr=DEVNULL)

# collect all Entry objects in entries
entries = {}

if ( (is_stdout == False) and (flush == False) ):
    try:
        f=open(output,'w+')
        if (onlyprobes == False):
            f.write(formatString.format("MAC", "SSID", "Last seen"))
    except Exception as e:
        print "An error has occurred: " + str(e)
        sys.exit(1)

for line in iter(popen.stdout.readline, ''):
    line = line.rstrip()
    #    if verbose: 
    #        print 'line: "%s"' % (line,)
    if line.find(',') > 0:
        mac, ssid = line.split(',', 1)
        if line in entries:
            #if verbose:
            #    print "entry found (seen before): mac: '{0}', ssid: '{1}'".format(mac,ssid)
            entry = entries[line]
            entry.timeLastSeen = time.localtime()
        else:
            localtime=time.localtime()
            if ( (is_stdout == False) and (flush == False) ):
                if (onlyprobes == False):
                    f.write(formatString.format(mac, ssid, time.strftime("%Y%m%d-%H:%M:%S", localtime)))
                else:
                    f.write(ssid + "\n")
                f.flush()
            print "New entry found: mac: '{0}', ssid: '{1}'".format(mac,ssid)
            entries[line] = Entry(mac, ssid, localtime)

I need a small orientation about code optimization. I think that many lines of code can be grouped to avoid code duplication in a clever way.

My main style troubles are about code size reduction and readability. How can I simplify this code?

if not ( (output != '') and (output != '-') ):

Should I use myClassName, MyClassName, myclassname or my_class_name?
Should I use myFunctionName, MyFunctionName, myfunctionname or my_function_name?

\$\endgroup\$
  • 1
    \$\begingroup\$ That is not the right way around - please read and follow the Meta post I linked to about creating follow-up questions. \$\endgroup\$ – jonrsharpe Mar 13 '15 at 16:15
7
\$\begingroup\$

Some points on style.
Python has an official style guide which has a lot of information on how to properly layout your code.

  1. Don't put if, elif, while, etc. conditions in brackets (i.e if a == 2 and not if (a == 2)) This mostly comes from other languages which require this.
  2. Use newlines and indentations after if, else, etc. (Don't do if condition: code(), do

    if condition:
        code()
    

    in all circumstances.)

  3. Don't compare to Booleans. (if condition == True to if condition and if condition == False to if not condition)
  4. Use lower_case_with_underscores for variable names (see this)
  5. Boolean values should not be negative and be verbs or verb+noun (is_stopping = True or stopping = True, not dont_stop = False or not_stopping = False, then do if not stopping rather than if not_stopping)
  6. Have an entry point to your program. Right now, your code has no entry point, it is just in the global namespace. Do this:

    #!shebang
    
    imports
    
    classes
    
    function definitions
    
    def main():
         everything that is not an import, class or function
    
    if __name__ == "__main__":   # i.e. The module is not being imported
        main()
    
| improve this answer | |
\$\endgroup\$
  • 3
    \$\begingroup\$ @Helio Answers need not address every aspect of the code. Also, if you had those concerns, it would have been wise to mention them in your question. (It's not too late to edit the question, as long as you don't change the code.) \$\endgroup\$ – 200_success Mar 12 '15 at 15:03

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.