#! /usr/bin/env python
#
# Copyright 2001-2002 by Vinay Sajip. All Rights Reserved.
#
# Permission to use, copy, modify, and distribute this software and its
# documentation for any purpose and without fee is hereby granted,
# provided that the above copyright notice appear in all copies and that
# both that copyright notice and this permission notice appear in
# supporting documentation, and that the name of Vinay Sajip
# not be used in advertising or publicity pertaining to distribution
# of the software without specific, written prior permission.
# VINAY SAJIP DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING
# ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL
# VINAY SAJIP BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR
# ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN
# AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR
# IN CONNECTION WITH THE USE OR PERFORMANCE    OF THIS SOFTWARE.
#

"""
Simple socket-based logging event receiver for use with "logging.py"
logging module.

Should work under Python versions >= 1.5.2, except that source line
information is not available unless 'inspect' is.

Copyright (C) 2001-2002 Vinay Sajip. All Rights Reserved.
"""

from SocketServer import ThreadingTCPServer, ThreadingUDPServer, \
    StreamRequestHandler, DatagramRequestHandler
from select import select
import sys, string, types, cPickle, thread, socket
import logging

TIMEOUT         = 10
if sys.platform == "win32":
    RESET_ERROR    = 10054
else:
    RESET_ERROR = 0

class LogRecordStreamHandler(StreamRequestHandler):
    """
    Handler for a streaming logging request. It basically logs the record
    using whatever logging policy is configured locally.
    """
    def recv(self, msglen):
        """
        Receive a specific number of bytes.
        """
        data = ""
        while len(data) < msglen:
            ready = select([self.connection], [], [], TIMEOUT)
            if not len(ready[0]):
                print "timed out!"
                break
            chunk = self.connection.recv(1)
            data = data + chunk
        return data

    def finish(self):
        pass

    def handle(self):
        """
        Handle multiple requests - each expected to be a 2-byte length,
        followed by the LogRecord in pickle format. Logs the record
        according to whatever policy is configured locally.
        """
        while 1:
            try:
                chunk = self.recv(2)
                slen = (ord(chunk[0]) << 8) | ord(chunk[1])
                chunk = self.recv(slen)
                obj = cPickle.loads(chunk)
                record = logging.LogRecord(None, None, "", 0, "", (), None)
                record.__dict__.update(obj)

            except socket.error, e:
                if type(e.args) != types.TupleType:
                    raise
                else:
                    errcode = e.args[0]
                    if errcode != RESET_ERROR:
                        raise
                    break
            self.handleLogRecord(record)

    def handleLogRecord(self, record):
        logger = logging.getLogger(record.name)
        logger._logRecord(record)

class LogRecordDatagramHandler(DatagramRequestHandler):
    """
    Handler for a datagram logging request. It basically logs the record
    using whatever logging policy is configured locally.
    """
    def handle(self):
        chunk = self.packet
        slen = (ord(chunk[0]) << 8) | ord(chunk[1])
        chunk = chunk[2:]
        assert len(chunk) == slen
        obj = cPickle.loads(chunk)
        record = logging.LogRecord(None, None, "", 0, "", (), None)
        record.__dict__.update(obj)
        self.handleLogRecord(record)

    def handleLogRecord(self, record):
        logger = logging.getLogger(record.name)
        logger._logRecord(record)

    def finish(self):
        pass

class LogRecordSocketReceiver(ThreadingTCPServer):
    """
    A simple-minded TCP socket-based logging receiver suitable for test
    purposes.
    """
    allow_reuse_address = 1
    def __init__(self, host='localhost', port=logging.DEFAULT_TCP_LOGGING_PORT,
            handler=LogRecordStreamHandler):
        ThreadingTCPServer.__init__(self, (host, port), handler)

class LogRecordDatagramReceiver(ThreadingUDPServer):
    """
    A simple-minded UDP datagram-based logging receiver suitable for test
    purposes.
    """
    allow_reuse_address = 1
    def __init__(self, host='localhost', port=logging.DEFAULT_UDP_LOGGING_PORT,
            handler=LogRecordDatagramHandler):
        ThreadingUDPServer.__init__(self, (host, port), handler)

def runTCP():
    tcpserver = LogRecordSocketReceiver()
    print "About to start TCP server..."
    tcpserver.serve_forever()

def runUDP():
    udpserver = LogRecordDatagramReceiver()
    print "About to start UDP server..."
    udpserver.serve_forever()

if __name__ == "__main__":
    if (len(sys.argv) < 2) or not (string.lower(sys.argv[1]) in ["udp",
            "tcp", "both"]):
        print "usage: logrecv.py [UDP|TCP|BOTH]"
    else:
        both = string.lower(sys.argv[1]) == "both"
        logging.basicConfig()
        hdlr = logging.FileHandler("test.log")
        hdlr.setFormatter(logging.Formatter("%(asctime)s %(name)-19s %(level)-5s - %(message)s"))
        logging.getRootLogger().addHandler(hdlr)
        if both:
            import threading
            tcpthread = threading.Thread(target=runTCP)
            udpthread = threading.Thread(target=runUDP)
            tcpthread.start()
            udpthread.start()
            tcpthread.join()
            udpthread.join()
        else:
            tcp = string.lower(sys.argv[1]) == "tcp"
            if tcp:
                runTCP()
            else:
                runUDP()