From a29589adda00b320e8ec89d8e5a36305a240e6da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bruno=20Gon=C3=A7alves?= Date: Mon, 11 Nov 2024 18:30:51 +0000 Subject: [PATCH] Applied auto-formatting --- pycti/api/opencti_api_client.py | 12 ++++++--- pycti/connector/opencti_connector_helper.py | 14 ++++++++--- pycti/utils/opencti_logger.py | 28 +++++++++++++++------ 3 files changed, 41 insertions(+), 13 deletions(-) diff --git a/pycti/api/opencti_api_client.py b/pycti/api/opencti_api_client.py index 57282284..1d048c55 100644 --- a/pycti/api/opencti_api_client.py +++ b/pycti/api/opencti_api_client.py @@ -143,9 +143,15 @@ def __init__( raise ValueError("A TOKEN must be set") # Configure logger - self.logger_class = logger(log_level.upper(), json_logging, graylog_host, graylog_port, graylog_adapter, - log_shipping_level.upper() if log_shipping_level is not None else None, - log_shipping_env_var_prefix) + self.logger_class = logger( + log_level.upper(), + json_logging, + graylog_host, + graylog_port, + graylog_adapter, + log_shipping_level.upper() if log_shipping_level is not None else None, + log_shipping_env_var_prefix, + ) self.app_logger = self.logger_class("api") # Define API diff --git a/pycti/connector/opencti_connector_helper.py b/pycti/connector/opencti_connector_helper.py index a6411b66..4cb9c6d9 100644 --- a/pycti/connector/opencti_connector_helper.py +++ b/pycti/connector/opencti_connector_helper.py @@ -885,13 +885,21 @@ def __init__(self, config: Dict, playbook_compatible=False) -> None: "CONNECTOR_GRAYLOG_PORT", ["connector", "graylog_port"], config, True, 12201 ) self.graylog_adapter = get_config_variable( - "CONNECTOR_GRAYLOG_ADAPTER", ["connector", "graylog_adapter"], config, default="udp" + "CONNECTOR_GRAYLOG_ADAPTER", + ["connector", "graylog_adapter"], + config, + default="udp", ) self.log_shipping_level = get_config_variable( - "CONNECTOR_LOG_SHIPPING_LEVEL", ["connector", "log_shipping_level"], config, default="INFO" + "CONNECTOR_LOG_SHIPPING_LEVEL", + ["connector", "log_shipping_level"], + config, + default="INFO", ).upper() self.log_shipping_env_var_prefix = get_config_variable( - "CONNECTOR_LOG_SHIPPING_ENV_VAR_PREFIX", ["connector", "log_shipping_env_var_prefix"], config + "CONNECTOR_LOG_SHIPPING_ENV_VAR_PREFIX", + ["connector", "log_shipping_env_var_prefix"], + config, ) self.connect_run_and_terminate = get_config_variable( "CONNECTOR_RUN_AND_TERMINATE", diff --git a/pycti/utils/opencti_logger.py b/pycti/utils/opencti_logger.py index f0e36a2a..8d8bfaaf 100644 --- a/pycti/utils/opencti_logger.py +++ b/pycti/utils/opencti_logger.py @@ -2,7 +2,7 @@ import logging import os -from pygelf import GelfUdpHandler, GelfTcpHandler +from pygelf import GelfTcpHandler, GelfUdpHandler from pythonjsonlogger import jsonlogger @@ -34,8 +34,15 @@ def filter(self, record): return True -def logger(level, json_logging=True, graylog_host=None, graylog_port=None, graylog_adapter=None, - log_shipping_level=None, log_shipping_env_var_prefix=None): +def logger( + level, + json_logging=True, + graylog_host=None, + graylog_port=None, + graylog_adapter=None, + log_shipping_level=None, + log_shipping_env_var_prefix=None, +): # Exceptions logging.getLogger("urllib3").setLevel(logging.WARNING) logging.getLogger("pika").setLevel(logging.ERROR) @@ -51,14 +58,21 @@ def logger(level, json_logging=True, graylog_host=None, graylog_port=None, grayl if graylog_host is not None: if graylog_adapter == "tcp": - shipping_handler = GelfTcpHandler(host=graylog_host, port=graylog_port, include_extra_fields=True) + shipping_handler = GelfTcpHandler( + host=graylog_host, port=graylog_port, include_extra_fields=True + ) else: - shipping_handler = GelfUdpHandler(host=graylog_host, port=graylog_port, include_extra_fields=True) + shipping_handler = GelfUdpHandler( + host=graylog_host, port=graylog_port, include_extra_fields=True + ) shipping_handler.setLevel(log_shipping_level) if log_shipping_env_var_prefix is not None: - filtered_env = [(k.removeprefix(log_shipping_env_var_prefix), v) for k, v in os.environ.items() - if k.startswith(log_shipping_env_var_prefix)] + filtered_env = [ + (k.removeprefix(log_shipping_env_var_prefix), v) + for k, v in os.environ.items() + if k.startswith(log_shipping_env_var_prefix) + ] shipping_filter = ContextFilter(filtered_env) shipping_handler.addFilter(shipping_filter)