llm-ingestion-pipeline-security/docs/redos-sweep.py

"""Sweep v3 — arm-by-arm. v2 found only the FIRST quadratic run in a pattern
because its units were generic. The `[system](` arm of markdown:link-anchor-
injection is separately quadratic and v2 missed it, so v2 is not evidence for
the other 82 patterns either.

Fix: synthesise an almost-match sample string from each regex's own skeleton
(classes -> a member, alternations -> each branch, quantifiers -> one copy),
then use every token-boundary PREFIX of that sample as a repeating unit. That
generates `[`, `[system]`, `[system](` ... automatically, one per run.

Found with this: markdown:link-anchor-injection (both arms) and
markdown:link-ref-comment, fixed in 0.3.3. A generic-payload pass found only the
first of the two.

READ BEFORE TRUSTING A FLAG. The ratio is computed from TWO points, so a hit
near the noise floor is a coin flip, not a finding. Post-fix this script flags
`sub-agent:delegate-bypass` at x2.8 -- measured properly over four doublings its
exponent is 0.96-1.05, i.e. LINEAR (0.2s at the 1M cap). Machine load inflates
the small measurements. Always re-measure a flagged arm across 4+ doublings and
read the exponent before concluding anything; x4 on doubling is quadratic, x2 is
linear. The two real findings above sat at 297s and 55s, not near the floor.
"""
from __future__ import annotations

import json
import re
import sys
import time
from pathlib import Path

SRC = Path(__file__).resolve().parent.parent / "src" / "llm_ingestion_guard"
sys.path.insert(0, str(SRC.parent))

from llm_ingestion_guard.lexicon import load_lexicon  # noqa: E402

N1, N2 = 4_000, 8_000
RATIO_FLAG = 2.6
NOISE_FLOOR = 0.0015
HARD_CAP = 20.0

_CLASS_SAMPLE = {
    "\\s": " ", "\\S": "a", "\\w": "a", "\\W": "-",
    "\\d": "1", "\\D": "a", "\\b": "", "\\B": "",
}


def sample_class(body: str) -> str:
    """Pick one character a class accepts (negated -> something not listed)."""
    negated = body.startswith("^")
    inner = body[1:] if negated else body
    if negated:
        for cand in "a1 <>[](){}:/@.-x":
            if cand not in inner:
                return cand
        return "\x01"
    m = re.search(r"([A-Za-z0-9])-([A-Za-z0-9])", inner)
    if m:
        return m.group(1)
    for ch in inner:
        if ch not in "\\^-":
            return ch
    return "a"


def sample(src: str) -> list[str]:
    """Emit prefix samples at token boundaries. Returns cumulative prefixes."""
    prefixes: list[str] = []
    buf = ""
    i, n = 0, len(src)
    while i < n:
        ch = src[i]
        if ch == "\\" and i + 1 < n:
            tok = src[i:i + 2]
            buf += _CLASS_SAMPLE.get(tok, tok[1])
            i += 2
        elif ch == "[":
            j = i + 1
            if j < n and src[j] == "^":
                j += 1
            if j < n and src[j] == "]":
                j += 1
            while j < n and src[j] != "]":
                j += 2 if src[j] == "\\" else 1
            buf += sample_class(src[i + 1:j])
            i = j + 1
        elif ch == "(":
            depth, j = 1, i + 1
            while j < n and depth:
                if src[j] == "\\":
                    j += 2
                    continue
                depth += (src[j] == "(") - (src[j] == ")")
                j += 1
            body = src[i + 1:j - 1]
            body = re.sub(r"^\?(:|P<[^>]*>|=|!|<=|<!)", "", body)
            branch = body.split("|")[0]
            inner = sample(branch)
            buf += inner[-1] if inner else ""
            i = j
        elif ch in "?*+":
            i += 1
        elif ch == "{":
            j = src.find("}", i)
            i = (j + 1) if j != -1 else i + 1
        elif ch in "^$|":
            i += 1
        else:
            buf += ch
            i += 1
        prefixes.append(buf)
    seen, out = set(), []
    for p in prefixes:
        if p and p not in seen and len(p) <= 32:
            seen.add(p)
            out.append(p)
    return out


def build(unit: str, n: int) -> str:
    return (unit * (n // len(unit) + 1))[:n]


def t(rx: re.Pattern[str], text: str) -> float:
    start = time.monotonic()
    rx.search(text)
    return time.monotonic() - start


def main() -> None:
    raw = json.loads((SRC / "injection_lexicon.json").read_text())["patterns"]
    src_by_id = {e["id"]: e["regex"] for e in raw}
    flagged = []
    print(f"== arm-by-arm sweep, N={N1}->{N2}, flag ratio>={RATIO_FLAG} ==")
    for p in load_lexicon():
        src = src_by_id[p.id]
        units = sample(src) + ["[", "<a:", "<a ", "![", "a://:"]
        hits = []
        for u in units:
            t1 = t(p.regex, build(u, N1))
            if t1 > HARD_CAP:
                hits.append((u, t1, float("inf"), float("inf")))
                continue
            t2 = t(p.regex, build(u, N2))
            ratio = t2 / t1 if t1 > 0 else 0.0
            if t2 >= NOISE_FLOOR and ratio >= RATIO_FLAG:
                hits.append((u, t1, t2, ratio))
        for u, t1, t2, r in sorted(hits, key=lambda h: -h[2])[:3]:
            flagged.append((p.id, u, t1, t2, r))
            print(f"   {p.id:40} unit={u!r:14} {t1:.4f}->{t2:.4f}s x{r:.1f}")
    ids = {f[0] for f in flagged}
    print(f"\ncandidates: {len(ids)} patterns / 83, {len(flagged)} arms")


if __name__ == "__main__":
    main()