app.py

import spaces  # MUST come before torch / any CUDA-touching import
import os
import torch
import torch.nn as nn
import numpy as np
import gradio as gr
from transformers import PreTrainedTokenizerFast

# ---------------------------------------------------------------------------
# Model definition (ported from the official DRIFT repository)
# https://github.com/snsec-net/2026-DSN-DRIFT  ·  model.py
# ---------------------------------------------------------------------------
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, padding_idx):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_idx)
        torch.nn.init.xavier_normal_(self.embedding.weight)

    def forward(self, input):
        return self.embedding(input)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len):
        super().__init__()
        self.pos_embed = nn.Embedding(max_len, d_model)
        torch.nn.init.xavier_normal_(self.pos_embed.weight)

    def forward(self, x):
        B, L, _ = x.size()
        device = x.device
        pos_ids = torch.arange(L, device=device).unsqueeze(0).expand(B, L)
        return x + self.pos_embed(pos_ids)


class Transformer(nn.Module):
    def __init__(self, d_model, n_heads, dim_feedforward, num_layers, dropout=0.1):
        super().__init__()
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        if mask is not None and mask.any():
            out = self.encoder(x, src_key_padding_mask=mask)
        else:
            out = self.encoder(x)
        return self.dropout(out)


class PretrainedModel(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, dim_feedforward,
                 num_layers, max_len, dropout=0.1, padding_idx=0, tov_norm='pool'):
        super().__init__()
        self.d_model = d_model
        self.padding_idx = padding_idx
        self.max_len = max_len
        self.tov_norm = tov_norm
        self.embedding = TokenEmbedding(vocab_size, d_model, padding_idx)
        self.positional_encoding = PositionalEncoding(d_model, max_len)
        self.transformer = Transformer(d_model, n_heads, dim_feedforward, num_layers, dropout)

    def create_padding_mask(self, input_ids):
        return (input_ids == self.padding_idx)


class FinetuningHead(nn.Module):
    def __init__(self, input_dim, d_model, dropout):
        super().__init__()
        self.input_dim = input_dim
        self.d_model = d_model
        self.dense1 = nn.Linear(input_dim, d_model * 2)
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(d_model * 2, 2)

    def forward(self, encoder_output):
        x = self.dropout(encoder_output)
        x = self.dense1(x)
        x = torch.relu(x)
        x = self.dropout(x)
        logits = self.classifier(x)
        return logits


class FineTuningModel(nn.Module):
    def __init__(self, pretrain_model_t=None, pretrain_model_c=None,
                 dropout=0.1, padding_idx=0, clf_norm='pool', freeze_backbone=False):
        super().__init__()
        self.padding_idx = padding_idx
        self.clf_norm = clf_norm
        self.use_token = pretrain_model_t is not None
        self.use_char = pretrain_model_c is not None
        sample_model = pretrain_model_t if self.use_token else pretrain_model_c
        d_model = sample_model.d_model
        num_active_paths = sum([self.use_token, self.use_char])
        dim_per_path = d_model * 2 if clf_norm == 'pool' else d_model
        total_input_dim = dim_per_path * num_active_paths

        if self.use_token:
            self.transformer_encoder_t = pretrain_model_t.transformer
            self.embedding_t = pretrain_model_t.embedding
            self.positional_encoding_t = pretrain_model_t.positional_encoding
        if self.use_char:
            self.transformer_encoder_c = pretrain_model_c.transformer
            self.embedding_c = pretrain_model_c.embedding
            self.positional_encoding_c = pretrain_model_c.positional_encoding

        self.classifier_head = FinetuningHead(
            input_dim=total_input_dim, d_model=d_model, dropout=dropout
        )

    def create_padding_mask(self, input_ids):
        return (input_ids == self.padding_idx).to(input_ids.device)

    def forward(self, input_ids_t=None, input_ids_c=None):
        features = []

        if self.use_token and input_ids_t is not None:
            t_embed = self.embedding_t(input_ids_t)
            t_x = self.positional_encoding_t(t_embed)
            t_mask = self.create_padding_mask(input_ids_t)
            t_out = self.transformer_encoder_t(t_x, mask=t_mask)
            if self.clf_norm == 'pool':
                valid_mask = (~t_mask).float().unsqueeze(-1)
                t_sum = (t_out * valid_mask).sum(dim=1)
                t_len = valid_mask.sum(dim=1).clamp(min=1)
                t_mean = t_sum / t_len
                t_max = (t_out.masked_fill(valid_mask == 0, -1e9)).max(dim=1).values
                t_feat = torch.cat([t_max, t_mean], dim=1)
            else:
                t_feat = t_out[:, 0, :]
            features.append(t_feat)

        if self.use_char and input_ids_c is not None:
            c_embed = self.embedding_c(input_ids_c)
            c_x = self.positional_encoding_c(c_embed)
            c_mask = self.create_padding_mask(input_ids_c)
            c_out = self.transformer_encoder_c(c_x, mask=c_mask)
            if self.clf_norm == 'pool':
                valid_mask = (~c_mask).float().unsqueeze(-1)
                c_sum = (c_out * valid_mask).sum(dim=1)
                c_len = valid_mask.sum(dim=1).clamp(min=1)
                c_mean = c_sum / c_len
                c_max = (c_out.masked_fill(valid_mask == 0, -1e9)).max(dim=1).values
                c_feat = torch.cat([c_max, c_mean], dim=1)
            else:
                c_feat = c_out[:, 0, :]
            features.append(c_feat)

        combined_output = torch.cat(features, dim=1) if len(features) > 1 else features[0]
        return self.classifier_head(combined_output)


# ---------------------------------------------------------------------------
# Configuration constants (from utility/config.py)
# ---------------------------------------------------------------------------
D_MODEL = 256
N_HEADS = 8
NUM_LAYERS = 12
DIM_FEEDFORWARD = 768
MAX_LEN_SUBWORD = 30
MAX_LEN_CHAR = 77
VOCAB_SIZE_SUBWORD = 30522
VOCAB_SIZE_CHAR = 43
PADDING_IDX = 0
CLF_NORM = 'pool'

# Special token IDs (from preprocessing.py SpecialIDs)
PAD_ID = 0
UNK_ID = 1
CLS_ID = 2
SEP_ID = 3
MASK_ID = 4

CHAR_LIST = list("abcdefghijklmnopqrstuvwxyz0123456789-.")
SPECIAL_TOKENS = ['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]']
ALL_TOKENS = SPECIAL_TOKENS + CHAR_LIST
CHAR2ID = {char: idx for idx, char in enumerate(ALL_TOKENS)}

MODEL_ID = "snsec-net/dga-detector-drift26dsn"

# ---------------------------------------------------------------------------
# Load tokenizer and model at module scope
# ---------------------------------------------------------------------------
TOKENIZER_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                              "tokenizer-0-30522-both.json")
tokenizer = PreTrainedTokenizerFast(tokenizer_file=TOKENIZER_PATH)

pt_model_c = PretrainedModel(
    vocab_size=VOCAB_SIZE_CHAR,
    d_model=D_MODEL,
    n_heads=N_HEADS,
    dim_feedforward=DIM_FEEDFORWARD,
    num_layers=NUM_LAYERS,
    max_len=MAX_LEN_CHAR,
)
pt_model_t = PretrainedModel(
    vocab_size=VOCAB_SIZE_SUBWORD,
    d_model=D_MODEL,
    n_heads=N_HEADS,
    dim_feedforward=DIM_FEEDFORWARD,
    num_layers=NUM_LAYERS,
    max_len=MAX_LEN_SUBWORD,
)
model = FineTuningModel(pt_model_t, pt_model_c, clf_norm=CLF_NORM)

CKPT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "finetuning.pt")
state = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)
model.load_state_dict(state, strict=False)
model = model.to("cuda").eval()


# ---------------------------------------------------------------------------
# Preprocessing (ported from preprocessing.py FineTuningDataset)
# ---------------------------------------------------------------------------
def domain_to_char_ids(domain: str) -> np.ndarray:
    """Character-level encoding: [CLS] + chars + [SEP], padded to MAX_LEN_CHAR."""
    domain = domain.lower()
    token_indices = [CHAR2ID.get(c, UNK_ID) for c in domain]
    if len(token_indices) > MAX_LEN_CHAR - 2:
        token_indices = token_indices[:MAX_LEN_CHAR - 2]
    ids = [CLS_ID] + token_indices + [SEP_ID]
    if len(ids) < MAX_LEN_CHAR:
        ids += [PAD_ID] * (MAX_LEN_CHAR - len(ids))
    return np.array(ids, dtype=np.int64)


def domain_to_subword_ids(domain: str) -> np.ndarray:
    """Subword-level encoding: [CLS] + subwords + [SEP], padded to MAX_LEN_SUBWORD."""
    domain = domain.lower()
    encoded = tokenizer(domain, add_special_tokens=False)
    token_indices = encoded["input_ids"]
    if len(token_indices) > MAX_LEN_SUBWORD - 2:
        token_indices = token_indices[:MAX_LEN_SUBWORD - 2]
    ids = [CLS_ID] + token_indices + [SEP_ID]
    if len(ids) < MAX_LEN_SUBWORD:
        ids += [PAD_ID] * (MAX_LEN_SUBWORD - len(ids))
    return np.array(ids, dtype=np.int64)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
@spaces.GPU(duration=30)
def detect_dga(domain: str):
    """Classify a domain name as Benign or DGA-generated.

    Args:
        domain: A domain name (e.g. "google.com" or "xkqjhfwiwbfw.info").

    Returns:
        A tuple of (label_dict, confidence, benign_prob, dga_prob).
    """
    domain = domain.strip()
    if not domain:
        return {"—": 1.0}, 0.0, 0.0, 0.0

    char_ids = domain_to_char_ids(domain)
    subword_ids = domain_to_subword_ids(domain)

    x_t = torch.tensor(np.array([subword_ids]), dtype=torch.long, device="cuda")
    x_c = torch.tensor(np.array([char_ids]), dtype=torch.long, device="cuda")

    with torch.no_grad():
        logits = model(x_t, x_c)
        probs = torch.softmax(logits, dim=1)
        pred = torch.argmax(logits, dim=1).item()

    benign_prob = probs[0, 0].item()
    dga_prob = probs[0, 1].item()
    label = "DGA" if pred == 1 else "Benign"
    confidence = max(benign_prob, dga_prob)

    return {label: confidence}, confidence, benign_prob, dga_prob


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# DRIFT: Drift-Resilient Invariant-Feature Transformer for DGA Detection\n"
        "Enter a domain name (effective second-level domain) to classify it as "
        "**Benign** or **DGA-generated**. DRIFT uses a dual-branch Transformer "
        "(character + subword) to learn invariant structural features that remain "
        "robust against concept drift.\n\n"
        "[Paper](https://huggingface.co/papers/2605.10436) · "
        "[GitHub](https://github.com/snsec-net/2026-DSN-DRIFT) · "
        "[Model](https://huggingface.co/snsec-net/dga-detector-drift26dsn)"
    )

    with gr.Column(elem_id="col-container"):
        domain_input = gr.Textbox(
            label="Domain name (effective second-level domain)",
            placeholder="e.g. google or xkqjhfwiwbfw",
            info="Enter the effective second-level domain (eSLD) — TLD stripped, "
                 "lowercased, characters a–z, 0–9, '-', '.' only.",
        )
        run_btn = gr.Button("Classify", variant="primary")

        with gr.Row():
            label_out = gr.Label(label="Classification")
            confidence_out = gr.Number(label="Confidence", precision=4)

        with gr.Accordion("Class probabilities", open=False):
            benign_bar = gr.Number(label="P(Benign)", precision=4)
            dga_bar = gr.Number(label="P(DGA)", precision=4)

    gr.Examples(
        examples=[
            ["google"],
            ["github"],
            ["xkqjhfwiwbfw"],
            ["my-secure-login-portal"],
            ["qpalzm-xnkvjf"],
        ],
        inputs=[domain_input],
        fn=detect_dga,
        outputs=[label_out, confidence_out, benign_bar, dga_bar],
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=detect_dga,
        inputs=[domain_input],
        outputs=[label_out, confidence_out, benign_bar, dga_bar],
        api_name="detect_dga",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)