#!/usr/bin/env python
"""
Write a NetLogger-formatted message to the console, or
a TCP or UDP socket.

Optionally add a syslog header to the message. Combining this
with UDP allows this tool to write directly to syslog.
"""
__author__ = "Dan Gunter <dkgunter@lbl.gov>"
__rcsid__ = "$Id: nl_write 25122 2010-08-16 19:22:49Z mgoode $"

import socket
import sys
import syslog
import time
#
from netlogger.nllog import OptionParser, get_logger
from netlogger import nlapi
from netlogger.nlapi import Level
from netlogger.amqp.connection import Produce, Connect, ConnectionException

## Globals

# some defaults
_P = { 'event' : 'nlwrite.event',
       'level' : Level.INFO,
       'udp_port' : 514,
       'tcp_port' : 14380,
       'amqp_port' : 5672,
}

## Functions

def parseNVP(args, option_parser):
    d = { }
    for arg in args:
        try:
            name, value = arg.split('=')
        except ValueError:
            option_parser.error("argument '%s' not in form name=value" % arg)
        # eval values in curly braces
        if value[0] == '{' and value[-1] == '}':
            value = eval(value[1:-1])
        d[name] = value
    return d

def syslogHeader(host, program="cedps-logger"):
    """Return string to be used as a message header for syslogging it
    """
    # build syslog 'priority'
    pri =   "<%d>" % (syslog.LOG_USER + syslog.LOG_INFO)
    # build syslog formatted date
    localtime = time.localtime()
    day = time.strftime("%d", localtime)
    if day[0] == "0":
        day = " " + day[1:]  # syslog RFC says this MUST be a space
    val = time.strftime("%b %%s %H:%M:%S", localtime)
    date = val % day
    return "%s%s %s %s: " % (pri, date, host, program)

def main():
    """ Program entry point """
    usage = "%prog [options] name=value.."
    desc = ' '.join(__doc__.split())
    parser = OptionParser(usage=usage, description=desc)
    parser.add_option('-g', '--guid',
                      action='store_true', dest='guid',
                      help="add guid=GUID to message. " +
                      "This is overridden by an explicit guid=GUID argument.")
    parser.add_option('-i', '--ip',
                      action='store_true', dest='ipaddr', metavar='IP',
                      help="add 'host=IP' to message. " +
                      "This is overridden by an explicit host=HOST argument.")
    parser.add_option('-n', '--num', metavar='NUM',
                      action='store', type='int', dest='num', default=1,
                      help="Write NUM messages, each with n=<1..NUM> in them"
                      " (default=%default)")
    parser.add_option('-H', '--host',
                      action='store', dest='host', default="localhost",
                      help="for UDP/TCP/AMQP, the remote host (default=%default)")
    parser.add_option('-P', '--port', action="store", type="int",
                      dest="port", metavar="PORT", default=-1,
                      help="For UDP/TCP/AMQP, the port to write to " +
                      "(default=UDP %d, TCP %d, AMQP %s)" %
                      (_P['udp_port'], _P['tcp_port'], _P['amqp_port']))
    parser.add_option('-S', '--syslog', action="store_true", default=None,
                      dest="syslog_hdr",
                      help="add a header for syslog (default=False unless " +
                      "-U is given, then True)")
    parser.add_option('-T', '--tcp', action="store_true", default=False,
                      dest="tcp",
                      help="write message to TCP (default port=%d)" %
                      _P['tcp_port'])
    parser.add_option('-U', '--udp', action="store_true", default=False,
                      dest="udp",
                      help="write message to UDP (default port=%d)" %
                      _P['udp_port'])
    parser.add_option('-A', '--amqp', action="store_true", default=False,
                      dest="amqp",
                      help="write message to AMQP server (default port=%d)" %
                      _P['amqp_port'])
    parser.add_option('-D', '--amqp_disconnect', action="store_true", default=False,
                    dest="amqp_disconnect",
                    help="send disconnect message to AMQP server when done. " +
                    "No effect if not used with -A.")
    parser.add_option('-O', '--amqp_option', action='append', 
                    dest='amqp_option', metavar="name=val",
                    help="optional arg to feed name/value options to amqp " +
                    "connection/producer.  (repeatable: -O user=foo -O pw=bar)")
    options, args = parser.parse_args()
    log = get_logger(__file__)  # Should be first done, just after parsing args
    if options.tcp and options.udp:
        parser.error("both TCP and UDP options given, can only do one")
    # make argument into dictionary
    d = parseNVP(args, parser)
    # add values from options
    if options.guid:
        if d.has_key('guid'):
            log.warn("opt.ignored", value="guid", msg="guid keyword given")
        else:
            d['guid'] = nlapi.getGuid()
    if options.ipaddr:
        if d.has_key('host'):
            log.warn("opt.ignored", value="ip", msg="host keyword given")
        else:
            d['host'] = nlapi.getHost()
    # pull out required fields event and level
    if d.has_key('event'):
        event = d['event']
        del d['event']
    else:
        event = _P['event']
    if d.has_key('level'):
        levelstr = d['level']
        del d['level']
        # check level
        try:
            level = Level.getLevel(levelstr.upper())
        except ValueError:
            parser.error("bad level name '%s'" % levelstr)
    else:
        level = _P['level']
    # write message(s)
    log.info("run.start", num=options.num)
    olog = nlapi.Log()
    if options.udp:
        port = (options.port, _P['udp_port'])[options.port < 0]
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    elif options.tcp:
        port = (options.port, _P['tcp_port'])[options.port < 0]
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            sock.connect((options.host, port))
        except socket.error, serr:
            log.error("connect.error", host=options.host, port=port, msg=serr)
            return 1
    elif options.amqp:
        port = (options.port, _P['amqp_port'])[options.port < 0]
        dd = {}
        if options.amqp_option:
            dd = parseNVP(options.amqp_option, parser)
        # remove kwargs that potentially conflict with flags/defaults
        for badarg in ['port', 'host', 'connection']:
            if dd.has_key(badarg):
                del dd[badarg]
        if not Produce:
            log.error('amqp.error', msg='py-amqplib support not enabled')
            return 1
        try:
            conn = Connect(host=options.host, port=port, **dd)
            prod = Produce(connection=conn, **dd)
        except ConnectionException, e:
            log.error('amqp.error', msg=e)
            return 1
    for i in xrange(options.num):
        if options.num > 1:
            d['n'] = i
        logstr = olog.write(event=event, level=level, **d)
        if options.syslog_hdr or (options.syslog_hdr is None and options.udp):
            logstr = syslogHeader(d.get('host', 'localhost')) + logstr
        if options.udp:
            sock.sendto(logstr, (options.host, port))
        elif options.tcp:
            sock.send(logstr)
        elif options.amqp:
            prod.send(logstr.strip())
        else:
            sys.stdout.write(logstr)
    if options.udp or options.tcp:
        sock.close()
    elif options.amqp:
        if options.amqp_disconnect:
            prod.send_disconnect()
        prod.close()
    log.info("run.end", status=0)


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