analysis.py

"""
Stage 4 - Analysing the output.

Per-window verdicts in, two reports out: one row per window, and one consolidated document
verdict.

The consolidation is where the honest work happens. Three things have to be got right:

1. **Overlapping windows are not independent findings.** Windows overlap by 50%, so one payload
   normally lands in two of them. Counting flagged windows would report a single injection as
   two. Adjacent flagged windows are merged into regions, and regions are what "in how many parts"
   counts.
2. **A window verdict is not a document verdict.** Part B measured Gemma on one window per
   document at a 3% false-alarm rate. A 40-window document asks the same question 40 times, so
   isolated single-window hits deserve less confidence than a region two or more windows wide.
   That is reported rather than silently thresholded away.
3. **The family is a suggestion.** Gemma's family accuracy is 0.630 overall and 0.00 on `ssrf`.
   The consolidated report says so, every time, next to the family it is suggesting.
"""

from collections import Counter

# Straight from the evaluation card, so the app cannot quote a score without its baseline.
MEASURED = {
    "gemma": {"f1": 0.969, "precision": 0.993, "recall": 0.946,
              "false_alarm_rate": 0.03, "family_acc": 0.630},
    "qwen":  {"f1": 0.957, "precision": 0.995, "recall": 0.921,
              "false_alarm_rate": 0.02, "family_acc": 0.524},
    "mimo":  {"f1": 0.945, "precision": 0.988, "recall": 0.906,
              "false_alarm_rate": 0.05, "family_acc": 0.433},
    "phi":   {"f1": 0.900, "precision": 0.824, "recall": 0.991,
              "false_alarm_rate": 0.95, "family_acc": 0.330},
}
BASELINE_F1 = 0.900     # "always malicious" - the bar every number above must be read against
WEAK_FAMILIES = {"ssrf"}    # families the winner catches but cannot name (0.00 in Part B)


def merge_regions(findings: list) -> list:
    """
    Collapse flagged windows into regions of the document.

    Two flagged windows belong to the same region when their character ranges touch or overlap.
    Each region carries the windows that produced it, the family they voted for, and the single
    best evidence string among them.
    """
    flagged = [f for f in findings if f["pred_injected"] == 1]
    if not flagged:
        return []

    flagged.sort(key=lambda f: f["start"])
    regions, current = [], [flagged[0]]

    for f in flagged[1:]:
        if f["start"] <= current[-1]["end"]:        # touching or overlapping
            current.append(f)
        else:
            regions.append(current)
            current = [f]
    regions.append(current)

    out = []
    for group in regions:
        votes = Counter(f["pred_family"] for f in group if f["pred_family"] != "none")
        family, family_votes = (votes.most_common(1)[0] if votes else ("none", 0))
        # The longest quoted evidence in the region: the model that had most to point at.
        best = max(group, key=lambda f: len(f["evidence"] or ""))
        out.append({
            "start": group[0]["start"],
            "end": max(f["end"] for f in group),
            "windows": [f["index"] for f in group],
            "n_windows": len(group),
            "family": family,
            "family_agreement": round(family_votes / len(group), 2) if group else 0.0,
            "evidence": best["evidence"],
            "evidence_offset": _locate(best),
            "reasoning": best["reasoning"],
        })
    return out


def _locate(finding: dict):
    """Character offset of the quoted evidence in the whole document, or None if it is not there."""
    ev = finding.get("evidence") or ""
    if not ev:
        return None
    i = (finding.get("text") or "").find(ev)
    return None if i < 0 else finding["start"] + i


def consolidate(findings: list, doc: dict, model: str = "gemma") -> dict:
    """The document-level verdict, built only from what the windows actually said."""
    regions = merge_regions(findings)
    n_flagged = sum(f["pred_injected"] for f in findings)
    unreadable = sum(1 for f in findings if not f["parse_ok"])
    families = Counter(r["family"] for r in regions if r["family"] != "none")

    # A region seen in two or more overlapping windows is a stronger signal than a lone hit: the
    # model was shown the same bytes twice, framed differently, and said yes both times.
    corroborated = [r for r in regions if r["n_windows"] >= 2]

    return {
        "injected": bool(regions),
        "model": model,
        "n_regions": len(regions),
        "n_corroborated": len(corroborated),
        "n_windows_flagged": n_flagged,
        "n_windows": len(findings),
        "unreadable": unreadable,
        "families": families.most_common(),
        "regions": regions,
        "seconds": round(sum(f.get("seconds", 0) for f in findings), 1),
        "document": doc,
    }


# --------------------------------------------------------------------------------------------
# Rendering
# --------------------------------------------------------------------------------------------

def window_rows(findings: list) -> list:
    """
    One row per window, for the per-window table in the GUI.

    The last column is what the model actually said, trimmed. It is here because "unreadable" is
    the one verdict a reader cannot check: every other row can be compared against the document,
    while a window that failed to parse offers nothing to look at unless the raw answer is shown.
    With it, a run of unreadable windows says which failure it is - an unfinished <think> block,
    an answer cut off mid-quote, or plain prose that never reached a verdict.
    """
    return [[f["index"],
             f"{f['start']:,}-{f['end']:,}",
             "INJECTED" if f["pred_injected"] else "clean",
             f["pred_family"] if f["pred_injected"] else "-",
             (f["evidence"] or "")[:90].replace("\n", " "),
             f["parsed_by"],
             (f.get("raw") or "")[:160].replace("\n", " ⏎ ")]
            for f in findings]


WINDOW_COLUMNS = ["window", "characters", "verdict", "family", "evidence", "answer read as",
                  "what the model said"]


def report_markdown(summary: dict) -> str:
    """The consolidated report."""
    doc = summary["document"]
    m = MEASURED.get(summary["model"], {})
    lines = []

    if summary["injected"]:
        lines.append(f"# Injection found - {summary['n_regions']} "
                     f"{'region' if summary['n_regions'] == 1 else 'regions'} of the document")
    else:
        lines.append("# No injection found")

    lines.append("")
    lines.append(f"**{summary['n_windows_flagged']} of {summary['n_windows']} windows flagged** "
                 f"| {doc['chars']:,} characters extracted from {doc['file_bytes']:,} bytes "
                 f"| {summary['seconds']}s on {summary['model']}")

    if doc["was_truncated"]:
        lines.append("")
        lines.append("> **This document was truncated to 120,000 characters** (head and tail kept, "
                     "middle elided) - the same budget the corpus was built with. An injection in "
                     "the elided middle would not be seen.")

    if summary["unreadable"]:
        lines.append("")
        lines.append(f"> **{summary['unreadable']} of {summary['n_windows']} windows produced an "
                     f"answer that could not be read**, and each counts as *clean*. Roughly one in "
                     f"seven is normal for this model; a much higher rate means this verdict rests "
                     f"on less evidence than the window count suggests.")

    if summary["injected"]:
        lines += ["", "## What was found", "",
                  "| region | characters | windows | family | agreement | evidence |",
                  "|---|---|---|---|---|---|"]
        for i, r in enumerate(summary["regions"], 1):
            ev = (r["evidence"] or "-")[:70].replace("\n", " ").replace("|", "\\|")
            where = f"@{r['evidence_offset']:,}" if r["evidence_offset"] is not None else ""
            lines.append(f"| {i} | {r['start']:,}-{r['end']:,} | {r['n_windows']} | "
                         f"`{r['family']}` | {r['family_agreement']:.0%} | `{ev}` {where} |")

        if summary["families"]:
            named = ", ".join(f"`{f}` x{n}" for f, n in summary["families"])
            lines += ["", f"**Families suggested:** {named}"]

        lines += ["", f"**Corroboration:** {summary['n_corroborated']} of "
                      f"{summary['n_regions']} regions were flagged in two or more overlapping "
                      f"windows. A region seen once is a weaker signal than one seen twice."]

        weak = [f for f, _ in summary["families"] if f in WEAK_FAMILIES]
        if weak:
            lines += ["", f"> The suggested family {', '.join('`'+w+'`' for w in weak)} is one this "
                          f"model **never names correctly** (0.00 in the evaluation). Treat the "
                          f"detection as real and the label as unreliable."]
    else:
        lines += ["", "Every window was scored and none reported an injection. This is a negative "
                      "result from a model with recall "
                      f"{m.get('recall', 0):.3f} on the evaluation corpus - it misses roughly "
                      f"{1 - m.get('recall', 0):.0%} of the attacks it is shown."]

    lines += ["", "---", "", "## How much this verdict is worth", "",
              f"On 1,100 held-out PDFs this model scored **F1 {m.get('f1', 0):.3f}**, against "
              f"**{BASELINE_F1:.3f} for a detector that flags everything without reading it**. "
              f"That is a {(m.get('f1', 0) / BASELINE_F1 - 1):.1%} relative improvement, not an "
              f"order of magnitude.",
              "",
              f"- **Precision {m.get('precision', 0):.3f}**, false alarms on "
              f"{m.get('false_alarm_rate', 0):.0%} of clean files.",
              f"- **Family named correctly {m.get('family_acc', 0):.3f} of the time** - the family "
              f"above is a suggestion, not a verdict.",
              "- Measured on synthetic injections of harmless test markers (EICAR/AMTSO/WICAR/"
              "RANSIM) into ordinary PDFs. It is **not** a general malware scanner.",
              ]
    return "\n".join(lines)