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

"""entropy — high-entropy / encoded-blob detection with decode-and-rescan.

A ``text -> findings`` detector (design principle 3): pure, no I/O, no mutation.
It flags runs of base64/hex-alphabet characters that look like encoded or
encrypted payloads — the carriers that smuggle instructions past a reader as an
opaque blob. Two independent mechanisms, ported from the ``llm-security``
entropy-scanner:

  1. **Length-calibrated Shannon entropy.** Random-looking text has high
     entropy, but the achievable maximum is length-dependent (a short base64
     string cannot reach the entropy of a long one), so thresholds pair an
     entropy floor with a minimum length: CRITICAL 5.4/128, HIGH 5.1/64,
     MEDIUM 4.7/40 (bits-per-char / chars).
  2. **Shape floor.** A base64-like blob (len > 100) or a hex blob (len > 64)
     is at least MEDIUM even when entropy alone would not trigger. This is the
     *only* path that catches hex: a 16-symbol alphabet caps entropy at
     log2(16) = 4.0 bits/char, below the 4.7 MEDIUM floor, so a hex blob never
     classifies on entropy.

**Decode-and-rescan** (the must-have): base64 blobs that decode to printable
text are exposed on :class:`EntropyResult.decoded` so a later stage (``lexicon``)
can rescan the *decoded plaintext* for injection strings the encoding hid.
Entropy is a weak, evadable signal on its own; the decoded rescan is where the
real detection happens.

Scope note — false-positive suppression is the *text-relevant* subset of the
seed: known base64 media-URI prefixes (embedded images/fonts) and SRI hashes
(``sha384-…``). The seed's source-code-specific rules (GLSL, CSS-in-JS,
lockfiles, ffmpeg, import specifiers, …) are intentionally omitted — this
library scans ingested content, not source trees. The shared input-size /
decompression guard (OWASP LLM10 self-safety) lands with ``lexicon`` per
docs/PLAN.md; the extraction regex here is a bounded character class and is
linear-time (ReDoS-safe) on its own.
"""
from __future__ import annotations

import base64
import math
import re
from dataclasses import dataclass, field

from .report import Finding, Report, Severity, Source

# --- length-calibrated entropy thresholds (bits/char, min length) -----------
# Empirically calibrated against real distributions in the seed scanner:
# plaintext prose H ~3.5-4.2; base64 len64 H ~5.2; base64 len128 H ~5.6.
_CRITICAL_H, _CRITICAL_LEN = 5.4, 128
_HIGH_H, _HIGH_LEN = 5.1, 64
_MEDIUM_H, _MEDIUM_LEN = 4.7, 40

# Shape-floor lengths: structured encodings that are suspicious by size even
# when their entropy sits below the classification floor.
_BASE64_FLOOR_LEN = 100
_HEX_FLOOR_LEN = 64

# Candidate extraction: maximal runs of the base64 alphabet (hex is a subset),
# with optional trailing padding. Bounded class, no nested quantifier -> linear.
_BLOB_RE = re.compile(r"[A-Za-z0-9+/]{20,}={0,3}")

# Known base64 media data-URI magic prefixes -> benign embedded image/font/av.
_DATA_URI_PREFIXES = (
    "iVBORw0KGgo",   # PNG
    "/9j/",           # JPEG
    "R0lGOD",        # GIF
    "PHN2Zy",        # SVG
    "AAABAA",        # ICO
    "T2dnUw",        # OGG
    "AAAAFGZ0",      # MP4
    "UklGR",         # WebP / RIFF
    "d09G",          # WOFF font
    "AAEAAAALAAI",   # TTF font
)

# Subresource-Integrity marker immediately before the blob: `sha256-`, etc. The
# `-` breaks the base64 run, so a matched blob starts right after this prefix.
_SRI_PREFIX_BEFORE = re.compile(r"sha(?:256|384|512)-\Z")
_SRI_LOOKBEHIND = 8  # chars of preceding context to inspect (len("sha512-") = 7)

_PRINTABLE_EXTRA = frozenset("\n\r\t")


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


# --- ported primitives ------------------------------------------------------

def shannon_entropy(s: str) -> float:
    """Shannon entropy of ``s`` in bits per character (0.0 for empty)."""
    if not s:
        return 0.0
    freq: dict[str, int] = {}
    for ch in s:
        freq[ch] = freq.get(ch, 0) + 1
    n = len(s)
    entropy = 0.0
    for count in freq.values():
        p = count / n
        entropy -= p * math.log2(p)
    return entropy


def is_base64_like(s: str) -> bool:
    """True if ``s`` looks like a base64-encoded blob (>= 20 chars)."""
    if len(s) < 20:
        return False
    return re.fullmatch(r"[A-Za-z0-9+/]{20,}={0,3}", s) is not None


def is_hex_blob(s: str) -> bool:
    """True if ``s`` looks like a hex-encoded blob (>= 32 hex chars)."""
    if len(s) < 32:
        return False
    return re.fullmatch(r"(?:0x)?[0-9a-fA-F]{32,}", s) is not None


def try_decode_base64(s: str) -> str | None:
    """Decode ``s`` as base64 to text, or ``None`` if it is not printable text.

    Mirrors the seed: only base64-like input is attempted, and the result is
    kept only when it is >= 80% printable ASCII (so binary/media blobs, which
    decode to bytes, are rejected as rescan candidates).
    """
    if not is_base64_like(s):
        return None
    padded = s + "=" * (-len(s) % 4)
    try:
        raw = base64.b64decode(padded, validate=False)
    except ValueError:  # binascii.Error is a ValueError subclass
        return None
    if not raw:
        return None
    decoded = raw.decode("utf-8", errors="replace")
    if not decoded:
        return None
    printable = sum(
        1 for ch in decoded if 0x20 <= ord(ch) <= 0x7E or ch in _PRINTABLE_EXTRA
    )
    if printable / len(decoded) < 0.8:
        return None
    return decoded


# --- result types -----------------------------------------------------------

@dataclass(frozen=True)
class DecodedBlob:
    """A base64 blob decoded to printable text, ready for lexicon rescan.

    ``offset`` locates the encoded blob in the original text; ``decoded`` is the
    plaintext to re-scan; ``evidence`` is a redacted view of the encoded form.
    """

    offset: int
    decoded: str
    evidence: str


@dataclass(frozen=True)
class EntropyResult:
    """Findings plus the decoded blobs a caller should feed back to the lexicon."""

    report: Report
    decoded: list[DecodedBlob] = field(default_factory=list)


# --- classification ---------------------------------------------------------

def _classify(entropy: float, length: int) -> Severity | None:
    """Length-calibrated entropy tier, or ``None`` if below all thresholds."""
    if entropy >= _CRITICAL_H and length >= _CRITICAL_LEN:
        return Severity.CRITICAL
    if entropy >= _HIGH_H and length >= _HIGH_LEN:
        return Severity.HIGH
    if entropy >= _MEDIUM_H and length >= _MEDIUM_LEN:
        return Severity.MEDIUM
    return None


def _floor_medium(current: Severity | None) -> Severity:
    """Raise ``None`` to MEDIUM; leave an already-classified severity as-is."""
    return current if current is not None else Severity.MEDIUM


def _is_suppressed(token: str, text: str, offset: int) -> bool:
    """True for benign high-entropy blobs (embedded media, SRI hashes)."""
    if any(token.startswith(prefix) for prefix in _DATA_URI_PREFIXES):
        return True
    before = text[max(0, offset - _SRI_LOOKBEHIND):offset]
    if _SRI_PREFIX_BEFORE.search(before):
        return True
    return False


def scan_entropy(text: str, source: Source = Source.INPUT) -> EntropyResult:
    """Scan ``text`` for encoded/high-entropy blobs; return findings + decodes."""
    report = Report()
    decoded: list[DecodedBlob] = []

    for match in _BLOB_RE.finditer(text):
        token = match.group()
        offset = match.start()

        # Decode-and-rescan runs BEFORE suppression: it is independent of both
        # the entropy verdict and false-positive suppression. An attacker can
        # prefix an injection blob with an SRI/media marker to suppress the
        # entropy *finding* (below), but the hidden plaintext must still reach
        # the lexicon. Real media/SRI blobs decode to binary -> try_decode_base64
        # returns None, so this adds no false rescan candidates.
        if is_base64_like(token):
            plain = try_decode_base64(token)
            if plain is not None:
                decoded.append(
                    DecodedBlob(offset=offset, decoded=plain, evidence=_redact(token))
                )

        # Suppression gates only the entropy finding below, not the decode above.
        if _is_suppressed(token, text, offset):
            continue

        length = len(token)
        entropy = shannon_entropy(token)
        severity = _classify(entropy, length)
        if is_base64_like(token) and length > _BASE64_FLOOR_LEN:
            severity = _floor_medium(severity)
        if is_hex_blob(token) and length > _HEX_FLOOR_LEN:
            severity = _floor_medium(severity)
        if severity is None:
            continue

        label = "entropy:hex-blob" if is_hex_blob(token) else "entropy:base64-blob"
        report.add(
            Finding(
                label=label,
                severity=severity,
                source=source,
                detector="entropy",
                offset=offset,
                # Encoded content that can carry instructions -> injection carrier.
                owasp="LLM01",
                evidence=f"H={entropy:.2f}, len={length}: {_redact(token)}",
            )
        )

    return EntropyResult(report=report, decoded=decoded)