#!/usr/bin/env python

import os
import sys
import errno
import subprocess
import time
import syslog

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
       Failure possibilitites
       NOTE 'iwlist wlan0 scan' must run as root. Otherwise, results are incorrect
       NOTE 'iwlist wlan0 scan' will fail if WF111 is already in AP mode (ie. AS_BSS_START)

    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
        try:
            output = subprocess.check_output(['iwlist', 'wlan0', 'scan'])
            output.decode()
            lines = output.split('\n')
        except subprocess.CalledProcessError as e:
            err_str = "'iwlist wlan0 scan' failure: " + e.__str__()
            syslog.syslog(syslog.LOG_ERR, err_str)
            print err_str
            return []
        
    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 (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 "\nselected-channel:%s"%channel

    line = "Avg Channel Quality: "
    for item in valid_channels:
        line += "%s:%s " % (item, stats[item]['quality'])
    line +=" Selected Channel:%s" % channel
    syslog.syslog(syslog.LOG_INFO, line)

    return channel


def run_iwlist_script(channel, filename):
    # This invoked script is meant to run "iwpriv wlan0 AP_Set_CFG..." to configure
    # the WF111 for AP mode (Note: The 'AP_SET_CFG' configure the WF111 for AP mode
    # and 'AS_BSS_START' puts it in AP mode)
    # Note: iwpriv must run at root level

    invoke_str = filename + ' ' + channel
    rtn = subprocess.call([invoke_str], shell=True)
    if rtn:
        err_str = "Error invoking script '%s'. Return: %d" % (invoke_str, rtn)
        print err_str
        syslog.syslog(syslog.LOG_ERR, err_str)
        
    else:
        print "Success invoking script '%s'" % invoke_str


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


def main():

    # Must run script as root.
    # 'iwlist wlan0 scan' will only give static results or no results as non-root
    # 'iwpriv' (invoked via -a) will not run at all as non-root
    if os.geteuid() != 0:
        print "Must run as root"
        syslog.syslog(syslog.LOG_ERR, "Must run as root. Exiting.")
        sys.exit(errno.EPERM)

    # 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()
