Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 76 additions & 0 deletions src/trailmark/query/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,40 @@ def callees_of(self, name: str) -> list[dict[str, Any]]:
return []
return [_unit_to_dict(u) for u in self._store.callees_of(node_id)]

def ancestors_of(self, name: str) -> list[dict[str, Any]]:
"""Find every function/method that can transitively reach ``name``.

The dual of ``callees_of`` extended transitively: given a sensitive
sink, this surfaces every function that could eventually call it,
directly or indirectly. Useful for upward slicing during audits.
"""
node_id = self._store.find_node_id(name)
if node_id is None:
return []
ancestor_ids = self._store.ancestors_of(node_id)
result: list[dict[str, Any]] = []
for aid in ancestor_ids:
unit = self._store._graph.nodes.get(aid) # noqa: SLF001
if unit is not None:
result.append(_unit_to_dict(unit))
return result

def reachable_from(self, name: str) -> list[dict[str, Any]]:
"""Find every function/method transitively reachable from ``name``.

The transitive closure of ``callees_of``.
"""
node_id = self._store.find_node_id(name)
if node_id is None:
return []
reachable_ids = self._store.reachable_from(node_id)
result: list[dict[str, Any]] = []
for rid in reachable_ids:
unit = self._store._graph.nodes.get(rid) # noqa: SLF001
if unit is not None:
result.append(_unit_to_dict(unit))
return result

def paths_between(
self,
src: str,
Expand All @@ -121,6 +155,48 @@ def paths_between(
return []
return self._store.paths_between(src_id, dst_id)

def entrypoint_paths_to(
self,
name: str,
max_depth: int = 20,
) -> list[list[str]]:
"""Find call paths from any entrypoint to ``name``.

Answers the canonical attack-surface question: "given this sink,
what concrete entrypoint paths can reach it?" Returns a list of
id-path lists, one per reachable entrypoint.
"""
node_id = self._store.find_node_id(name)
if node_id is None:
return []
return self._store.entrypoint_paths_to(node_id, max_depth=max_depth)

def nodes_with_annotation(
self,
kind: AnnotationKind,
) -> list[dict[str, Any]]:
"""Return every node tagged with the given annotation kind."""
return [_unit_to_dict(u) for u in self._store.nodes_with_annotation(kind)]

def functions_that_raise(
self,
exception_name: str,
) -> list[dict[str, Any]]:
"""Return functions/methods whose parser-detected exception list
includes the named exception.

Looks at the ``exception_types`` field parsers populate when they
extract ``raise``/``throw`` statements. Match is by type name
(``TypeRef.name``) — modules/generics are ignored.
"""
result: list[dict[str, Any]] = []
for unit in self._store._graph.nodes.values(): # noqa: SLF001
for exc in unit.exception_types:
if exc.name == exception_name:
result.append(_unit_to_dict(unit))
break
return result

def attack_surface(self) -> list[dict[str, Any]]:
"""List all entrypoints with their trust levels."""
return [
Expand Down
125 changes: 125 additions & 0 deletions tests/test_hidden_apis.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
"""Tests for the newly exposed QueryEngine methods.

Covers ``ancestors_of``, ``reachable_from``, ``entrypoint_paths_to``,
``nodes_with_annotation``, and ``functions_that_raise`` — methods that
existed in ``GraphStore`` but had no public-facing analogue.
"""

from __future__ import annotations

from pathlib import Path

from trailmark.models.annotations import AnnotationKind
from trailmark.query.api import QueryEngine


class TestAncestorsOf:
def test_direct_caller_is_ancestor(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text(
"def sink():\n pass\n\ndef caller():\n sink()\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
ancestors = engine.ancestors_of("sink")
names = {a["name"] for a in ancestors}
assert "caller" in names

def test_transitive_ancestor_surfaces(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text(
"def sink():\n pass\ndef middle():\n sink()\ndef top():\n middle()\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
names = {a["name"] for a in engine.ancestors_of("sink")}
assert {"middle", "top"}.issubset(names)

def test_unknown_node_returns_empty(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text("def only():\n pass\n")
engine = QueryEngine.from_directory(str(tmp_path))
assert engine.ancestors_of("does_not_exist") == []

def test_node_with_no_callers_has_only_module_ancestor(self, tmp_path: Path) -> None:
"""A function with no callers still has the module as an ancestor
via the `contains` edge. Function-level ancestors should be empty.
"""
(tmp_path / "app.py").write_text("def solo():\n pass\n")
engine = QueryEngine.from_directory(str(tmp_path))
ancestor_kinds = {a["kind"] for a in engine.ancestors_of("solo")}
assert "function" not in ancestor_kinds
assert "method" not in ancestor_kinds


class TestReachableFrom:
def test_direct_callee_is_reachable(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text(
"def a():\n b()\n\ndef b():\n pass\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
names = {n["name"] for n in engine.reachable_from("a")}
assert "b" in names

def test_transitive_reachability(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text(
"def a():\n b()\ndef b():\n c()\ndef c():\n pass\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
names = {n["name"] for n in engine.reachable_from("a")}
assert {"b", "c"}.issubset(names)


class TestEntrypointPathsTo:
def test_path_from_entrypoint_to_sink(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text(
"def main():\n validate()\n"
"def validate():\n sensitive()\n"
"def sensitive():\n pass\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
paths = engine.entrypoint_paths_to("sensitive")
assert paths, "Expected at least one entrypoint->sink path"
assert any(p[0] == "app:main" and p[-1] == "app:sensitive" for p in paths), paths

def test_unreachable_sink_returns_empty(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text(
"def main():\n pass\n\ndef orphan():\n pass\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
assert engine.entrypoint_paths_to("orphan") == []


class TestNodesWithAnnotation:
def test_manual_annotation_retrievable(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text("def risky():\n pass\n")
engine = QueryEngine.from_directory(str(tmp_path))
engine.annotate(
"risky",
AnnotationKind.FINDING,
"manually flagged",
source="test",
)
finds = engine.nodes_with_annotation(AnnotationKind.FINDING)
names = {n["name"] for n in finds}
assert "risky" in names

def test_empty_when_no_annotations(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text("def ok():\n pass\n")
engine = QueryEngine.from_directory(str(tmp_path))
assert engine.nodes_with_annotation(AnnotationKind.FINDING) == []


class TestFunctionsThatRaise:
def test_raises_detected_from_python_source(self, tmp_path: Path) -> None:
"""Python parser populates exception_types from `raise` statements."""
(tmp_path / "app.py").write_text(
"def bad():\n raise ValueError('nope')\n\ndef good():\n return 1\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
raisers = engine.functions_that_raise("ValueError")
names = {n["name"] for n in raisers}
assert "bad" in names
assert "good" not in names

def test_unknown_exception_returns_empty(self, tmp_path: Path) -> None:
(tmp_path / "app.py").write_text(
"def bad():\n raise ValueError('nope')\n",
)
engine = QueryEngine.from_directory(str(tmp_path))
assert engine.functions_that_raise("KeyError") == []
Loading