#!/usr/bin/python

"""
    Analyze WLAN STA logs.

        Jul  3 23:05:04 buffalo-g300n-647682 daemon.info hostapd: wlan0-1: STA aa:bb:cc:dd:ee:ff WPA: group key handshake completed (RSN)

"""

__version__ = '0.1'

import pvl.args
import pvl.syslog.args
import pvl.web.args
import pvl.rrd.hosts
import pvl.verkko.wlan

import optparse
import logging; log = logging.getLogger('main')

WLAN_STA_PROG = 'hostapd'

def parse_argv (argv, doc = __doc__) :
    """
        Parse command-line argv, returning (options, args).
    """

    prog = argv.pop(0)
    args = argv

    # optparse
    parser = optparse.OptionParser(
        prog        = prog,
        usage       = '%prog: [options] [<input.txt> [...]]',
        version     = __version__,
        description = doc,
    )

    # common
    parser.add_option_group(pvl.args.parser(parser))
    parser.add_option_group(pvl.syslog.args.parser(parser, prog=WLAN_STA_PROG))
    parser.add_option_group(pvl.verkko.db.parser(parser, table=db.wlan_sta))
    parser.add_option_group(pvl.web.args.parser(parser))

    parser.add_option('--interfaces', metavar='PATH',
            help="Load interface/node names from mapping file")

    # parse
    options, args = parser.parse_args(args)

    # apply
    pvl.args.apply(options)

    return options, args

import re
from pvl.verkko import db

class KeyTimestampDatabase (object) :
    """
        A pvl.verkko.db table that tracks events by key/timestamp.
    """

    DB_TABLE = None
    DB_LAST_SEEN = None
    DB_COUNT = None
    DB_DISTINCT = None

    def __init__ (self, db) :
        """
            db      - pvl.verkko.db.Database
        """

        self.db = db

    def select (self, interval=None) :
        """
            SELECT unique gw/ip hosts, for given interval.
        """

        query = db.select(self.DB_DISTINCT, distinct=True)

        if interval :
            # timedelta
            query = query.where(db.func.now() - self.DB_LAST_SEEN < interval)

        return self.db.select(query)

    def insert (self, key, timestamp, update) :
        """
            INSERT new row.
        """

        query = self.DB_TABLE.insert().values(**key).values(**update).values(
                first_seen  = timestamp,
                last_seen   = timestamp,
        )

        if self.DB_COUNT is not None :
            query = query.values(count=1)
        
        # -> id
        return self.db.insert(query)

    def update (self, key, timestamp, update) :
        """
            UPDATE existing row, or return False if not found.
        """

        table = self.DB_TABLE
        query = table.update()

        for col, value in key.iteritems() :
            query = query.where(table.c[col] == value)

        query = query.values(last_seen=timestamp)

        if self.DB_COUNT is not None :
            query = query.values(count=db.func.coalesce(self.DB_COUNT, 0) + 1)

        query = query.values(**update)
        
        # -> any matched rows?
        return self.db.update(query)

    def touch (self, key, timestamp, update, **opts) :
        # update existing?
        if self.update(key, timestamp, update, **opts) :
            log.info("Update: %s: %s: %s", key, timestamp, update)
        else :
            log.info("Insert: %s: %s: %s", key, timestamp, update)
            self.insert(key, timestamp, update, **opts)


class WlanStaDatabase (KeyTimestampDatabase) :
    HOSTAPD_STA_RE = re.compile(r'(?P<wlan>.+?): STA (?P<sta>.+?) (?P<msg>.+)')

    DB_TABLE = db.wlan_sta
    DB_LAST_SEEN = db.wlan_sta.c.last_seen
    DB_COUNT = db.wlan_sta.c.count

    DB_DISTINCT = (db.wlan_sta.c.sta, )

    def __init__ (self, db, interface_map=None) :
        """
            interface_map       - {(hostname, interface): (nodename, wlan)}
        """
        KeyTimestampDatabase.__init__(self, db)
        self.interface_map = interface_map

    def parse (self, item) :
        """
            Parse fields from a hostapd syslog message.
        """
        
        match = self.HOSTAPD_STA_RE.match(item['msg'])

        if not match :
            return None

        return match.groupdict()

    def lookup_wlan (self, host, iface) :
        """
            Lookup ap/ssid by host/iface.
        """
        mapping = None
        if self.interface_map :
            mapping = self.interface_map.get((host, iface))
        
        if mapping :
            return mapping
        else :
            # as-is
            log.warning("Unknown host/iface: %s/%s", host, iface)
            return host, iface

    def __call__ (self, item) :
        match = self.parse(item)

        if not match :
            return

        # lookup?
        ap, ssid = self.lookup_wlan(item['host'], match['wlan'])
        
        # update/insert
        self.touch(
            dict(ap=ap, wlan=ssid, sta=match['sta']),
            item['timestamp'],
            dict(msg=match['msg']),
        )

def main (argv) :
    options, args = parse_argv(argv)
    
    # database
    db = pvl.verkko.db.apply(options)

    if options.interfaces :
        interfaces = dict(pvl.rrd.hosts.map_interfaces(options, open(options.interfaces)))
    else :
        interfaces = None

    # syslog
    log.info("Open up syslog...")
    syslog = pvl.syslog.args.apply(options, optional=True)
        
    if syslog :
        # handler
        handler = WlanStaDatabase(db, interface_map=interfaces)

        log.info("Enter mainloop...")
        for source in syslog.main() :
            for item in source:
                handler(item)
    else :
        # run web
        application = pvl.web.args.apply(options, pvl.verkko.wlan.Application, db)
        return pvl.web.args.main(options, application)

    return 0
    
if __name__ == '__main__':
    import sys

    sys.exit(main(sys.argv))