src/cti_attack/baselines.py

"""Non-neural baselines.

These exist to make the fine-tuned encoder falsifiable. A large share of
published cyber-ML results do not clear a TF-IDF + logistic regression floor,
and a model that fails to beat these has no reason to be downloaded.

  * ``frequency``  — always predict the k most common training techniques.
    Establishes what a model that has learned nothing about the input scores.
  * ``keyword``    — substring match on ATT&CK technique names. Zero training.
  * ``tfidf_lr``   — TF-IDF word+char n-grams, one-vs-rest logistic regression.
    The real floor. Usually much stronger than people expect.
"""

from __future__ import annotations

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import FeatureUnion

from . import attack_meta


def frequency_scores(train_Y: np.ndarray, n_eval: int) -> np.ndarray:
    """Score every technique by its training prevalence, identically for all rows."""
    prior = train_Y.mean(axis=0)
    return np.tile(prior, (n_eval, 1))


def keyword_scores(texts: list[str], labels: list[str]) -> np.ndarray:
    """1.0 where the technique's name appears verbatim in the sentence, else 0."""
    kw = attack_meta.technique_keywords(labels)
    scores = np.zeros((len(texts), len(labels)), dtype=np.float32)
    lowered = [t.lower() for t in texts]
    for j, tid in enumerate(labels):
        forms = kw.get(tid, [])
        if not forms:
            continue
        for i, text in enumerate(lowered):
            if any(f in text for f in forms):
                scores[i, j] = 1.0
    return scores


def _vectorizer() -> FeatureUnion:
    return FeatureUnion([
        ("word", TfidfVectorizer(
            ngram_range=(1, 2), min_df=2, sublinear_tf=True,
            strip_accents="unicode", lowercase=True)),
        ("char", TfidfVectorizer(
            analyzer="char_wb", ngram_range=(3, 5), min_df=3, sublinear_tf=True,
            lowercase=True)),
    ])


def tfidf_lr_scores(
    train_texts: list[str], train_Y: np.ndarray, eval_sets: dict[str, list[str]]
) -> dict[str, np.ndarray]:
    """One-vs-rest logistic regression over TF-IDF word + character n-grams.

    Classes with a single training polarity (all-negative after splitting) are
    handled explicitly rather than allowed to raise, and score a constant 0.
    """
    vec = _vectorizer()
    Xtr = vec.fit_transform(train_texts)
    Xev = {name: vec.transform(texts) for name, texts in eval_sets.items()}

    out = {name: np.zeros((X.shape[0], train_Y.shape[1]), dtype=np.float32)
           for name, X in Xev.items()}

    for j in range(train_Y.shape[1]):
        y = train_Y[:, j]
        if y.sum() == 0 or y.sum() == len(y):
            continue  # degenerate class: leave scores at 0
        clf = LogisticRegression(
            C=4.0, max_iter=2000, class_weight="balanced", solver="liblinear")
        clf.fit(Xtr, y)
        for name, X in Xev.items():
            out[name][:, j] = clf.predict_proba(X)[:, 1]
    return out