from src.detector.detector import DetectorBase
import math
import numpy as np
from src.base.log_config import get_logger
module_name = "data_analysis.detector"
logger = get_logger(module_name)
[docs]
class DGADetector(DetectorBase):
"""
Detector implementation for identifying Domain Generation Algorithm (DGA) domains.
This class extends the DetectorBase to provide specific functionality for detecting
malicious domains generated by domain generation algorithms. It uses a machine learning
model to analyze domain name characteristics and identify potential DGA activity.
The detector extracts various statistical and structural features from domain names
to make predictions about whether a domain is likely generated by a DGA.
"""
def __init__(
self,
detector_config,
consume_topic,
produce_topics=None,
downstream_detector_topics=None,
):
"""
Initialize the DGA detector with configuration parameters.
Sets up the detector with the model base URL and passes configuration to the
base class for standard detector initialization.
Args:
detector_config (dict): Configuration dictionary containing detector-specific
parameters including base_url, model, checksum, and threshold.
consume_topic (str): Kafka topic from which the detector will consume messages.
"""
self.model_base_url = detector_config["base_url"]
super().__init__(
detector_config, consume_topic, produce_topics, downstream_detector_topics
)
def predict(self, message):
"""
Process a message and predict if the domain is likely generated by a DGA.
Extracts features from the domain name in the message and uses the loaded
machine learning model to generate prediction probabilities.
Args:
message (dict): A dictionary containing message data, expected to have
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).
"""
y_pred = self.model.predict_proba(self._get_features(message["domain_name"]))
return y_pred
def _get_features(self, query: str) -> np.ndarray:
"""Extracts feature vector from domain name for ML model inference.
Computes various statistical and linguistic features from the domain name
including label lengths, character frequencies, entropy measures, and
counts of different character types across domain name levels.
Args:
query (str): Domain name string to extract features from.
Returns:
numpy.ndarray: Feature vector ready for ML model prediction.
"""
# Splitting by dots to calculate label length and max length
query = query.strip(".")
label_parts = query.split(".")
levels = {
"fqdn": query,
"secondleveldomain": label_parts[-2] if len(label_parts) >= 2 else "",
"thirdleveldomain": (
".".join(label_parts[:-2]) if len(label_parts) > 2 else ""
),
}
label_length = len(label_parts)
parts = query.split(".")
label_max = len(max(parts, key=str)) if parts else 0
label_average = len(query)
basic_features = np.array(
[label_length, label_max, label_average], dtype=np.float64
)
alc = "abcdefghijklmnopqrstuvwxyz"
query_len = len(query)
freq = np.array(
[query.lower().count(c) / query_len if query_len > 0 else 0.0 for c in alc],
dtype=np.float64,
)
logger.debug("Get full, alpha, special, and numeric count.")
def calculate_counts(level: str) -> np.ndarray:
if not level:
return np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float64)
full_count = len(level) / len(level)
alpha_ratio = sum(c.isalpha() for c in level) / len(level)
numeric_ratio = sum(c.isdigit() for c in level) / len(level)
special_ratio = sum(
not c.isalnum() and not c.isspace() for c in level
) / len(level)
return np.array(
[full_count, alpha_ratio, numeric_ratio, special_ratio],
dtype=np.float64,
)
fqdn_counts = calculate_counts(levels["fqdn"])
third_counts = calculate_counts(levels["thirdleveldomain"])
second_counts = calculate_counts(levels["secondleveldomain"])
level_features = np.hstack([third_counts, second_counts, fqdn_counts])
def calculate_entropy(s: str) -> float:
if len(s) == 0:
return 0.0
probs = [s.count(c) / len(s) for c in dict.fromkeys(s)]
return -sum(p * math.log(p, 2) for p in probs)
logger.debug("Start entropy calculation")
entropy_features = np.array(
[
calculate_entropy(levels["fqdn"]),
calculate_entropy(levels["thirdleveldomain"]),
calculate_entropy(levels["secondleveldomain"]),
],
dtype=np.float64,
)
logger.debug("Entropy features calculated")
all_features = np.concatenate(
[basic_features, freq, level_features, entropy_features]
)
logger.debug("Finished data transformation")
return all_features.reshape(1, -1)
def detect(self) -> None:
"""
Process messages to detect malicious requests.
This method applies the detection model to each message in the current batch,
identifies potential threats based on the model's predictions, and collects
warnings for further processing.
The detection uses a threshold to determine if a prediction indicates
malicious activity, and only warnings exceeding this threshold are retained.
Note:
This method relies on the implementation of ``predict``of the rspective subclass
"""
for message in self.messages:
y_pred = self.predict(message)
logger.info(f"Prediction: {y_pred}")
if np.argmax(y_pred, axis=1) == 1 and y_pred[0][1] > self.threshold:
logger.debug("Append malicious request to warning.")
warning = {
"request": message,
"probability": float(y_pred[0][1]),
"name": self.name,
"sha256": self.checksum,
}
self.warnings.append(warning)
