@murk-security/runner-linux-x64 0.0.0-stage → 0.7.1

package/runners/remediate.py

@@ -0,0 +1,429 @@
+#!/usr/bin/env python3
+"""Generate a code patch for one confirmed finding.
+
+The exploit/assessment harness is strictly READ-ONLY, so remediation is a separate,
+isolated step. Rather than a single one-shot model call, this generates a HIGH-QUALITY
+patch via a small pipeline:
+
+  1. CONTEXT: read the whole target file (not a tiny window) plus a few files it references,
+     within a budget, so the model sees enough to produce hunks that actually match.
+  2. BEST-OF-N candidates, via diverse strategies: a minimal unified diff (with a repair
+     retry that feeds back the exact `git apply` error), a full-file rewrite we diff
+     locally (robust against mis-formatted hunks), and a structural-fix variant. Only
+     candidates that pass `git apply --check` survive.
+  3. REVIEWER pass: a stronger model judges each surviving candidate — does it actually FIX
+     the vuln, is it MINIMAL, any REGRESSION risk — and we pick the best-reviewed one.
+
+The control plane takes the chosen validated diff and opens a PR via the GitHub App;
+nothing here writes to the repo or runs git beyond the read-only `apply --check`.
+
+Importable (the product runner calls generate_patch); model calls are injectable so the
+logic is testable without the network.
+"""
+
+from __future__ import annotations
+
+import difflib
+import json
+import os
+import re
+import subprocess
+import urllib.request
+from pathlib import Path
+
+OPENROUTER_URL = os.environ.get("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1").rstrip("/") + "/chat/completions"
+MAX_SOURCE_CHARS = 24000    # windowed view of the target file sent in the diff prompt
+MAX_CONTEXT_CHARS = 60000   # total context budget (target view + related files)
+MAX_FULLFILE_CHARS = 40000  # only attempt a full-file rewrite when the file is at/under this
+
+
+def load_key(key_file: str) -> str:
+    key = os.environ.get("OPENROUTER_API_KEY")
+    if key:
+        return key
+    try:
+        for line in Path(key_file).read_text().splitlines():
+            if line.startswith("OPENROUTER_API_KEY="):
+                return line.split("=", 1)[1].strip().strip("\"'")
+    except OSError:
+        pass
+    return ""
+
+
+# ── context ────────────────────────────────────────────────────────────────────────────
+
+def _referenced_files(checkout: str, file: str, text: str) -> list[str]:
+    """Best-effort: repo-local files the target references (Python imports, JS/TS relative
+    imports). Gives the model the blast radius, not just the one file. Bounded to a few."""
+    root = Path(checkout)
+    base = Path(file).parent
+    out: list[str] = []
+    for m in re.finditer(r"^\s*(?:from|import)\s+([.\w]+)", text, re.M):
+        mod = m.group(1).lstrip(".").split(".")[0]
+        if not mod:
+            continue
+        for cand in (base / f"{mod}.py", Path(f"{mod}.py"), base / mod / "__init__.py"):
+            if (root / cand).is_file():
+                out.append(str(cand))
+                break
+    for m in re.finditer(r"""(?:from|require\()\s*['"](\.[^'"]+)['"]""", text):
+        rel = (base / m.group(1)).as_posix()
+        for ext in ("", ".js", ".ts", ".jsx", ".tsx", "/index.js", "/index.ts"):
+            if (root / f"{rel}{ext}").is_file():
+                out.append(f"{rel}{ext}")
+                break
+    seen: set[str] = set()
+    res: list[str] = []
+    for p in out:
+        if p != file and p not in seen:
+            seen.add(p)
+            res.append(p)
+    return res[:3]
+
+
+def _gather_context(checkout: str, file: str, line: int | None) -> tuple[dict | None, str | None]:
+    """Read the target file (whole, or windowed around `line` if huge) plus a few referenced
+    files, confined to the repo root. Returns (context, error)."""
+    root = Path(checkout).resolve()
+    target = (root / file).resolve()
+    try:
+        target.relative_to(root)
+    except ValueError:
+        return None, "source path escapes the repository"
+    if not target.is_file():
+        return None, f"source file not found in the checkout: {file}"
+    full = target.read_text(errors="replace")
+    if len(full) > MAX_SOURCE_CHARS and line:
+        lines = full.splitlines()
+        lo, hi = max(0, line - 200), min(len(lines), line + 200)
+        view = "\n".join(lines[lo:hi])[:MAX_SOURCE_CHARS]
+    else:
+        view = full[:MAX_SOURCE_CHARS]
+    related: list[tuple[str, str]] = []
+    budget = max(0, MAX_CONTEXT_CHARS - len(view))
+    for rel in _referenced_files(checkout, file, full):
+        if budget <= 0:
+            break
+        p = (root / rel).resolve()
+        try:
+            p.relative_to(root)
+            t = p.read_text(errors="replace")
+        except (OSError, ValueError):
+            continue
+        t = t[:budget]
+        related.append((rel, t))
+        budget -= len(t)
+    return {"target_full": full, "target_view": view, "related": related}, None
+
+
+def _related_blob(ctx: dict) -> str:
+    return "".join(
+        f"\n\nRELATED FILE {p} (context only — do NOT edit it):\n```\n{t}\n```"
+        for p, t in ctx.get("related", [])
+    )
+
+
+# ── prompts ────────────────────────────────────────────────────────────────────────────
+
+_DIFF_SYSTEM = (
+    "You are a security remediation engineer. You are given a CONFIRMED vulnerability and the "
+    "current contents of the source file that contains it. Produce the MINIMAL code change that "
+    "fixes the vulnerability without altering unrelated behavior. Output EXACTLY two things and "
+    "nothing else:\n"
+    "1) a single fenced ```diff block containing a valid unified diff that `git apply -p1` accepts: "
+    "it MUST start with `diff --git a/<path> b/<path>`, use `--- a/<path>` / `+++ b/<path>` and `@@` "
+    "hunks, with paths relative to the repo root, and its context lines MUST match the file EXACTLY.\n"
+    "2) after the block, one line: `RATIONALE: <one sentence on why this fixes it>`.\n"
+    "Keep the diff tight; change only what the fix requires."
+)
+_STRUCTURAL_HINT = (
+    "\nPrefer a robust STRUCTURAL fix — cryptographic signing/verification, server-side validation, "
+    "parameterized/safe APIs, correct escaping, access checks — over blocklists or superficial guards."
+)
+
+
+def _diff_messages(finding: dict, file: str, line: int | None, ctx: dict, variant: int = 0) -> list[dict]:
+    system = _DIFF_SYSTEM + (_STRUCTURAL_HINT if variant else "")
+    user = (
+        f"FINDING: {finding.get('title', '')}\n"
+        f"DETAILS: {finding.get('description', '')}\n"
+        f"FILE: {file}" + (f" (around line {line})" if line else "") + "\n\n"
+        f"CURRENT CONTENTS OF {file}:\n```\n{ctx['target_view']}\n```" + _related_blob(ctx)
+    )
+    return [{"role": "system", "content": system}, {"role": "user", "content": user}]
+
+
+def _fullfile_messages(finding: dict, file: str, full: str) -> list[dict]:
+    system = (
+        "You are a security remediation engineer. Fix the vulnerability with the MINIMAL change and "
+        "do not alter unrelated behavior. Return the COMPLETE corrected contents of the file inside a "
+        "single fenced code block, and nothing else except a final line `RATIONALE: <one sentence>`."
+    )
+    user = (
+        f"FINDING: {finding.get('title', '')}\nDETAILS: {finding.get('description', '')}\n"
+        f"FILE: {file}\n\nCURRENT CONTENTS:\n```\n{full}\n```"
+    )
+    return [{"role": "system", "content": system}, {"role": "user", "content": user}]
+
+
+_REVIEW_SYSTEM = (
+    "You are a strict, skeptical security code reviewer. You are given a vulnerability, the original "
+    "file, and a proposed patch (unified diff). Judge the patch on three axes and respond with ONLY a "
+    "JSON object, no prose:\n"
+    '{"fixes": true|false, "minimal": true|false, "regression_risk": "none"|"low"|"med"|"high", '
+    '"score": <0-100>, "approved": true|false, "notes": "<=2 sentences"}\n'
+    "fixes = does it actually close THIS vulnerability (trace the fixed path). minimal = no unrelated "
+    "changes. approved = you would merge it. Be honest; a patch that applies but does not fix the vuln "
+    "is NOT approved."
+)
+
+
+def _review_messages(finding: dict, file: str, original: str, diff: str) -> list[dict]:
+    user = (
+        f"VULNERABILITY: {finding.get('title', '')}\nDETAILS: {finding.get('description', '')}\n"
+        f"FILE: {file}\n\nORIGINAL FILE:\n```\n{original[:MAX_SOURCE_CHARS]}\n```\n\n"
+        f"PROPOSED PATCH:\n```diff\n{diff}\n```"
+    )
+    return [{"role": "system", "content": _REVIEW_SYSTEM}, {"role": "user", "content": user}]
+
+
+# ── model + diff plumbing ────────────────────────────────────────────────────────────────
+
+def call_model(key: str, model: str, messages: list[dict], timeout: int = 150) -> str:
+    body = json.dumps({"model": model, "messages": messages, "temperature": 0}).encode()
+    req = urllib.request.Request(
+        OPENROUTER_URL, data=body, method="POST",
+        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
+    )
+    with urllib.request.urlopen(req, timeout=timeout) as resp:
+        data = json.loads(resp.read().decode("utf-8", "replace"))
+    return data["choices"][0]["message"]["content"]
+
+
+def extract_diff(text: str) -> str:
+    """Pull a unified diff out of the model output: prefer a fenced block, then trim to the
+    first real diff marker so surrounding prose never reaches `git apply`."""
+    m = re.search(r"```(?:diff|patch)?\s*\n(.*?)```", text or "", re.S)
+    body = m.group(1) if m else (text or "")
+    idx = None
+    for marker in ("diff --git", "--- a/", "--- "):
+        i = body.find(marker)
+        if i != -1 and (idx is None or i < idx):
+            idx = i
+    if idx is None:
+        return ""
+    return body[idx:].rstrip() + "\n"
+
+
+def extract_rationale(text: str) -> str:
+    m = re.search(r"RATIONALE:\s*(.+)", text or "")
+    return m.group(1).strip()[:500] if m else ""
+
+
+def _extract_code_block(text: str) -> str:
+    """Pull FULL-FILE contents out of a fenced block for the rewrite strategy. Skips
+    ```diff/```patch fences (those are diffs, not file contents)."""
+    for m in re.finditer(r"```([A-Za-z0-9_+-]*)\s*\n(.*?)```", text or "", re.S):
+        if m.group(1).lower() in ("diff", "patch"):
+            continue
+        return m.group(2)
+    return ""
+
+
+def _unified_diff(file: str, old: str, new: str) -> str:
+    """Build a git-apply-able unified diff between two full file contents, computed locally
+    (robust against a model mis-formatting hunk headers)."""
+    a = old.splitlines(keepends=True)
+    b = new.splitlines(keepends=True)
+    body = "".join(difflib.unified_diff(a, b, fromfile=f"a/{file}", tofile=f"b/{file}", n=3))
+    if not body.strip():
+        return ""
+    return f"diff --git a/{file} b/{file}\n{body}"
+
+
+def validate_diff(checkout: str, diff: str) -> tuple[bool, str | None]:
+    """Does the diff apply cleanly to the checkout? Read-only (`--check`): applies nothing."""
+    if not diff.strip():
+        return False, "no diff produced"
+    try:
+        p = subprocess.run(
+            ["git", "apply", "--check", "-p1", "-"], input=diff, text=True,
+            cwd=checkout, capture_output=True, timeout=30,
+        )
+        return (p.returncode == 0), (p.stderr.strip() or None)
+    except Exception as e:  # noqa: BLE001
+        return False, f"git apply --check failed to run: {e}"
+
+
+# ── candidates ─────────────────────────────────────────────────────────────────────────
+
+def _diff_candidate(checkout: str, finding: dict, file: str, line: int | None, ctx: dict,
+                    key: str, model: str, model_fn, max_attempts: int, variant: int) -> dict | None:
+    """One diff-strategy candidate, with a repair retry that feeds back the git-apply error.
+    Returns the applying candidate if found, else the last non-applying one, else None."""
+    messages = _diff_messages(finding, file, line, ctx, variant)
+    last: dict | None = None
+    for attempt in range(max(1, max_attempts)):
+        try:
+            out = model_fn(key, model, messages)
+        except Exception:  # noqa: BLE001
+            return last
+        diff = extract_diff(out)
+        if diff:
+            applies, aerr = validate_diff(checkout, diff)
+            cand = {"diff": diff, "rationale": extract_rationale(out), "applies": applies,
+                    "strategy": "diff", "error": None if applies else (aerr or "did not apply cleanly")}
+            if applies:
+                return cand
+            last = cand
+            feedback = (f"That diff did not apply (`git apply --check`: {aerr}). Produce a corrected "
+                        "unified diff whose @@ hunk headers and context lines EXACTLY match the file "
+                        "shown above. Output only the ```diff block and the RATIONALE line.")
+        else:
+            feedback = ("You did not return a unified diff. Output only a single ```diff block that "
+                        "`git apply -p1` accepts, then the RATIONALE line.")
+        if attempt + 1 < max(1, max_attempts):
+            messages = messages + [{"role": "assistant", "content": out},
+                                   {"role": "user", "content": feedback}]
+    return last
+
+
+def _fullfile_candidate(checkout: str, finding: dict, file: str, full: str,
+                        key: str, model: str, model_fn) -> dict | None:
+    """Full-file-rewrite strategy: ask for the whole corrected file, diff it ourselves. Robust
+    against hunk-format mistakes. Only for files small enough to rewrite without truncation."""
+    if not full or len(full) > MAX_FULLFILE_CHARS:
+        return None
+    try:
+        out = model_fn(key, model, _fullfile_messages(finding, file, full))
+    except Exception:  # noqa: BLE001
+        return None
+    new = _extract_code_block(out)
+    if not new or new.strip() == full.strip():
+        return None
+    if re.match(r"\s*(diff --git|--- |\+\+\+ |@@ )", new):
+        return None  # a diff slipped through, not a file rewrite
+    diff = _unified_diff(file, full, new)
+    if not diff:
+        return None
+    applies, aerr = validate_diff(checkout, diff)
+    return {"diff": diff, "rationale": extract_rationale(out), "applies": applies,
+            "strategy": "full-file", "error": None if applies else (aerr or "did not apply cleanly")}
+
+
+def generate_candidates(checkout: str, finding: dict, file: str, line: int | None, ctx: dict,
+                        key: str, model: str, model_fn, n: int, max_attempts: int) -> tuple[list[dict], list[dict]]:
+    """Produce up to `n` distinct candidate patches via diverse strategies. Returns
+    (applying_unique, all_produced)."""
+    cands: list[dict] = []
+    seen: set[str] = set()
+
+    def add(c: dict | None) -> None:
+        if c and c.get("diff") and c["diff"] not in seen:
+            seen.add(c["diff"])
+            cands.append(c)
+
+    add(_diff_candidate(checkout, finding, file, line, ctx, key, model, model_fn, max_attempts, variant=0))
+    add(_fullfile_candidate(checkout, finding, file, ctx["target_full"], key, model, model_fn))
+    variant = 1
+    while len([c for c in cands if c.get("applies")]) < n and variant <= 2:
+        add(_diff_candidate(checkout, finding, file, line, ctx, key, model, model_fn, 1, variant=variant))
+        variant += 1
+    applying = [c for c in cands if c.get("applies")]
+    return applying, cands
+
+
+# ── reviewer ──────────────────────────────────────────────────────────────────────────
+
+def _parse_review(text: str) -> dict | None:
+    m = re.search(r"\{.*\}", text or "", re.S)
+    if not m:
+        return None
+    try:
+        v = json.loads(m.group(0))
+    except (json.JSONDecodeError, TypeError):
+        return None
+    if not isinstance(v, dict):
+        return None
+    score = v.get("score", 0)
+    try:
+        score = max(0, min(100, int(score)))
+    except (ValueError, TypeError):
+        score = 0
+    return {
+        "fixes": bool(v.get("fixes")),
+        "minimal": bool(v.get("minimal", True)),
+        "regression_risk": str(v.get("regression_risk", "unknown"))[:16],
+        "score": score,
+        "approved": bool(v.get("approved", v.get("fixes"))),
+        "notes": str(v.get("notes", ""))[:400],
+    }
+
+
+def review_patch(finding: dict, file: str, original: str, diff: str,
+                 key: str, reviewer_model: str, reviewer_fn) -> dict | None:
+    try:
+        out = reviewer_fn(key, reviewer_model, _review_messages(finding, file, original, diff))
+    except Exception:  # noqa: BLE001
+        return None
+    return _parse_review(out)
+
+
+# ── orchestration ────────────────────────────────────────────────────────────────────────
+
+def _failed(error: str, file: str | None = None) -> dict:
+    r = {"status": "failed", "error": error}
+    if file:
+        r["targetFile"] = file
+    return r
+
+
+def generate_patch(checkout: str, finding: dict, key: str, model: str, model_fn=call_model,
+                   reviewer_model: str | None = None, reviewer_fn=None,
+                   candidates: int = 3, max_attempts: int = 2) -> dict:
+    """Produce {diff, rationale, targetFile, applies, status, ...} for one finding.
+
+    Pipeline: gather context -> generate best-of-N candidates (diverse strategies, each
+    validated with `git apply --check`) -> optional reviewer pass picks the best. status is
+    'proposed' when a diff was produced (whether or not it applies), 'failed' when none was.
+    When a reviewer runs, the result carries `review` and `verified`."""
+    src = finding.get("source") or {}
+    file = src.get("file")
+    if not file:
+        return _failed("finding has no source file to patch")
+    ctx, err = _gather_context(checkout, file, src.get("line"))
+    if err or ctx is None:
+        return _failed(err or "could not read source", file)
+
+    applying, produced = generate_candidates(
+        checkout, finding, file, src.get("line"), ctx, key, model, model_fn, candidates, max_attempts)
+
+    if not applying:
+        last = next((c for c in produced if c.get("diff")), None)
+        if last:
+            return {"diff": last["diff"], "rationale": last.get("rationale", ""), "targetFile": file,
+                    "applies": False, "status": "proposed",
+                    "error": last.get("error") or "patch did not apply cleanly after retries"}
+        return _failed("model did not return a unified diff", file)
+
+    reviewer_fn = reviewer_fn or model_fn
+    chosen = applying[0]
+    if reviewer_model:
+        for c in applying:
+            c["review"] = review_patch(finding, file, ctx["target_full"], c["diff"], key, reviewer_model, reviewer_fn)
+        reviewed = [c for c in applying if c.get("review")]
+        if reviewed:
+            approved = [c for c in reviewed if c["review"]["approved"]]
+            pool = approved or reviewed
+            chosen = max(pool, key=lambda c: c["review"]["score"])
+
+    result = {
+        "diff": chosen["diff"], "rationale": chosen.get("rationale", ""), "targetFile": file,
+        "applies": True, "status": "proposed",
+        "strategy": chosen.get("strategy"), "candidates": len(applying),
+    }
+    review = chosen.get("review")
+    if review is not None:
+        result["review"] = review
+        result["verified"] = bool(review["approved"])
+    return result