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
5from __future__ import absolute_import
6
7from six.moves import socketserver
8import json
9import socket
10
11
12class LogMessageServer(socketserver.TCPServer):
13
14    def __init__(self, server_address, logger, message_callback=None, timeout=3):
15        socketserver.TCPServer.__init__(self, server_address, LogMessageHandler)
16        self._logger = logger
17        self._message_callback = message_callback
18        self.timeout = timeout
19
20
21class LogMessageHandler(socketserver.BaseRequestHandler):
22    """Processes output from a connected log source, logging to an
23    existing logger upon receipt of a well-formed log messsage."""
24
25    def handle(self):
26        """Continually listens for log messages."""
27        self._partial_message = ''
28        self.request.settimeout(self.server.timeout)
29
30        while True:
31            try:
32                data = self.request.recv(1024)
33                if not data:
34                    return
35                self.process_message(data)
36            except socket.timeout:
37                return
38
39    def process_message(self, data):
40        """Processes data from a connected log source. Messages are assumed
41        to be newline delimited, and generally well-formed JSON."""
42        for part in data.split('\n'):
43            msg_string = self._partial_message + part
44            try:
45                msg = json.loads(msg_string)
46                self._partial_message = ''
47                self.server._logger.log_structured(msg.get('action', 'UNKNOWN'), msg)
48                if self.server._message_callback:
49                    self.server._message_callback()
50
51            except ValueError:
52                self._partial_message = msg_string
53