security hardening: F-002/F-005/F-007/F-008/F-009/F-010/F-016/F-031/F-035/F-038/PR747-NEW-2
This commit is contained in:
+32
-5
@@ -1735,17 +1735,33 @@ def main() -> None:
|
||||
|
||||
elif cmd == "merge-driver":
|
||||
# git merge driver for graph.json — takes (base, current, other) and writes
|
||||
# the union of current+other nodes/edges back to current. Always exits 0
|
||||
# so git never marks graph.json as conflicted.
|
||||
# the union of current+other nodes/edges back to current. Exits 1 on
|
||||
# corrupt input so git surfaces the conflict instead of silently
|
||||
# accepting a poisoned merge (see F-005).
|
||||
# Usage: graphify merge-driver %O %A %B (set in .git/config merge driver)
|
||||
if len(sys.argv) < 5:
|
||||
print("Usage: graphify merge-driver <base> <current> <other>", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
_base_path, _current_path, _other_path = sys.argv[2], sys.argv[3], sys.argv[4]
|
||||
# Hard caps so a malicious or corrupted graph.json cannot exhaust memory
|
||||
# at parse time. 50 MB / 100k nodes are well above any realistic graph
|
||||
# (typical graphs are <5 MB / <50k nodes); anything larger should fail
|
||||
# the merge so a human can investigate.
|
||||
_MERGE_MAX_BYTES = 50 * 1024 * 1024
|
||||
_MERGE_MAX_NODES = 100_000
|
||||
import networkx as _nx
|
||||
from networkx.readwrite import json_graph as _jg
|
||||
def _load_graph(p: str):
|
||||
data = json.loads(Path(p).read_text(encoding="utf-8"))
|
||||
path_obj = Path(p)
|
||||
try:
|
||||
size = path_obj.stat().st_size
|
||||
except OSError as exc:
|
||||
raise RuntimeError(f"cannot stat {p}: {exc}") from exc
|
||||
if size > _MERGE_MAX_BYTES:
|
||||
raise RuntimeError(
|
||||
f"graph.json {p} is {size} bytes, exceeds {_MERGE_MAX_BYTES}-byte cap"
|
||||
)
|
||||
data = json.loads(path_obj.read_text(encoding="utf-8"))
|
||||
try:
|
||||
return _jg.node_link_graph(data, edges="links"), data
|
||||
except TypeError:
|
||||
@@ -1755,8 +1771,15 @@ def main() -> None:
|
||||
G_oth, _ = _load_graph(_other_path)
|
||||
except Exception as exc:
|
||||
print(f"[graphify merge-driver] error loading graphs: {exc}", file=sys.stderr)
|
||||
sys.exit(0) # exit 0 so git doesn't block the merge
|
||||
sys.exit(1) # surface the conflict so git doesn't accept a corrupt merge
|
||||
merged = _nx.compose(G_cur, G_oth)
|
||||
if merged.number_of_nodes() > _MERGE_MAX_NODES:
|
||||
print(
|
||||
f"[graphify merge-driver] merged graph has {merged.number_of_nodes()} nodes, "
|
||||
f"exceeds {_MERGE_MAX_NODES}-node cap; aborting merge.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
try:
|
||||
out_data = _jg.node_link_data(merged, edges="links")
|
||||
except TypeError:
|
||||
@@ -1839,6 +1862,7 @@ def main() -> None:
|
||||
print(" svg [--graph PATH] [--labels PATH]", file=sys.stderr)
|
||||
print(" graphml [--graph PATH]", file=sys.stderr)
|
||||
print(" neo4j [--graph PATH] [--push URI] [--user U] [--password P]", file=sys.stderr)
|
||||
print(" (or set NEO4J_PASSWORD instead of --password to keep it off argv)", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Parse shared args
|
||||
@@ -1851,7 +1875,10 @@ def main() -> None:
|
||||
obsidian_dir = Path(_GRAPHIFY_OUT) / "obsidian"
|
||||
neo4j_uri: str | None = None
|
||||
neo4j_user = "neo4j"
|
||||
neo4j_password: str | None = None
|
||||
# F-031: prefer the NEO4J_PASSWORD env var so the password never
|
||||
# appears on argv (visible in `ps` output / shell history). The
|
||||
# explicit --password flag still overrides it for compatibility.
|
||||
neo4j_password: str | None = os.environ.get("NEO4J_PASSWORD") or None
|
||||
i = 0
|
||||
while i < len(args):
|
||||
a = args[i]
|
||||
|
||||
+8
-1
@@ -297,7 +297,14 @@ def _llm_tiebreak(
|
||||
|
||||
try:
|
||||
from graphify.llm import _call_llm
|
||||
except ImportError:
|
||||
except ImportError as exc:
|
||||
# F-038: previously this silent fallback hid the fact that `_call_llm`
|
||||
# didn't exist in `graphify.llm` at all, so `--dedup-llm` was a no-op.
|
||||
# Surface the import failure so future regressions are visible.
|
||||
print(
|
||||
f"[graphify] --dedup-llm: cannot import _call_llm ({exc}); skipping LLM tiebreaker.",
|
||||
flush=True,
|
||||
)
|
||||
return
|
||||
|
||||
for batch_start in range(0, len(ambiguous), batch_size):
|
||||
|
||||
+6
-1
@@ -247,7 +247,12 @@ def xlsx_extract_structure(path: Path) -> dict:
|
||||
except Exception:
|
||||
return {"nodes": [], "edges": []}
|
||||
|
||||
stem = _re.sub(r"[^a-z0-9]", "_", path.stem.lower())
|
||||
# F-035: typo fix — was `_re.sub` (NameError, but unreachable because the
|
||||
# whole xlsx codepath is currently behind a feature flag / not yet wired
|
||||
# into the dispatcher). Before re-enabling this path, re-audit it for
|
||||
# zip/XML bombs (openpyxl is built on top of zipfile and lxml-style XML
|
||||
# parsing — a malicious .xlsx can blow up memory at load_workbook time).
|
||||
stem = re.sub(r"[^a-z0-9]", "_", path.stem.lower())
|
||||
str_path = str(path)
|
||||
file_nid = _nid(str_path)
|
||||
nodes: list[dict] = [{"id": file_nid, "label": path.name, "file_type": "document",
|
||||
|
||||
+95
-10
@@ -26,6 +26,44 @@ def _strip_diacritics(text: str) -> str:
|
||||
return "".join(c for c in nfkd if not unicodedata.combining(c))
|
||||
|
||||
|
||||
def _yaml_str(s: str) -> str:
|
||||
"""Escape a value for safe embedding in a YAML double-quoted scalar (F-009).
|
||||
|
||||
See `graphify.ingest._yaml_str` for the full rationale; duplicated here to
|
||||
avoid pulling the URL-fetching `ingest` module into export's dependency
|
||||
graph. Handles backslash, double-quote, all line breaks (\\n, \\r,
|
||||
U+2028, U+2029), tab, NUL, and other C0/DEL control characters that
|
||||
would otherwise let a hostile `source_file` / `community` / etc. break
|
||||
out of the YAML scalar and inject sibling keys.
|
||||
"""
|
||||
if s is None:
|
||||
return ""
|
||||
out: list[str] = []
|
||||
for ch in str(s):
|
||||
cp = ord(ch)
|
||||
if ch == "\\":
|
||||
out.append("\\\\")
|
||||
elif ch == '"':
|
||||
out.append('\\"')
|
||||
elif ch == "\n":
|
||||
out.append("\\n")
|
||||
elif ch == "\r":
|
||||
out.append("\\r")
|
||||
elif ch == "\t":
|
||||
out.append("\\t")
|
||||
elif ch == "\0":
|
||||
out.append("\\0")
|
||||
elif cp == 0x2028:
|
||||
out.append("\\L")
|
||||
elif cp == 0x2029:
|
||||
out.append("\\P")
|
||||
elif cp < 0x20 or cp == 0x7F:
|
||||
out.append(f"\\x{cp:02x}")
|
||||
else:
|
||||
out.append(ch)
|
||||
return "".join(out)
|
||||
|
||||
|
||||
COMMUNITY_COLORS = [
|
||||
"#4E79A7", "#F28E2B", "#E15759", "#76B7B2", "#59A14F",
|
||||
"#EDC948", "#B07AA1", "#FF9DA7", "#9C755F", "#BAB0AC",
|
||||
@@ -426,8 +464,48 @@ def prune_dangling_edges(graph_data: dict) -> tuple[dict, int]:
|
||||
|
||||
|
||||
def _cypher_escape(s: str) -> str:
|
||||
"""Escape a string for safe embedding in a Cypher single-quoted literal."""
|
||||
return s.replace("\\", "\\\\").replace("'", "\\'")
|
||||
"""Escape a string for safe embedding in a Cypher single-quoted literal.
|
||||
|
||||
Handles all characters that could prematurely terminate the literal or
|
||||
inject control sequences:
|
||||
- `\\` and `'` (literal terminators)
|
||||
- newlines/CRs (would break the per-line statement framing)
|
||||
- NUL/control bytes (defensive — Neo4j errors on raw NULs)
|
||||
|
||||
Also strips any leading/trailing whitespace that would let an attacker
|
||||
break the `;`-terminated statement boundary used by `cypher-shell`.
|
||||
Closing `}` and `)` are NOT special inside a single-quoted Cypher string,
|
||||
so escaping the quote and backslash correctly is sufficient (a `}` inside
|
||||
a properly-closed `'...'` literal is just a character) — but we previously
|
||||
missed `\\n` / `\\r` which DO let a payload break out of the statement
|
||||
line and inject a fresh MATCH/DELETE on the following line. See F-008.
|
||||
"""
|
||||
# First normalise: drop NUL and other C0 control chars except tab.
|
||||
s = "".join(ch for ch in s if ch >= " " or ch == "\t")
|
||||
return (
|
||||
s.replace("\\", "\\\\")
|
||||
.replace("'", "\\'")
|
||||
.replace("\n", "\\n")
|
||||
.replace("\r", "\\r")
|
||||
)
|
||||
|
||||
|
||||
# Restrict identifier-position values (labels and relationship types are NOT
|
||||
# quoted in Cypher and so cannot be safely escaped — they must be allowlisted).
|
||||
_CYPHER_IDENT_RE = re.compile(r"[^A-Za-z0-9_]")
|
||||
|
||||
|
||||
def _cypher_label(raw: str, fallback: str) -> str:
|
||||
"""Sanitise a value used in identifier position (node label / rel type).
|
||||
|
||||
Cypher does not provide a way to escape `:Foo` label syntax, so we must
|
||||
strip everything except `[A-Za-z0-9_]` and require the result to start
|
||||
with a letter; otherwise we fall back to a safe constant.
|
||||
"""
|
||||
cleaned = _CYPHER_IDENT_RE.sub("", raw or "")
|
||||
if not cleaned or not cleaned[0].isalpha():
|
||||
return fallback
|
||||
return cleaned
|
||||
|
||||
|
||||
def to_cypher(G: nx.Graph, output_path: str) -> None:
|
||||
@@ -435,12 +513,17 @@ def to_cypher(G: nx.Graph, output_path: str) -> None:
|
||||
for node_id, data in G.nodes(data=True):
|
||||
label = _cypher_escape(data.get("label", node_id))
|
||||
node_id_esc = _cypher_escape(node_id)
|
||||
_ft = re.sub(r"[^A-Za-z0-9_]", "", data.get("file_type", "unknown").capitalize())
|
||||
ftype = (_ft if _ft and _ft[0].isalpha() else "Entity")
|
||||
ftype = _cypher_label(
|
||||
(data.get("file_type", "unknown") or "unknown").capitalize(),
|
||||
"Entity",
|
||||
)
|
||||
lines.append(f"MERGE (n:{ftype} {{id: '{node_id_esc}', label: '{label}'}});")
|
||||
lines.append("")
|
||||
for u, v, data in G.edges(data=True):
|
||||
rel = re.sub(r"[^A-Za-z0-9_]", "_", data.get("relation", "RELATES_TO").upper())
|
||||
rel = _cypher_label(
|
||||
(data.get("relation", "RELATES_TO") or "RELATES_TO").upper(),
|
||||
"RELATES_TO",
|
||||
)
|
||||
conf = _cypher_escape(data.get("confidence", "EXTRACTED"))
|
||||
u_esc = _cypher_escape(u)
|
||||
v_esc = _cypher_escape(v)
|
||||
@@ -697,15 +780,17 @@ def to_obsidian(
|
||||
|
||||
lines: list[str] = []
|
||||
|
||||
# YAML frontmatter - readable in Obsidian's properties panel
|
||||
# YAML frontmatter - readable in Obsidian's properties panel.
|
||||
# All scalars pass through _yaml_str so a hostile source_file or
|
||||
# community label cannot break out and inject sibling keys (F-009).
|
||||
lines += [
|
||||
"---",
|
||||
f'source_file: "{data.get("source_file", "")}"',
|
||||
f'type: "{ftype}"',
|
||||
f'community: "{community_name}"',
|
||||
f'source_file: "{_yaml_str(data.get("source_file", ""))}"',
|
||||
f'type: "{_yaml_str(ftype)}"',
|
||||
f'community: "{_yaml_str(community_name)}"',
|
||||
]
|
||||
if data.get("source_location"):
|
||||
lines.append(f'location: "{data["source_location"]}"')
|
||||
lines.append(f'location: "{_yaml_str(str(data["source_location"]))}"')
|
||||
# Add tags list to frontmatter
|
||||
lines.append("tags:")
|
||||
for tag in node_tags:
|
||||
|
||||
+8
-1
@@ -2816,6 +2816,13 @@ def _cpp_preprocess(path: Path) -> bytes:
|
||||
Falls back to raw file bytes if cpp is not available. Capital-F extensions
|
||||
conventionally require C preprocessor expansion (#ifdef MPI, #define REAL8, etc.)
|
||||
before parsing.
|
||||
|
||||
Security (F-007): we pass `-nostdinc` and `-I /dev/null` so a malicious
|
||||
source file containing `#include "/home/victim/.ssh/id_rsa"` (or any other
|
||||
include directive) cannot inline arbitrary host files into the output that
|
||||
we then ship to an LLM. Without these flags `cpp` happily resolves any
|
||||
relative or absolute include path it can read, which is a corpus-side
|
||||
file-exfiltration vector.
|
||||
"""
|
||||
import shutil
|
||||
import subprocess
|
||||
@@ -2823,7 +2830,7 @@ def _cpp_preprocess(path: Path) -> bytes:
|
||||
return path.read_bytes()
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["cpp", "-w", "-P", str(path)],
|
||||
["cpp", "-w", "-P", "-nostdinc", "-I", "/dev/null", str(path)],
|
||||
capture_output=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
+11
-2
@@ -2,6 +2,7 @@
|
||||
from __future__ import annotations
|
||||
import configparser
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_HOOK_MARKER = "# graphify-hook-start"
|
||||
@@ -168,8 +169,16 @@ def _hooks_dir(root: Path) -> Path:
|
||||
else:
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
return p
|
||||
except Exception:
|
||||
pass
|
||||
except (configparser.Error, OSError) as exc:
|
||||
# Narrow the exception (PR747-NEW-2): a bare `except Exception: pass`
|
||||
# was hiding tampering signals (corrupt .git/config, permission flips
|
||||
# by another tool). Surface them on stderr instead of silently
|
||||
# falling through to the default hooks directory.
|
||||
print(
|
||||
f"[graphify hooks] could not read core.hooksPath from "
|
||||
f"{root / '.git' / 'config'}: {exc}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
d = root / ".git" / "hooks"
|
||||
d.mkdir(exist_ok=True)
|
||||
return d
|
||||
|
||||
+39
-2
@@ -11,8 +11,45 @@ from graphify.security import safe_fetch, safe_fetch_text, validate_url
|
||||
|
||||
|
||||
def _yaml_str(s: str) -> str:
|
||||
"""Escape a string for embedding in a YAML double-quoted scalar."""
|
||||
return s.replace("\\", "\\\\").replace('"', '\\"').replace("\n", " ").replace("\r", " ")
|
||||
"""Escape a string for embedding in a YAML double-quoted scalar.
|
||||
|
||||
Handles every YAML 1.1/1.2 line-break and control character that could
|
||||
let a hostile value (e.g. a fetched page title) break out of the quoted
|
||||
scalar and inject sibling YAML keys (F-009 / F-019). The previous
|
||||
implementation missed `\\t`, `\\0`, the unicode line-separator U+2028 and
|
||||
paragraph-separator U+2029 — all of which YAML treats as line breaks.
|
||||
|
||||
We intentionally do not depend on PyYAML (not in pyproject deps) and
|
||||
instead emit safely-escaped double-quoted scalars by hand: the YAML
|
||||
double-quoted form recognises `\\\\`, `\\"`, `\\n`, `\\r`, `\\t`, `\\0`,
|
||||
`\\L` (U+2028), `\\P` (U+2029), and `\\xNN`/`\\uNNNN` numeric escapes.
|
||||
"""
|
||||
if s is None:
|
||||
return ""
|
||||
out: list[str] = []
|
||||
for ch in str(s):
|
||||
cp = ord(ch)
|
||||
if ch == "\\":
|
||||
out.append("\\\\")
|
||||
elif ch == '"':
|
||||
out.append('\\"')
|
||||
elif ch == "\n":
|
||||
out.append("\\n")
|
||||
elif ch == "\r":
|
||||
out.append("\\r")
|
||||
elif ch == "\t":
|
||||
out.append("\\t")
|
||||
elif ch == "\0":
|
||||
out.append("\\0")
|
||||
elif cp == 0x2028:
|
||||
out.append("\\L")
|
||||
elif cp == 0x2029:
|
||||
out.append("\\P")
|
||||
elif cp < 0x20 or cp == 0x7F:
|
||||
out.append(f"\\x{cp:02x}")
|
||||
else:
|
||||
out.append(ch)
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _safe_filename(url: str, suffix: str) -> str:
|
||||
|
||||
+151
-11
@@ -142,8 +142,22 @@ def _read_files(paths: list[Path], root: Path) -> str:
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
_LLM_JSON_MAX_BYTES = 10 * 1024 * 1024 # 10 MB hard cap before json.loads (F-016)
|
||||
|
||||
|
||||
def _parse_llm_json(raw: str) -> dict:
|
||||
"""Strip optional markdown fences and parse JSON. Returns empty fragment on failure."""
|
||||
"""Strip optional markdown fences and parse JSON. Returns empty fragment on failure.
|
||||
|
||||
Caps the input at `_LLM_JSON_MAX_BYTES` so a hostile or runaway model
|
||||
response cannot exhaust memory inside `json.loads` (F-016).
|
||||
"""
|
||||
if len(raw) > _LLM_JSON_MAX_BYTES:
|
||||
print(
|
||||
f"[graphify] LLM response exceeds {_LLM_JSON_MAX_BYTES} bytes "
|
||||
f"({len(raw)} bytes); refusing to parse and dropping chunk.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return {"nodes": [], "edges": [], "hyperedges": []}
|
||||
if raw.startswith("```"):
|
||||
raw = raw.split("```", 2)[1]
|
||||
if raw.startswith("json"):
|
||||
@@ -334,7 +348,18 @@ def extract_files_direct(
|
||||
cfg = BACKENDS[backend]
|
||||
key = api_key or _get_backend_api_key(backend)
|
||||
if not key and backend == "ollama":
|
||||
key = "ollama" # Ollama ignores auth but openai client requires non-empty
|
||||
# Ollama ignores auth but the OpenAI client library requires a non-empty
|
||||
# string. Use a placeholder and surface a visible warning so this never
|
||||
# silently routes traffic without the user realising — see F-029.
|
||||
ollama_url = os.environ.get("OLLAMA_BASE_URL", cfg.get("base_url", ""))
|
||||
_validate_ollama_base_url(ollama_url)
|
||||
print(
|
||||
"[graphify] WARNING: ollama backend selected with no OLLAMA_API_KEY set; "
|
||||
f"sending corpus to {ollama_url}. Set OLLAMA_API_KEY (any non-empty value) "
|
||||
"to suppress this warning.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
key = "ollama"
|
||||
if not key and backend != "bedrock":
|
||||
raise ValueError(
|
||||
f"No API key for backend '{backend}'. "
|
||||
@@ -609,6 +634,83 @@ def _merge_into(merged: dict, result: dict) -> None:
|
||||
merged["output_tokens"] += result.get("output_tokens", 0)
|
||||
|
||||
|
||||
def _call_llm(prompt: str, *, backend: str, max_tokens: int = 200) -> str:
|
||||
"""Send a plain-text prompt to `backend` and return the model's text reply.
|
||||
|
||||
Used by lightweight callers (e.g. `graphify.dedup` LLM tiebreaker) that
|
||||
don't need the full extraction prompt or JSON-shaped output. Mirrors the
|
||||
backend dispatch logic of `extract_files_direct` but skips the
|
||||
`_EXTRACTION_SYSTEM` prompt and JSON parsing.
|
||||
|
||||
Previously `graphify.dedup` imported a `_call_llm` symbol that did not
|
||||
exist in this module, so the LLM tiebreaker silently no-op'd on
|
||||
`ImportError` (F-038). Adding the function here re-enables it.
|
||||
"""
|
||||
if backend not in BACKENDS:
|
||||
raise ValueError(f"Unknown backend {backend!r}")
|
||||
cfg = BACKENDS[backend]
|
||||
key = _get_backend_api_key(backend)
|
||||
if not key and backend == "ollama":
|
||||
ollama_url = os.environ.get("OLLAMA_BASE_URL", cfg.get("base_url", ""))
|
||||
_validate_ollama_base_url(ollama_url)
|
||||
key = "ollama"
|
||||
if not key and backend != "bedrock":
|
||||
raise ValueError(
|
||||
f"No API key for backend '{backend}'. Set {_format_backend_env_keys(backend)}."
|
||||
)
|
||||
mdl = _default_model_for_backend(backend)
|
||||
|
||||
if backend == "claude":
|
||||
try:
|
||||
import anthropic
|
||||
except ImportError as exc:
|
||||
raise ImportError("anthropic package required for claude backend") from exc
|
||||
client = anthropic.Anthropic(api_key=key)
|
||||
resp = client.messages.create(
|
||||
model=mdl,
|
||||
max_tokens=max_tokens,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
)
|
||||
return resp.content[0].text if resp.content else ""
|
||||
|
||||
if backend == "bedrock":
|
||||
try:
|
||||
import boto3
|
||||
except ImportError as exc:
|
||||
raise ImportError("boto3 required for bedrock backend") from exc
|
||||
region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-east-1"
|
||||
profile = os.environ.get("AWS_PROFILE")
|
||||
session = boto3.Session(profile_name=profile, region_name=region)
|
||||
client = session.client("bedrock-runtime")
|
||||
resp = client.converse(
|
||||
modelId=mdl,
|
||||
messages=[{"role": "user", "content": [{"text": prompt}]}],
|
||||
inferenceConfig={"maxTokens": max_tokens, "temperature": 0},
|
||||
)
|
||||
return resp.get("output", {}).get("message", {}).get("content", [{}])[0].get("text", "")
|
||||
|
||||
# OpenAI-compatible (kimi, openai, gemini, ollama)
|
||||
try:
|
||||
from openai import OpenAI
|
||||
except ImportError as exc:
|
||||
raise ImportError("openai package required for this backend") from exc
|
||||
client = OpenAI(api_key=key, base_url=cfg["base_url"])
|
||||
kwargs: dict = {
|
||||
"model": mdl,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_completion_tokens": max_tokens,
|
||||
}
|
||||
temperature = cfg.get("temperature", 0)
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
if cfg.get("reasoning_effort"):
|
||||
kwargs["reasoning_effort"] = cfg["reasoning_effort"]
|
||||
if "moonshot" in cfg["base_url"]:
|
||||
kwargs["extra_body"] = {"thinking": {"type": "disabled"}}
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
return resp.choices[0].message.content or ""
|
||||
|
||||
|
||||
def estimate_cost(backend: str, input_tokens: int, output_tokens: int) -> float:
|
||||
"""Estimate USD cost for a given token count using published pricing."""
|
||||
if backend not in BACKENDS:
|
||||
@@ -617,21 +719,59 @@ def estimate_cost(backend: str, input_tokens: int, output_tokens: int) -> float:
|
||||
return (input_tokens * p["input"] + output_tokens * p["output"]) / 1_000_000
|
||||
|
||||
|
||||
def _validate_ollama_base_url(url: str) -> None:
|
||||
"""Warn (do not raise) if OLLAMA_BASE_URL looks unsafe.
|
||||
|
||||
Sending an entire corpus to a non-loopback http:// endpoint silently leaks
|
||||
proprietary code; we surface a visible stderr warning instead of failing
|
||||
closed (some users genuinely run Ollama on a LAN host they trust).
|
||||
"""
|
||||
try:
|
||||
from urllib.parse import urlparse
|
||||
parsed = urlparse(url)
|
||||
except Exception:
|
||||
print(
|
||||
f"[graphify] WARNING: OLLAMA_BASE_URL={url!r} is not a parseable URL.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
print(
|
||||
f"[graphify] WARNING: OLLAMA_BASE_URL has unexpected scheme {parsed.scheme!r}; "
|
||||
"expected http or https.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return
|
||||
host = (parsed.hostname or "").lower()
|
||||
is_loopback = host in ("localhost", "127.0.0.1", "::1") or host.startswith("127.")
|
||||
if not is_loopback:
|
||||
scheme_note = " (UNENCRYPTED)" if parsed.scheme == "http" else ""
|
||||
print(
|
||||
f"[graphify] WARNING: OLLAMA_BASE_URL points to non-loopback host {host!r}{scheme_note}. "
|
||||
"Your full corpus will be sent to that endpoint. "
|
||||
"Set OLLAMA_BASE_URL=http://localhost:11434/v1 to keep extraction local.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
def detect_backend() -> str | None:
|
||||
"""Return the name of whichever backend has an API key set, or None.
|
||||
|
||||
Priority: gemini → kimi → ollama (opt-in via OLLAMA_BASE_URL) → claude → openai.
|
||||
Ollama is opt-in via env var — never auto-probed without OLLAMA_BASE_URL set.
|
||||
Claude is the default for the skill.md subagent pipeline and is never forced here.
|
||||
Priority: gemini → kimi → claude → openai → bedrock → ollama (last, opt-in).
|
||||
|
||||
Ollama is intentionally checked LAST so a paid API key (Anthropic/OpenAI/etc.)
|
||||
is never silently shadowed by an incidental OLLAMA_BASE_URL in the environment
|
||||
— see security finding F-002/F-029. Setting OLLAMA_BASE_URL alongside a paid
|
||||
key now keeps you on the paid backend; remove the paid key (or pass
|
||||
--backend ollama explicitly) to route to the local model.
|
||||
"""
|
||||
for backend in ("gemini", "kimi"):
|
||||
for backend in ("gemini", "kimi", "claude", "openai"):
|
||||
if _get_backend_api_key(backend):
|
||||
return backend
|
||||
if os.environ.get("OLLAMA_BASE_URL"):
|
||||
return "ollama"
|
||||
if os.environ.get("AWS_PROFILE") or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION"):
|
||||
return "bedrock"
|
||||
for backend in ("claude", "openai"):
|
||||
if _get_backend_api_key(backend):
|
||||
return backend
|
||||
ollama_url = os.environ.get("OLLAMA_BASE_URL")
|
||||
if ollama_url:
|
||||
_validate_ollama_base_url(ollama_url)
|
||||
return "ollama"
|
||||
return None
|
||||
|
||||
+30
-11
@@ -172,17 +172,28 @@ def _subgraph_to_text(G: nx.Graph, nodes: set[str], edges: list[tuple], token_bu
|
||||
sorted(nodes - seed_set, key=lambda n: G.degree(n), reverse=True)
|
||||
for nid in ordered:
|
||||
d = G.nodes[nid]
|
||||
line = f"NODE {sanitize_label(d.get('label', nid))} [src={d.get('source_file', '')} loc={d.get('source_location', '')} community={d.get('community', '')}]"
|
||||
# Every LLM-derived field passes through sanitize_label before being
|
||||
# concatenated into MCP tool output (F-010): an attacker who controls a
|
||||
# corpus document can otherwise inject ANSI escapes, fake graphify-out
|
||||
# log lines, or prompt-injection markup into the model's context via
|
||||
# source_file / source_location / community.
|
||||
line = (
|
||||
f"NODE {sanitize_label(d.get('label', nid))} "
|
||||
f"[src={sanitize_label(str(d.get('source_file', '')))} "
|
||||
f"loc={sanitize_label(str(d.get('source_location', '')))} "
|
||||
f"community={sanitize_label(str(d.get('community', '')))}]"
|
||||
)
|
||||
lines.append(line)
|
||||
for u, v in edges:
|
||||
if u in nodes and v in nodes:
|
||||
raw = G[u][v]
|
||||
d = next(iter(raw.values()), {}) if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)) else raw
|
||||
context = d.get("context")
|
||||
context_suffix = f" context={context}" if context else ""
|
||||
context_suffix = f" context={sanitize_label(str(context))}" if context else ""
|
||||
line = (
|
||||
f"EDGE {sanitize_label(G.nodes[u].get('label', u))} "
|
||||
f"--{d.get('relation', '')} [{d.get('confidence', '')}{context_suffix}]--> "
|
||||
f"--{sanitize_label(str(d.get('relation', '')))} "
|
||||
f"[{sanitize_label(str(d.get('confidence', '')))}{context_suffix}]--> "
|
||||
f"{sanitize_label(G.nodes[v].get('label', v))}"
|
||||
)
|
||||
lines.append(line)
|
||||
@@ -373,12 +384,13 @@ def serve(graph_path: str = "graphify-out/graph.json") -> None:
|
||||
if not matches:
|
||||
return f"No node matching '{label}' found."
|
||||
nid, d = matches[0]
|
||||
# Sanitise every LLM-derived field before concatenation (F-010).
|
||||
return "\n".join([
|
||||
f"Node: {d.get('label', nid)}",
|
||||
f" ID: {nid}",
|
||||
f" Source: {d.get('source_file', '')} {d.get('source_location', '')}",
|
||||
f" Type: {d.get('file_type', '')}",
|
||||
f" Community: {d.get('community', '')}",
|
||||
f"Node: {sanitize_label(d.get('label', nid))}",
|
||||
f" ID: {sanitize_label(nid)}",
|
||||
f" Source: {sanitize_label(str(d.get('source_file', '')))} {sanitize_label(str(d.get('source_location', '')))}",
|
||||
f" Type: {sanitize_label(str(d.get('file_type', '')))}",
|
||||
f" Community: {sanitize_label(str(d.get('community', '')))}",
|
||||
f" Degree: {G.degree(nid)}",
|
||||
])
|
||||
|
||||
@@ -389,13 +401,16 @@ def serve(graph_path: str = "graphify-out/graph.json") -> None:
|
||||
if not matches:
|
||||
return f"No node matching '{label}' found."
|
||||
nid = matches[0]
|
||||
lines = [f"Neighbors of {G.nodes[nid].get('label', nid)}:"]
|
||||
lines = [f"Neighbors of {sanitize_label(G.nodes[nid].get('label', nid))}:"]
|
||||
for neighbor in G.neighbors(nid):
|
||||
d = G.edges[nid, neighbor]
|
||||
rel = d.get("relation", "")
|
||||
if rel_filter and rel_filter not in rel.lower():
|
||||
continue
|
||||
lines.append(f" --> {G.nodes[neighbor].get('label', neighbor)} [{rel}] [{d.get('confidence', '')}]")
|
||||
lines.append(
|
||||
f" --> {sanitize_label(G.nodes[neighbor].get('label', neighbor))} "
|
||||
f"[{sanitize_label(str(rel))}] [{sanitize_label(str(d.get('confidence', '')))}]"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
def _tool_get_community(arguments: dict) -> str:
|
||||
@@ -406,7 +421,11 @@ def serve(graph_path: str = "graphify-out/graph.json") -> None:
|
||||
lines = [f"Community {cid} ({len(nodes)} nodes):"]
|
||||
for n in nodes:
|
||||
d = G.nodes[n]
|
||||
lines.append(f" {d.get('label', n)} [{d.get('source_file', '')}]")
|
||||
# Sanitise label and source_file (F-010).
|
||||
lines.append(
|
||||
f" {sanitize_label(d.get('label', n))} "
|
||||
f"[{sanitize_label(str(d.get('source_file', '')))}]"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
def _tool_god_nodes(arguments: dict) -> str:
|
||||
|
||||
@@ -25,11 +25,15 @@ def test_detect_backend_kimi_beats_ollama(monkeypatch):
|
||||
assert detect_backend() == "kimi"
|
||||
|
||||
|
||||
def test_detect_backend_ollama_beats_claude(monkeypatch):
|
||||
def test_detect_backend_claude_beats_ollama(monkeypatch):
|
||||
# ANTHROPIC_API_KEY (paid, intentional) should win over OLLAMA_BASE_URL
|
||||
# (env-driven, easy to set accidentally) -- security fix F-002/F-029.
|
||||
monkeypatch.delenv("MOONSHOT_API_KEY", raising=False)
|
||||
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
|
||||
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
|
||||
monkeypatch.setenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
assert detect_backend() == "ollama"
|
||||
assert detect_backend() == "claude"
|
||||
|
||||
|
||||
def test_detect_backend_none_without_envvars(monkeypatch):
|
||||
|
||||
Reference in New Issue
Block a user