From ae29a6637beb53b72742b88555f535e4b03dd424 Mon Sep 17 00:00:00 2001 From: maldwg Date: Tue, 30 Jun 2026 15:55:28 +0200 Subject: [PATCH 1/3] fix inspector issue --- src/inspector/inspector.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/inspector/inspector.py b/src/inspector/inspector.py index f63127a1..b5366123 100644 --- a/src/inspector/inspector.py +++ b/src/inspector/inspector.py @@ -34,18 +34,18 @@ module_name = "data_inspection.inspector" logger = get_logger(module_name) -APP_CONFIG = setup_config() -PRODUCE_TOPIC_PREFIX = APP_CONFIG["environment"]["kafka_topics_prefix"]["pipeline"][ +config = setup_config() +PRODUCE_TOPIC_PREFIX = config["environment"]["kafka_topics_prefix"]["pipeline"][ "inspector_to_detector" ] -CONSUME_TOPIC_PREFIX = APP_CONFIG["environment"]["kafka_topics_prefix"]["pipeline"][ +CONSUME_TOPIC_PREFIX = config["environment"]["kafka_topics_prefix"]["pipeline"][ "prefilter_to_inspector" ] -SENSOR_PROTOCOLS = get_zeek_sensor_topic_base_names(APP_CONFIG) -PREFILTERS = APP_CONFIG["pipeline"]["log_filtering"] -INSPECTORS = APP_CONFIG["pipeline"]["data_inspection"] -COLLECTORS = APP_CONFIG["pipeline"]["log_collection"]["collectors"] -DETECTORS = APP_CONFIG["pipeline"]["data_analysis"] +SENSOR_PROTOCOLS = get_zeek_sensor_topic_base_names(config) +PREFILTERS = config["pipeline"]["log_filtering"] +INSPECTORS = config["pipeline"]["data_inspection"] +COLLECTORS = config["pipeline"]["log_collection"]["collectors"] +DETECTORS = config["pipeline"]["data_analysis"] PLUGIN_PATH = "src.inspector.plugins" @@ -99,7 +99,7 @@ def __init__(self, consume_topic, produce_topics, config) -> None: self.time_range = config["time_range"] self.name = config["name"] self.acceleration = resolve_acceleration_config( - APP_CONFIG["pipeline"], + config["pipeline"], config, component_name=f"{module_name}.{self.name}", logger=logger, From 0c6156e5387ad05b1d00b9e004930b89d5d83cd2 Mon Sep 17 00:00:00 2001 From: maldwg Date: Tue, 30 Jun 2026 15:56:09 +0200 Subject: [PATCH 2/3] Appease the linting gods --- config.yaml | 2 +- docker/create_tables/alerts.sql | 2 +- docker/create_tables/logline_to_batches.sql | 2 - .../create_tables/server_log_to_logline.sql | 1 - docker/create_tables/server_logs.sql | 2 +- .../create_tables/server_logs_timestamps.sql | 2 +- .../suspicious_batch_timestamps.sql | 2 +- docs/api/base.rst | 1 - docs/api/detector.rst | 1 - docs/api/inspector.rst | 1 - docs/api/logcollector.rst | 1 - docs/api/logserver.rst | 1 - docs/api/monitoring.rst | 1 - docs/api/prefilter.rst | 1 - docs/api/train.rst | 1 - docs/api/version.rst | 1 - docs/index.rst | 2 +- src/alerter/alerter.py | 4 +- src/base/execution.py | 4 +- src/base/kafka_handler.py | 41 ++++++++++++++----- src/base/retry.py | 6 +-- src/detector/detector.py | 32 +++++++-------- src/detector/plugins/dga_detector.py | 3 +- src/detector/plugins/domainator_attributor.py | 10 ++--- src/detector/plugins/domainator_utils.py | 7 +--- src/monitoring/clickhouse_batch_sender.py | 11 +++-- tests/alerter/__init__.py | 1 - tests/alerter/test_alerter.py | 1 - tests/detector/test_domainator_attributor.py | 10 +++-- tests/detector/test_domainator_detector.py | 6 ++- tests/kafka/test_kafka_consume_handler.py | 15 ++++--- tests/logcollector/test_collector.py | 4 +- tests/logserver/test_server.py | 12 +++--- tests/miscellaneous/test_acceleration.py | 5 +-- tests/miscellaneous/test_monitoring_agent.py | 4 +- 35 files changed, 97 insertions(+), 103 deletions(-) diff --git a/config.yaml b/config.yaml index 8bd8ae20..5aba2f3d 100644 --- a/config.yaml +++ b/config.yaml @@ -200,7 +200,7 @@ pipeline: threshold: 0.05 consume_from: inspector inspector_name: "domainator_inspector" - next_detectors: + next_detectors: - "domainator_attributor_behaviour" - "domainator_attributor_identification_behaviour" - "domainator_attributor_identification" diff --git a/docker/create_tables/alerts.sql b/docker/create_tables/alerts.sql index 16b11a59..4e218868 100644 --- a/docker/create_tables/alerts.sql +++ b/docker/create_tables/alerts.sql @@ -9,4 +9,4 @@ CREATE TABLE IF NOT EXISTS alerts ( ENGINE = MergeTree PARTITION BY toYYYYMM(alert_timestamp) ORDER BY (alert_timestamp, src_ip, suspicious_batch_id) -TTL toDateTime(alert_timestamp) + INTERVAL 60 DAY; \ No newline at end of file +TTL toDateTime(alert_timestamp) + INTERVAL 60 DAY; diff --git a/docker/create_tables/logline_to_batches.sql b/docker/create_tables/logline_to_batches.sql index 8b306171..8d68b65c 100644 --- a/docker/create_tables/logline_to_batches.sql +++ b/docker/create_tables/logline_to_batches.sql @@ -7,5 +7,3 @@ ENGINE = MergeTree ORDER BY (timestamp, batch_id, logline_id) PARTITION BY toYYYYMM(timestamp) TTL toDateTime(timestamp) + INTERVAL 1 DAY; - - diff --git a/docker/create_tables/server_log_to_logline.sql b/docker/create_tables/server_log_to_logline.sql index 96fe645f..e96ce7ab 100644 --- a/docker/create_tables/server_log_to_logline.sql +++ b/docker/create_tables/server_log_to_logline.sql @@ -7,4 +7,3 @@ ENGINE = MergeTree ORDER BY (timestamp, message_id, logline_id) PARTITION BY toYYYYMM(timestamp) TTL toDateTime(timestamp) + INTERVAL 1 DAY; - diff --git a/docker/create_tables/server_logs.sql b/docker/create_tables/server_logs.sql index d166d9a5..8f803cc8 100644 --- a/docker/create_tables/server_logs.sql +++ b/docker/create_tables/server_logs.sql @@ -6,4 +6,4 @@ CREATE TABLE IF NOT EXISTS server_logs ( ENGINE = MergeTree PARTITION BY toYYYYMM(timestamp_in) ORDER BY (timestamp_in, message_id) -TTL toDateTime(timestamp_in) + INTERVAL 1 DAY; \ No newline at end of file +TTL toDateTime(timestamp_in) + INTERVAL 1 DAY; diff --git a/docker/create_tables/server_logs_timestamps.sql b/docker/create_tables/server_logs_timestamps.sql index 81a6b972..81e5a71b 100644 --- a/docker/create_tables/server_logs_timestamps.sql +++ b/docker/create_tables/server_logs_timestamps.sql @@ -6,4 +6,4 @@ CREATE TABLE IF NOT EXISTS server_logs_timestamps ( ENGINE = MergeTree PARTITION BY toYYYYMM(event_timestamp) ORDER BY (event, event_timestamp, message_id) -TTL toDateTime(event_timestamp) + INTERVAL 1 DAY; \ No newline at end of file +TTL toDateTime(event_timestamp) + INTERVAL 1 DAY; diff --git a/docker/create_tables/suspicious_batch_timestamps.sql b/docker/create_tables/suspicious_batch_timestamps.sql index 4bc857b9..c1dc908d 100644 --- a/docker/create_tables/suspicious_batch_timestamps.sql +++ b/docker/create_tables/suspicious_batch_timestamps.sql @@ -11,4 +11,4 @@ CREATE TABLE IF NOT EXISTS suspicious_batch_timestamps ( ENGINE = MergeTree PARTITION BY toYYYYMM(timestamp) ORDER BY (stage, status, timestamp, instance_name, suspicious_batch_id, src_ip) -TTL toDateTime(timestamp) + INTERVAL 1 DAY; \ No newline at end of file +TTL toDateTime(timestamp) + INTERVAL 1 DAY; diff --git a/docs/api/base.rst b/docs/api/base.rst index f494dcfe..653730a7 100644 --- a/docs/api/base.rst +++ b/docs/api/base.rst @@ -2,4 +2,3 @@ Base ==== .. automodule:: src.base - diff --git a/docs/api/detector.rst b/docs/api/detector.rst index 03c31da2..f5110bee 100644 --- a/docs/api/detector.rst +++ b/docs/api/detector.rst @@ -2,4 +2,3 @@ Detector ======== .. automodule:: src.detector - diff --git a/docs/api/inspector.rst b/docs/api/inspector.rst index 7ca22f1e..3faaec4d 100644 --- a/docs/api/inspector.rst +++ b/docs/api/inspector.rst @@ -2,4 +2,3 @@ Inspector ========= .. automodule:: src.inspector - diff --git a/docs/api/logcollector.rst b/docs/api/logcollector.rst index f7c3ee09..69ca0c71 100644 --- a/docs/api/logcollector.rst +++ b/docs/api/logcollector.rst @@ -2,4 +2,3 @@ Log Collector ============= .. automodule:: src.logcollector - diff --git a/docs/api/logserver.rst b/docs/api/logserver.rst index a8394239..939646e2 100644 --- a/docs/api/logserver.rst +++ b/docs/api/logserver.rst @@ -2,4 +2,3 @@ Log Server ========== .. automodule:: src.logserver - diff --git a/docs/api/monitoring.rst b/docs/api/monitoring.rst index 652c988c..8242f8b6 100644 --- a/docs/api/monitoring.rst +++ b/docs/api/monitoring.rst @@ -2,4 +2,3 @@ Monitoring ========== .. automodule:: src.monitoring - diff --git a/docs/api/prefilter.rst b/docs/api/prefilter.rst index f806c639..9e1976e3 100644 --- a/docs/api/prefilter.rst +++ b/docs/api/prefilter.rst @@ -2,4 +2,3 @@ Prefilter ========= .. automodule:: src.prefilter - diff --git a/docs/api/train.rst b/docs/api/train.rst index 836ba206..744d6191 100644 --- a/docs/api/train.rst +++ b/docs/api/train.rst @@ -2,4 +2,3 @@ Train ===== .. automodule:: src.train - diff --git a/docs/api/version.rst b/docs/api/version.rst index e3c1f69f..c0416bc0 100644 --- a/docs/api/version.rst +++ b/docs/api/version.rst @@ -2,4 +2,3 @@ Version ======= .. automodule:: src.version - diff --git a/docs/index.rst b/docs/index.rst index 06875993..93cc0c53 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -2,7 +2,7 @@ Welcome to HAMSTRING's documentation! ====================================== **HAMSTRING** is a CIDS framework to run signature and machine learning-based IDS classifiers. It employs several string and anomaly-based filtering techniques -to maximize detection efficiency. Currently the tool focueses on DNS attacks, as it incorporates heiDGAF (DGA detection) & Domainator (DNS Tunneling Detector) +to maximize detection efficiency. Currently the tool focueses on DNS attacks, as it incorporates heiDGAF (DGA detection) & Domainator (DNS Tunneling Detector) :cite:p:`petrov_domainator_2025` :cite:p:`machmeier_heidgaf`. Check out the :doc:`usage` section for further information on how to use the software, including how to diff --git a/src/alerter/alerter.py b/src/alerter/alerter.py index a04f5fca..38f27886 100644 --- a/src/alerter/alerter.py +++ b/src/alerter/alerter.py @@ -348,9 +348,7 @@ async def start(self): def build_alerter_worker(alerter_config, consume_topic, worker_id=None): class_name = alerter_config.get("alerter_class_name", "GenericAlerter") - alerter_module_name = alerter_config.get( - "alerter_module_name", "generic_alerter" - ) + alerter_module_name = alerter_config.get("alerter_module_name", "generic_alerter") plugin_module_name = f"{PLUGIN_PATH}.{alerter_module_name}" plugin_module = importlib.import_module(plugin_module_name) alerter_class = getattr(plugin_module, class_name) diff --git a/src/base/execution.py b/src/base/execution.py index 317012b0..30a9bfe9 100644 --- a/src/base/execution.py +++ b/src/base/execution.py @@ -97,7 +97,9 @@ async def start_pipeline_worker_replicas( return if process_entrypoint is None: - raise ValueError("process_entrypoint is required for process and hybrid scaling") + raise ValueError( + "process_entrypoint is required for process and hybrid scaling" + ) loop = asyncio.get_running_loop() executor = ProcessPoolExecutor(max_workers=executor_config.processes) diff --git a/src/base/kafka_handler.py b/src/base/kafka_handler.py index 007f9553..f906a991 100644 --- a/src/base/kafka_handler.py +++ b/src/base/kafka_handler.py @@ -53,8 +53,8 @@ ) KAFKA_TOPIC_STAGE_CONFIG = KAFKA_TOPIC_CONFIG.get("stages", {}) KAFKA_TOPIC_EXACT_CONFIG = KAFKA_TOPIC_CONFIG.get("topics", {}) -KAFKA_PIPELINE_TOPIC_PREFIXES = config["environment"].get("kafka_topics_prefix", {}).get( - "pipeline", {} +KAFKA_PIPELINE_TOPIC_PREFIXES = ( + config["environment"].get("kafka_topics_prefix", {}).get("pipeline", {}) ) @@ -390,17 +390,23 @@ def _reset_producer(self) -> None: if self.producer: self.producer.flush(5) except Exception as exception: - logger.warning("Ignoring Kafka producer flush failure during reconnect: %s", exception) + logger.warning( + "Ignoring Kafka producer flush failure during reconnect: %s", exception + ) self.producer = self._new_producer() - def _with_producer_retry(self, description: str, operation: Callable[[], None]) -> None: + def _with_producer_retry( + self, description: str, operation: Callable[[], None] + ) -> None: def attempt(): try: operation() except Exception as exception: if not _is_retriable_kafka_exception(exception): raise - logger.warning("%s failed, recreating Kafka producer: %s", description, exception) + logger.warning( + "%s failed, recreating Kafka producer: %s", description, exception + ) self._reset_producer() raise @@ -599,7 +605,9 @@ def operation(): try: self.producer.abort_transaction() except Exception as abort_exception: - logger.warning("Kafka transaction abort failed: %s", abort_exception) + logger.warning( + "Kafka transaction abort failed: %s", abort_exception + ) logger.error("Transaction aborted.") logger.error(e) raise @@ -711,7 +719,9 @@ def connect(): ) target_partitions_by_topic = ensure_topics(admin_client, self.topics) - if not self._all_topics_created(self.topics, target_partitions_by_topic, consumer): + if not self._all_topics_created( + self.topics, target_partitions_by_topic, consumer + ): try: consumer.close() except Exception: @@ -721,14 +731,18 @@ def connect(): consumer.subscribe(self.topics) return consumer - self.consumer = retry_forever(connect, f"Kafka consumer setup for {self.topics}") + self.consumer = retry_forever( + connect, f"Kafka consumer setup for {self.topics}" + ) def _reset_consumer(self) -> None: try: if self.consumer: self.consumer.close() except Exception as exception: - logger.warning("Ignoring Kafka consumer close failure during reconnect: %s", exception) + logger.warning( + "Ignoring Kafka consumer close failure during reconnect: %s", exception + ) self._last_consumed_message = None self._connect_consumer() @@ -842,7 +856,9 @@ def _poll_message(self): try: msg = self.consumer.poll(timeout=1.0) except (KafkaException, RuntimeError, OSError) as exception: - logger.warning("Kafka consumer poll failed, reconnecting: %s", exception) + logger.warning( + "Kafka consumer poll failed, reconnecting: %s", exception + ) self._reset_consumer() continue @@ -853,7 +869,10 @@ def _poll_message(self): if msg.error().code() == KafkaError._PARTITION_EOF: return None if _is_retriable_kafka_error(msg.error()): - logger.warning("Kafka consumer error is retriable, reconnecting: %s", msg.error()) + logger.warning( + "Kafka consumer error is retriable, reconnecting: %s", + msg.error(), + ) self._reset_consumer() return None diff --git a/src/base/retry.py b/src/base/retry.py index 779798e5..2b4a2c11 100644 --- a/src/base/retry.py +++ b/src/base/retry.py @@ -21,11 +21,7 @@ def resilience_config() -> dict[str, Any]: config = setup_config() - retry_config = ( - config.get("pipeline", {}) - .get("resilience", {}) - .get("retry", {}) - ) + retry_config = config.get("pipeline", {}).get("resilience", {}).get("retry", {}) merged = dict(_DEFAULT_CONFIG) merged.update(retry_config) return merged diff --git a/src/detector/detector.py b/src/detector/detector.py index 41901f7a..dd121f36 100644 --- a/src/detector/detector.py +++ b/src/detector/detector.py @@ -61,7 +61,9 @@ def _normalize_topic_suffixes(topic_suffixes) -> list[str]: if topic_suffixes is None: return [] if isinstance(topic_suffixes, str): - return [suffix.strip() for suffix in topic_suffixes.split(",") if suffix.strip()] + return [ + suffix.strip() for suffix in topic_suffixes.split(",") if suffix.strip() + ] if isinstance(topic_suffixes, (list, tuple, set)): return [str(suffix).strip() for suffix in topic_suffixes if str(suffix).strip()] return [str(topic_suffixes).strip()] @@ -85,7 +87,9 @@ def build_alerter_topics(detector_config: dict) -> list[str]: def build_downstream_detector_topics(detector_config: dict) -> list[str]: detector_names = [] for config_key in ("next_detectors", "produce_detector_topics"): - detector_names.extend(_normalize_topic_suffixes(detector_config.get(config_key))) + detector_names.extend( + _normalize_topic_suffixes(detector_config.get(config_key)) + ) return [ f"{DETECTOR_TO_DETECTOR_TOPIC_PREFIX}-{detector_name}" @@ -352,7 +356,6 @@ def _sha256sum(self, file_path: str) -> str: return h.hexdigest() - def get_model_download_url(self): """ Generate the complete URL for downloading the Domainator detection model. @@ -387,8 +390,6 @@ def get_scaler_download_url(self): ) return f"{self.model_base_url}/files/?p=%2F{self.model_name}%2F{self.checksum}%2Fscaler.pickle&dl=1" - - def _get_model(self): """ Download and validate the detection model. @@ -410,7 +411,9 @@ def _get_model(self): requests.HTTPError: If there's an error downloading the model. """ logger.info(f"Get model: {self.model_name} with checksum {self.checksum}") - scaler_download_url = self.get_scaler_download_url() if self.use_scaler else None + scaler_download_url = ( + self.get_scaler_download_url() if self.use_scaler else None + ) if not os.path.isfile(self.model_path): model_download_url = self.get_model_download_url() @@ -460,7 +463,7 @@ def detect(self) -> None: then just overwrite the method and append to `self.warnings` a warning if a given messag is to be regarded malicious. Note: - + """ logger.info("general") logger.info("Start detecting malicious requests.") @@ -479,7 +482,7 @@ def detect(self) -> None: "sha256": self.checksum, } self.warnings.append(warning) - + def clear_data(self): """Clears the data in the internal data structures.""" self.messages = [] @@ -646,8 +649,7 @@ def _send_detector_batch(self, parent_row_id, messages) -> None: def _build_persisted_warnings(self) -> list[dict]: return [ - self._normalize_warning_for_storage(warning) - for warning in self.warnings + self._normalize_warning_for_storage(warning) for warning in self.warnings ] def _normalize_warning_for_storage(self, warning: dict) -> dict: @@ -658,11 +660,7 @@ def _normalize_warning_for_storage(self, warning: dict) -> dict: request_messages = ( self._flatten_messages(request) if request is not None else [] ) - detector_name = ( - warning.get("detector_name") - or warning.get("name") - or self.name - ) + detector_name = warning.get("detector_name") or warning.get("name") or self.name score = warning.get("score", warning.get("probability")) domains = self._extract_warning_domains(warning) logline_ids = self._extract_message_values(request_messages, "logline_id") @@ -753,9 +751,7 @@ def _extract_warning_domains(self, warning) -> list[str]: field_value = item.get(field_name) if isinstance(field_value, list): domains.extend( - domain - for domain in field_value - if isinstance(domain, str) + domain for domain in field_value if isinstance(domain, str) ) return sorted(set(domains)) diff --git a/src/detector/plugins/dga_detector.py b/src/detector/plugins/dga_detector.py index c3cebfcd..ac3c784c 100644 --- a/src/detector/plugins/dga_detector.py +++ b/src/detector/plugins/dga_detector.py @@ -152,7 +152,6 @@ def calculate_entropy(s: str) -> float: return all_features.reshape(1, -1) - def detect(self) -> None: """ Process messages to detect malicious requests. @@ -178,4 +177,4 @@ def detect(self) -> None: "name": self.name, "sha256": self.checksum, } - self.warnings.append(warning) \ No newline at end of file + self.warnings.append(warning) diff --git a/src/detector/plugins/domainator_attributor.py b/src/detector/plugins/domainator_attributor.py index 87f097fb..80be9dbc 100644 --- a/src/detector/plugins/domainator_attributor.py +++ b/src/detector/plugins/domainator_attributor.py @@ -19,17 +19,17 @@ class DomainatorAttributor(DetectorBase): """ - Detector implementation for the attribution of a tool or malware that was used in the + Detector implementation for the attribution of a tool or malware that was used in the data exfiltration and command and control, using the subdomain level labels. This class extends the DetectorBase to provide specific functionality for identifying - the tool sending the malicious queries. It analyzes subdomain similarity characteristics + the tool sending the malicious queries. It analyzes subdomain similarity characteristics based on grouping of the queries in windows of fixed size, similar to the Domainator detector approach. It can be used both as a standalone detector or as a next stage in a pipeline of detectors, - dependent on the provided model. + dependent on the provided model. The identity detector extracts various statistical similarity features from windows of subdomains - to make predictions about what tool likely sent the malicious query or what job/behaviour was observed. + to make predictions about what tool likely sent the malicious query or what job/behaviour was observed. """ def __init__( @@ -82,7 +82,7 @@ def predict(self, messages): def detect(self): logger.info("Start detecting malicious requests.") - + for message in self.messages: if isinstance(message, list): # This assumes the current example for the request structure. diff --git a/src/detector/plugins/domainator_utils.py b/src/detector/plugins/domainator_utils.py index 80671e70..3fe4db00 100644 --- a/src/detector/plugins/domainator_utils.py +++ b/src/detector/plugins/domainator_utils.py @@ -40,6 +40,7 @@ def strip_domain(query: str): return domain + def get_domainator_features(queries: list) -> pd.DataFrame: """Extracts feature vector from domain name for ML model inference. @@ -69,10 +70,7 @@ def get_domainator_features(queries: list) -> pd.DataFrame: [Levenshtein.jaro(product[0], product[1]) for product in cartesian] ) metrics["jaro_reversed"] = np.mean( - [ - Levenshtein.jaro(product[0][::-1], product[1][::-1]) - for product in cartesian - ] + [Levenshtein.jaro(product[0][::-1], product[1][::-1]) for product in cartesian] ) metrics["jaro_winkler"] = np.mean( [ @@ -116,4 +114,3 @@ def get_domainator_features(queries: list) -> pd.DataFrame: [[metrics[column] for column in DOMAINATOR_FEATURE_COLUMNS]], columns=DOMAINATOR_FEATURE_COLUMNS, ) - diff --git a/src/monitoring/clickhouse_batch_sender.py b/src/monitoring/clickhouse_batch_sender.py index e4f73c4d..8bb4ae54 100644 --- a/src/monitoring/clickhouse_batch_sender.py +++ b/src/monitoring/clickhouse_batch_sender.py @@ -236,7 +236,10 @@ def _reset_client(self) -> None: if self._client: self._client.close() except Exception as exception: - logger.warning("Ignoring ClickHouse client close failure during reconnect: %s", exception) + logger.warning( + "Ignoring ClickHouse client close failure during reconnect: %s", + exception, + ) self._client = self._connect_client() def __del__(self): @@ -298,10 +301,10 @@ def insert_batch(): self._reset_client() raise - retry_forever(insert_batch, f"ClickHouse insert for table '{table_name}'") - logger.debug( - f"Inserted {table_name=},{pending_rows=},{column_names=}" + retry_forever( + insert_batch, f"ClickHouse insert for table '{table_name}'" ) + logger.debug(f"Inserted {table_name=},{pending_rows=},{column_names=}") self.batch[table_name] = [] def insert_all(self): diff --git a/tests/alerter/__init__.py b/tests/alerter/__init__.py index 8b137891..e69de29b 100644 --- a/tests/alerter/__init__.py +++ b/tests/alerter/__init__.py @@ -1 +0,0 @@ - diff --git a/tests/alerter/test_alerter.py b/tests/alerter/test_alerter.py index 01bfd01d..574d6311 100644 --- a/tests/alerter/test_alerter.py +++ b/tests/alerter/test_alerter.py @@ -81,4 +81,3 @@ def test_log_to_file_writes_to_rotated_file(self): log_file = Path(temp_dir) / "alerts-2026-06-25.txt" self.assertEqual('{"src_ip": "192.0.2.1"}\n', log_file.read_text()) - diff --git a/tests/detector/test_domainator_attributor.py b/tests/detector/test_domainator_attributor.py index f076c76c..bcaee6cc 100644 --- a/tests/detector/test_domainator_attributor.py +++ b/tests/detector/test_domainator_attributor.py @@ -13,7 +13,7 @@ from src.base.data_classes.batch import Batch from src.detector.plugins.domainator_utils import ( DOMAINATOR_FEATURE_COLUMNS, - get_domainator_features + get_domainator_features, ) @@ -83,7 +83,7 @@ def test_detect(self): mock_kafka = MagicMock() mock_ch = MagicMock() sut = self._create_detector(mock_kafka, mock_ch) - sut.labels = ['tool-A', 'tool-B'] + sut.labels = ["tool-A", "tool-B"] for _ in range(0, 4, 1): sut.messages.append((DEFAULT_DATA)) with patch( @@ -136,7 +136,7 @@ def test_detect_message_list(self): mock_kafka = MagicMock() mock_ch = MagicMock() sut = self._create_detector(mock_kafka, mock_ch) - sut.labels = ['tool-A', 'tool-B'] + sut.labels = ["tool-A", "tool-B"] for _ in range(0, 2, 1): sut.messages.append([DEFAULT_DATA, DEFAULT_DATA]) with patch( @@ -229,7 +229,9 @@ def test_get_features_empty_domains(self): self.assertEqual( features.iloc[0, 0], 1.0 ) # Levenshtein distance of empty strings is 1 - self.assertEqual(features.iloc[0, 1], 1.0) # Jaro distance of empty strings is 1 + self.assertEqual( + features.iloc[0, 1], 1.0 + ) # Jaro distance of empty strings is 1 self.assertEqual( features.iloc[0, 2], 1.0 ) # Jaro distance on the reverse empty strings is 1 diff --git a/tests/detector/test_domainator_detector.py b/tests/detector/test_domainator_detector.py index 0a11c5a1..106a232f 100644 --- a/tests/detector/test_domainator_detector.py +++ b/tests/detector/test_domainator_detector.py @@ -13,7 +13,7 @@ from src.base.data_classes.batch import Batch from src.detector.plugins.domainator_utils import ( DOMAINATOR_FEATURE_COLUMNS, - get_domainator_features + get_domainator_features, ) @@ -175,7 +175,9 @@ def test_get_features_empty_domains(self): self.assertEqual( features.iloc[0, 0], 1.0 ) # Levenshtein distance of empty strings is 1 - self.assertEqual(features.iloc[0, 1], 1.0) # Jaro distance of empty strings is 1 + self.assertEqual( + features.iloc[0, 1], 1.0 + ) # Jaro distance of empty strings is 1 self.assertEqual( features.iloc[0, 2], 1.0 ) # Jaro distance on the reverse empty strings is 1 diff --git a/tests/kafka/test_kafka_consume_handler.py b/tests/kafka/test_kafka_consume_handler.py index ac61cd23..039e4306 100644 --- a/tests/kafka/test_kafka_consume_handler.py +++ b/tests/kafka/test_kafka_consume_handler.py @@ -50,12 +50,10 @@ def test_missing_topic_is_created_with_target_partitions(self, mock_new_topic): _metadata({"test_topic": 4}), ] admin_client.create_topics.return_value = {"test_topic": MagicMock()} - mock_new_topic.side_effect = ( - lambda topic, partitions, replication_factor: ( - topic, - partitions, - replication_factor, - ) + mock_new_topic.side_effect = lambda topic, partitions, replication_factor: ( + topic, + partitions, + replication_factor, ) target_partitions_by_topic = ensure_topics( @@ -79,8 +77,9 @@ def test_existing_topic_with_too_few_partitions_is_expanded( _metadata({"test_topic": 2}), ] admin_client.create_partitions.return_value = {"test_topic": MagicMock()} - mock_new_partitions.side_effect = ( - lambda topic, total_count: (topic, total_count) + mock_new_partitions.side_effect = lambda topic, total_count: ( + topic, + total_count, ) ensure_topics( diff --git a/tests/logcollector/test_collector.py b/tests/logcollector/test_collector.py index 1b4a5cec..afba936e 100644 --- a/tests/logcollector/test_collector.py +++ b/tests/logcollector/test_collector.py @@ -86,9 +86,7 @@ async def test_start_successful_execution( mock_loop.run_in_executor.assert_awaited_once_with( mock_executor, self.sut.fetch ) - mock_executor.shutdown.assert_called_once_with( - wait=False, cancel_futures=True - ) + mock_executor.shutdown.assert_called_once_with(wait=False, cancel_futures=True) class _StopFetching(RuntimeError): diff --git a/tests/logserver/test_server.py b/tests/logserver/test_server.py index c59fc44f..05d8d4e3 100644 --- a/tests/logserver/test_server.py +++ b/tests/logserver/test_server.py @@ -72,9 +72,7 @@ async def test_start( mock_loop.run_in_executor.assert_awaited_once_with( mock_executor, self.sut.fetch_from_kafka ) - mock_executor.shutdown.assert_called_once_with( - wait=False, cancel_futures=True - ) + mock_executor.shutdown.assert_called_once_with(wait=False, cancel_futures=True) class TestSend(unittest.TestCase): @@ -212,7 +210,9 @@ def fetch_wrapper(*args, **kwargs): class TestMain(unittest.IsolatedAsyncioTestCase): @patch("src.logserver.server.logger") - @patch("src.logserver.server.start_pipeline_worker_replicas", new_callable=AsyncMock) + @patch( + "src.logserver.server.start_pipeline_worker_replicas", new_callable=AsyncMock + ) @patch("asyncio.create_task") @patch("asyncio.run") @patch("src.logserver.server.SENSOR_PROTOCOLS", ["dns"]) @@ -243,7 +243,9 @@ async def test_main( @patch("src.logserver.server.SENSOR_PROTOCOLS", ["dns", "http"]) @patch("src.logserver.server.logger") - @patch("src.logserver.server.start_pipeline_worker_replicas", new_callable=AsyncMock) + @patch( + "src.logserver.server.start_pipeline_worker_replicas", new_callable=AsyncMock + ) @patch("asyncio.create_task") @patch("asyncio.run") async def test_main_multiple_protocols( diff --git a/tests/miscellaneous/test_acceleration.py b/tests/miscellaneous/test_acceleration.py index 7fac4eed..527d725f 100644 --- a/tests/miscellaneous/test_acceleration.py +++ b/tests/miscellaneous/test_acceleration.py @@ -27,9 +27,7 @@ def test_auto_uses_cuda_when_available(self, mock_cuda_available): self.assertTrue(result.uses_gpu) @patch("src.base.acceleration.is_cuda_available", return_value=False) - def test_auto_falls_back_to_cpu_when_cuda_is_unavailable( - self, mock_cuda_available - ): + def test_auto_falls_back_to_cpu_when_cuda_is_unavailable(self, mock_cuda_available): config = { "acceleration": { "enabled": True, @@ -129,4 +127,3 @@ def set_params(self, **params): self.assertEqual({"device": "cuda:0"}, model.params) self.assertIs(model, result) - diff --git a/tests/miscellaneous/test_monitoring_agent.py b/tests/miscellaneous/test_monitoring_agent.py index 90cbfd47..c989b7c6 100644 --- a/tests/miscellaneous/test_monitoring_agent.py +++ b/tests/miscellaneous/test_monitoring_agent.py @@ -158,9 +158,7 @@ async def test_successful(self): mock_loop.run_in_executor.assert_any_await( mock_executor, self.sut.kafka_consumer.consume ) - mock_executor.shutdown.assert_called_once_with( - wait=False, cancel_futures=True - ) + mock_executor.shutdown.assert_called_once_with(wait=False, cancel_futures=True) self.sut.batch_sender.add.assert_called_once_with( "server_logs", { From 01e47da6e71027b36060fc9410e2f88c3b378bdb Mon Sep 17 00:00:00 2001 From: maldwg Date: Tue, 30 Jun 2026 16:11:37 +0200 Subject: [PATCH 3/3] Fix test issues --- src/base/retry.py | 6 +++++- src/inspector/inspector.py | 2 +- src/monitoring/monitoring_agent.py | 12 +++++------- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/src/base/retry.py b/src/base/retry.py index 2b4a2c11..c10affab 100644 --- a/src/base/retry.py +++ b/src/base/retry.py @@ -21,7 +21,11 @@ def resilience_config() -> dict[str, Any]: config = setup_config() + if not isinstance(config, dict): + return dict(_DEFAULT_CONFIG) retry_config = config.get("pipeline", {}).get("resilience", {}).get("retry", {}) + if not isinstance(retry_config, dict): + retry_config = {} merged = dict(_DEFAULT_CONFIG) merged.update(retry_config) return merged @@ -33,7 +37,7 @@ def retry_forever( retry_config: dict[str, Any] | None = None, retryable: tuple[type[BaseException], ...] = (Exception,), ) -> T: - config = retry_config or resilience_config() + config = retry_config if retry_config is not None else resilience_config() initial_delay = _float_setting(config, "initial_delay_seconds") max_delay = _float_setting(config, "max_delay_seconds") multiplier = max(1.0, _float_setting(config, "backoff_multiplier")) diff --git a/src/inspector/inspector.py b/src/inspector/inspector.py index b5366123..a3e5d8a6 100644 --- a/src/inspector/inspector.py +++ b/src/inspector/inspector.py @@ -99,7 +99,7 @@ def __init__(self, consume_topic, produce_topics, config) -> None: self.time_range = config["time_range"] self.name = config["name"] self.acceleration = resolve_acceleration_config( - config["pipeline"], + globals()["config"].get("pipeline", {}), config, component_name=f"{module_name}.{self.name}", logger=logger, diff --git a/src/monitoring/monitoring_agent.py b/src/monitoring/monitoring_agent.py index 5000d351..f69ec9cc 100644 --- a/src/monitoring/monitoring_agent.py +++ b/src/monitoring/monitoring_agent.py @@ -22,6 +22,7 @@ CLICKHOUSE_HOSTNAME = CONFIG["environment"]["monitoring"]["clickhouse_server"][ "hostname" ] +RETRY_CONFIG = CONFIG.get("pipeline", {}).get("resilience", {}).get("retry", {}) def prepare_all_tables(): @@ -49,17 +50,14 @@ def _iter_statements(sql_content: str): file_path = os.path.join(CREATE_TABLES_DIRECTORY, filename) sql_content = _load_contents(file_path) - client = retry_forever( + with retry_forever( lambda: clickhouse_connect.get_client(host=CLICKHOUSE_HOSTNAME), "ClickHouse table preparation connection", - ) - with client: + retry_config=RETRY_CONFIG, + ) as client: for statement in _iter_statements(sql_content): try: - retry_forever( - lambda statement=statement: client.command(statement), - "ClickHouse CREATE TABLE statement", - ) + client.command(statement) except Exception as e: logger.critical("Error in CREATE TABLE statement") raise e