| """ |
| DGA-CNN: Character-level CNN for DGA detection. |
| Architecture from Patton et al. (adapted), trained on 54 DGA families. |
| """ |
| import string |
| import torch |
| import torch.nn as nn |
|
|
| CHARS = string.ascii_lowercase + string.digits + "-._" |
| CHAR2IDX = {c: i + 1 for i, c in enumerate(CHARS)} |
| VOCAB_SIZE = len(CHARS) + 1 |
| MAXLEN = 75 |
|
|
|
|
| def encode_domain(domain: str) -> list: |
| domain = str(domain).lower().strip() |
| encoded = [CHAR2IDX.get(c, 0) for c in domain[:MAXLEN]] |
| return encoded + [0] * (MAXLEN - len(encoded)) |
|
|
|
|
| class DGACNN(nn.Module): |
| def __init__(self, vocab_size=VOCAB_SIZE, embedding_dim=32, num_classes=2): |
| super().__init__() |
| self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0) |
| self.conv1 = nn.Conv1d(embedding_dim, 64, kernel_size=3, padding=1) |
| self.relu = nn.ReLU() |
| self.pool = nn.MaxPool1d(2) |
| self.dropout = nn.Dropout(0.3) |
| self.fc = nn.Linear(64 * (MAXLEN // 2), num_classes) |
|
|
| def forward(self, x): |
| x = self.embedding(x).transpose(1, 2) |
| x = self.pool(self.relu(self.conv1(x))) |
| x = x.view(x.size(0), -1) |
| x = self.dropout(x) |
| return self.fc(x) |
|
|
|
|
| def load_model(weights_path: str, device: str = None): |
| """Load trained model from a local weights path.""" |
| if device is None: |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| model = DGACNN() |
| model.load_state_dict(torch.load(weights_path, map_location=device)) |
| model.to(device) |
| model.eval() |
| return model |
|
|
|
|
| def predict(model, domains, device: str = None, batch_size: int = 256): |
| """ |
| Predict DGA vs legit for a list of domain strings. |
| Returns list of dicts: [{"domain": ..., "label": "dga"/"legit", "score": float}] |
| """ |
| if device is None: |
| device = next(model.parameters()).device |
| if isinstance(domains, str): |
| domains = [domains] |
|
|
| results = [] |
| for i in range(0, len(domains), batch_size): |
| batch = domains[i : i + batch_size] |
| encoded = [encode_domain(d) for d in batch] |
| x = torch.tensor(encoded, dtype=torch.long).to(device) |
| with torch.no_grad(): |
| logits = model(x) |
| probs = torch.softmax(logits, dim=1) |
| preds = logits.argmax(dim=1).cpu().tolist() |
| scores = probs[:, 1].cpu().tolist() |
| for domain, pred, score in zip(batch, preds, scores): |
| results.append({ |
| "domain": domain, |
| "label": "dga" if pred == 1 else "legit", |
| "score": round(score, 4), |
| }) |
| return results |
|
|
