Fizcko/argus_sentinel

HTTP Request Text
     |
     v
[BPE Tokenizer (vocab=8192, max_len=128)]
     |
     +---> [Embedding (128-dim)]
     |           |
     |     [Conv1D (128 ch, k=3) + BatchNorm + ReLU] x2
     |           |
     |     [AdaptiveMaxPool1d → 128-dim]
     |
     +---> [6 Numeric Features]
               |
         [Linear 6→32 + ReLU]
               |
         [Concatenate (128 + 32 = 160)]
               |
         [Linear 160→128→64 + ReLU + Dropout(0.1)]
               |
     +---------+---------+
     |                   |
[Label Head → 7]   [Risk Head → 1]
     |                   |
[Sigmoid]            [Sigmoid]
     |                   |
label_probs [7]    risk_score [1]

from tokenizers import Tokenizer
tok = Tokenizer.from_file("tokenizer.json")
result = tok.encode("GET /search?q=test HTTP/1.1")
# result.ids → [2, 546, 287, ...]  (starts with [CLS]=2)
# result.attention_mask → [1, 1, 1, ...]

def extract_numeric_features(request: dict) -> list[float]:
    body = request.get("body") or ""
    headers = request.get("headers") or {}
    return [
        float(len(body)),                                    # content_length
        float(len(headers)),                                 # num_headers
        1.0 if body else 0.0,                                # has_body
        float(request.get("session_request_count") or 0),    # session_request_count
        float(request.get("session_duration") or 0),         # session_duration
        float(request.get("session_pattern_score") or 0),    # session_pattern_score
    ]

fn extract_numeric_features(request: &HttpRequest) -> [f32; 6] {
    let body_len = request.body.as_ref().map_or(0, |b| b.len());
    [
        body_len as f32,
        request.headers.len() as f32,
        if body_len > 0 { 1.0 } else { 0.0 },
        request.session_request_count.unwrap_or(0) as f32,
        request.session_duration.unwrap_or(0.0),
        request.session_pattern_score.unwrap_or(0.0),
    ]
}

import onnxruntime as ort
import numpy as np
import json
from tokenizers import Tokenizer

# Load model and tokenizer
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
tokenizer = Tokenizer.from_file("tokenizer.json")
thresholds = json.load(open("thresholds.json"))["thresholds"]
label_names = ["clean", "xss", "sqli", "path_traversal",
               "command_injection", "scanner", "spam_bot"]

def classify_request(method, path, query, headers, body):
    # 1. Build text
    text = f"{method} {path}"
    if query: text += f"?{query}"
    if body: text += f" {body[:200]}"
    text = text[:500]

    # 2. Tokenize
    enc = tokenizer.encode(text)
    input_ids = np.array([enc.ids], dtype=np.int32)
    attention_mask = np.array([enc.attention_mask], dtype=np.int32)

    # 3. Numeric features (RAW values)
    numeric = np.array([[
        float(len(body or "")),
        float(len(headers)),
        1.0 if body else 0.0,
        0.0, 0.0, 0.0,  # session features (0 if unavailable)
    ]], dtype=np.float32)

    # 4. Inference
    probs, risk = session.run(None, {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "numeric_features": numeric,
    })

    # 5. Apply per-label thresholds
    detections = {
        name: float(probs[0][i])
        for i, name in enumerate(label_names)
        if name != "clean" and probs[0][i] >= thresholds[name]
    }

    return {
        "risk_score": float(risk[0][0]),
        "detections": detections,
        "is_clean": len(detections) == 0,
    }

# Example
result = classify_request("GET", "/search", "q=' OR 1=1--", {"Host": "example.com"}, None)
print(result)
# {'risk_score': 0.87, 'detections': {'sqli': 0.94}, 'is_clean': False}

use ort::{Session, Value};
use ndarray::Array2;

fn main() -> anyhow::Result<()> {
    let session = Session::builder()?
        .with_model_from_file("model.onnx")?;

    let input_ids = Array2::<i32>::zeros((1, 128));       // from tokenizer
    let attention_mask = Array2::<i32>::zeros((1, 128));   // from tokenizer
    let numeric_features = Array2::<f32>::zeros((1, 6));   // extract_numeric_features()

    let outputs = session.run(ort::inputs![
        "input_ids" => &input_ids,
        "attention_mask" => &attention_mask,
        "numeric_features" => &numeric_features,
    ]?)?;

    let label_probs: Vec<f32> = outputs[0].extract_tensor::<f32>()?.view().iter().copied().collect();
    let risk_score: f32 = *outputs[1].extract_tensor::<f32>()?.view().first().unwrap();

    // Apply thresholds from thresholds.json
    let thresholds = [0.20, 0.50, 0.74, 0.68, 0.66, 0.70, 0.72];
    let labels = ["clean", "xss", "sqli", "path_traversal",
                  "command_injection", "scanner", "spam_bot"];

    for (i, (prob, thr)) in label_probs.iter().zip(thresholds.iter()).enumerate() {
        if i > 0 && prob >= thr {
            println!("DETECTED: {} ({:.3})", labels[i], prob);
        }
    }
    println!("Risk score: {:.4}", risk_score);

    Ok(())
}

thresholds = json.load(open("thresholds.json"))["thresholds"]

# Per-label detection
triggered = [name for i, name in enumerate(label_names)
             if name != "clean" and probs[0][i] >= thresholds[name]]

# Risk-score action
score = float(risk[0][0])
if score >= 0.8:   action = "BLOCK"
elif score >= 0.5:  action = "CHALLENGE"
elif score >= 0.2:  action = "LOG"
else:               action = "ALLOW"

@misc{argus_sentinel_2026,
  title        = {Argus Sentinel: A Low-Latency CNN-Based WAF Classifier},
  author       = {Fizcko},
  year         = {2026},
  howpublished = {Hugging Face Model Hub},
  note         = {V3, 1.17M params, 0.24ms latency, Macro F1 0.866, FPR 0.83\%}
}