detector.py

"""
Stage 3 - Prompting the model.

The prompt, the parser, and the model registry. **The prompt and the parser are lifted verbatim
from Part B** (Final_Project_Evaluation_V1.ipynb, cells 48-49): same system message, same closed
family list in the same order, same brace-counting parser, same greedy decoding, same 4-bit NF4
quantisation. A reworded prompt is a different experiment, and the measured F1 would no longer
apply to this app.

The registry is the only place a model is named. Every field in a `ModelSpec` is a per-model
difference Part B actually ran into.
"""

import gc
import json
import os
import re
import sys
import time
from dataclasses import dataclass, field

FAMILIES = ["cross_site_scripting", "dde_template_injection", "javascript_injection",
            "llm_prompt_injection", "object_action_injection", "polyglot_file",
            "ransomware_simulation", "shellcode_embedded_exe", "ssrf",
            "steganographic_payload", "uri_redirect_phishing", "xfa_acroform_injection"]

# A short paragraph per family: how it works, then why that matters. Each is written against the
# generator's own object shapes rather than the general security literature, so a reader who opens
# a flagged file finds the thing the paragraph describes.
#
# Ordered by how directly the attack reaches a person, not alphabetically - these are read one at a
# time while something downloads, so the first ones a user meets should be the ones that explain
# why any of this is worth detecting.
#
# These are for the *reader*. They are deliberately kept out of SYSTEM: the prompt is Part B's, byte
# for byte, and adding twelve definitions to it would make the measured scores describe a different
# prompt from the one running.
FAMILY_NOTES = {
    "llm_prompt_injection":
        "LLM prompt injections work by hiding a written instruction inside the PDF - in the "
        "metadata, or in an annotation sized so small that no human ever sees it. When someone "
        "feeds the file to an AI assistant, the assistant reads that hidden text along with "
        "everything else and follows it. This is dangerous because the attacker is now giving "
        "orders to a tool that the victim trusts and has given access to their files, their inbox "
        "or their code - and nothing looked wrong on the page.",
    "javascript_injection":
        "JavaScript injections put code in a `/JavaScript` action, which a PDF viewer runs the "
        "moment the file is opened. The code is usually obfuscated, so the file carries a string "
        "of hex or character codes rather than anything readable. This is dangerous because it "
        "needs no click and no mistake by the reader: opening the document is the whole attack, "
        "and the code runs with whatever the viewer is allowed to do.",
    "object_action_injection":
        "Object and action injections use a `/Launch` action pointing at a program - typically "
        "`cmd.exe` with a command attached. Opening the document asks the operating system to run "
        "something rather than to display a page. This is dangerous because it steps outside the "
        "PDF viewer entirely: whatever the command does happens with the user's own permissions, "
        "and a PDF is the last place most people expect a program to start.",
    "shellcode_embedded_exe":
        "Embedded-executable payloads carry a real Windows program inside the PDF, recognisable "
        "by the `MZ` header that starts every .exe, stored as an attachment named something "
        "ordinary. This is dangerous because the PDF becomes a delivery envelope: it passes the "
        "mail filter as a document, and the malware only has to be extracted and double-clicked "
        "by someone who thinks they are opening a file the sender meant to attach.",
    "ransomware_simulation":
        "Ransomware payloads combine code that runs on open with an on-page notice announcing "
        "that files have been encrypted and demanding payment. This is dangerous because of what "
        "the code does before the notice appears - by the time the message is read, the work it "
        "describes is finished. The samples in this corpus are harmless simulations built from "
        "public test material, shaped exactly like the real thing so a detector can be measured.",
    "cross_site_scripting":
        "Cross-site scripting hides a `<script>` tag where a web address belongs, in a link "
        "annotation. This is dangerous because PDFs are increasingly opened inside browsers and "
        "web applications rather than desktop readers: a viewer that renders that link as HTML "
        "runs the attacker's script inside the page, where it can reach the session, the cookies "
        "and anything else that page is trusted with.",
    "uri_redirect_phishing":
        "Redirect phishing places a link annotation over the entire page and points it at an "
        "attacker's site, usually with an identifying token attached. This is dangerous because "
        "there is nothing to avoid clicking - any click anywhere in the document opens the link, "
        "and the destination is a login page that looks like one the reader already trusts.",
    "ssrf":
        "Server-side request forgery aims a URI action at an address only the server can reach - "
        "classically `169.254.169.254`, the cloud metadata service. This is dangerous when the "
        "PDF is processed by a machine rather than a person: a thumbnail generator or preview "
        "service that follows the link fetches its own host's credentials and hands them back, "
        "and the attacker never needed access to the network at all.",
    "dde_template_injection":
        "DDE and template injections carry a spreadsheet formula such as `=cmd|' /c ...'!A1` in a "
        "form field, alongside a link to a remote Office template. This is dangerous because the "
        "payload activates once the content moves into Office - copied into a spreadsheet, or "
        "opened through the linked template - so the PDF itself looks inert and the attack "
        "happens in an application the reader thinks is unrelated.",
    "xfa_acroform_injection":
        "XFA and AcroForm injections bury the payload in the XFA form definition - a whole XML "
        "document living inside the PDF. This is dangerous because it is a second format inside "
        "the first: many scanners parse PDF objects and never parse the form XML, so the payload "
        "sits in a part of the file that was never inspected while still reaching the viewer.",
    "polyglot_file":
        "Polyglot files are valid as two formats at once - here a ZIP archive embedded so that "
        "the same bytes read as both a PDF and an archive. This is dangerous because every "
        "security tool decides what a file is before deciding how to check it: the scanner reads "
        "the document, the operating system reads the archive, and the contents of whichever "
        "format was not chosen are never examined.",
    "steganographic_payload":
        "Steganographic payloads hide data inside an image stream, often a 1x1 pixel nobody will "
        "ever see rendered. This is dangerous because it defeats inspection rather than "
        "execution: the file carries the payload past filters that are looking for code or links, "
        "and the hidden data is retrieved later by something that already knows where to look.",
}

SYSTEM = (
    "You are a PDF security analyst. You are given the raw extracted text of a PDF file - object "
    "definitions, stream contents and metadata, exactly as they appear in the file. Some of these "
    "files have had a malicious payload injected into them; most, but not all, have. Your job is to "
    "say which, and to point at your evidence.\n\n"
    "Answer with a single JSON object and nothing else:\n"
    '{"injected": true or false, '
    '"injection_type": one of ' + json.dumps(FAMILIES) + ' or "none", '
    '"evidence": the exact substring from the input that convinced you, at most 200 characters, '
    'or "" if none, '
    '"reasoning": one short sentence}\n\n'
    "If the file looks clean, answer injected=false and injection_type=\"none\". Do not guess a "
    "family when you do not believe there is an injection."
)

MAX_CHARS = 3000        # a window is already this size; belt and braces
MAX_NEW = 200
BATCH = 8               # quartered automatically on CUDA OOM, same policy as Parts A and B

# What a 7-9B model in 4-bit NF4 actually costs on the card: about 5 GB of weights, plus room for
# the activations of a batch of 8 at 2,048 tokens. Checked before loading rather than discovered
# during it, so a full GPU is a sentence rather than a page of allocator arithmetic.
NEEDED_BYTES = 7_000_000_000


@dataclass
class ModelSpec:
    repo: str
    label: str
    prefill: str = ""
    trust_remote_code: bool = False
    kwargs: dict = field(default_factory=dict)
    notes: str = ""
    gated: bool = False                 # the weights need an accepted licence, not just a token
    seconds_per_window: float = 10.95   # measured in Part B on a T4, 4-bit, batch 8
    summary: str = ""                   # one line, shown next to the picker


DETECTORS = {
    "mimo": ModelSpec(
        repo="XiaomiMiMo/MiMo-7B-RL", seconds_per_window=4.2, label="MiMo-7B - fast, no token needed",
        prefill='<think>\n\n</think>\n\n{"injected":', trust_remote_code=True,
        summary="F1 0.945 · names the family 43% of the time · 2.6x faster than the others · "
                "downloads without a Hugging Face account.",
        notes="F1 0.945, family 0.433, 2.6x faster. Reasoning-trained: without the prefilled "
              "think-block it never reaches the JSON within 200 tokens."),
    "gemma": ModelSpec(
        repo="google/gemma-2-9b-it", seconds_per_window=10.95, label="Gemma-2-9B - most accurate, gated",
        kwargs={"attn_implementation": "eager"}, gated=True,
        summary="F1 0.969 · names the family 63% of the time · 3% false alarms. The best of the "
                "four, and the only one Google gates: needs your own token.",
        notes="F1 0.969, family 0.630, 6 false alarms of 200. Gated repo - needs HF_TOKEN. "
              "No system role (handled in render_prompt); eager attention per Google."),
    "qwen": ModelSpec(
        repo="Qwen/Qwen2.5-7B-Instruct", seconds_per_window=10.94, label="Qwen2.5-7B - most cautious",
        summary="F1 0.957 · names the family 52% of the time · almost never cries wolf "
                "(precision 0.995), but leaves the most answers unreadable.",
        notes="F1 0.957, family 0.524. Most precise (0.995), worst parse rate (159 unreadable)."),
    "phi": ModelSpec(
        repo="microsoft/Phi-4-mini-instruct", seconds_per_window=2.12, label="Phi-4-mini - not recommended",
        summary="F1 0.900, which is exactly what calling every file malicious scores. Flags 95% "
                "of clean files. Here because Part B measured it, not because you should use it.",
        notes="F1 0.900 - ties the always-malicious constant. 95% false-alarm rate. Included "
              "because Part B measured it, not because it should be used."),
}

# The app opens on MiMo, not on the winner. Gemma is gated, so defaulting to it means the first
# thing a new user meets is a licence refusal; MiMo downloads for anyone, gives up 0.024 F1, and
# is the only one of the four that can be pre-fetched before a token exists. The picker is one
# click away for anyone who has Gemma access.
ACTIVE = "mimo"

# Where a gated model's licence is accepted. Built from the repo id rather than written out, so
# it cannot drift from the repo the code actually downloads.
def gate_url(name: str = None) -> str:
    return f"https://huggingface.co/{DETECTORS[name or ACTIVE].repo}"

# name -> (model, tokenizer). One at a time; two do not fit on one GPU.
#
# Parked on `sys` rather than held as an ordinary module global, because this module gets reloaded.
# `reload_stages()` drops "detector" from sys.modules so an edited cell takes effect, and the fresh
# module then starts with an empty cache while the previous incarnation's model is still sitting on
# the GPU - reachable by nothing, releasable by nobody. That is how a 22 GB card ends up with 0.7 GB
# free before anything has been asked to load. One dict, shared across reloads, makes release()
# able to free a model this module did not personally load.
_loaded = getattr(sys, "_pdf_injection_detectors", None)
if _loaded is None:
    _loaded = sys._pdf_injection_detectors = {}


def hf_token():
    """
    The Hugging Face token, from wherever this runtime keeps it.

    The licence for a gated repo is granted once, to an account. What differs between machines is
    only how the token reaches the process, so all four places are checked instead of making the
    user paste it again:

      1. HF_TOKEN / HUGGING_FACE_HUB_TOKEN in the environment - this is how a Space injects a
         repository secret, so nothing extra is needed there;
      2. Colab's secret store (the key icon in the sidebar), for a hosted notebook;
      3. the `huggingface-cli login` cache, for an ordinary machine;
      4. nothing, in which case ungated models still work and gated ones 401.

    Whatever is found is also exported to the environment, because `transformers` reads it from
    there on its own in a few code paths.
    """
    tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

    if not tok:
        try:                                     # Colab secrets
            from google.colab import userdata
            tok = userdata.get("HF_TOKEN")
        except Exception:                        # not Colab, or the secret is not enabled
            tok = None

    if not tok:
        try:                                     # huggingface-cli login cache
            from huggingface_hub import get_token
            tok = get_token()
        except Exception:
            tok = None

    if tok:
        os.environ["HF_TOKEN"] = tok
    return tok



# --------------------------------------------------------------------------------------------
# Prompt and parser - verbatim from Part B
# --------------------------------------------------------------------------------------------

def build_messages(text: str):
    """The chat turns for one window. Identical for all models - only the template differs."""
    return [{"role": "system", "content": SYSTEM},
            {"role": "user", "content": "PDF extract:\n\n```\n" + text[:MAX_CHARS] + "\n```"}]


def render_prompt(tokenizer, text: str, model_name: str = "") -> str:
    """
    Apply a model's own chat template to one window.

    Gemma-2 has no `system` role - Google's template raises TemplateError and expects system
    instructions folded into the first user turn. The same words are moved into the user turn
    only for the templates that refuse them, so every model reads identical text.
    """
    messages = build_messages(text)
    try:
        rendered = tokenizer.apply_chat_template(messages, tokenize=False,
                                                 add_generation_prompt=True)
    except Exception:                       # jinja2.TemplateError
        merged = [{"role": "user", "content": SYSTEM + "\n\n" + messages[1]["content"]}]
        rendered = tokenizer.apply_chat_template(merged, tokenize=False, add_generation_prompt=True)
    prefill = DETECTORS[model_name].prefill if model_name in DETECTORS else ""
    return rendered + prefill


def scan_objects(raw: str):
    r"""
    Every balanced {...} in the text, counting braces and skipping string literals.

    A regex cannot do this. The naive `\{[^{}]*\}` this replaced was actively harmful: every
    injected file carries a marker containing a `}`, so the moment a model quoted its evidence the
    match truncated and the verdict was discarded - the bug fired exactly when the model was RIGHT.
    """
    objs, depth, start, in_str, esc = [], 0, None, False, False
    for i, ch in enumerate(raw or ""):
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
            continue
        if ch == '"':
            in_str = True
        elif ch == "{":
            if depth == 0:
                start = i
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0 and start is not None:
                objs.append(raw[start:i + 1])
                start = None
            depth = max(depth, 0)
    return objs


# The salvage patterns, for answers that are not valid JSON.
#
# Part B's version of this was `"injected"\s*:\s*(true|false)` - double quotes, bare boolean - and
# every other way a model writes the same verdict was thrown away as unreadable. Widened here for
# what these four models actually emit:
#   ["'*]*   the key wrapped in double quotes, single quotes, markdown bold, or nothing
#   [:=]     a colon, or the equals sign models sometimes use in prose
#   "?(...)  the value quoted or bare, and true/yes/1 as well as True (the pattern is case-blind)
# Evidence is salvaged too, greedily up to the last quote before the next known key, which is what
# lets a quote containing its own quotes and backslashes - i.e. most real PDF text - survive.
#
# This changes verdicts: an answer Part B scored as unreadable-therefore-clean may now be read as
# injected. The reported F1 was measured with the narrow pattern, so it describes the old parser,
# and any number quoted beside this one needs re-deriving.
VERDICT_RE = re.compile(r'''["'*]*injected["'*]*\s*[:=]\s*["']?(true|false|yes|no|1|0)''', re.I)
FAMILY_RE = re.compile(r'''["'*]*injection_type["'*]*\s*[:=]\s*["']?([a-z_]+)''', re.I)
#
# Evidence needs two patterns, because the two ways it arrives are different problems. A complete
# answer is closed by the next key, so the quote can be found by looking for that key - which lets
# the value itself contain quotes, as PDF text usually does. A truncated one has no closing
# anything, so whatever follows the opening quote is taken as a partial quote and marked as one.
EVIDENCE_RE = re.compile(r'''["'*]*evidence["'*]*\s*[:=]\s*["'](.*?)["']\s*[,}]\s*'''
                         r'''["'*]*(?:reasoning|injection_type|injected)''', re.S | re.I)
EVIDENCE_TAIL_RE = re.compile(r'''["'*]*evidence["'*]*\s*[:=]\s*["'](.*)$''', re.S | re.I)


def _salvage_evidence(raw: str) -> str:
    """The quoted evidence from an answer that did not parse as JSON, or ""."""
    m = EVIDENCE_RE.search(raw)
    if m:
        return m.group(1).strip()[:200]
    m = EVIDENCE_TAIL_RE.search(raw)
    if m:                                       # cut off mid-quote; say so rather than imply more
        tail = m.group(1).strip().rstrip('"\'').strip()
        return (tail[:200] + " …[cut off]") if tail else ""
    return ""


def parse_response(raw: str) -> dict:
    """
    Pull the verdict out of whatever the model said.

    The LAST balanced object wins, not the first: reasoning-trained models restate the schema
    while thinking. If nothing parses, the two fields that decide the answer are lifted out with a
    field-level regex. Only a response with no recoverable verdict is parse_ok=False, and that
    reads as 'not injected' - a detector that cannot make itself understood has caught nothing.
    """
    for m in reversed(scan_objects(raw)):
        try:
            obj = json.loads(m)
        except json.JSONDecodeError:
            continue
        if "injected" in obj:
            inj = obj["injected"]
            inj = inj if isinstance(inj, bool) else str(inj).strip().lower() in {"true", "yes", "1"}
            fam = str(obj.get("injection_type", "none") or "none").strip().lower()
            return {"parse_ok": True, "parsed_by": "balanced JSON",
                    "pred_injected": int(inj),
                    "pred_family": fam if fam in FAMILIES else "none",
                    "evidence": str(obj.get("evidence", ""))[:200],
                    "reasoning": str(obj.get("reasoning", ""))[:300]}

    m = re.search(VERDICT_RE, raw or "")
    if m:
        f = re.search(FAMILY_RE, raw)
        fam = f.group(1).lower() if f else "none"
        return {"parse_ok": True, "parsed_by": "field regex",
                "pred_injected": int(m.group(1).lower() in {"true", "yes", "1"}),
                "pred_family": fam if fam in FAMILIES else "none",
                "evidence": _salvage_evidence(raw), "reasoning": ""}

    return {"parse_ok": False, "parsed_by": "unrecoverable", "pred_injected": 0,
            "pred_family": "none", "evidence": "", "reasoning": ""}


# --------------------------------------------------------------------------------------------
# Loading and generation
# --------------------------------------------------------------------------------------------

def load(name: str = None):
    """
    Load one model in 4-bit NF4 - the same quantisation Part B measured.

    Cached per name. Loading a second model releases the first: a T4 holds one of these at a time.
    """
    name = name or ACTIVE
    if name in _loaded:
        return _loaded[name]

    spec = DETECTORS[name]
    for other in list(_loaded):
        release(other)

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    # bitsandbytes cannot quantise a layer that is not on a GPU. With device_map="auto" and no
    # device to put things on, `accelerate` silently assigns layers to the CPU and the load then
    # dies deep inside bitsandbytes with "Some modules are dispatched on the CPU or the disk" -
    # a message about offloading that never mentions the actual problem, which is that this
    # runtime has no GPU. Say so here instead, before anything is downloaded onto nothing.
    if not torch.cuda.is_available():
        raise RuntimeError(
            "this runtime has no GPU, and 4-bit loading needs one. In Colab: Runtime > Change "
            "runtime type > T4 GPU, then press Start again. (Running these models on CPU is not "
            "an option worth offering - a single window would take minutes.)")

    # Reclaim first, then look. release() dropped the other entries above; this is what actually
    # returns their memory to the driver, and it is worth doing even when the cache was empty -
    # a failed load leaves allocations behind that only gc plus empty_cache will free.
    free, total = reclaim()
    if free < NEEDED_BYTES:
        # Not enough after the polite attempt. Push strays onto the CPU and look again: a model
        # stranded by a module reload is unreachable by name but still movable by object.
        before = free
        free, total = reclaim(aggressive=True)
        print(f"reclaimed {(free - before) / 1e9:.1f} GB from stray modules")

    if free < NEEDED_BYTES:
        holders = gpu_holders()
        who = (", ".join(f"{name} {size / 1e9:.1f} GB" for name, size in holders)
               if holders else "no live torch module - the memory is held by something this "
                               "process can no longer name")
        raise RuntimeError(
            f"only {free / 1e9:.1f} GB of this GPU's {total / 1e9:.1f} GB is free, and "
            f"{spec.label.split(' - ')[0]} in 4-bit needs about {NEEDED_BYTES / 1e9:.0f} GB. "
            f"Still resident: {who}. A 4-bit model cannot be moved off the card, so if that is "
            f"what is listed, restart the runtime (in Colab: Runtime > Restart session) and press "
            f"Start again without running the earlier model cells.")

    q = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                           bnb_4bit_compute_dtype=torch.float16,
                           bnb_4bit_use_double_quant=True)
    t0 = time.time()
    tok = AutoTokenizer.from_pretrained(spec.repo, token=hf_token(),
                                        trust_remote_code=spec.trust_remote_code)

    # device_map={"": 0} rather than "auto": everything on one GPU or nothing. "auto" is free to
    # spill the layers that do not fit onto the CPU, which is exactly the state bitsandbytes then
    # refuses to work in, so the failure arrives as an offload complaint instead of "it does not
    # fit". Pinning it turns that into a plain out-of-memory error, reported below with the
    # numbers that explain it.
    try:
        model = AutoModelForCausalLM.from_pretrained(
            spec.repo, quantization_config=q, device_map={"": 0},
            token=hf_token(), trust_remote_code=spec.trust_remote_code, **spec.kwargs)
    except (ValueError, torch.cuda.OutOfMemoryError) as e:
        free, total = torch.cuda.mem_get_info()
        raise RuntimeError(
            f"{spec.label} would not fit on this GPU: {free / 1e9:.1f} GB free of "
            f"{total / 1e9:.1f} GB, and a 7-9B model in 4-bit needs roughly 6 GB. "
            f"If another model is still resident, restart the runtime; otherwise use a larger "
            f"GPU. (Underlying error: {type(e).__name__}: {e})") from e
    model.eval()
    print(f"{spec.repo}: loaded in {time.time() - t0:.0f}s")

    _loaded[name] = (model, tok)
    return _loaded[name]


def gpu_holders(limit: int = 4) -> list:
    """
    What is actually sitting on the GPU right now: [(class name, bytes), ...], largest first.

    "Something is holding the card" is not a useful thing to tell someone. This walks the garbage
    collector's object list for live torch modules with CUDA parameters and adds up what each one
    costs, so the message can name the thing instead of describing its shadow. Sizes are summed per
    top-level module, and submodules are skipped so a model is not counted once per layer.
    """
    import warnings

    import torch

    # Walking gc.get_objects() touches every attribute of every live object, which makes torch's
    # deprecated aliases fire FutureWarnings that have nothing to do with anything here.
    modules, children = [], set()
    warnings.simplefilter("ignore")
    for obj in gc.get_objects():
        try:
            if isinstance(obj, torch.nn.Module):
                modules.append(obj)
                children.update(id(m) for m in obj.children())
        except Exception:                        # a half-built module mid-construction
            continue

    sizes = {}
    for mod in modules:
        if id(mod) in children:                  # a submodule; its parent already counts it
            continue
        try:
            total = sum(p.numel() * p.element_size()
                        for p in mod.parameters(recurse=True) if p.is_cuda)
        except Exception:
            continue
        if total:
            sizes[type(mod).__name__] = sizes.get(type(mod).__name__, 0) + total
    return sorted(sizes.items(), key=lambda kv: -kv[1])[:limit]


def reclaim(aggressive: bool = False):
    """
    Give the GPU back, as far as Python allows.

    Dropping our own cache and collecting is the whole of the safe part. `aggressive` additionally
    pushes stray CUDA modules - ones this module never loaded, left behind by a reloaded copy of
    itself - onto the CPU, which is the only way to free memory that something else still
    references. A 4-bit model cannot be moved and raises; that one genuinely needs a restart.
    """
    import torch

    import warnings

    release()
    if aggressive:
        # This can also push the embedding model off the card, which is a fair trade: it is small
        # and runs acceptably on the CPU, whereas the detector does not run at all without a GPU.
        warnings.simplefilter("ignore")
        for obj in gc.get_objects():
            try:
                if isinstance(obj, torch.nn.Module) and any(p.is_cuda for p in obj.parameters()):
                    obj.to("cpu")
            except Exception:                    # quantised weights refuse to move; skip them
                continue
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return torch.cuda.mem_get_info() if torch.cuda.is_available() else (0, 0)


def release(name: str = None):
    """Free the GPU. ZeroGPU reclaims the device between calls, so this is called on the way out."""
    import torch
    for key in ([name] if name else list(_loaded)):
        if key in _loaded:
            del _loaded[key]
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def scan(windows: list, name: str = None, batch: int = BATCH, progress=None) -> list:
    """
    Score every window. Returns one dict per window, in order.

    **No early stop.** The old design stopped at the first hit, which is faster but cannot answer
    "in how many parts was malware found" - and cannot honestly say "clean" either, since that
    claim requires having looked everywhere. Every window is scored and reported.

    Windows are batched the way Part B batched documents, which is the difference between seconds
    and minutes on a document of 30-80 windows.
    """
    import torch

    name = name or ACTIVE
    model, tok = load(name)
    prompts = [render_prompt(tok, w["text"], name) for w in windows]

    tok.padding_side = "left"               # decoder-only: right padding starts generation after it
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    out, i, t0 = [], 0, time.time()
    while i < len(prompts):
        chunk = prompts[i:i + batch]
        try:
            enc = tok(chunk, return_tensors="pt", padding=True,
                      truncation=True, max_length=2048).to(model.device)
            with torch.inference_mode():
                gen = model.generate(**enc, max_new_tokens=MAX_NEW, do_sample=False,
                                     pad_token_id=tok.pad_token_id)
            texts = tok.batch_decode(gen[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            if batch == 1:
                raise
            batch = max(1, batch // 4)
            print(f"    OOM -> batch {batch}")
            continue

        # generate() returns only new tokens, so a prefill was given to the model but never
        # generated. Put it back before parsing or the opening brace is missing.
        pre = DETECTORS[name].prefill
        for w, raw in zip(windows[i:i + len(chunk)], texts):
            raw = pre + raw
            out.append({**w, "raw": raw.strip()[:2000], **parse_response(raw)})

        i += len(chunk)
        if progress is not None:
            progress(min(i, len(prompts)) / len(prompts))

    elapsed = time.time() - t0
    for r in out:
        r["seconds"] = round(elapsed / max(len(out), 1), 2)
    return out