#!/usr/bin/env python
"""
Process NetLogger best-practices log files or messages from
a log or broker, using the specified loader \"module\".

For log files, extra arguments that are NOT in the form name=value
are interpreted as event prefixes which form a filter.
"""
__rcsid__ = "$Id: nl_loader 23968 2009-10-15 14:30:44Z dang $"
__author__ = "Dan Gunter"

import imp
import os
import re
import socket
import sys
import time
#
from netlogger import info_broker
from netlogger import util, module_util
from netlogger.nllog import OptionParser, get_logger
from netlogger.nlapi import EVENT_FIELD
from netlogger.parsers.base import NLSimpleParser
from netlogger.analysis.modules import _base as analysis_base
from netlogger.amqp import connection
from netlogger.amqp.connection import Consume, Connect, ConnectionException
from netlogger.amqp.scriptutil import AMQPOptionParser

# Variables

_P = {
    'broker_port' : info_broker.FANOUT_PORT,
    'amqp_port' : 5672,
}

g_progress_meter = None

#
# Exceptions
#

class OptionError(Exception):
    pass

#
# Signal handlers
#

# Stop things that are in a loop
g_stop = False
# loader visible to sig handlers
loader = None
consumer = None

def on_kill(signo, frame):
    """Signal handler for a graceful exit.
    """
    global g_stop, loader, consumer
    log = get_logger(__file__)
    log.warn("killed", signo=signo)
    g_stop = True
    if loader:
        time.sleep(1)
        log.info('waiting for loader module to finish')
        loader.finish()
        log.info('loader done')
    if consumer:
        consumer.close()
    time.sleep(1)
    sys.exit(0)

def on_hup(signo, frame):
    
    return

#
# Module/class loading
#

def load_class(module_name):
    log = get_logger(__file__)
    clazz = None
    try:
        mod = module_util.load_module("netlogger.analysis.modules." + 
                                      module_name)
        clazz = vars(mod)['Analyzer']
    except module_util.ModuleLoadError, error:
        log.critical("load_class.error", msg=error)
    return clazz

#
# Utility functions
#

def _notify(loader, data, mode='logfile', **context):
    """Handle exceptions from loader.notify(data)
    in a standard way: log them.

    Args:
      - loader: Loader module, should have a notify(data) method.
      - data (dict): NetLogger event dictionary
    Kwargs:
      - mode: run-mode.
      - **context:  added to the error message
            if an exception occurs. Keywords to avoid: msg, event.
    Returns: None
    """
    g_progress_meter.advance()
    try:
        loader.notify(data)
    except analysis_base.AnalyzerException, err:
        log = get_logger(__file__)
        log.error("notify.error", mode=mode,
                  msg=err, event_name=data.get(EVENT_FIELD,"<None>"),
                  **context)

#
# Run functions for each major mode
#

def infobroker_callback(data, loader=None):
    _notify(loader, data, mode="info_broker")
    
def run_client(loader, host, port, event_list=[ ], reconnect=0,
               progress=False):
    """Run as a client to the info-broker running at 'host':'port'.
    Filter using the list of events in 'event-list'.
    """
    global g_stop
    log = get_logger(__file__)
    # query expression is just the list of events
    expr = ' '.join(event_list)
    # Run until stopped
    consumer = info_broker.StreamConsumer(infobroker_callback,
                                          loader=loader)
    while not g_stop:
        try:
            # Create a client to the info. broker, which
            # passes parsed data on to the loader
            info_broker.Client(host, port, query_text=expr,
                               ostrm=consumer)
            # run the clients
            info_broker.run_clients()
        except socket.error, error:
            if reconnect > 0:
                log.warn("run_client.error", type="socket", msg=error)
                time.sleep(reconnect)
            else:
                g_stop = 1
                
def amqp_callback(msg, loader, parser):
    data = parser.parseLine(msg.body)
    _notify(loader, data, mode="AMQP")
    msg.channel.basic_ack(msg.delivery_tag)
                
def run_amqp(loader, host, port, opts, log):
    global consumer
    try:
        conn = Connect(host=host, port=port, **opts)
        consumer = Consume(connection=conn, no_ack=False, **opts)
    except ConnectionException, e:
        log.error('amqp.error', msg=e)
        return
    parser = NLSimpleParser()
    consumer.register(amqp_callback, loader, parser)
    consumer.wait()
    consumer.close()
 
def run_logfile(loader, filename, tail=False, event_list=[ ]):
    """Run by reading the logfile in 'filename'
    """
    from netlogger.parsers.base import NLSimpleParser
    stdin = False
    log = get_logger(__file__)
    #                    
    if filename == "":
        infile = sys.stdin
        stdin = True
    else:
        infile = file(filename)
    parser = NLSimpleParser()
    while 1:
        where = None
        if not stdin:
            where = infile.tell()
        line = infile.readline()
        if not line:
            if tail and not stdin:
                time.sleep(0.1)
                infile.seek(where)
                continue
            else:
                break
        data = parser.parseLine(line)
        # Only process events that match one of the prefixes
        # in the event_list, if it is non-empty.
        if event_list:
            event = data[EVENT_FIELD]
            for prefix in event_list:
                if event.startswith(prefix):
                    _notify(loader, data, mode="logfile",
                            event_prefix=prefix)
                    break
        # Empty event list = all
        else:
            _notify(loader, data, mode="logfile", event_prefix="ALL")

#
# Main
#

def main():
    """Program entry point.
    """
    global g_progress_meter
    usage = "%prog {-a HOST | -c HOST | -f FILE} module_name " \
        "[option=value..] [prefix1 prefix2 ..]"
    desc = ' '.join(__doc__.split())
    parser = AMQPOptionParser(usage=usage, description=desc)    
    parser.add_option('-c', '--host', dest='host', metavar='HOST', default=None,
                      help='Connect to NetLogger info-broker at HOST ' +
                      '(default=localhost)')
    parser.add_option('-f', '--infile', dest='filename', metavar='FILE', default="",
                      help='Read NetLogger logs from FILE (default=stdin)')
    parser.add_option("-g", "--progress", action="store_true",
                      dest="progress", default=False,
                      help="report progress to stderr")
    parser.add_option('-i', '--info', dest='info', action='store_true',
                      default=False,
                      help="Print information on selected module")
    parser.add_option('-l', '--list', dest='listmod', action='store_true',
                      help="List available modules")
    parser.add_option('-M', '--module-opt', dest="mopt", default="",
                       metavar="FILE",
                       help="Read module options from a file with one"
                       " name=value pair per line (default=No file; use command-line)")
    parser.add_option('-p', '--port', dest='port', type='int', metavar='PORT',
                      default=-1,
                      help="For info_broker or amqp server, the port to connect to"
                      " (default=info_broker %d, amqp broker %d)" \
                      % (_P['broker_port'], _P['amqp_port']))
    parser.add_option('-r', '--reconnect', dest='reconnect', type='int', metavar='SEC',
                      default=10,
                      help="If connection to broker at HOST fails, "
                      "try again every SEC seconds (default=%default). "
                      "0=don't retry")
    parser.add_option('-t', '--tail', dest='tail', action='store_true',
                      help="With -f, tail the file instead of stopping "
                      "at EOF")
    options, args = parser.parse_args(sys.argv[1:])
    log = get_logger(__file__)  # Must come after parsing args
    # Check for 'list' mode
    if options.listmod:
        print 'Available modules:'
        avail = [ ]
        try:
            avail = module_util.list_modules('analysis', 'modules')
        except ImportError, error:
            parser.error('No module found: %s' % error)
        print ', '.join(avail)
        return 0
    # Load the class
    if len(args) == 0:
        parser.error("A module name is required")
    module_name = args[0]
    clazz = load_class(module_name)
    if clazz is None:
        return -1
    # Check for 'info' mode:
    if options.info:
        print(module_util.module_info('loader', module_name, clazz))
        return 0
    # Initialize the class
    init_kw, event_list = { }, [ ]
    # ~ read values from a file if given
    if options.mopt:
        try:
            infile = open(options.mopt)
        except IOError, err:
            parser.error("Cannot open module options file '%s': %s" % (
                options.mopt, err))
        for i, line in enumerate(infile):
            line = line.strip()
            # skip comments and blank lines
            if line == '' or line.startswith('#'):
                continue
            try:                
                key, value = util.process_kvp(line)
            except ValueError, err:
                parser.error("%s:%d => %s." % (options.mopt, i, err,))
            init_kw[key] = value
    # ~ override file values with cmdline
    for init_arg in args[1:]:        
        parts = init_arg.split('=', 1)
        if len(parts) != 2:
            event_list.append(init_arg)
            continue
        key, value = parts[0], parts[1]
        init_kw[key] = value
    try:
        global loader
        log.info("init.class.begin")
        loader = clazz(**init_kw)
        log.info("init.class.end", status=0)
    except Exception, error:
        log.error("init.class.end", status=-1, msg=error)
        return -1
    # Set up signal handlers
    util.handleSignals(
        (on_kill, ('SIGTERM', 'SIGINT', 'SIGUSR2')),
        (on_hup, ('SIGHUP',)) )
    # Run in selected mode
    num_modes = (0,1)[bool(options.filename)] + (0,1)[bool(options.host)] + \
                (0,1)[bool(options.amqp_host)]
    if num_modes > 1:
        parser.error("Cannot read from more than one input source. "
                     "If this is what you want to do, run multiple instances.")
    # Progress meter
    if options.progress:
        g_progress_meter = util.ProgressMeter(sys.stderr, units="records")
    else:
        g_progress_meter = util.NullProgressMeter()
    # Run in appropriate mode
    try:
        # netlogger homebrewed TCP broker
        if options.host:
            port = (options.port, _P['broker_port'])[options.port < 0]
            status = run_client(loader, options.host, port,
                                reconnect=options.reconnect,
                                event_list=event_list)
        # AMQP server
        elif (options.amqp_host or options.amqp_option):
            if not Consume:
                log.error('amqp.error', msg='py-amqplib support not enabled')
                status = -1
                return -1
            if not options.amqp_host:
                default_host = 'localhost'
                options.ensure_value('amqp_host', default_host)
            port = (options.port, _P['amqp_port'])[options.port < 0]
            # add all other options
            amqp_opt_vals = parser.get_amqp_options(options)
            status = run_amqp(loader, options.amqp_host,
                              port, amqp_opt_vals, log)
        # Log file
        else:
            status = run_logfile(loader,
                                 options.filename, tail=options.tail, 
                                 event_list=event_list)
        # call finish in case using threaded buffer loader
        log.info('loader.finishing')
        loader.finish()
        log.info('loader.finished')
    except Exception, error:
        #log.fatal("run.error")
        raise
        status = -1
    return status

if __name__ == '__main__':
    sys.exit(main())
