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.

TABLES (`python docs/redos-sweep.py [table ...]`, default: all). The lexicon is
one of six regex surfaces; 0.3.2 swept the other five with HAND-WRITTEN rows,
which is the class of sweep this script exists because it misses arms. The
collector is therefore MECHANICAL on both axes -- it walks each module's
namespace for compiled patterns (so a pattern added later is swept without
anyone remembering to list it) and derives each one's CALL MODE by grepping the
module source for `NAME.<method>(`. Mode matters: `.match()`/`.fullmatch()` are
anchored at position 0 and cannot pay the per-start rescan cost that makes a run
quadratic, so timing them with `search` would manufacture unreachable flags,
while `.sub()`/`.finditer()` scan every position and must be timed as such.
"""
from __future__ import annotations

import importlib
import json
import re
import sys
import time
from dataclasses import dataclass
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, mode: str = "search") -> float:
    """Time one scan of ``text`` in the mode the production code actually uses."""
    mode = mode.rstrip("*")
    start = time.monotonic()
    if mode == "finditer":
        for _ in rx.finditer(text):
            pass
    elif mode == "sub":
        rx.sub("", text)
    elif mode == "match":
        rx.match(text)
    elif mode == "fullmatch":
        rx.fullmatch(text)
    else:
        rx.search(text)
    return time.monotonic() - start


# --- targets ----------------------------------------------------------------

@dataclass(frozen=True)
class Target:
    table: str
    id: str
    regex: re.Pattern[str]
    src: str
    mode: str


_MODE_ORDER = ("sub", "finditer", "search", "match", "fullmatch")


_CALL_RE = r"\b%s\b(?:\[[^\]]*\]|\.\w+)*\s*\.\s*(search|match|fullmatch|finditer|subn?)\("


def _mode_for(names: tuple[str, ...], module_src: str) -> str:
    """Derive the call mode from the module source: worst scanning mode wins.

    ``names`` are the identifiers a pattern is reachable through (its own name
    plus, for a pattern living in a table, the table's name). A pattern in a
    table is usually never named directly -- it is reached through the loop
    variable of ``for p in TABLE:`` -- so those bindings are resolved too, else
    the whole egress table reads as ``search`` when it is really ``finditer``.
    A pattern with NO direct reference at all is reached indirectly some other
    way -- ``active_content`` passes each construct regex into a local ``_scan``
    helper that calls ``pattern.sub`` on the parameter -- so it falls back to the
    worst scanning mode rather than the mildest. ``sub``/``finditer`` visit every
    start position where ``search`` stops at the first match, so the fallback can
    only over-measure, never miss. Fallback rows print with a ``*``.
    """
    idents = set(names)
    for name in names:
        for m in re.finditer(r"for\s+([\w,\s]+?)\s+in\s+%s\b" % re.escape(name),
                             module_src):
            idents.update(v.strip() for v in m.group(1).split(",") if v.strip())
    found = set()
    for ident in idents:
        for m in re.finditer(_CALL_RE % re.escape(ident), module_src):
            found.add("sub" if m.group(1).startswith("sub") else m.group(1))
    for mode in _MODE_ORDER:
        if mode in found:
            return mode
    return "sub*"


def _walk(value, path: str, depth: int = 0):
    """Yield (path, compiled_pattern) for patterns nested in tables/dataclasses."""
    if isinstance(value, re.Pattern):
        yield path, value
    elif depth < 3 and isinstance(value, (list, tuple, frozenset, set)):
        for i, item in enumerate(value):
            yield from _walk(item, f"{path}[{i}]", depth + 1)
    elif (depth < 3 and not isinstance(value, type)
            and hasattr(value, "__dataclass_fields__")):
        ident = getattr(value, "id", None)
        for f in value.__dataclass_fields__:
            sub = getattr(value, f)
            if isinstance(sub, re.Pattern):
                yield (f"{ident}" if ident else f"{path}.{f}"), sub


def collect_module(mod_name: str, table: str, only=None, skip=()) -> list[Target]:
    """Every compiled pattern reachable from a module's namespace."""
    mod = importlib.import_module(f"llm_ingestion_guard.{mod_name}")
    module_src = (SRC / f"{mod_name}.py").read_text()
    out: list[Target] = []
    for name, value in vars(mod).items():
        if name.startswith("__") or name in skip:
            continue
        if only is not None and name not in only:
            continue
        for path, rx in _walk(value, name):
            out.append(Target(table, path, rx, rx.pattern,
                              _mode_for((name, path), module_src)))
    return out


def collect_lexicon() -> list[Target]:
    raw = json.loads((SRC / "injection_lexicon.json").read_text())["patterns"]
    src_by_id = {e["id"]: e["regex"] for e in raw}
    # scan_lexicon drives every entry through `.search()` (lexicon.py:289,352).
    return [Target("lexicon", p.id, p.regex, src_by_id[p.id], "search")
            for p in load_lexicon()]


def collect_egress() -> list[Target]:
    return collect_module("output", "egress", only={"_SECRET_PATTERNS"})


def collect_output() -> list[Target]:
    return collect_module("output", "output", skip={"_SECRET_PATTERNS"})


TABLES = {
    "lexicon": collect_lexicon,
    # The normalizers run `.sub()` over EVERY input before any pattern matches;
    # 0.3.3 swept the 83 patterns and not these. `_LEXICON_CACHE` is skipped: it
    # is empty until `load_lexicon()` runs, so counting it would make this
    # table's size depend on whether `lexicon` was swept first.
    "normalize": lambda: collect_module("lexicon", "normalize",
                                        skip={"_LEXICON_CACHE"}),
    "active_content": lambda: collect_module("active_content", "active_content"),
    "entropy": lambda: collect_module("entropy", "entropy") + [
        # Inline literals in is_base64_like / is_hex_blob (entropy.py:111,118),
        # invisible to a namespace walk.
        Target("entropy", "is_base64_like", re.compile(r"[A-Za-z0-9+/]{20,}={0,3}"),
               r"[A-Za-z0-9+/]{20,}={0,3}", "fullmatch"),
        Target("entropy", "is_hex_blob", re.compile(r"(?:0x)?[0-9a-fA-F]{32,}"),
               r"(?:0x)?[0-9a-fA-F]{32,}", "fullmatch"),
    ],
    "egress": collect_egress,
    "output": collect_output,
    # Not "detector tables", but the same regex surface on the same paths.
    # Leaving them out would reproduce the 0.3.2 mistake at module granularity.
    "sanitize": lambda: collect_module("sanitize", "sanitize"),
    "neutralize": lambda: collect_module("neutralize", "neutralize"),
    "okf": lambda: collect_module("okf", "okf"),
    "contract": lambda: collect_module("contract", "contract"),
    "fence": lambda: collect_module("fence", "fence"),
}

EXTRA_UNITS = ["[", "<a:", "<a ", "![", "a://:"]


def sweep(targets: list[Target]) -> list[tuple]:
    flagged = []
    for tgt in targets:
        units = sample(tgt.src) + EXTRA_UNITS
        hits = []
        for u in units:
            t1 = t(tgt.regex, build(u, N1), tgt.mode)
            if t1 > HARD_CAP:
                hits.append((u, t1, float("inf"), float("inf")))
                continue
            t2 = t(tgt.regex, build(u, N2), tgt.mode)
            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((tgt, u, t1, t2, r))
            print(f"   {tgt.table}/{tgt.id:34} [{tgt.mode:9}] "
                  f"unit={u!r:14} {t1:.4f}->{t2:.4f}s x{r:.1f}")
    return flagged


def main() -> None:
    argv = sys.argv[1:]
    listing = "--list" in argv
    wanted = [a for a in argv if a != "--list"] or list(TABLES)
    unknown = [w for w in wanted if w not in TABLES]
    if unknown:
        sys.exit(f"unknown table(s): {unknown}; known: {list(TABLES)}")
    print(f"== arm-by-arm sweep, N={N1}->{N2}, flag ratio>={RATIO_FLAG} ==")
    total, flagged = 0, []
    for name in wanted:
        targets = TABLES[name]()
        total += len(targets)
        print(f"-- {name}: {len(targets)} patterns")
        if listing:
            for tgt in targets:
                print(f"   {tgt.id:34} [{tgt.mode:9}] {tgt.src[:70]}")
            continue
        flagged += sweep(targets)
    if listing:
        print(f"\n{total} patterns across {len(wanted)} table(s)")
        return
    ids = {f[0].id for f in flagged}
    print(f"\ncandidates: {len(ids)} patterns / {total}, {len(flagged)} arms")


if __name__ == "__main__":
    main()