pipeline.py

"""
Stage 1 - PDF processing.

PDF bytes in, the text a model is allowed to read out, cut into windows.

**Everything above `window_document` is lifted verbatim from the EDA notebook that built the
corpus** (Final_project_V7_EDA.ipynb, cells 85-90). That is not tidiness - it is the whole
correctness argument of this application. Part B measured Gemma on `payload_window` strings
produced by exactly this code. If the app extracted text even slightly differently, the published
F1 of 0.969 would stop describing this program. `tests/test_fidelity.py` re-derives the corpus
columns from the original PDFs and fails on a single differing character.
"""

import re
import zlib

import numpy as np

# --------------------------------------------------------------------------------------------
# Lifted verbatim from the corpus build. Do not "improve" anything in this block.
# --------------------------------------------------------------------------------------------

STREAM_RE = re.compile(rb"(stream\r?\n)(.*?)(endstream)", re.S)

INJECTION_MARKERS = {
    "javascript_injection":     [rb"/S\s*/JavaScript", rb"/JS\s*\("],
    "cross_site_scripting":     [rb"<script>", rb"fetch\("],
    "ssrf":                     [rb"169\.254\.169\.254", rb"/latest/meta-data"],
    "object_action_injection":  [rb"/S\s*/Launch", rb"cmd\.exe"],
    "llm_prompt_injection":     [rb"IGNORE PREVIOUS INSTRUCTIONS", rb"LLM-INJECT"],
    "shellcode_embedded_exe":   [rb"application#2Fx-msdownload", rb"MZ.{0,20}\\x90\\x90"],
    "polyglot_file":            [rb"PK\\x03\\x04", rb"POLYGLOT ZIP\+PDF"],
    "dde_template_injection":   [rb"remote\.dotm\?dde=", rb"=cmd\|"],
    "steganographic_payload":   [rb"LSB-STEGO"],
    "ransomware_simulation":    [rb"RANSIM TEST"],
    "xfa_acroform_injection":   [rb"/Type\s*/XFA", rb"xfa:data"],
    "uri_redirect_phishing":    [rb"phishing\.\w+\.test"],
}
FRAMEWORK_MARKERS = ["AMTSO", "WICAR", "OWASP", "AtomicRedTeam", "Metasploit",
                     "Glasswall", "mindcrypt", "RanSim", "RANSIM", "custom"]
BASE_PAYLOAD_MARKERS = ["EICAR-STANDARD-ANTIVIRUS-TEST-FILE", "AMTSO-POTENTIALLY-UNWANTED-TEST-FILE",
                        "WICAR-BROWSER-TEST-PAYLOAD", "RANSIM-KNOWBE4-ENCRYPTION-SIMULATION"]

# Brand identifiers a model could string-match on without reasoning about PDF structure. The
# structural markers above are deliberately NOT masked: those ARE the attack, not a giveaway label.
LEAK_STRINGS = sorted(
    set(FRAMEWORK_MARKERS) | set(BASE_PAYLOAD_MARKERS)
    | {"EICAR", "RANSIM", "eicar-standard-antivirus-test-file",
       r"X5O!P%@AP[4\PZX54(P^)7CC)7}$"},
    key=len, reverse=True)            # longest first: "EICAR-STANDARD-..." masks before "EICAR"
LEAK_RE = re.compile("|".join(re.escape(s) for s in LEAK_STRINGS), re.IGNORECASE)

STREAM_BODY_CAP      = 4_096      # chars kept from any single stream body
SKELETON_CHAR_BUDGET = 120_000    # ~30k tokens
HEAD_SHARE           = 0.45       # of the budget; the rest is the tail window
SCAN_WINDOW          = 8_000_000  # bytes scanned from each end for markers

CLEAN_RE  = re.compile(r"[^\x20-\x7e\n]")
SPACES_RE = re.compile(r"[ ]{4,}")


def _printable_frac(chunk: bytes, sample: int = 200_000) -> float:
    """Fraction of bytes that are ordinary printable ASCII."""
    if not chunk:
        return 1.0
    arr = np.frombuffer(chunk[:sample], dtype=np.uint8)
    ok = ((arr >= 32) & (arr < 127)) | (arr == 9) | (arr == 10) | (arr == 13)
    return float(ok.mean())


def build_skeleton(data: bytes):
    """Render a PDF as payload-preserving text. Returns (skeleton, was_truncated, n_binary_dropped)."""
    dropped = 0

    def replace(match):
        nonlocal dropped
        opener, body, closer = match.group(1), match.group(2), match.group(3)
        try:                                            # most streams are FlateDecode
            inflated = zlib.decompress(body)
            if _printable_frac(inflated) > 0.6:
                return opener + inflated[:STREAM_BODY_CAP] + b"\n" + closer
        except zlib.error:
            pass
        if _printable_frac(body) > 0.6:                 # already plain text
            return opener + body[:STREAM_BODY_CAP] + closer
        dropped += 1                                    # genuinely binary (an image)
        return opener + b"<<BINARY %d bytes>>" % len(body) + closer

    text = STREAM_RE.sub(replace, data).decode("latin-1")

    truncated = len(text) > SKELETON_CHAR_BUDGET
    if truncated:
        # Head AND tail: the two insertion strategies put payloads at opposite ends of the file,
        # so a plain head truncation would lose most of them.
        head_n = int(SKELETON_CHAR_BUDGET * HEAD_SHARE)
        tail_n = SKELETON_CHAR_BUDGET - head_n
        elided = len(text) - SKELETON_CHAR_BUDGET
        text = (text[:head_n * 2]
                + f"\n<<... {elided} characters elided ...>>\n"
                + text[-tail_n * 2:])

    text = CLEAN_RE.sub(" ", text)                      # drop control/binary residue
    text = SPACES_RE.sub("   ", text)

    if len(text) > SKELETON_CHAR_BUDGET:                # enforce the budget after cleanup
        head_n = int(SKELETON_CHAR_BUDGET * HEAD_SHARE)
        text = text[:head_n] + text[-(SKELETON_CHAR_BUDGET - head_n):]
        truncated = True
    return text, truncated, dropped


def mask_leaks(skeleton: str) -> str:
    """Blank the brand identifiers, keep the structural shape of the injection intact."""
    return LEAK_RE.sub(lambda m: "X" * len(m.group(0)), skeleton)


def scan_window(data: bytes, window: int = SCAN_WINDOW) -> bytes:
    """Head and tail of a file, for the regex marker scans."""
    if len(data) <= 2 * window:
        return data
    return data[:window] + data[-window:]


def detect_markers(blob) -> list:
    """Which injection types are structurally present. Works on raw bytes or on a skeleton."""
    data = blob.encode("latin-1", errors="replace") if isinstance(blob, str) else blob
    data = scan_window(data)
    return [name for name, patterns in INJECTION_MARKERS.items()
            if any(re.search(p, data, re.IGNORECASE) for p in patterns)]


ANY_MARKER_RE = re.compile(
    b"|".join([p for ps in INJECTION_MARKERS.values() for p in ps]
              + [re.escape(s).encode() for s in FRAMEWORK_MARKERS + BASE_PAYLOAD_MARKERS]),
    re.IGNORECASE)

# --------------------------------------------------------------------------------------------
# The application's own code starts here.
# --------------------------------------------------------------------------------------------

WINDOW = 1_500                 # characters either side; a window is 2*WINDOW = 3,000 chars
STRIDE = WINDOW                # 50% overlap, so no payload can straddle a boundary unseen

# 120,000-char skeleton budget / 1,500 stride caps this at ~80 windows for any document.
MAX_WINDOWS = SKELETON_CHAR_BUDGET // STRIDE + 1


def payload_window(text, half=WINDOW):
    """
    The evaluation's window: the neighbourhood of the *known* injection.

    THIS IS THE ANSWER KEY. It locates the payload by searching for the markers the corpus was
    generated with, so it only works on corpus files. On a PDF a user uploads nothing matches and
    every document falls through to the head - which the EDA measured as a bad detector, since
    only 40.5% of injected files carry their payload in the first 3,000 characters.

    Kept because `tests/test_fidelity.py` needs it to reproduce the published column, and for
    scoring corpus files. Never call it on user input; call `window_document` instead.
    """
    m = ANY_MARKER_RE.search(text.encode("latin-1", errors="replace"))
    if not m:
        return text[:2 * half]
    return text[max(0, m.start() - half): m.start() + half]


def pdf_to_text(data: bytes) -> dict:
    """One uploaded PDF, rendered the way the corpus was rendered."""
    skeleton, truncated, dropped = build_skeleton(data)
    return {
        "skeleton": skeleton,
        "masked": mask_leaks(skeleton),
        "chars": len(skeleton),
        "was_truncated": truncated,
        "binary_streams_dropped": dropped,
        "file_bytes": len(data),
    }


def window_document(text: str, size: int = 2 * WINDOW, stride: int = STRIDE) -> list:
    """
    Cut a document into the input shape Part B measured, covering all of it.

    Returns [{"index", "start", "end", "text"}]. Windows are `size` characters at `stride`
    intervals, so consecutive windows overlap by half. The overlap is deliberate: a payload
    sitting across a boundary would otherwise be split into two halves, neither of which reads
    as an attack. It also means **one payload normally appears in two windows**, which is why
    `analysis.py` merges overlapping hits into regions before counting them.

    The final window is anchored to the end of the text rather than left short, so the tail -
    where the EDA found 63.4% of injections - is always covered at full width.
    """
    text = text or ""
    if not text:
        return []
    if len(text) <= size:
        return [{"index": 0, "start": 0, "end": len(text), "text": text}]

    starts = list(range(0, len(text) - size + 1, stride))
    if starts[-1] + size < len(text):
        starts.append(len(text) - size)          # anchored tail window

    return [{"index": i, "start": s, "end": s + size, "text": text[s:s + size]}
            for i, s in enumerate(starts[:MAX_WINDOWS])]