[client] Implement logger

Author: SamuelHassine

pyobas/helpers.py

@@ -12,7 +12,7 @@
 import pika
 import yaml
 
-from pyobas import OpenBAS
+from pyobas import OpenBAS, utils
 
 TRUTHY: List[str] = ["yes", "true", "True"]
 FALSY: List[str] = ["no", "false", "False"]
@@ -152,6 +152,7 @@ def __init__(
         self,
         config: Dict,
         injector_config: Dict,
+        logger,
         callback,
     ) -> None:
         threading.Thread.__init__(self)
@@ -162,6 +163,7 @@ def __init__(
 
         self.callback = callback
         self.config = config
+        self.logger = logger
         self.host = injector_config.connection["host"]
         self.vhost = injector_config.connection["vhost"]
         self.use_ssl = injector_config.connection["use_ssl"]
@@ -197,10 +199,10 @@ def _data_handler(self, json_data) -> None:
         self.callback(json_data)
 
     def run(self) -> None:
-        print("Starting ListenQueue thread")
+        self.logger.info("Starting ListenQueue thread")
         while not self.exit_event.is_set():
             try:
-                print("ListenQueue connecting to rabbitMq.")
+                self.logger.info("ListenQueue connecting to RabbitMQ.")
                 # Connect the broker
                 self.pika_credentials = pika.PlainCredentials(self.user, self.password)
                 self.pika_parameters = pika.ConnectionParameters(
@@ -221,7 +223,7 @@ def run(self) -> None:
                     # when not in cluster mode this line raise an exception
                     self.channel.confirm_delivery()
                 except Exception as err:  # pylint: disable=broad-except
-                    print(str(err))
+                    self.logger.error(str(err))
                 self.channel.basic_qos(prefetch_count=1)
                 assert self.channel is not None
                 self.channel.basic_consume(
@@ -232,25 +234,26 @@ def run(self) -> None:
                 try:
                     self.pika_connection.close()
                 except Exception as errInException:
-                    print(str(errInException))
+                    self.logger.error(str(errInException))
                 traceback.print_exc()
                 # Wait some time and then retry ListenQueue again.
                 time.sleep(10)
 
     def stop(self):
-        print("Preparing ListenQueue for clean shutdown")
+        self.logger.info("Preparing ListenQueue for clean shutdown")
         self.exit_event.set()
         self.pika_connection.close()
         if self.thread:
             self.thread.join()
 
 
 class PingAlive(threading.Thread):
-    def __init__(self, api, config, ping_type) -> None:
+    def __init__(self, api, config, logger, ping_type) -> None:
         threading.Thread.__init__(self)
         self.ping_type = ping_type
         self.api = api
         self.config = config
+        self.logger = logger
         self.in_error = False
         self.exit_event = threading.Event()
 
@@ -262,15 +265,15 @@ def ping(self) -> None:
                 else:
                     self.api.collector.create(self.config, False)
             except Exception as e:  # pylint: disable=broad-except
-                print(str(e))
+                self.logger.error(str(e))
             self.exit_event.wait(40)
 
     def run(self) -> None:
-        print("Starting PingAlive thread")
+        self.logger.info("Starting PingAlive thread")
         self.ping()
 
     def stop(self) -> None:
-        print("Preparing PingAlive for clean shutdown")
+        self.logger.info("Preparing PingAlive for clean shutdown")
         self.exit_event.set()
 
 
@@ -309,12 +312,20 @@ def __init__(self, config: OpenBASConfigHelper, icon) -> None:
             url=config.get_conf("openbas_url"),
             token=config.get_conf("openbas_token"),
         )
+
         self.config = {
             "collector_id": config.get_conf("collector_id"),
             "collector_name": config.get_conf("collector_name"),
             "collector_type": config.get_conf("collector_type"),
             "collector_period": config.get_conf("collector_period"),
         }
+
+        self.logger_class = utils.logger(
+            config.get_conf("collector_log_level", default="info").upper(),
+            config.get_conf("collector_json_logging", default=True),
+        )
+        self.collector_logger = self.logger_class(config.get_conf("collector_name"))
+
         icon_name = config.get_conf("collector_id") + ".png"
         collector_icon = (icon_name, icon, "image/png")
         self.api.collector.create(self.config, collector_icon)
@@ -323,7 +334,9 @@ def __init__(self, config: OpenBASConfigHelper, icon) -> None:
         self.scheduler = sched.scheduler(time.time, time.sleep)
         # Start ping thread
         if not self.connect_run_and_terminate:
-            self.ping = PingAlive(self.api, self.config, "collector")
+            self.ping = PingAlive(
+                self.api, self.config, self.collector_logger, "collector"
+            )
             self.ping.start()
         self.listen_queue = None
 
@@ -371,21 +384,28 @@ def __init__(self, config: OpenBASConfigHelper, icon) -> None:
                 "injector_executor_clear_commands", default=None
             ),
         }
+
+        self.logger_class = utils.logger(
+            config.get_conf("injector_log_level", default="info").upper(),
+            config.get_conf("injector_json_logging", default=True),
+        )
+        self.injector_logger = self.logger_class(config.get_conf("injector_name"))
+
         icon_name = config.get_conf("injector_type") + ".png"
         injector_icon = (icon_name, icon, "image/png")
         self.injector_config = self.api.injector.create(self.config, injector_icon)
         self.connect_run_and_terminate = False
         self.scheduler = sched.scheduler(time.time, time.sleep)
         # Start ping thread
         if not self.connect_run_and_terminate:
-            self.ping = PingAlive(self.api, self.config, "injector")
+            self.ping = PingAlive(
+                self.api, self.config, self.injector_logger, "injector"
+            )
             self.ping.start()
         self.listen_queue = None
 
     def listen(self, message_callback: Callable[[Dict], None]) -> None:
         self.listen_queue = ListenQueue(
-            self.config,
-            self.injector_config,
-            message_callback,
+            self.config, self.injector_config, self.injector_logger, message_callback
         )
         self.listen_queue.start()

pyobas/utils.py

@@ -1,11 +1,13 @@
 import dataclasses
+import datetime
 import email.message
 import json
 import logging
 import urllib.parse
 from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple, Union
 
 import requests
+from pythonjsonlogger import jsonlogger
 
 
 class _StdoutStream:
@@ -20,31 +22,6 @@ def get_content_type(content_type: Optional[str]) -> str:
     return message.get_content_type()
 
 
-class MaskingFormatter(logging.Formatter):
-    """A logging formatter that can mask credentials"""
-
-    def __init__(
-        self,
-        fmt: Optional[str] = logging.BASIC_FORMAT,
-        datefmt: Optional[str] = None,
-        style: Literal["%", "{", "$"] = "%",
-        validate: bool = True,
-        masked: Optional[str] = None,
-    ) -> None:
-        super().__init__(fmt, datefmt, style, validate)
-        self.masked = masked
-
-    def _filter(self, entry: str) -> str:
-        if not self.masked:
-            return entry
-
-        return entry.replace(self.masked, "[MASKED]")
-
-    def format(self, record: logging.LogRecord) -> str:
-        original = logging.Formatter.format(self, record)
-        return self._filter(original)
-
-
 def response_content(
     response: requests.Response,
     streamed: bool,
@@ -136,3 +113,60 @@ def validate_attrs(
                     f"Must provide one of these attributes: "
                     f"{', '.join(self.exclusive)}"
                 )
+
+
+class CustomJsonFormatter(jsonlogger.JsonFormatter):
+    def add_fields(self, log_record, record, message_dict):
+        super(CustomJsonFormatter, self).add_fields(log_record, record, message_dict)
+        if not log_record.get("timestamp"):
+            # This doesn't use record.created, so it is slightly off
+            now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+            log_record["timestamp"] = now
+        if log_record.get("level"):
+            log_record["level"] = log_record["level"].upper()
+        else:
+            log_record["level"] = record.levelname
+
+
+def logger(level, json_logging=True):
+    # Exceptions
+    logging.getLogger("urllib3").setLevel(logging.WARNING)
+    logging.getLogger("pika").setLevel(logging.ERROR)
+    # Exceptions
+    if json_logging:
+        log_handler = logging.StreamHandler()
+        log_handler.setLevel(level)
+        formatter = CustomJsonFormatter("%(timestamp)s %(level)s %(name)s %(message)s")
+        log_handler.setFormatter(formatter)
+        logging.basicConfig(handlers=[log_handler], level=level, force=True)
+    else:
+        logging.basicConfig(level=level)
+
+    class AppLogger:
+        def __init__(self, name):
+            self.local_logger = logging.getLogger(name)
+
+        @staticmethod
+        def prepare_meta(meta=None):
+            return None if meta is None else {"attributes": meta}
+
+        @staticmethod
+        def setup_logger_level(lib, log_level):
+            logging.getLogger(lib).setLevel(log_level)
+
+        def debug(self, message, meta=None):
+            self.local_logger.debug(message, extra=AppLogger.prepare_meta(meta))
+
+        def info(self, message, meta=None):
+            self.local_logger.info(message, extra=AppLogger.prepare_meta(meta))
+
+        def warning(self, message, meta=None):
+            self.local_logger.warning(message, extra=AppLogger.prepare_meta(meta))
+
+        def error(self, message, meta=None):
+            # noinspection PyTypeChecker
+            self.local_logger.error(
+                message, exc_info=1, extra=AppLogger.prepare_meta(meta)
+            )
+
+    return AppLogger