Merge pull request #50 from odin-detector/graylog

Add optional logging to a graylog server

.github/workflows/test_odin_control.yml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [2.7, 3.6, 3.7, 3.8, 3.9]
+        python-version: [3.7, 3.8, 3.9]
 
     steps:
     - uses: actions/checkout@v2

setup.cfg

@@ -34,6 +34,8 @@ dev =
     tox
     pytest-asyncio
     pytest-cov
+graylog =
+    pygelf
 
 [options.packages.find]
 where = src

src/odin/logconfig.py

@@ -0,0 +1,44 @@
+import os
+import sys
+import logging
+import getpass
+from typing import Optional
+
+
+def add_graylog_handler(
+    log_server: str, log_level: int = logging.INFO, static_fields: Optional[str] = None
+) -> None:
+    """Add a graylog handler to the root logger
+
+    Args:
+        log_server: Graylog server endpoint, e.g. "127.0.0.1:12201"
+        log_level: Log level to filter messages by in handler
+        static_fields: Comma-separated string of extra fields to include in log message
+            metadata, e.g. "_field1=value1,_field2=value2" (fields should have a
+            leading underscore).
+    """
+    try:
+        from pygelf import GelfUdpHandler
+    except ImportError:
+        logging.error("Cannot add graylog handler - pygelf is not installed")
+        return
+
+    host, port = log_server.split(":")
+    config = {
+        "host": host,
+        "port": int(port),
+        "debug": True,  # Include file, line, module, func, logger_name
+        # Add custom fields
+        "include_extra_fields": True,
+        "_username": getpass.getuser(),
+        "_process_id": os.getpid(),
+        "_application_name": os.path.split(sys.argv[0])[1]
+    }
+
+    if static_fields is not None:
+        static_fields = dict(entry.split("=") for entry in static_fields.split(","))
+        config.update(static_fields)
+
+    handler: logging.Handler = GelfUdpHandler(**config)
+    handler.setLevel(log_level)
+    logging.getLogger().addHandler(handler)

src/odin/main.py

@@ -14,6 +14,7 @@
 
 from odin.http.server import HttpServer
 from odin.config.parser import ConfigParser, ConfigError
+from odin.logconfig import add_graylog_handler
 
 
 def shutdown_handler():  # pragma: no cover
@@ -43,6 +44,10 @@ def main(argv=None):
     config.define('enable_cors', default=False,
                   option_help='Enable cross-origin resource sharing (CORS)')
     config.define('cors_origin', default='*', option_help='Specify allowed CORS origin')
+    config.define('graylog_server', default=None, option_help="Graylog server address and :port")
+    config.define('graylog_logging_level', default=logging.INFO, option_help="Graylog logging level")
+    config.define('graylog_static_fields', default=None,
+                  option_help="Comma separated list of key=value pairs to add to every log message metadata")
 
     # Parse configuration options and any configuration file specified
     try:
@@ -51,6 +56,13 @@ def main(argv=None):
         logging.error('Failed to parse configuration: %s', e)
         return 2
 
+    if config.graylog_server is not None:
+        add_graylog_handler(
+            config.graylog_server,
+            config.graylog_logging_level,
+            config.graylog_static_fields
+        )
+
     # Launch the HTTP server with the parsed configuration
     http_server = HttpServer(config)
     http_server.listen(config.http_port, config.http_addr)

tests/test_server.py

@@ -5,9 +5,9 @@
 
 import pytest
 if sys.version_info[0] == 3:  # pragma: no cover
-    from unittest.mock import Mock
+    from unittest import mock
 else:                         # pragma: no cover
-    from mock import Mock
+    import mock
 
 from odin.http.server import HttpServer
 from odin import main
@@ -139,6 +139,39 @@ def test_background_task_in_adapter(self, odin_test_server):
         count = result.json()['response']['background_task_count']
         assert count > 0
 
+def test_graylog_handler_pygelf():
+    """Test that gelf handler is added if pygelf is available"""
+    try:
+        import pygelf
+        del pygelf
+    except ImportError:
+        return  # Cannot test pygelf functionality
+
+    with mock.patch.object(logging.getLogger(), "addHandler") as mock_add_handler, \
+        mock.patch("pygelf.GelfUdpHandler") as mock_gelf:
+
+        server = OdinTestServer(
+            graylog_server="127.0.0.1:12210",
+            graylog_static_fields="key1=val1,key2=val2"
+        )
+        server.stop()
+
+        mock_add_handler.assert_called_with(mock_gelf.return_value)
+
+
+def test_graylog_handler_no_pygelf(caplog):
+    """Test that an error is logged if called without pygelf available"""
+    with mock.patch.dict(sys.modules, {"pygelf": None}):
+        server = OdinTestServer(graylog_server="127.0.0.1:12210")
+        server.stop()
+
+    assert log_message_seen(
+        caplog,
+        logging.ERROR,
+        "Cannot add graylog handler - pygelf is not installed"
+    )
+
+
 class TestBadServerConfig(object):
     """Class for testing a server with a bad configuration argument."""
 

tests/utils.py

@@ -37,7 +37,14 @@ class OdinTestServer(object):
     server_addr = '127.0.0.1'
     server_api_version = 0.1
 
-    def __init__(self, server_port=server_port, adapter_config=None, access_logging=None):
+    def __init__(
+        self,
+        server_port=server_port,
+        adapter_config=None,
+        access_logging=None,
+        graylog_server=None,
+        graylog_static_fields=None,
+    ):
 
         self.server_thread = None
         self.server_event_loop = None
@@ -61,6 +68,11 @@ def __init__(self, server_port=server_port, adapter_config=None, access_logging=
         if access_logging is not None:
             parser.set("server", 'access_logging', access_logging)
 
+        if graylog_server is not None:
+            parser.set("server", 'graylog_server', graylog_server)
+            if graylog_static_fields is not None:
+                parser.set("server", 'graylog_static_fields', graylog_static_fields)
+
         parser.add_section('tornado')
         parser.set('tornado', 'logging', 'debug')
 

tox.ini

@@ -1,7 +1,7 @@
 # tox test configuration for odin-control
 
 [tox]
-envlist = clean,py27-tornado{4,5},py{36,37,38,39}-tornado{5,6},report
+envlist = clean,py27-tornado{4,5},py{36,37,38,39}-tornado{5,6},py{37}-tornado{6}-pygelf,report
 
 [gh-actions]
 python =
@@ -21,6 +21,7 @@ deps =
     tornado4: tornado>=4.0,<5.0
     tornado5: tornado>=5.0,<6.0
     tornado6: tornado>=6.0
+    py37: pygelf
 setenv =
     py{27,36,37,38,39}: COVERAGE_FILE=.coverage.{envname}
 commands =