pwnguard 0.2.4__py3-none-any.whl

pwnguard/__init__.py

@@ -0,0 +1,8 @@
+"""PwnGuard: AI-powered security review for git workflows.
+
+Implementation modules. The legacy ``audit`` module at the repo root
+re-exports the public surface so existing callers (the pre-commit
+hook, the README CLI examples, the test suite) keep working unchanged.
+"""
+
+__version__ = "0.2.4"

pwnguard/__main__.py

@@ -0,0 +1,4 @@
+from pwnguard.cli import main
+
+if __name__ == "__main__":
+    main()

pwnguard/anchors.py

@@ -0,0 +1,192 @@
+"""Opaque anchor token tagging and resolution.
+
+Replaces fuzzy line-number / snippet matching: each ``+`` (added) and
+context (`` ``) line in the diff is prefixed with ``[a<N>]``. The
+model is told to echo the bare token back in each finding's
+``anchor`` field; the host resolves the token via the table built
+during tagging.
+
+Why opaque tokens beat plain line numbers: they have no numeric
+semantics the model can drift on, so they can only be copied
+verbatim, not regenerated from "where this probably is in the file."
+"""
+
+import re
+from typing import Optional
+
+from pwnguard.constants import DIFF_WRAPPER_OPEN, DIFF_WRAPPER_CLOSE
+
+
+def wrap_diff(diff: str) -> tuple:
+    """Wrap diff in delimiters, tag content lines with anchor tokens, and
+    return both the wrapped text and an anchor lookup table.
+
+    Each ``+`` (added) and context (`` ``) content line is prefixed
+    with an opaque token of the form ``[a<N>]`` (e.g. ``[a1]``,
+    ``[a42]``). The model is told to echo the bare token back in each
+    finding's ``anchor`` field; the returned table maps each token to
+    ``(file, line, content, kind, hunk_context)`` so post-parse
+    resolution is a single dict lookup. No counting, no fuzzy quote
+    matching, no function-name regex fallback.
+
+    Opaque tokens beat the old 5-digit line-number prefix because
+    they have no numeric semantics the model can drift on - they can
+    only be copied verbatim, not regenerated from "where this
+    probably is in the file". The namespace resets per call, which
+    is fine because each AI request is parsed against its own table.
+
+    Diff metadata (file headers, hunk headers, removed lines) is left
+    untagged: those don't correspond to a new-file position.
+
+    Returns ``(wrapped_text, anchor_table)`` where ``anchor_table`` is
+    a ``dict[str, dict]`` keyed by the bare token (no brackets).
+    """
+    body, anchors = _anchor_diff_lines(diff)
+    wrapped = f"{DIFF_WRAPPER_OPEN}\n{body}\n{DIFF_WRAPPER_CLOSE}"
+    return wrapped, anchors
+
+
+def _anchor_diff_lines(diff: str) -> tuple:
+    """Walk the diff once, emitting ``[a<N>]`` tokens and building the
+    anchor table.
+
+    Replaces the old ``_number_diff_lines`` line-number prefixer. The
+    token namespace is a simple incrementing counter (``a1``, ``a2``,
+    ...) reset on every call. Removed (``-``) lines, file headers,
+    hunk headers, and the ``diff --git`` line stay untagged - the
+    model can't anchor a finding to them, so giving them a token
+    would only invite hallucinated references.
+    """
+    out = []
+    anchors = {}
+    next_id = 1
+    current_file: Optional[str] = None
+    current_lineno = 0
+    current_hunk_context: Optional[str] = None
+
+    # ``git diff`` always emits a trailing newline; ``split("\n")`` then
+    # yields a trailing "" that isn't a real content line. Drop it so we
+    # don't tag a phantom context anchor past the last real line (which
+    # the model could then pick, resolving to an empty-content row).
+    lines = diff.split("\n")
+    if lines and lines[-1] == "":
+        lines = lines[:-1]
+
+    for line in lines:
+        if line.startswith("+++ b/"):
+            current_file = line[6:]
+            current_lineno = 0
+            current_hunk_context = None
+            out.append(line)
+            continue
+        if line.startswith("@@"):
+            m = re.match(
+                r"@@ -\d+(?:,\d+)? \+(\d+)(?:,\d+)? @@\s*(.*)",
+                line,
+            )
+            if m:
+                current_lineno = int(m.group(1)) - 1
+                ctx = m.group(2).strip()
+                current_hunk_context = ctx or None
+            out.append(line)
+            continue
+        # Until we've seen a `+++ b/<file>` header we don't know what
+        # file a line belongs to. Skip tagging in that state - emitting
+        # tokens with file="" would let the model anchor findings to a
+        # blank file and produce findings with no path / no preview
+        # (the typical symptom of feeding a non-diff to --diff-file).
+        if current_file is None:
+            out.append(line)
+            continue
+        if line.startswith("+") and not line.startswith("+++"):
+            current_lineno += 1
+            tok = f"a{next_id}"
+            next_id += 1
+            anchors[tok] = {
+                "file": current_file,
+                "line": current_lineno,
+                "content": line[1:],
+                "kind": "added",
+                "hunk_context": current_hunk_context,
+            }
+            out.append(f"[{tok}] {line}")
+            continue
+        if line.startswith(" ") or line == "":
+            current_lineno += 1
+            tok = f"a{next_id}"
+            next_id += 1
+            anchors[tok] = {
+                "file": current_file,
+                "line": current_lineno,
+                "content": line[1:] if line else "",
+                "kind": "context",
+                "hunk_context": current_hunk_context,
+            }
+            out.append(f"[{tok}] {line}")
+            continue
+        # `-` removed lines, `---` headers, `diff --git ...`, etc.
+        out.append(line)
+
+    return "\n".join(out), anchors
+
+
+def resolve_anchors(result, anchor_table: dict) -> int:
+    """Resolve each finding's / observation's ``anchor`` token to a real
+    location, rewriting ``file``, ``line``, and ``hunk_context``.
+
+    Returns the number of items dropped because their anchor was
+    unknown - "loud failure" rather than fuzzy recovery. The host
+    should surface the count on stderr so the user sees that a
+    fabricated anchor occurred.
+
+    Resolution rules per finding/observation:
+
+    - Anchor present AND in table: rewrite ``file`` / ``line`` /
+      ``hunk_context`` from the table entry. Trust the table.
+    - Anchor present but NOT in table: model fabricated an anchor.
+      Drop the item; do not try fuzzy matching - that's exactly the
+      pre-anchor fallback chain we're replacing.
+    - Anchor absent: file-level carve-out. Keep the item, leave
+      ``file`` as the model supplied it, ``line`` stays None.
+
+    The previous pipeline (``_anchor_findings_by_snippet`` etc.) is
+    gone; one O(1) lookup replaces the whole repair chain.
+    """
+    dropped = 0
+    kept_findings = []
+    for f in result.findings:
+        if f.anchor is None:
+            # File-level carve-out: model deliberately omitted anchor
+            # and supplied "file" directly. Keep if file is set;
+            # otherwise drop (no way to render it).
+            if f.file:
+                kept_findings.append(f)
+            else:
+                dropped += 1
+            continue
+        entry = anchor_table.get(f.anchor)
+        if not entry:
+            dropped += 1
+            continue
+        f.file = entry["file"]
+        f.line = entry["line"]
+        f.hunk_context = entry.get("hunk_context")
+        kept_findings.append(f)
+    result.findings = kept_findings
+
+    kept_obs = []
+    for o in result.observations:
+        if o.anchor is None:
+            kept_obs.append(o)
+            continue
+        entry = anchor_table.get(o.anchor)
+        if not entry:
+            # Drop the observation silently - we don't count these in
+            # the "dropped" tally because they're opt-in informational.
+            continue
+        o.file = entry["file"]
+        o.line = entry["line"]
+        kept_obs.append(o)
+    result.observations = kept_obs
+
+    return dropped

pwnguard/backends/__init__.py

@@ -0,0 +1,73 @@
+"""AI backend dispatch.
+
+Each backend lives in its own submodule (one ``query_X`` function plus
+any backend-specific helpers). ``dispatch_backend`` is the single
+entry point used by ``scan`` and ``monitor`` so the wrap-diff /
+anchor-table contract is enforced in one place.
+"""
+
+from typing import Optional
+
+from pwnguard import runtime
+from pwnguard.anchors import wrap_diff
+from pwnguard.prompts import build_system_prompt
+
+from pwnguard.backends.claude_api import query_claude_api
+from pwnguard.backends.claude_code import claude_code_available, query_claude_code
+from pwnguard.backends.ollama import query_ollama
+from pwnguard.backends.openai_compat import query_openai_compat
+
+
+__all__ = [
+    "claude_code_available",
+    "dispatch_backend",
+    "query_claude_api",
+    "query_claude_code",
+    "query_ollama",
+    "query_openai_compat",
+]
+
+
+def dispatch_backend(
+    backend: str,
+    diff: str,
+    config: dict,
+    system_prompt: Optional[str] = None,
+    pre_wrapped: bool = False,
+) -> tuple:
+    """Run the requested backend. Centralizes the dispatch logic.
+
+    Wraps ``diff`` (assigning anchor tokens) and ships only the
+    wrapped text to the backend; returns ``(response, anchor_table)``
+    so the caller can resolve each finding's ``anchor`` field with a
+    single dict lookup.
+
+    Set ``pre_wrapped=True`` when the caller has already embedded a
+    wrapped diff inside its own custom ``system_prompt`` (the
+    --explain path is the only current case). In that mode the
+    anchor table comes back empty - explain is a re-query for one
+    already-resolved finding, so it doesn't need anchors.
+
+    When no system_prompt is given, builds one from the current code-preview
+    setting: if the caller won't render fix_example, we don't ask the model
+    to generate it (saves a few hundred tokens of prompt + output time on
+    7B local models).
+    """
+    if system_prompt is None:
+        system_prompt = build_system_prompt(
+            include_preview_fields=runtime.show_code_preview,
+            include_observations=runtime.show_observations,
+        )
+    if pre_wrapped:
+        wrapped, anchors = diff, {}
+    else:
+        wrapped, anchors = wrap_diff(diff)
+    if backend == "claude-api":
+        response = query_claude_api(wrapped, config, system_prompt)
+    elif backend == "claude-code":
+        response = query_claude_code(wrapped, config, system_prompt)
+    elif backend == "openai-compat":
+        response = query_openai_compat(wrapped, config, system_prompt)
+    else:
+        response = query_ollama(wrapped, config, system_prompt)
+    return response, anchors

pwnguard/backends/claude_api.py

@@ -0,0 +1,74 @@
+"""Anthropic API backend (requires ``ANTHROPIC_API_KEY``).
+
+This is the only paid path PwnGuard supports directly, so the
+"large prompt" confirmation gate lives here.
+"""
+
+import os
+import sys
+
+from pwnguard import ui
+from pwnguard.constants import LARGE_PROMPT_TOKEN_THRESHOLD
+from pwnguard.diff import estimate_tokens
+from pwnguard.prompts import SYSTEM_PROMPT
+
+
+def maybe_confirm_large_prompt(prompt: str, backend: str) -> None:
+    """Warn before sending a very large prompt to a paid backend.
+
+    Only prompts for confirmation on the claude-api backend (the one with
+    direct per-token cost) and only when the prompt clearly crosses our
+    arbitrary "this is big" threshold. The user can disable the prompt by
+    setting PWNGUARD_NO_PROMPT=1 (e.g. for non-interactive scripts).
+    """
+    if backend != "claude-api":
+        return
+    if os.environ.get("PWNGUARD_NO_PROMPT") == "1":
+        return
+    tokens = estimate_tokens(prompt)
+    if tokens < LARGE_PROMPT_TOKEN_THRESHOLD:
+        return
+    if not sys.stdin.isatty():
+        # Non-interactive (CI) - log the size but don't block.
+        print(
+            ui.dim(f"PwnGuard: large prompt (~{tokens:,} tokens)."),
+            file=sys.stderr,
+        )
+        return
+    print(
+        f"\nPwnGuard: estimated ~{tokens:,} input tokens for this scan.",
+        file=sys.stderr,
+    )
+    answer = input("Send to claude-api anyway? [y/N] ").strip().lower()
+    if answer != "y":
+        sys.exit("Aborted by user.")
+
+
+def query_claude_api(diff: str, config: dict, system_prompt: str = SYSTEM_PROMPT) -> str:
+    """Send diff to Claude API for analysis (requires ANTHROPIC_API_KEY)."""
+    try:
+        import anthropic
+    except ImportError:
+        sys.exit("Error: 'anthropic' package not installed. Run: pip install anthropic")
+
+    api_key = os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        sys.exit("Error: ANTHROPIC_API_KEY environment variable not set")
+
+    client = anthropic.Anthropic(api_key=api_key)
+    claude_config = config.get("claude_api", {})
+
+    user_content = f"Review this git diff for security vulnerabilities:\n\n{diff}"
+    maybe_confirm_large_prompt(system_prompt + user_content, backend="claude-api")
+
+    message = client.messages.create(
+        model=claude_config.get("model", "claude-opus-4-7"),
+        max_tokens=claude_config.get("max_tokens", 4096),
+        system=system_prompt,
+        messages=[{"role": "user", "content": user_content}],
+    )
+
+    # Guard against an empty content list (refusals, safety stops).
+    if not message.content:
+        return '{"findings": []}'
+    return message.content[0].text

pwnguard/backends/claude_code.py

@@ -0,0 +1,73 @@
+"""Claude Code CLI backend (uses the user's Pro subscription).
+
+The CLI is shelled out and the prompt is fed via stdin. Detection is
+cached so we don't probe ``claude --version`` on every invocation.
+"""
+
+import subprocess
+import sys
+from typing import Optional
+
+from pwnguard.prompts import SYSTEM_PROMPT
+
+
+# Cache so we don't shell out to `claude --version` more than once per run.
+_claude_code_available: Optional[bool] = None
+
+
+def claude_code_available() -> bool:
+    """Detect whether the `claude` CLI is installed. Cached after first call."""
+    global _claude_code_available
+    if _claude_code_available is not None:
+        return _claude_code_available
+    try:
+        result = subprocess.run(
+            ["claude", "--version"],
+            capture_output=True, text=True, timeout=5,
+        )
+        _claude_code_available = result.returncode == 0
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        _claude_code_available = False
+    return _claude_code_available
+
+
+def query_claude_code(diff: str, config: dict, system_prompt: str = SYSTEM_PROMPT) -> str:
+    """Send diff to Claude Code CLI for analysis (uses Pro subscription)."""
+    if not claude_code_available():
+        sys.exit(
+            "Error: 'claude' CLI not found. "
+            "Install Claude Code: https://docs.anthropic.com/en/docs/claude-code"
+        )
+
+    cc_config = config.get("claude_code", {})
+    timeout = cc_config.get("timeout", 120)
+
+    # Combine system prompt and user prompt for -p mode. ``diff`` arrives
+    # pre-wrapped (in <diff_to_review>...</diff_to_review> with anchor
+    # tokens) from dispatch_backend; the system prompt's input-format
+    # rules already cover that envelope.
+    full_prompt = (
+        f"{system_prompt}\n\n"
+        f"Review this git diff for security vulnerabilities:\n\n{diff}"
+    )
+
+    try:
+        result = subprocess.run(
+            ["claude", "-p", "--output-format", "text"],
+            input=full_prompt,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+        )
+    except subprocess.TimeoutExpired:
+        sys.exit(f"Error: Claude Code timed out after {timeout}s")
+    except FileNotFoundError:
+        sys.exit("Error: 'claude' command not found in PATH")
+
+    if result.returncode != 0:
+        msg = f"Error: Claude Code returned exit code {result.returncode}"
+        if result.stderr:
+            msg += f"\nstderr: {result.stderr[:500]}"
+        sys.exit(msg)
+
+    return result.stdout

pwnguard/backends/ollama.py

@@ -0,0 +1,201 @@
+"""Local Ollama backend.
+
+Refuses to ship the diff to a non-loopback host unless the user has
+explicitly opted in (``ollama.allow_remote: true``) - otherwise an
+attacker-controlled ``pwnguard.yaml`` could redirect every diff to
+an external endpoint.
+"""
+
+import json
+import socket
+import sys
+import urllib.error
+import urllib.parse
+import urllib.request
+
+from pwnguard import runtime, ui
+from pwnguard.constants import SAFE_OLLAMA_HOSTS
+from pwnguard.prompts import SYSTEM_PROMPT
+from pwnguard.security import _sanitize
+
+
+def query_ollama(diff: str, config: dict, system_prompt: str = SYSTEM_PROMPT) -> str:
+    """Send diff to local Ollama instance for analysis."""
+    ollama_config = config.get("ollama", {})
+    url = ollama_config.get("url", "http://localhost:11434")
+    model = ollama_config.get("model", "qwen2.5-coder:7b")
+    allow_remote = ollama_config.get("allow_remote", False)
+    timeout = ollama_config.get("timeout", 600)
+
+    # Refuse to send the diff to a non-local host unless explicitly opted in.
+    # Otherwise a committed pwnguard.yaml pointing ollama.url at an external
+    # endpoint would silently exfiltrate every diff at the next CI run.
+    host = urllib.parse.urlparse(url).hostname or ""
+    if host not in SAFE_OLLAMA_HOSTS and not allow_remote:
+        sys.exit(
+            f"Refusing to send diff to non-local Ollama host: {host!r}.\n"
+            f"Set ollama.allow_remote: true in pwnguard.yaml to override "
+            f"(you accept that diffs leave the local machine)."
+        )
+
+    payload_dict = {
+        "model": model,
+        "stream": False,
+        "messages": [
+            {"role": "system", "content": system_prompt},
+            {
+                "role": "user",
+                "content": f"Review this git diff for security vulnerabilities:\n\n{diff}",
+            },
+        ],
+    }
+    # Ollama's "format": "json" mode constrains every generated token to
+    # produce valid JSON. Reliable but slow (~2x on 7B). When the user
+    # opts out, parse_response's regex-extract + escape-fix fallbacks
+    # are what catch any non-JSON preamble / markdown fences.
+    if runtime.ollama_json_mode:
+        payload_dict["format"] = "json"
+
+    # Forward optional model tunables. Each is opt-in via config so the
+    # default behaviour matches whatever the model's own defaults are.
+    keep_alive = ollama_config.get("keep_alive")
+    if keep_alive:
+        payload_dict["keep_alive"] = keep_alive
+    options = {}
+    for opt_key in ("num_ctx", "num_predict", "temperature", "seed"):
+        if opt_key in ollama_config:
+            options[opt_key] = ollama_config[opt_key]
+    if options:
+        payload_dict["options"] = options
+
+    # Debug mode uses streaming so the user sees the model's output as it
+    # arrives. Non-debug mode keeps the single-blob request (works with
+    # the spinner and stays quieter for normal use).
+    payload_dict["stream"] = bool(runtime.debug_mode)
+    payload = json.dumps(payload_dict).encode()
+
+    req = urllib.request.Request(
+        f"{url}/api/chat",
+        data=payload,
+        headers={"Content-Type": "application/json"},
+    )
+
+    try:
+        if runtime.debug_mode:
+            return _query_ollama_stream(req, timeout, url)
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            data = json.loads(resp.read().decode())
+    except (urllib.error.URLError, socket.timeout, TimeoutError) as e:
+        msg = str(e)
+        if "timed out" in msg.lower():
+            sys.exit(
+                f"Ollama request timed out after {timeout}s. For large "
+                f"diffs, raise ollama.timeout in pwnguard.yaml or switch "
+                f"to --backend claude-code (much faster on large inputs)."
+            )
+        sys.exit(
+            f"Error: cannot reach Ollama at {url}: {e}\n"
+            f"Is Ollama running (ollama serve)?"
+        )
+
+    # Ollama returns {"message": {"content": ...}} on success, but error
+    # responses or unexpected shapes would otherwise raise a bare KeyError.
+    message = data.get("message") if isinstance(data, dict) else None
+    if not isinstance(message, dict) or "content" not in message:
+        err = data.get("error") if isinstance(data, dict) else None
+        # Sanitize before printing: an Ollama error blob is attacker-shaped
+        # data and would otherwise inject ANSI escapes straight into the
+        # user's terminal when sys.exit prints the message.
+        snippet = _sanitize(err or str(data)[:300]) or "(empty)"
+        sys.exit(f"Error: unexpected Ollama response: {snippet}")
+    return message["content"]
+
+
+def _query_ollama_stream(req: urllib.request.Request, timeout: int, url: str) -> str:
+    """Stream Ollama's chat response, echoing each token chunk to stderr.
+
+    Used in --debug mode so the user can watch the model's output in
+    real time and spot truncation, refusals, or stalls. Returns the
+    full accumulated content string after the stream is done so the
+    rest of the pipeline (parse_response, etc.) can treat it the same
+    as a non-streamed response.
+
+    Shows a "waiting for first token" spinner during prompt processing
+    so the user sees something is happening (large diffs on local 7B
+    models can sit silent for many seconds while Ollama tokenises and
+    runs prompt eval before any output token is emitted).
+    """
+    full_content = ""
+    final_meta: dict = {}
+    waiting = ui.Spinner("Waiting for response (ollama)")
+    waiting.__enter__()
+    first_token_seen = False
+    any_chunk_seen = False
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            for raw_line in resp:
+                line = raw_line.strip()
+                if not line:
+                    continue
+                try:
+                    chunk = json.loads(line.decode())
+                except (json.JSONDecodeError, UnicodeDecodeError):
+                    continue
+                # First sign of life from Ollama: prompt eval is running.
+                # Update the label so the user sees activity rather than
+                # a static "waiting" message while the server processes.
+                if not any_chunk_seen:
+                    any_chunk_seen = True
+                    waiting.label = "Model responding (ollama)"
+                content = chunk.get("message", {}).get("content", "")
+                if content:
+                    if not first_token_seen:
+                        first_token_seen = True
+                        waiting.__exit__(None, None, None)
+                        print(ui.dim("--- begin model output ---"), file=sys.stderr)
+                    full_content += content
+                    sys.stderr.write(ui.dim(content))
+                    sys.stderr.flush()
+                if chunk.get("done"):
+                    final_meta = chunk
+                    break
+    except (urllib.error.URLError, socket.timeout, TimeoutError) as e:
+        if not first_token_seen:
+            waiting.__exit__(None, None, None)
+        msg = str(e)
+        if "timed out" in msg.lower():
+            sys.exit(
+                f"\nOllama request timed out after {timeout}s. For large "
+                f"diffs, raise ollama.timeout in pwnguard.yaml or switch "
+                f"to --backend claude-code."
+            )
+        sys.exit(f"\nError: cannot reach Ollama at {url}: {e}")
+    finally:
+        # Stream ended without ever producing a token (server closed
+        # cleanly, empty response, etc.): make sure the spinner thread
+        # is stopped so the program can exit.
+        if not first_token_seen:
+            waiting.__exit__(None, None, None)
+
+    sys.stderr.write("\n")
+    print(ui.dim("--- end model output ---"), file=sys.stderr)
+    # Surface Ollama's per-request metrics when present (handy diagnostic
+    # for why a run was slow or stopped early).
+    if final_meta:
+        eval_count = final_meta.get("eval_count")
+        eval_duration = final_meta.get("eval_duration")
+        prompt_eval = final_meta.get("prompt_eval_count")
+        done_reason = final_meta.get("done_reason")
+        bits = []
+        if prompt_eval:
+            bits.append(f"prompt: {prompt_eval} tokens")
+        if eval_count:
+            bits.append(f"output: {eval_count} tokens")
+        if eval_count and eval_duration:
+            t_per_s = eval_count / (eval_duration / 1e9)
+            bits.append(f"{t_per_s:.1f} t/s")
+        if done_reason:
+            bits.append(f"stop: {done_reason}")
+        if bits:
+            print(ui.dim("PwnGuard: " + "  ·  ".join(bits)), file=sys.stderr)
+    return full_content