llm-ingestion-pipeline-security/src/llm_ingestion_guard/sanitize.py

"""sanitize — carrier stripping for untrusted content.

Removes the invisible/steganographic carrier classes that smuggle instructions
past a human reader but reach the model: zero-width characters, BIDI overrides,
Unicode-tag steganography, HTML comments, and ``data:`` URIs.

Contract (BRIEF §5, §9): this function only ever *removes* — never rewrites.
Clean input returns byte-identical with an empty report, and the output is
always a subsequence of the input. Per-class counts are reported so the caller
can gate (WARN / block) on them. Ported from the ``llm-security`` unicode-scanner
and string-utils primitives.
"""
from __future__ import annotations

import re
from dataclasses import dataclass

from .calibration import MAX_INPUT_CHARS
from .contract import assert_within_input_cap
from .report import Finding, Report, Severity, Source

# Invisible / steganographic character classes (codepoints).
_ZERO_WIDTH = frozenset({0x200B, 0x200C, 0x200D, 0xFEFF, 0x00AD})
_BIDI = frozenset({0x202A, 0x202B, 0x202C, 0x202D, 0x202E, 0x2066, 0x2067, 0x2068, 0x2069})
_TAG_LO, _TAG_HI = 0xE0000, 0xE007F  # Unicode Tags block (U+E0000–U+E007F)

# ZWJ (U+200D) is the one zero-width codepoint with a legitimate, extremely
# common use: it composes emoji. 👩‍💻 is WOMAN + ZWJ + PERSONAL COMPUTER, and
# families, skin-tone professions and flag variants are all built the same way.
# Testing membership alone therefore flagged — and *stripped* — the joiner in any
# first-party document containing such an emoji, which `disposition` grades as an
# any-tier carrier (FAIL_SECURE, no appeal). The strip was the worse half: it
# silently decomposed 👩‍💻 into two unrelated emoji, so `sanitize` corrupted
# content it was only ever supposed to remove carriers from.
#
# The fix is the one our own lexicon row `unicode:zero-width-in-word`
# (`\w[ZW]\w`) already used: judge the joiner by CONTEXT, not identity. A ZWJ is
# exempt only when BOTH neighbours are emoji-context codepoints — half-context is
# not context, so `a<ZWJ>👩` stays a carrier and an attacker cannot buy exemption
# with a single emoji.
#
# The ranges are blocks, not an emoji table. Measured against Unicode 17.0's
# `emoji-zwj-sequences.txt`: 1614 RGI sequences use 122 distinct codepoints
# adjacent to a ZWJ, and these five ranges cover 122/122. Blocks are the point —
# shipping the RGI sequence list itself would be precise the day it landed and
# stale at the next Unicode release, reintroducing this exact false positive for
# every new emoji until someone bumped the file. Whole blocks include the
# unassigned headroom (458 Cn codepoints here) that future emoji are allocated
# into, so the table does not age.
#
# Residual, documented in LIMITATIONS: a ZWJ hidden *between two emoji* is
# exempt and could carry a covert channel. It costs an emoji per bit and cannot
# split a word, which is the shape the word-splitting attack actually needs.
_EMOJI_CTX_RANGES = (
    (0x2190, 0x21FF),    # Arrows — ↔ ↕ (U+2194/2195)
    (0x2600, 0x27BF),    # Misc Symbols + Dingbats — ❤ ☠ ⚕ ⚧ ✈ ❄ ♀ ♂ ➡
    (0x2B00, 0x2BFF),    # Misc Symbols & Arrows — ⬛ (U+2B1B)
    (0xFE0F, 0xFE0F),    # VARIATION SELECTOR-16, the emoji presentation selector
    (0x1F000, 0x1FAFF),  # Emoji planes, incl. skin-tone modifiers U+1F3FB–FF
)


def _is_emoji_context(cp: int) -> bool:
    """True when ``cp`` may legitimately sit adjacent to an emoji-composing ZWJ."""
    return any(lo <= cp <= hi for lo, hi in _EMOJI_CTX_RANGES)


def _is_joiner_in_emoji_sequence(text: str, i: int) -> bool:
    """True when ``text[i]`` is a ZWJ composing an emoji rather than a carrier.

    Requires an emoji-context codepoint on BOTH sides. A joiner at either edge of
    the document has no neighbour and so is never exempt.
    """
    if i == 0 or i + 1 >= len(text):
        return False
    return _is_emoji_context(ord(text[i - 1])) and _is_emoji_context(ord(text[i + 1]))

# Span carriers.
#
# ReDoS note (OWASP LLM10). The comment stripper used to be `<!--.*?-->` with a
# comment that absence of nesting meant no catastrophic backtracking. That claim
# was wrong in the same way `output`'s was: a lazy run in front of a REQUIRED
# literal costs a full tail rescan at *every* start position when the literal
# never arrives, so `<!--` repeated to 100_000 chars measured 20.1s (exponent
# 1.96–2.14 over four doublings) — and at the time this module, unlike
# `scan_lexicon` / `scan_output`, applied no input cap, so nothing bounded that
# above. MAX_INPUT_CHARS now does, but as a second line only: the cap bounds a
# *future* quadratic pattern's damage, it does not make a quadratic one safe.
#
# Neither of the two fixes used elsewhere fits here. Excluding the opener (`<`)
# from the run would drop every comment containing markup — `<!-- <b>x</b> -->`
# is the ordinary case, not an edge one. Bounding the run would be a one-line
# carrier bypass: a comment padded past the bound is exactly what this stripper
# exists to remove. So the scan is done with `str.find`, which is linear and
# semantically identical to the lazy regex — leftmost opener, nearest following
# terminator, unterminated trailer left in place.
_COMMENT_OPEN, _COMMENT_CLOSE = "<!--", "-->"
# `data:` not preceded by a letter (so "metadata:" / "userdata:" do not match),
# consuming up to the next whitespace / quote / angle bracket / closing paren.
_DATA_URI_RE = re.compile(r"(?<![A-Za-z])data:[^\s'\"<>)]+", re.IGNORECASE)


@dataclass(frozen=True)
class SanitizeResult:
    """The cleaned text plus a report of what was stripped."""

    text: str
    report: Report


def _redact(s: str, show_start: int = 12, show_end: int = 4) -> str:
    if len(s) <= show_start + show_end + 3:
        return s
    return f"{s[:show_start]}...{s[-show_end:]}"


def _strip_html_comments(text: str) -> tuple[str, int]:
    """Remove ``<!-- ... -->`` spans; return the cleaned text and the count.

    Linear replacement for the quantifier form (see the ReDoS note above). An
    unterminated ``<!--`` is left verbatim, matching the regex it replaces.
    """
    if _COMMENT_OPEN not in text:
        return text, 0
    out: list[str] = []
    pos = count = 0
    while True:
        start = text.find(_COMMENT_OPEN, pos)
        if start == -1:
            break
        end = text.find(_COMMENT_CLOSE, start + len(_COMMENT_OPEN))
        if end == -1:  # unterminated — not a comment, keep the rest verbatim
            break
        out.append(text[pos:start])
        pos = end + len(_COMMENT_CLOSE)
        count += 1
    if not count:
        return text, 0
    out.append(text[pos:])
    return "".join(out), count


def _decode_tags(codepoints: list[int]) -> str:
    """Decode Unicode-tag codepoints to their hidden ASCII (cp - 0xE0000)."""
    out = []
    for cp in codepoints:
        ascii_cp = cp - 0xE0000
        out.append(chr(ascii_cp) if 0x20 <= ascii_cp <= 0x7E else "?")
    return "".join(out)


def sanitize(
    text: str,
    source: Source = Source.INPUT,
    max_input_chars: int = MAX_INPUT_CHARS,
) -> SanitizeResult:
    """Strip carrier classes from ``text`` and report per-class counts.

    Raises :class:`~llm_ingestion_guard.contract.OversizeInputError` above
    ``max_input_chars``: this is step 1 of the input path, so the refusal bounds
    the whole path, and a *partially* sanitized document is worse than none —
    the unstripped tail is where a carrier would be placed.
    """
    assert_within_input_cap(text, surface="sanitize", max_input_chars=max_input_chars)
    report = Report()

    # Character-class carriers: single pass, keep everything else verbatim.
    zero_width = 0
    bidi = 0
    tag_cps: list[int] = []
    kept: list[str] = []
    for i, ch in enumerate(text):
        cp = ord(ch)
        if cp == 0x200D and _is_joiner_in_emoji_sequence(text, i):
            kept.append(ch)  # composing an emoji, not carrying a payload
        elif cp in _ZERO_WIDTH:
            zero_width += 1
        elif cp in _BIDI:
            bidi += 1
        elif _TAG_LO <= cp <= _TAG_HI:
            tag_cps.append(cp)
        else:
            kept.append(ch)
    # Preserve object identity (and byte-identity) when nothing was stripped.
    cleaned = "".join(kept) if (zero_width or bidi or tag_cps) else text

    # Span carriers.
    cleaned, n_comments = _strip_html_comments(cleaned)
    cleaned, n_data = _DATA_URI_RE.subn("", cleaned)

    if zero_width:
        report.add(Finding(label="sanitize:zero-width", severity=Severity.HIGH,
                           source=source, detector="sanitize", count=zero_width, owasp="LLM01"))
    if bidi:
        report.add(Finding(label="sanitize:bidi-override", severity=Severity.HIGH,
                           source=source, detector="sanitize", count=bidi, owasp="LLM01"))
    if tag_cps:
        report.add(Finding(label="sanitize:unicode-tag", severity=Severity.CRITICAL,
                           source=source, detector="sanitize", count=len(tag_cps),
                           evidence=_redact(_decode_tags(tag_cps)), owasp="LLM01"))
    if n_comments:
        report.add(Finding(label="sanitize:html-comment", severity=Severity.MEDIUM,
                           source=source, detector="sanitize", count=n_comments, owasp="LLM01"))
    if n_data:
        report.add(Finding(label="sanitize:data-uri", severity=Severity.MEDIUM,
                           source=source, detector="sanitize", count=n_data, owasp="LLM01"))

    return SanitizeResult(text=cleaned, report=report)