#!/usr/bin/env python

import sys
import subprocess
import time


def _get_name(cell):
    """
    For the supplied AP (cell) from the iwlist scan, parse the ESSID property
    """
    name = _matching_line(cell, "ESSID:")
    # keep only data between quotes
    return name[name.index("\"") + 1: name.rindex("\"")]


def _get_quality(cell):
    """
    For the supplied AP (cell) from the iwlist scan, parse the quality property
    """
    quality = _matching_line(cell, "Quality=").split()[0].split('/')
    return str(int(round(float(quality[0]) / float(quality[1]) * 100))).rjust(3)


def _get_signal_level(cell):
    """
    For the supplied AP (cell) from the iwlist scan, parse the signal level property
    """
    return _matching_line(cell, "Signal level=").split()[0]


def _get_noise_level(cell):
    """
    For the supplied AP (cell) from the iwlist scan, parse the noise level property
    """
    return _matching_line(cell, "Noise level=").split()[0]


def _get_channel(cell):
    """
    For the supplied AP (cell) from the iwlist scan, parse the channel property
    """
    return _matching_line(cell, "Channel ")[:-1]


def _get_encryption(cell):
    """
    For the supplied AP (cell) from the iwlist scan, parse the encryption property
    """
    enc = ""
    if _matching_line(cell, "Encryption key:") == "off":
        enc = "Open"
    else:
        for line in cell:
            matching = _match(line, "IE:")
            if matching is not None:
                wpa = _match(matching, "WPA Version ")
                if wpa is not None:
                    enc = "WPA v." + wpa
        if enc == "":
            enc = "WEP"
    return enc


def _get_address(cell):
    """
    For the supplied AP (cell) from the iwlist scan, parse the address  property
    """
    return _matching_line(cell, "Address: ")


# List of dictionary rules that will be applied to the description of each
# cell. The key will be the name of the column in the table. The value is a
# function defined above.
rules = {"Name": _get_name,
         "Quality": _get_quality,
         "Channel": _get_channel,
         "Encryption": _get_encryption,
         "Address": _get_address,
         "SignalLvl": _get_signal_level,
         "NoiseLvl": _get_noise_level}

# Selects which columns to display here, and most importantly in what order. Of
# course, they must exist as keys in the dict rules.
columns = ["Name", "Address", "Quality", "SignalLvl", "NoiseLvl", "Channel", "Encryption"]


def _sort_cells(cells, sortby):
    """
    Sorts the table where sortby must be a key of the dictionary rules
    """
    reverse = True
    cells.sort(None, lambda el: el[sortby], reverse)


def _get_max_widths(tables):
    widths = []
    all_widths = []
    for table in tables:
        for item in table:
            widths.append(len(str(item)))
        all_widths.append(widths)
        widths = []

    return(map(max, zip(*all_widths)))


def _print_table(table):
    widths = _get_max_widths(table)

    justified_table = []
    for line in table:
        justified_line = []
        for i, el in enumerate(line):
            # 'offet' makes everything aligned (look pretty)
            offset = widths[i] - len(str(el))
            x = ' ' * offset
            justified_line.append(str(el) + x)

        justified_table.append(justified_line)

    for line in justified_table:
        for el in line:
            print el,
        print


def _print_cells(cells):

    # sort cells by quality
    _sort_cells(cells, "Quality")

    table = [columns]
    for cell in cells:
        cell_properties = []
        for column in columns:
            cell_properties.append(cell[column])
        table.append(cell_properties)
    _print_table(table)


def _matching_line(lines, keyword):
    """
    Returns the first matching line in a list of lines. See match()
    """
    for line in lines:
        matching = _match(line, keyword)
        if matching is not None:
            return matching
    return None


def _match(line, keyword):
    """If the first part of line (modulo blanks) matches keyword,
    returns the end of that line. Otherwise returns None"""

    x = line.find(keyword)
    if x != -1:
        x = line[line.find(keyword) + len(keyword):]
        return x
    else:
        return None


def _parse_cell(cell):
    """Applies the rules to the bunch of text describing a cell and returns the
    corresponding dictionary"""
    parsed_cell = {}
    for key in rules:
        rule = rules[key]
        parsed_cell.update({key: rule(cell)})
    return parsed_cell


def _parse_cells(cells):
    """Parse cells into list of dicts (example as shown)
    parsed_cells
     [
       {'Encryption': 'WPA v.1', 'Quality': ' 25 %', 'Name': 'McWane_Main', 'Channel': 1, 'Address': '5C:A4:8A:2C:3C:90'i},
       {'Encryption': 'WPA v.1', 'Quality': ' 25 %', 'Name': 'McWane_Main', 'Channel': 1, 'Address': '5C:A4:8A:2C:3C:90'i},
     ]
    """
    parsed_cells = []
    for cell in cells:
        parsed_cells.append(_parse_cell(cell))
    return parsed_cells


def get_iwlist_input(use_stdin):
    """Execute "iwlist wlan0 scan" and parse the output into cell list (example as shown) or
       read parse input from stdin

    cells
     [
       ['Address: 5C:A4:8A:2C:3C:90', 'ESSID:"McWane_Main"', 'Mode:Managed', 'Frequency:2.412 GHz (Channel 1)',
        'Quality=10/40  Signal level=-91 dBm  Noise level=-101 dBm', 'Encryption key:on',...],
       ['Address: 5C:A4:8A:2C:3C:90', 'ESSID:"McWane_Main"', 'Mode:Managed', 'Frequency:2.412 GHz (Channel 1)',
        'Quality=10/40  Signal level=-91 dBm  Noise level=-101 dBm', 'Encryption key:on',...],
     ]
    """
    cells = [[]]
    lines = []
    if use_stdin:
        # read input from stdin
        for line in sys.stdin:
            lines.append(line)
    else:
        # execute 'wilist wlan0 scan' and get its output
        output = subprocess.check_output(['iwlist', 'wlan0', 'scan'])
        output.decode()
        lines = output.split('\n')

    for line in lines:
        cell_line = _match(line, "Cell ")
        if cell_line is not None:
            cells.append([])
        cells[-1].append(line.rstrip())

    return cells[1:]


def get_args():
    import argparse
    parser = argparse.ArgumentParser(description='Selects an optimum wifi channel by parsing the '
                                     'output of an "iwlist scan" commmand from Bluegiga\'s WF111 wifi module. '
                                     'Optionally, configures the WF111 module for an AP using the selected channel. '
                                     'Note: Should run with root privileges')
                                     #formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-s", "--stdin", action="store_true", dest="use_stdin", default=False,
                        help="Uses stdin to read \'iwlist wlan0 scan\' output. (Default executes \'iwlist scan\' and parses "
                             "its output. Note: \'iwlist scan\' cannot be invoked if WF111 is already an AP)")
    parser.add_argument("-a", "--ap_config", action="store", type=str, dest="ap_config_fname", default=None,
                        help="Executes bash script (to configure AP) in specified filename passing selected channel as command line parameter $1 (Default no script called)")
    parser.add_argument("-c", "--channel_file", action="store", type=str, dest="channel_fname", default=None,
                        help="Write selected channel into specified filename. (Default no channel is written to file)")
    parser.add_argument("-r", "--repeat", action="store", type=int, dest="avg_count", default=1,
                        help="Number of times to repeat the scan for averaging. (Default is 1)")
    parser.add_argument("-p", "--pause", action="store", type=int, dest="pause_time", default=1,
                        help="Amount of time to sleep between repeated scans. (Default is 1 sec)")
    # Parse command-line options
    args = parser.parse_args()
    return args


def select_channel(parsed_cells):
    """
    Select the best channel from the statistics
    stats =
        {
         chan: {'quality': xxx, 'count': xxx},
         chan: {'quality': xxx, 'count': xxx},
         chan: {'quality': xxx, 'count': xxx}
        }

    - If channels 1, 6, 11 are not being used, one of them gets selected
    - Otherwise, select the channel with the lowest average signal quality
    """
    valid_channels = ('1', '6', '11')

    # gather statistics for channel 1, 6, 11
    stats = {}
    for cell in parsed_cells:
        if cell['Channel'] in valid_channels:
            chan = stats.setdefault(cell['Channel'], {'quality': 0, 'count': 0, 'signallvl': 0, 'noiselvl': 0})
            chan['quality'] = chan['quality'] + int(cell['Quality'])
            chan['signallvl'] = chan['signallvl'] + int(cell['SignalLvl'])
            chan['noiselvl'] = chan['noiselvl'] + int(cell['NoiseLvl'])
            chan['count'] = chan['count'] + 1

    # Average the signal quality
    for item in stats:
        stats[item]['quality'] = str(stats[item]['quality']/stats[item]['count'])
        stats[item]['signallvl'] = str(stats[item]['signallvl']/stats[item]['count'])
        stats[item]['noiselvl'] = str(stats[item]['noiselvl']/stats[item]['count'])

    # fill in with "---" missing channels from 1, 6, 11
    for chan_sel in valid_channels:
        if chan_sel not in stats:
            stats[chan_sel] = {'quality': '---', 'signallvl': '---', 'noiselvl': '---'}

    # print out the averages
    print "\n*** Averages for channel 1, 6, 11 ***"
    title = "Channel Quality SignalLvl NoiseLvl"
    titles = title.split()
    print title
    line = "%s" + " " * (len(titles[0])-2) + "%s" + " " * (len(titles[1])-2) + "%s" + " " * (len(titles[2])-2) + "%s" + " " * (len(titles[3])-2)

    for item in valid_channels:
        print line % (item.rjust(3), stats[item]['quality'].rjust(3),
                      stats[item]['signallvl'].rjust(3), stats[item]['noiselvl'].rjust(3))

    # For channel selection, select a unused channel first
    channel = None
    for chan_sel in valid_channels:
        if stats[chan_sel]['quality'] is '---':
            channel = chan_sel
            break

   # if all channels are being used, select channel with lowest signal quality
    if channel is None:
        current_quality = 100
        channel = '1'
        for item in stats:
            if int(stats[item]['quality']) < int(current_quality):
                current_quality = stats[item]['quality']
                channel = item

    print "\nOptimum channel: ", channel
    return channel


def run_iwlist_script(channel, filename):
    invoke_str = filename + ' ' + channel
    print "Invoking script: ", invoke_str
    rtn = subprocess.call([invoke_str], shell=True)
    if rtn:
        print "Error invoking script. Return: ", rtn
    else:
        print "Success invoking script'"


def store_channel(channel, filename):
    try:
        with open(filename, "w") as f:
            f.write(channel)
            f.close()
    except IOError:
        print "IOError opening %s for writing" % filename


def main():

    # get command line args
    args = get_args()

    total_parsed_cells = []
    for x in xrange(args.avg_count):
        # get the iwlist datai (parse into cells)
        cells = get_iwlist_input(args.use_stdin)
        if cells == []:
            continue
        # parse the cells (i.e. parse components of each cell)
        parsed_cells = _parse_cells(cells)

        #print the parsed components
        print "***Pass %d scan***" % (x+1)
        _print_cells(parsed_cells)

        total_parsed_cells.extend(parsed_cells)
        time.sleep(args.pause_time)

    # get the best channel to use for ap mode
    channel = select_channel(total_parsed_cells)

    if args.ap_config_fname is not None:
        run_iwlist_script(channel, args.ap_config_fname)

    if args.channel_fname is not None:
        store_channel(channel, args.channel_fname)


if __name__ == '__main__':
    main()
