1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this file,
3# You can obtain one at http://mozilla.org/MPL/2.0/.
4
5import SocketServer
6import socket
7import json
8
9
10class LogMessageServer(SocketServer.TCPServer):
11
12    def __init__(self, server_address, logger, message_callback=None, timeout=3):
13        SocketServer.TCPServer.__init__(self, server_address, LogMessageHandler)
14        self._logger = logger
15        self._message_callback = message_callback
16        self.timeout = timeout
17
18
19class LogMessageHandler(SocketServer.BaseRequestHandler):
20    """Processes output from a connected log source, logging to an
21    existing logger upon receipt of a well-formed log messsage."""
22
23    def handle(self):
24        """Continually listens for log messages."""
25        self._partial_message = ''
26        self.request.settimeout(self.server.timeout)
27
28        while True:
29            try:
30                data = self.request.recv(1024)
31                if not data:
32                    return
33                self.process_message(data)
34            except socket.timeout:
35                return
36
37    def process_message(self, data):
38        """Processes data from a connected log source. Messages are assumed
39        to be newline delimited, and generally well-formed JSON."""
40        for part in data.split('\n'):
41            msg_string = self._partial_message + part
42            try:
43                msg = json.loads(msg_string)
44                self._partial_message = ''
45                self.server._logger.log_structured(msg.get('action', 'UNKNOWN'), msg)
46                if self.server._message_callback:
47                    self.server._message_callback()
48
49            except ValueError:
50                self._partial_message = msg_string
51