1# The source code is from following Python documentation:
2# https://docs.python.org/2/howto/logging-cookbook.html#network-logging
3
4# Classes in this file are used to create a simple TCP socket-based logging
5# receiver. The receiver listens to default logging port (9020) and save log to
6# any given log configuration, e.g., a local file. Once the receiver is running,
7# client can add a logging handler to write log to the receiver with following
8# sample code:
9# socketHandler = logging.handlers.SocketHandler('localhost',
10#         logging.handlers.DEFAULT_TCP_LOGGING_PORT)
11# logging.getLogger().addHandler(socketHandler)
12
13import ctypes
14import pickle
15import logging
16import multiprocessing
17import select
18import SocketServer
19import struct
20import time
21
22import common
23from autotest_lib.client.common_lib import utils
24
25class LogRecordStreamHandler(SocketServer.StreamRequestHandler):
26    """Handler for a streaming logging request.
27
28    This basically logs the record using whatever logging policy is
29    configured locally.
30    """
31
32    def handle(self):
33        """
34        Handle multiple requests - each expected to be a 4-byte length,
35        followed by the LogRecord in pickle format. Logs the record
36        according to whatever policy is configured locally.
37        """
38        while True:
39            chunk = self.connection.recv(4)
40            if len(chunk) < 4:
41                return
42            slen = struct.unpack('>L', chunk)[0]
43            chunk = self.connection.recv(slen)
44            while len(chunk) < slen:
45                chunk = chunk + self.connection.recv(slen - len(chunk))
46            obj = self.unpickle(chunk)
47            record = logging.makeLogRecord(obj)
48            self.handle_log_record(record)
49
50
51    def unpickle(self, data):
52        """Unpickle data received.
53
54        @param data: Received data.
55        @returns: unpickled data.
56        """
57        return pickle.loads(data)
58
59
60    def handle_log_record(self, record):
61        """Process log record.
62
63        @param record: log record.
64        """
65        # if a name is specified, we use the named logger rather than the one
66        # implied by the record.
67        if self.server.logname is not None:
68            name = self.server.logname
69        else:
70            name = record.name
71        logger = logging.getLogger(name)
72        # N.B. EVERY record gets logged. This is because Logger.handle
73        # is normally called AFTER logger-level filtering. If you want
74        # to do filtering, do it at the client end to save wasting
75        # cycles and network bandwidth!
76        logger.handle(record)
77
78
79class LogRecordSocketReceiver(SocketServer.ThreadingTCPServer):
80    """Simple TCP socket-based logging receiver.
81    """
82
83    allow_reuse_address = 1
84
85    def __init__(self, host='localhost', port=None,
86                 handler=LogRecordStreamHandler):
87        if not port:
88            port = utils.get_unused_port()
89        SocketServer.ThreadingTCPServer.__init__(self, (host, port), handler)
90        self.abort = 0
91        self.timeout = 1
92        self.logname = None
93        self.port = port
94
95
96    def serve_until_stopped(self):
97        """Run the socket receiver until aborted."""
98        print 'Log Record Socket Receiver is started.'
99        abort = 0
100        while not abort:
101            rd, wr, ex = select.select([self.socket.fileno()], [], [],
102                                       self.timeout)
103            if rd:
104                self.handle_request()
105            abort = self.abort
106        print 'Log Record Socket Receiver is stopped.'
107
108
109class LogSocketServer:
110    """A wrapper class to start and stop a TCP server for logging."""
111
112    process = None
113    port = None
114
115    @staticmethod
116    def start(**kwargs):
117        """Start Log Record Socket Receiver in a new process.
118
119        @param kwargs: log configuration, e.g., format, filename.
120
121        @raise Exception: if TCP server is already running.
122        """
123        if LogSocketServer.process:
124            raise Exception('Log Record Socket Receiver is already running.')
125        server_started = multiprocessing.Value(ctypes.c_bool, False)
126        port = multiprocessing.Value(ctypes.c_int, 0)
127        LogSocketServer.process = multiprocessing.Process(
128                target=LogSocketServer._start_server,
129                args=(server_started, port),
130                kwargs=kwargs)
131        LogSocketServer.process.start()
132        while not server_started.value:
133            time.sleep(0.1)
134        LogSocketServer.port = port.value
135        print 'Log Record Socket Server is started at port %d.' % port.value
136
137
138    @staticmethod
139    def _start_server(server_started, port, **kwargs):
140        """Start the TCP server to receive log.
141
142        @param server_started: True if socket log server is started.
143        @param port: Port used by socket log server.
144        @param kwargs: log configuration, e.g., format, filename.
145        """
146        # Clear all existing log handlers.
147        logging.getLogger().handlers = []
148        if not kwargs:
149            logging.basicConfig(
150                    format='%(asctime)s - %(levelname)s - %(message)s')
151        else:
152            logging.basicConfig(**kwargs)
153
154        tcp_server = LogRecordSocketReceiver()
155        print('Starting TCP server...')
156        server_started.value = True
157        port.value = tcp_server.port
158        tcp_server.serve_until_stopped()
159
160
161    @staticmethod
162    def stop():
163        """Stop Log Record Socket Receiver.
164        """
165        if LogSocketServer.process:
166            LogSocketServer.process.terminate()
167            LogSocketServer.process = None
168            LogSocketServer.port = None
169