diff --git a/config.yaml b/config.yaml index 930b63e..d9a8b07 100644 --- a/config.yaml +++ b/config.yaml @@ -196,6 +196,8 @@ pipeline: detector_class_name: "DomainatorDetector" model: domainator use_scaler: false + domain_expiration_hours: 1 + expiry_check_interval_hours: 1 checksum: a4aac4c585f1e614c3cf0d737e80b960c5de6e87b253f7cdd07125d9ce486476 base_url: https://heibox.uni-heidelberg.de/d/0d5cbcbe16cd46a58021/ threshold: 0.5 diff --git a/docs/configuration.rst b/docs/configuration.rst index 468a14b..0718d70 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -540,6 +540,12 @@ To entirely skip the anomaly detection phase, you can set ``inspector_module_nam * - threshold - ``0.5`` - Threshold for the detector's classification. + * - domain_expiration_hours + - ``1`` + - Domainator only: discard source/domain windows that have been inactive for more than this many hours, measured using observed message timestamps rather than wall-clock time. + * - expiry_check_interval_hours + - ``1`` + - Domainator only: interval between stale-window checks, measured using observed message timestamps. * - produce_topics - ``(empty)`` - (Optional) Comma-separated list of alerter topic suffixes to produce alerts to. If left empty, defaults to the ``generic`` alerter topic. Use ``send_to_alerter: false`` or ``produce_topics: []`` for intermediary detectors that should not produce to an alerter. diff --git a/src/detector/plugins/domainator_detector.py b/src/detector/plugins/domainator_detector.py index 5d9fcd3..0bd6a77 100644 --- a/src/detector/plugins/domainator_detector.py +++ b/src/detector/plugins/domainator_detector.py @@ -1,6 +1,7 @@ from src.detector.detector import DetectorBase import numpy as np from collections import defaultdict +from datetime import datetime, timedelta from src.base.log_config import get_logger from src.detector.plugins.domainator_utils import ( @@ -45,11 +46,64 @@ def __init__( consume_topic (str): Kafka topic from which the detector will consume messages. """ self.model_base_url = detector_config["base_url"] - self.message_queues = defaultdict(list) + # How long a domain has to have been stale for before it gets removed from the dictionary + self.domain_expiration_hours = timedelta( + hours=detector_config.get("domain_expiration_hours", 1) + ) + # How often the detector checks whether entries have expired + self.expiry_check_interval = timedelta( + hours=detector_config.get("expiry_check_interval_hours", 1) + ) + + self.message_queues = defaultdict( + lambda: defaultdict(lambda: {"timestamp": None, "messages": []}) + ) + self.latest_message_timestamp = None + self.latest_expiry_timestamp = None + + self.requests_detect_threshold = 10 + self.max_window_requests = 10 + self.max_window_cut = 5 + super().__init__( detector_config, consume_topic, produce_topics, downstream_detector_topics ) + def expire_stale_domains(self): + """ + Drop windows for src_ip/domain pairs that have not been seen for a certain period of time. + + The method deliberately avoids using the current time and instead tracks the + timestamp of the observed data. This keeps the tracking correct for both live + traffic and historical data. + + Returns: + int: Number of ``(src_ip, domain)`` entries removed. + """ + if self.latest_message_timestamp is None: + return 0 + + removed_domains = 0 + + stale_entries = [] + for src, domains in self.message_queues.items(): + for domain, entry in domains.items(): + entry_timestamp = datetime.fromisoformat(entry["timestamp"]) + if ( + self.latest_message_timestamp - entry_timestamp + > self.domain_expiration_hours + ): + stale_entries.append((src, domain)) + + for src, domain in stale_entries: + del self.message_queues[src][domain] + removed_domains += 1 + + if not self.message_queues[src]: + del self.message_queues[src] + + return removed_domains + def predict(self, messages): """ Process a window of messages and predict if the domain is likely to be used @@ -63,19 +117,31 @@ def predict(self, messages): a "domain_name" key with the domain to analyze. Returns: - np.ndarray: Prediction probabilities for each class. Typically a 2D array - where the shape is (1, 2) for binary classification (benign/malicious). + dict: Prediction probabilities keyed by the model's class labels. """ queries = [message["domain_name"] for message in messages] y_pred = self.model.predict_proba(get_domainator_features(queries)) + y_pred = dict(zip(self.model.classes_, y_pred[0])) + return y_pred def detect(self): logger.info("Start detecting malicious requests.") for message in self.messages: + message_src = message["src_ip"] message_domain = strip_domain(message["domain_name"]) - self.message_queues[message_domain].append(message) + message_queue = self.message_queues[message_src][message_domain] + + message_queue["messages"].append(message) + message_queue["timestamp"] = message["ts"] + + datetime_timestamp = datetime.fromisoformat(message_queue["timestamp"]) + if ( + self.latest_message_timestamp is None + or datetime_timestamp > self.latest_message_timestamp + ): + self.latest_message_timestamp = datetime_timestamp # currently receives 1 message batch at a time iterates over it and adds to the queue # right now let's say we fetch 100 messages in a batch, we iterate over the messages to add to the dict, worst case: 97 predictions necessary (len of 3, sliding window over all 100 messsages) @@ -84,19 +150,35 @@ def detect(self): # improvement suggestion: try to predict only at the end for each domain seen in the batch IF len > 3. If prediction happened maybe clear the queue ? --> suboptimal for no sliding window approach # further improvement: kafka partitoin key shold not be src_ip but rather domain name --> make adjustable for other detectors as well !! - if len(self.message_queues[message_domain]) >= 3: - y_pred = self.predict(self.message_queues[message_domain]) + if len(message_queue["messages"]) >= self.requests_detect_threshold: + y_pred = self.predict(message_queue["messages"]) logger.info(f"Prediction: {y_pred}") - if np.argmax(y_pred, axis=1) == 1 and y_pred[0][1] > self.threshold: + if ( + y_pred["malicious"] == max(y_pred.values()) + and y_pred["malicious"] > self.threshold + ): logger.info("Append malicious request domain to warning.") warning = { - "request": self.message_queues[message_domain].copy(), - "probability": float(y_pred[0][1]), + "request": message_queue["messages"].copy(), + "probability": float(y_pred["malicious"]), "name": self.name, "sha256": self.checksum, } self.warnings.append(warning) - if len(self.message_queues[message_domain]) >= 5: - del self.message_queues[message_domain][0] + # The statement here can be removed if the minimum and maximum window sizes are kept the same + if len(message_queue["messages"]) >= self.max_window_requests: + del message_queue["messages"][0 : self.max_window_cut] + + # The check is done based on the message's timestamps to support historical data + if self.latest_message_timestamp is not None and ( + self.latest_expiry_timestamp is None + or self.latest_message_timestamp - self.latest_expiry_timestamp + > self.expiry_check_interval + ): + removed_domains = self.expire_stale_domains() + self.latest_expiry_timestamp = self.latest_message_timestamp + if removed_domains: + logger.info(f"Expired {removed_domains} stale window(s).") + logger.debug(f"Domainator Message queue length: {len(self.message_queues)}") diff --git a/src/detector/plugins/domainator_utils.py b/src/detector/plugins/domainator_utils.py index 3fe4db0..7c085a6 100644 --- a/src/detector/plugins/domainator_utils.py +++ b/src/detector/plugins/domainator_utils.py @@ -36,7 +36,7 @@ def strip_domain(query: str): domain = "" if len(query) >= 2: - domain = query[-2] + domain = ".".join(query[-2:]) return domain @@ -56,7 +56,7 @@ def get_domainator_features(queries: list) -> pd.DataFrame: """ queries = [query.strip(".") for query in queries] - subdomains = [".".join(domain.split(".")[:-2]) for domain in queries] + subdomains = ["".join(domain.split(".")[:-2]) for domain in queries] metrics = {column: [] for column in DOMAINATOR_FEATURE_COLUMNS} @@ -64,7 +64,15 @@ def get_domainator_features(queries: list) -> pd.DataFrame: cartesian = list(itertools.combinations(subdomains, 2)) metrics["levenshtein"] = np.mean( - [Levenshtein.ratio(product[0], product[1]) for product in cartesian] + [ + ( + Levenshtein.distance(product[0], product[1]) + / max(len(product[0]), len(product[1])) + if len(product[0]) and len(product[1]) + else 0.0 + ) + for product in cartesian + ] ) metrics["jaro"] = np.mean( [Levenshtein.jaro(product[0], product[1]) for product in cartesian] @@ -91,7 +99,7 @@ def get_domainator_features(queries: list) -> pd.DataFrame: [ ( pylcs.lcs_sequence_length(product[0], product[1]) - / ((len(product[0]) + len(product[1])) / 2) + / (max(len(product[0]), len(product[1]))) if len(product[0]) and len(product[1]) else 0.0 ) @@ -102,7 +110,7 @@ def get_domainator_features(queries: list) -> pd.DataFrame: [ ( pylcs.lcs_string_length(product[0], product[1]) - / ((len(product[0]) + len(product[1])) / 2) + / (max(len(product[0]), len(product[1]))) if len(product[0]) and len(product[1]) else 0.0 ) diff --git a/tests/detector/test_domainator_attributor.py b/tests/detector/test_domainator_attributor.py index ddf5b92..dd21e28 100644 --- a/tests/detector/test_domainator_attributor.py +++ b/tests/detector/test_domainator_attributor.py @@ -125,7 +125,7 @@ def test_detect_snapshots_each_warning_request(self): ) self.assertEqual( ["1", "2", "3", "4"], - [message["logline_id"] for message in sut.message_queues["example"]], + [message["logline_id"] for message in sut.message_queues["example.org"]], ) def test_detect_emits_for_attribution_class_other_than_index_one(self): @@ -248,7 +248,7 @@ def test_get_features_basic_attributes(self): jaro_dist = features.iloc[0, 1] # Jaro distance lcs = features.iloc[0, 6] # Longest common string - self.assertEqual(leven_dist, 0.75) + self.assertEqual(leven_dist, 0.25) self.assertAlmostEqual(jaro_dist, 0.833, 3) # Rounded to 3 decimal places self.assertEqual(lcs, 0.75) @@ -262,8 +262,8 @@ def test_get_features_empty_domains(self): # Basic features self.assertEqual( - features.iloc[0, 0], 1.0 - ) # Levenshtein distance of empty strings is 1 + features.iloc[0, 0], 0.0 + ) # Normalized Levenshtein distance is 0 for empty subdomains self.assertEqual( features.iloc[0, 1], 1.0 ) # Jaro distance of empty strings is 1 @@ -293,8 +293,8 @@ def test_get_features_single_same_character(self): # Basic features self.assertEqual( - features.iloc[0, 0], 1.0 - ) # Levenshtein distance of same strings is 1 + features.iloc[0, 0], 0.0 + ) # Identical subdomains have zero Levenshtein distance self.assertEqual(features.iloc[0, 1], 1.0) # Jaro distance of same strings is 1 self.assertEqual( features.iloc[0, 2], 1.0 diff --git a/tests/detector/test_domainator_detector.py b/tests/detector/test_domainator_detector.py index 5681c83..c5fbf33 100644 --- a/tests/detector/test_domainator_detector.py +++ b/tests/detector/test_domainator_detector.py @@ -3,6 +3,7 @@ import unittest from unittest.mock import MagicMock, patch, call from pandas.testing import assert_frame_equal +from datetime import datetime, timedelta import os import sys @@ -21,7 +22,7 @@ "src_ip": "192.168.0.167", "dns_ip": "10.10.0.10", "response_ip": "252.79.173.222", - "ts": "", + "ts": "2011-08-17T14:32:02", "status": "NXDOMAIN", "domain_name": "IF356gEnJHPdRxnkDId4RDUSgtqxx9I+pZ5n1V53MdghOGQncZWAQgAPRx3kswi.750jnH6iSqmiAAeyDUMX0W6SHGpVsVsKSX8ZkKYDs0GFh/9qU5N9cwl00XSD8ID.NNhBdHZIb7nc0hDQXFPlABDLbRwkJS38LZ8RMX4yUmR2Mb6YqTTJBn+nUcB9P+v.jBQdwdS53XV9W2p1BHjh.16.f.1.6037.tunnel.example.org", "record_type": "A", @@ -35,7 +36,12 @@ def setUp(self): self.mock_logger = patcher.start() self.addCleanup(patcher.stop) - def _create_detector(self, mock_kafka_handler=None, mock_clickhouse=None): + def _create_detector( + self, + mock_kafka_handler=None, + mock_clickhouse=None, + detector_config_overrides=None, + ): """Helper method to create a DomainatorDetector instance with proper mocks.""" if mock_kafka_handler is None: mock_kafka_handler = MagicMock() @@ -51,6 +57,8 @@ def _create_detector(self, mock_kafka_handler=None, mock_clickhouse=None): "base_url": "https://heibox.uni-heidelberg.de/d/0d5cbcbe16cd46a58021/", "threshold": 0.005, } + if detector_config_overrides: + detector_config.update(detector_config_overrides) with patch( "src.detector.detector.ExactlyOnceKafkaConsumeHandler", @@ -68,6 +76,35 @@ def _create_detector(self, mock_kafka_handler=None, mock_clickhouse=None): detector.scaler = MagicMock() return detector + def test_expiry_defaults_support_existing_config(self): + detector = self._create_detector() + + self.assertEqual(detector.domain_expiration_hours, timedelta(hours=1)) + self.assertEqual(detector.expiry_check_interval, timedelta(hours=1)) + + def test_expiry_settings_are_configurable(self): + detector = self._create_detector( + detector_config_overrides={ + "domain_expiration_hours": 2, + "expiry_check_interval_hours": 0.25, + } + ) + + self.assertEqual(detector.domain_expiration_hours, timedelta(hours=2)) + self.assertEqual(detector.expiry_check_interval, timedelta(minutes=15)) + + def test_detect_empty_batch_does_not_start_expiry_clock(self): + detector = self._create_detector() + + with patch.object(detector, "expire_stale_domains") as mock_expire: + detector.detect() + mock_expire.assert_not_called() + + self.assertIsNone(detector.latest_message_timestamp) + self.assertIsNone(detector.latest_expiry_timestamp) + self.assertEqual(detector.expire_stale_domains(), 0) + self.assertEqual(detector.message_queues, {}) + def test_get_model_download_url(self): """Test that the model download URL is correctly formatted.""" mock_kafka = MagicMock() @@ -83,30 +120,38 @@ def test_detect(self): mock_kafka = MagicMock() mock_ch = MagicMock() sut = self._create_detector(mock_kafka, mock_ch) - for _ in range(0, 4, 1): + window_size = 4 + sut.requests_detect_threshold = window_size + + for _ in range(0, window_size, 1): sut.messages.append((DEFAULT_DATA)) with patch( "src.detector.plugins.domainator_detector.DomainatorDetector.predict", - return_value=[[0.01, 0.99]], + return_value={"legitimate": 0.01, "malicious": 0.99}, ): sut.detect() self.assertNotEqual([], sut.warnings) def test_detect_snapshots_each_warning_request(self): sut = self._create_detector() + + domain_name = "example.org" sut.messages = [ { **DEFAULT_DATA, - "domain_name": f"subdomain-{index}.example.org", + "domain_name": f"subdomain-{index}.{domain_name}", "logline_id": str(index), } for index in range(5) ] + sut.requests_detect_threshold = 3 + sut.max_window_requests = 5 + sut.max_window_cut = 1 # [0:1] is the same as [0] with patch.object( sut, "predict", - return_value=np.array([[0.01, 0.99]]), + return_value={"legitimate": 0.01, "malicious": 0.99}, ): sut.detect() @@ -121,9 +166,12 @@ def test_detect_snapshots_each_warning_request(self): for warning in sut.warnings ], ) + + message_queues = sut.message_queues[DEFAULT_DATA["src_ip"]][domain_name] + self.assertEqual( ["1", "2", "3", "4"], - [message["logline_id"] for message in sut.message_queues["example"]], + [message["logline_id"] for message in message_queues["messages"]], ) def test_predict_calls_model(self): @@ -131,6 +179,7 @@ def test_predict_calls_model(self): mock_kafka = MagicMock() mock_ch = MagicMock() detector = self._create_detector(mock_kafka, mock_ch) + detector.model.classes_ = ["legitimate", "malicious"] # Mock model prediction mock_prediction = np.array([[0.2, 0.8]]) @@ -149,13 +198,15 @@ def test_predict_calls_model(self): assert_frame_equal(called_features, expected_features) # Verify prediction result - np.testing.assert_array_equal(result, mock_prediction) + np.testing.assert_array_equal(result["legitimate"], mock_prediction[0][0]) + np.testing.assert_array_equal(result["malicious"], mock_prediction[0][1]) def test_predict_uses_domainator_feature_order(self): """Test that predict passes features in Domainator's generated order.""" mock_kafka = MagicMock() mock_ch = MagicMock() detector = self._create_detector(mock_kafka, mock_ch) + detector.model.classes_ = ["legitimate", "malicious"] mock_prediction = np.array([[0.2, 0.8]]) detector.model.predict_proba.return_value = mock_prediction @@ -175,7 +226,8 @@ def test_predict_uses_domainator_feature_order(self): called_features.columns.tolist(), DOMAINATOR_FEATURE_COLUMNS, ) - np.testing.assert_array_equal(result, mock_prediction) + np.testing.assert_array_equal(result["legitimate"], mock_prediction[0][0]) + np.testing.assert_array_equal(result["malicious"], mock_prediction[0][1]) def test_get_features_basic_attributes(self): """Test basic label features calculation.""" @@ -193,7 +245,7 @@ def test_get_features_basic_attributes(self): jaro_dist = features.iloc[0, 1] # Jaro distance lcs = features.iloc[0, 6] # Longest common string - self.assertEqual(leven_dist, 0.75) + self.assertEqual(leven_dist, 0.25) self.assertAlmostEqual(jaro_dist, 0.833, 3) # Rounded to 3 decimal places self.assertEqual(lcs, 0.75) @@ -207,7 +259,7 @@ def test_get_features_empty_domains(self): # Basic features self.assertEqual( - features.iloc[0, 0], 1.0 + features.iloc[0, 0], 0.0 ) # Levenshtein distance of empty strings is 1 self.assertEqual( features.iloc[0, 1], 1.0 @@ -238,8 +290,8 @@ def test_get_features_single_same_character(self): # Basic features self.assertEqual( - features.iloc[0, 0], 1.0 - ) # Levenshtein distance of same strings is 1 + features.iloc[0, 0], 0.0 + ) # Levenshtein distance of same strings is 0 self.assertEqual(features.iloc[0, 1], 1.0) # Jaro distance of same strings is 1 self.assertEqual( features.iloc[0, 2], 1.0 @@ -291,3 +343,53 @@ def test_get_features_case_insensitivity(self): features_lower.to_numpy()[0], decimal=5, ) + + def test_expire_stale_domains_removes_old_entries(self): + """expire_stale_domains drops entries older than the threshold (strict >) and + removes any source bucket that becomes empty as a result.""" + sut = self._create_detector() + sut.latest_message_timestamp = datetime(2011, 8, 17, 16, 0, 0) + sut.message_queues["1.1.1.1"]["fresh.com"] = { + "timestamp": "2011-08-17T15:00:00", + "message_domain": ["a.fresh.com"], + } + sut.message_queues["2.2.2.2"]["stale.com"] = { + "timestamp": "2011-08-17T10:00:00", # 6h old -> stale + "message_domain": ["b.stale.com"], + } + + removed = sut.expire_stale_domains() + + self.assertEqual(removed, 1) + self.assertIn("1.1.1.1", sut.message_queues) + self.assertIn("fresh.com", sut.message_queues["1.1.1.1"]) + self.assertNotIn("2.2.2.2", sut.message_queues) + + def test_expire_on_interval(self): + """The check runs at most once every expiry_check_interval of data time.""" + sut = self._create_detector() + sut.message_queues["1.1.1.1"]["a.com"] = { + "timestamp": "2011-08-17T10:00:00", + "message_domain": ["x.a.com"], + } + sut.latest_message_timestamp = datetime(2011, 8, 17, 10, 10, 0) + + with patch.object(sut, "expire_stale_domains", return_value=0) as mock_expire: + + sut.detect() + # First run, latest_expiry_timestamp is still None + self.assertEqual(mock_expire.call_count, 1) + + sut.latest_message_timestamp = datetime(2011, 8, 17, 10, 15, 0) # +5m + sut.detect() # 5m < 15m -> skipped + self.assertEqual(mock_expire.call_count, 1) + self.assertEqual( + sut.latest_expiry_timestamp, datetime(2011, 8, 17, 10, 10, 0) + ) + + sut.latest_message_timestamp = datetime(2011, 8, 17, 11, 15, 0) # +1h15m + sut.detect() + self.assertEqual(mock_expire.call_count, 2) + self.assertEqual( + sut.latest_expiry_timestamp, datetime(2011, 8, 17, 11, 15, 0) + )