diff --git a/README.md b/README.md index a921f57..49503de 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,14 @@ A language-specific parser walks the directory, parses each file into a tree-sit | Objective-C | `.m`, `.mm`, `.h` | C functions, classes, methods (selector-based naming) | | Kotlin | `.kt`, `.kts` | functions, classes, interfaces, data classes, objects, methods | | Dart | `.dart` | functions, classes, abstract classes, methods, constructors | +| Move | `.move` | modules, functions, imports, direct calls | +| Tact | `.tact` | contracts, structs, receivers, functions | +| Func | `.fc`, `.func` | functions, includes, direct calls | +| Sway | `.sw` | ABI interfaces, structs, impl methods, functions | +| Rego | `.rego` | packages, imports, policy rules, rule calls | +| Proto | `.proto` | services, RPCs, messages, fields, enums | +| Thrift | `.thrift` | services, functions, structs, fields, enums | +| GraphQL | `.graphql`, `.gql` | object types, root operations, fields, enums | ```mermaid flowchart TD diff --git a/src/trailmark/analysis/entrypoints.py b/src/trailmark/analysis/entrypoints.py index df04f45..e389a62 100644 --- a/src/trailmark/analysis/entrypoints.py +++ b/src/trailmark/analysis/entrypoints.py @@ -330,6 +330,22 @@ def _detect_for_unit( return _detect_kotlin(cache, unit, path) if path.endswith(".dart"): return _detect_dart(cache, unit, path) + if path.endswith(".move"): + return _detect_move(cache, unit, path) + if path.endswith(".tact"): + return _detect_tact(unit) + if path.endswith((".fc", ".func")): + return _detect_func(unit) + if path.endswith(".sw"): + return _detect_sway(cache, unit, path) + if path.endswith(".rego"): + return _detect_rego(unit) + if path.endswith(".proto"): + return _detect_proto(cache, unit, path) + if path.endswith(".thrift"): + return _detect_thrift(cache, unit, path) + if path.endswith((".graphql", ".gql")): + return _detect_graphql(unit) if path.endswith(".go"): return _detect_go(cache, unit, path) if path.endswith(".rb"): @@ -901,6 +917,133 @@ def _detect_dart( return None +def _detect_move( + cache: _SourceCache, + unit: CodeUnit, + path: str, +) -> EntrypointTag | None: + signature = cache.signature_block(path, unit.location.start_line) or "" + if unit.kind.value == "function" and (" entry " in signature or "public" in signature): + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Move public/entry function", + asset_value=AssetValue.HIGH, + ) + return None + + +def _detect_tact(unit: CodeUnit) -> EntrypointTag | None: + tact_role = _unit_attr(unit, "tact_role") or unit.name + if unit.kind.value == "method" and tact_role in {"init", "receive", "external", "bounced"}: + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Tact contract receiver/initializer", + asset_value=AssetValue.HIGH, + ) + return None + + +def _detect_func(unit: CodeUnit) -> EntrypointTag | None: + if unit.name in {"recv_internal", "recv_external"} or unit.name.startswith("get_"): + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Func receiver/getter entrypoint", + asset_value=AssetValue.HIGH, + ) + return None + + +def _detect_sway( + cache: _SourceCache, + unit: CodeUnit, + path: str, +) -> EntrypointTag | None: + signature = cache.signature_block(path, unit.location.start_line) or "" + if unit.kind.value == "function" and "pub fn" in signature: + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Sway public function", + asset_value=AssetValue.HIGH, + ) + if ( + unit.kind.value == "method" + and re.search(r"\bfn\s+\w+\s*\(", signature) + and ";" in signature + ): + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Sway ABI method", + asset_value=AssetValue.HIGH, + ) + return None + + +def _detect_rego(unit: CodeUnit) -> EntrypointTag | None: + if unit.name in {"allow", "deny", "violation"}: + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Rego policy decision rule", + asset_value=AssetValue.HIGH, + ) + return None + + +def _detect_proto( + cache: _SourceCache, + unit: CodeUnit, + path: str, +) -> EntrypointTag | None: + del cache, path + if unit.kind.value == "method" and _unit_attr(unit, "schema_role") == "rpc": + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Protocol Buffers service RPC", + asset_value=AssetValue.HIGH, + ) + return None + + +def _detect_thrift( + cache: _SourceCache, + unit: CodeUnit, + path: str, +) -> EntrypointTag | None: + del cache, path + if unit.kind.value == "method" and _unit_attr(unit, "schema_role") == "service_function": + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="Thrift service function", + asset_value=AssetValue.HIGH, + ) + return None + + +def _detect_graphql(unit: CodeUnit) -> EntrypointTag | None: + if unit.kind.value == "method" and _unit_attr(unit, "schema_role") == "root_operation": + return EntrypointTag( + kind=EntrypointKind.API, + trust_level=TrustLevel.UNTRUSTED_EXTERNAL, + description="GraphQL root operation field", + asset_value=AssetValue.HIGH, + ) + return None + + +def _unit_attr(unit: CodeUnit, key: str) -> str | None: + for attr_key, attr_value in unit.attributes: + if attr_key == key and isinstance(attr_value, str): + return attr_value + return None + + def _detect_objc(unit: CodeUnit) -> EntrypointTag | None: """Detect Objective-C entrypoints: AppDelegate selectors and extern C. diff --git a/src/trailmark/parse.py b/src/trailmark/parse.py index 0901711..1f4173b 100644 --- a/src/trailmark/parse.py +++ b/src/trailmark/parse.py @@ -33,6 +33,14 @@ "objc": ("trailmark.parsers.objc", "ObjCParser"), "kotlin": ("trailmark.parsers.kotlin", "KotlinParser"), "dart": ("trailmark.parsers.dart", "DartParser"), + "move": ("trailmark.parsers.move", "MoveParser"), + "tact": ("trailmark.parsers.tact", "TactParser"), + "func": ("trailmark.parsers.func", "FuncParser"), + "sway": ("trailmark.parsers.sway", "SwayParser"), + "rego": ("trailmark.parsers.rego", "RegoParser"), + "proto": ("trailmark.parsers.proto", "ProtoParser"), + "thrift": ("trailmark.parsers.thrift", "ThriftParser"), + "graphql": ("trailmark.parsers.graphql", "GraphQLParser"), } # Extensions used for language auto-detection. Keep these aligned with each @@ -63,6 +71,14 @@ "objc": (".m", ".mm"), "kotlin": (".kt", ".kts"), "dart": (".dart",), + "move": (".move",), + "tact": (".tact",), + "func": (".fc", ".func"), + "sway": (".sw",), + "rego": (".rego",), + "proto": (".proto",), + "thrift": (".thrift",), + "graphql": (".graphql", ".gql"), } _SUPPORTED_LANGUAGES = tuple(_PARSER_MAP.keys()) diff --git a/src/trailmark/parsers/_simple.py b/src/trailmark/parsers/_simple.py new file mode 100644 index 0000000..5eae805 --- /dev/null +++ b/src/trailmark/parsers/_simple.py @@ -0,0 +1,255 @@ +"""Small helpers for declarative tree-sitter parsers.""" + +from __future__ import annotations + +import re +from collections.abc import Iterable + +from tree_sitter import Node + +from trailmark.models.edges import CodeEdge, EdgeConfidence, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import BranchInfo, Parameter, TypeRef +from trailmark.parsers._common import make_location, node_text + + +def named_children(node: Node) -> list[Node]: + """Return named children only.""" + return [child for child in node.children if child.is_named] + + +def descendants(node: Node) -> Iterable[Node]: + """Yield named descendants depth-first, excluding ``node`` itself.""" + stack = list(reversed(named_children(node))) + while stack: + child = stack.pop() + yield child + stack.extend(reversed(named_children(child))) + + +def first_child_type(node: Node, types: set[str] | frozenset[str]) -> Node | None: + """Return the first direct named child whose type is in ``types``.""" + for child in named_children(node): + if child.type in types: + return child + return None + + +def first_descendant_type(node: Node, types: set[str] | frozenset[str]) -> Node | None: + """Return the first named descendant whose type is in ``types``.""" + for child in descendants(node): + if child.type in types: + return child + return None + + +def child_text(node: Node, types: set[str] | frozenset[str]) -> str: + """Return text for the first matching direct child.""" + child = first_child_type(node, types) + return node_text(child) if child is not None else "" + + +def descendant_text(node: Node, types: set[str] | frozenset[str]) -> str: + """Return text for the first matching descendant.""" + child = first_descendant_type(node, types) + return node_text(child) if child is not None else "" + + +def extract_parameters( + node: Node, + *, + list_types: set[str] | frozenset[str] = frozenset({"parameters", "parameter_list"}), + item_types: set[str] | frozenset[str] = frozenset( + {"parameter", "parameter_declaration", "input_value_definition", "field"} + ), + name_types: set[str] | frozenset[str] = frozenset( + {"identifier", "parameter", "name", "var", "field_name"} + ), + type_types: set[str] | frozenset[str] = frozenset( + { + "primitive_type", + "type", + "type_identifier", + "qualified_type", + "message_or_enum_type", + "named_type", + "non_null_type", + "list_type", + "type_expression", + } + ), +) -> tuple[Parameter, ...]: + """Best-effort parameter extraction for simple grammars.""" + params_node = first_child_type(node, list_types) + if params_node is None: + return () + params: list[Parameter] = [] + for item in named_children(params_node): + if item.type not in item_types: + continue + name = _parameter_name(item, name_types) + if not name: + continue + type_ref = _parameter_type(item, type_types, name) + params.append(Parameter(name=name, type_ref=type_ref)) + return tuple(params) + + +def _parameter_name(node: Node, name_types: set[str] | frozenset[str]) -> str: + matches = [child for child in named_children(node) if child.type in name_types] + if not matches: + return "" + # Func parameters use a node named "parameter" for the identifier. + return node_text(matches[-1]) + + +def _parameter_type( + node: Node, + type_types: set[str] | frozenset[str], + name: str, +) -> TypeRef | None: + for child in named_children(node): + if child.type in type_types: + text = node_text(child) + if text and text != name: + return TypeRef(name=clean_type_name(text)) + return None + + +def return_type_after_params( + node: Node, + type_types: set[str] | frozenset[str] = frozenset( + { + "primitive_type", + "type", + "type_identifier", + "qualified_type", + "message_or_enum_type", + "named_type", + "non_null_type", + "list_type", + } + ), +) -> TypeRef | None: + """Return the first type-looking child after a parameter list.""" + seen_params = False + for child in named_children(node): + if child.type in {"parameters", "parameter_list"}: + seen_params = True + continue + if seen_params and child.type in type_types: + return TypeRef(name=clean_type_name(node_text(child))) + return None + + +def clean_type_name(text: str) -> str: + """Normalize compact type syntax into a stable TypeRef name.""" + stripped = text.strip() + stripped = stripped.removeprefix("bounced<").removesuffix(">") + stripped = stripped.replace("!", "") + return re.sub(r"\s+", " ", stripped) + + +def collect_branches_and_calls( + body: Node | None, + file_path: str, + *, + branch_types: set[str] | frozenset[str], + call_types: set[str] | frozenset[str], +) -> tuple[list[BranchInfo], list[tuple[str, Node]]]: + """Collect branch metadata and direct call expressions under ``body``.""" + branches: list[BranchInfo] = [] + calls: list[tuple[str, Node]] = [] + if body is None: + return branches, calls + for node in descendants(body): + if node.type in branch_types: + branches.append( + BranchInfo( + location=make_location(node, file_path), + condition=_condition_text(node), + ) + ) + if node.type in call_types: + call_name = call_expression_name(node) + if call_name: + calls.append((call_name, node)) + return branches, calls + + +def _condition_text(node: Node) -> str: + condition = node.child_by_field_name("condition") + if condition is not None: + return node_text(condition) + return node.type + + +def call_expression_name(node: Node) -> str: + """Extract a callable symbol from common call-expression shapes.""" + function = node.child_by_field_name("function") + if function is not None: + return _callable_text(function) + + for child in named_children(node): + if child.type in { + "identifier", + "function_name", + "func_identifier", + "field_identifier", + "scoped_identifier", + "method_call", + "field_access_expression", + "var", + "ref", + }: + return _callable_text(child) + text = node_text(node).strip() + return text.split("(", 1)[0].strip() + + +def _callable_text(node: Node) -> str: + text = node_text(node).strip() + text = text.split("(", 1)[0].strip() + text = text.removeprefix("self.").removeprefix("this.") + text = text.removeprefix("~") + if "::" in text: + return text.rsplit("::", 1)[-1] + if "." in text: + return text.rsplit(".", 1)[-1] + return text + + +def add_call_edges( + graph: CodeGraph, + calls: list[tuple[str, Node]], + source_id: str, + module_id: str, + file_path: str, + container_id: str | None = None, +) -> None: + """Add direct call edges with local best-effort target resolution.""" + for call_name, call_node in calls: + target_id = resolve_call_target(call_name, module_id, container_id) + graph.edges.append( + CodeEdge( + source_id=source_id, + target_id=target_id, + kind=EdgeKind.CALLS, + confidence=EdgeConfidence.CERTAIN, + location=make_location(call_node, file_path), + ) + ) + + +def resolve_call_target(call_name: str, module_id: str, container_id: str | None = None) -> str: + """Resolve a bare call into the current container/module namespace.""" + name = call_name.strip() + if not name: + return f"{module_id}:" + if ":" in name: + return name + if "." in name: + name = name.rsplit(".", 1)[-1] + if container_id is not None: + return f"{container_id}.{name}" + return f"{module_id}:{name}" diff --git a/src/trailmark/parsers/func/__init__.py b/src/trailmark/parsers/func/__init__.py new file mode 100644 index 0000000..2ce4ec6 --- /dev/null +++ b/src/trailmark/parsers/func/__init__.py @@ -0,0 +1,5 @@ +"""Func language parser for Trailmark.""" + +from trailmark.parsers.func.parser import FuncParser + +__all__ = ["FuncParser"] diff --git a/src/trailmark/parsers/func/parser.py b/src/trailmark/parsers/func/parser.py new file mode 100644 index 0000000..d3337f6 --- /dev/null +++ b/src/trailmark/parsers/func/parser.py @@ -0,0 +1,106 @@ +"""Func language parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.edges import CodeEdge, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind, TypeRef +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + compute_complexity, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import ( + add_call_edges, + collect_branches_and_calls, + extract_parameters, + first_child_type, + named_children, +) + +_EXTENSIONS = (".fc", ".func") +_BRANCH_TYPES = frozenset({"if_statement"}) +_CALL_TYPES = frozenset({"function_application", "method_call"}) + + +class FuncParser: + """Parses Func source files into CodeGraph using tree-sitter.""" + + @property + def language(self) -> str: + return "func" + + def __init__(self) -> None: + self._parser = Parser(get_language("func")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="func", root_path=file_path) + module_id = module_id_from_path(file_path) + _visit_root(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "func", dir_path, _EXTENSIONS) + + +def _visit_root(root: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + add_module_node(root, file_path, module_id, graph) + for child in named_children(root): + if child.type == "include_directive": + _extract_import(child, graph, module_id) + elif child.type == "function_definition": + _extract_function(child, file_path, module_id, graph) + + +def _extract_function(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name_node = first_child_type(node, frozenset({"function_name"})) + if name_node is None: + return + name = node_text(name_node) + func_id = f"{module_id}:{name}" + body = first_child_type(node, frozenset({"block_statement", "asm_function_body"})) + branches, calls = collect_branches_and_calls( + body, + file_path, + branch_types=_BRANCH_TYPES, + call_types=_CALL_TYPES, + ) + graph.nodes[func_id] = CodeUnit( + id=func_id, + name=name, + kind=NodeKind.FUNCTION, + location=make_location(node, file_path), + parameters=extract_parameters(node), + return_type=_return_type(node), + cyclomatic_complexity=compute_complexity(branches), + branches=tuple(branches), + ) + add_contains_edge(graph, module_id, func_id) + add_call_edges(graph, calls, func_id, module_id, file_path) + + +def _return_type(node: Node) -> TypeRef | None: + for child in named_children(node): + if child.type in {"unit_type", "primitive_type", "function_type", "type_identifier"}: + return TypeRef(name=node_text(child)) + if child.type == "function_name": + return None + return None + + +def _extract_import(node: Node, graph: CodeGraph, module_id: str) -> None: + target = node_text(node).removeprefix("#include").strip().strip(";").strip('"<>') + if target: + graph.dependencies.append(target) + graph.edges.append(CodeEdge(source_id=module_id, target_id=target, kind=EdgeKind.IMPORTS)) diff --git a/src/trailmark/parsers/graphql/__init__.py b/src/trailmark/parsers/graphql/__init__.py new file mode 100644 index 0000000..9727bcd --- /dev/null +++ b/src/trailmark/parsers/graphql/__init__.py @@ -0,0 +1,5 @@ +"""GraphQL schema parser for Trailmark.""" + +from trailmark.parsers.graphql.parser import GraphQLParser + +__all__ = ["GraphQLParser"] diff --git a/src/trailmark/parsers/graphql/parser.py b/src/trailmark/parsers/graphql/parser.py new file mode 100644 index 0000000..f37a3ca --- /dev/null +++ b/src/trailmark/parsers/graphql/parser.py @@ -0,0 +1,140 @@ +"""GraphQL schema parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind, Parameter, TypeRef +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import clean_type_name, descendants, first_child_type, named_children + +_EXTENSIONS = (".graphql", ".gql") + + +class GraphQLParser: + """Parses GraphQL schemas into CodeGraph.""" + + @property + def language(self) -> str: + return "graphql" + + def __init__(self) -> None: + self._parser = Parser(get_language("graphql")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="graphql", root_path=file_path) + module_id = module_id_from_path(file_path) + _visit_root(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "graphql", dir_path, _EXTENSIONS) + + +def _visit_root(root: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + add_module_node(root, file_path, module_id, graph) + for node in descendants(root): + if node.type in { + "object_type_definition", + "interface_type_definition", + "input_object_type_definition", + }: + _extract_object(node, file_path, module_id, graph) + elif node.type == "enum_type_definition": + _extract_enum(node, file_path, module_id, graph) + + +def _extract_object(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name_node = first_child_type(node, frozenset({"name"})) + if name_node is None: + return + name = node_text(name_node) + object_id = f"{module_id}:{name}" + graph.nodes[object_id] = CodeUnit( + id=object_id, + name=name, + kind=NodeKind.INTERFACE, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, object_id) + fields = first_child_type(node, frozenset({"fields_definition", "input_fields_definition"})) + if fields is None: + return + for child in named_children(fields): + if child.type in {"field_definition", "input_value_definition"}: + _extract_field(child, file_path, object_id, graph) + + +def _extract_field(node: Node, file_path: str, object_id: str, graph: CodeGraph) -> None: + name_node = first_child_type(node, frozenset({"name"})) + if name_node is None: + return + name = node_text(name_node) + field_id = f"{object_id}.{name}" + graph.nodes[field_id] = CodeUnit( + id=field_id, + name=name, + kind=NodeKind.METHOD, + location=make_location(node, file_path), + parameters=_arguments(node), + return_type=_type_ref(node), + attributes=(("schema_role", _field_role(object_id)),), + ) + add_contains_edge(graph, object_id, field_id) + + +def _field_role(object_id: str) -> str: + parent = object_id.rsplit(":", 1)[-1] + if parent in {"Query", "Mutation", "Subscription"}: + return "root_operation" + return "field" + + +def _arguments(node: Node) -> tuple[Parameter, ...]: + args = first_child_type(node, frozenset({"arguments_definition"})) + if args is None: + return () + params: list[Parameter] = [] + for arg in named_children(args): + if arg.type != "input_value_definition": + continue + name_node = first_child_type(arg, frozenset({"name"})) + if name_node is None: + continue + params.append(Parameter(name=node_text(name_node), type_ref=_type_ref(arg))) + return tuple(params) + + +def _type_ref(node: Node) -> TypeRef | None: + type_node = first_child_type(node, frozenset({"type"})) + if type_node is None: + return None + return TypeRef(name=clean_type_name(node_text(type_node))) + + +def _extract_enum(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name_node = first_child_type(node, frozenset({"name"})) + if name_node is None: + return + name = node_text(name_node) + enum_id = f"{module_id}:{name}" + graph.nodes[enum_id] = CodeUnit( + id=enum_id, + name=name, + kind=NodeKind.ENUM, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, enum_id) diff --git a/src/trailmark/parsers/move/__init__.py b/src/trailmark/parsers/move/__init__.py new file mode 100644 index 0000000..64536db --- /dev/null +++ b/src/trailmark/parsers/move/__init__.py @@ -0,0 +1,5 @@ +"""Move language parser for Trailmark.""" + +from trailmark.parsers.move.parser import MoveParser + +__all__ = ["MoveParser"] diff --git a/src/trailmark/parsers/move/parser.py b/src/trailmark/parsers/move/parser.py new file mode 100644 index 0000000..cc4a24d --- /dev/null +++ b/src/trailmark/parsers/move/parser.py @@ -0,0 +1,136 @@ +"""Move language parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.edges import CodeEdge, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + compute_complexity, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import ( + add_call_edges, + child_text, + collect_branches_and_calls, + extract_parameters, + first_child_type, + named_children, + return_type_after_params, +) + +_EXTENSIONS = (".move",) +_BRANCH_TYPES = frozenset({"if_expression"}) +_CALL_TYPES = frozenset({"call_expression"}) + + +class MoveParser: + """Parses Move source files into CodeGraph using tree-sitter.""" + + @property + def language(self) -> str: + return "move" + + def __init__(self) -> None: + self._parser = Parser(get_language("move")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="move", root_path=file_path) + module_id = module_id_from_path(file_path) + _visit_module(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "move", dir_path, _EXTENSIONS) + + +def _visit_module(root: Node, file_path: str, file_module_id: str, graph: CodeGraph) -> None: + add_module_node(root, file_path, file_module_id, graph) + move_module_id = _move_module_id(root, file_module_id) + if move_module_id != file_module_id: + graph.nodes[move_module_id] = CodeUnit( + id=move_module_id, + name=move_module_id.rsplit(":", 1)[-1], + kind=NodeKind.MODULE, + location=make_location(root, file_path), + ) + add_contains_edge(graph, file_module_id, move_module_id) + + body = first_child_type(root, frozenset({"module_body"})) or root + for child in named_children(body): + if child.type == "function_item": + _extract_function(child, file_path, move_module_id, graph) + elif child.type == "use_declaration": + _extract_import(child, graph, move_module_id) + elif child.type in {"struct_item", "struct"}: + _extract_struct(child, file_path, move_module_id, graph) + + +def _move_module_id(root: Node, fallback: str) -> str: + parts = [ + node_text(child) + for child in named_children(root) + if child.type in {"hex_address", "identifier"} + ] + if not parts: + return fallback + return f"{fallback}:{'::'.join(parts)}" + + +def _extract_function(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + func_id = f"{module_id}:{name}" + body = first_child_type(node, frozenset({"block"})) + branches, calls = collect_branches_and_calls( + body, + file_path, + branch_types=_BRANCH_TYPES, + call_types=_CALL_TYPES, + ) + graph.nodes[func_id] = CodeUnit( + id=func_id, + name=name, + kind=NodeKind.FUNCTION, + location=make_location(node, file_path), + parameters=extract_parameters(node), + return_type=return_type_after_params(node), + cyclomatic_complexity=compute_complexity(branches), + branches=tuple(branches), + ) + add_contains_edge(graph, module_id, func_id) + add_call_edges(graph, calls, func_id, module_id, file_path) + + +def _extract_struct(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + struct_id = f"{module_id}:{name}" + graph.nodes[struct_id] = CodeUnit( + id=struct_id, + name=name, + kind=NodeKind.STRUCT, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, struct_id) + + +def _extract_import(node: Node, graph: CodeGraph, module_id: str) -> None: + target = node_text(node).removeprefix("use").strip().rstrip(";") + if target: + graph.dependencies.append(target) + graph.edges.append(CodeEdge(source_id=module_id, target_id=target, kind=EdgeKind.IMPORTS)) diff --git a/src/trailmark/parsers/proto/__init__.py b/src/trailmark/parsers/proto/__init__.py new file mode 100644 index 0000000..b5547ab --- /dev/null +++ b/src/trailmark/parsers/proto/__init__.py @@ -0,0 +1,5 @@ +"""Protocol Buffers parser for Trailmark.""" + +from trailmark.parsers.proto.parser import ProtoParser + +__all__ = ["ProtoParser"] diff --git a/src/trailmark/parsers/proto/parser.py b/src/trailmark/parsers/proto/parser.py new file mode 100644 index 0000000..a91d57c --- /dev/null +++ b/src/trailmark/parsers/proto/parser.py @@ -0,0 +1,171 @@ +"""Protocol Buffers parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.edges import CodeEdge, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind, Parameter, TypeRef +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import child_text, clean_type_name, first_child_type, named_children + +_EXTENSIONS = (".proto",) + + +class ProtoParser: + """Parses Protocol Buffers schemas into CodeGraph.""" + + @property + def language(self) -> str: + return "proto" + + def __init__(self) -> None: + self._parser = Parser(get_language("proto")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="proto", root_path=file_path) + module_id = _package_id(tree.root_node) or module_id_from_path(file_path) + _visit_root(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "proto", dir_path, _EXTENSIONS) + + +def _visit_root(root: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + add_module_node(root, file_path, module_id, graph) + for child in named_children(root): + if child.type == "import": + _extract_import(child, graph, module_id) + elif child.type == "service": + _extract_service(child, file_path, module_id, graph) + elif child.type == "message": + _extract_message(child, file_path, module_id, module_id, graph) + elif child.type == "enum": + _extract_enum(child, file_path, module_id, module_id, graph) + + +def _package_id(root: Node) -> str: + for child in named_children(root): + if child.type == "package": + name = child_text(child, frozenset({"full_ident"})) + return name + return "" + + +def _extract_service(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"service_name"})) + if not name: + return + service_id = f"{module_id}:{name}" + graph.nodes[service_id] = CodeUnit( + id=service_id, + name=name, + kind=NodeKind.INTERFACE, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, service_id) + for child in named_children(node): + if child.type == "rpc": + _extract_rpc(child, file_path, service_id, graph) + + +def _extract_rpc(node: Node, file_path: str, service_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"rpc_name"})) + types = [ + node_text(child) for child in named_children(node) if child.type == "message_or_enum_type" + ] + params = (Parameter(name="request", type_ref=TypeRef(name=types[0])),) if types else () + return_type = TypeRef(name=types[1]) if len(types) > 1 else None + rpc_id = f"{service_id}.{name}" + graph.nodes[rpc_id] = CodeUnit( + id=rpc_id, + name=name, + kind=NodeKind.METHOD, + location=make_location(node, file_path), + parameters=params, + return_type=return_type, + attributes=(("schema_role", "rpc"),), + ) + add_contains_edge(graph, service_id, rpc_id) + + +def _extract_message( + node: Node, + file_path: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> None: + name = child_text(node, frozenset({"message_name"})) + if not name: + return + message_id = f"{module_id}:{name}" + graph.nodes[message_id] = CodeUnit( + id=message_id, + name=name, + kind=NodeKind.STRUCT, + location=make_location(node, file_path), + ) + add_contains_edge(graph, container_id, message_id) + body = first_child_type(node, frozenset({"message_body"})) + if body is None: + return + for child in named_children(body): + if child.type == "field": + _extract_field(child, file_path, message_id, graph) + + +def _extract_field(node: Node, file_path: str, message_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + type_node = first_child_type(node, frozenset({"type", "message_or_enum_type"})) + field_id = f"{message_id}.{name}" + graph.nodes[field_id] = CodeUnit( + id=field_id, + name=name, + kind=NodeKind.METHOD, + location=make_location(node, file_path), + return_type=TypeRef(name=clean_type_name(node_text(type_node))) if type_node else None, + attributes=(("schema_role", "field"),), + ) + add_contains_edge(graph, message_id, field_id) + + +def _extract_enum( + node: Node, + file_path: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> None: + name = child_text(node, frozenset({"enum_name", "identifier"})) + if not name: + return + enum_id = f"{module_id}:{name}" + graph.nodes[enum_id] = CodeUnit( + id=enum_id, + name=name, + kind=NodeKind.ENUM, + location=make_location(node, file_path), + ) + add_contains_edge(graph, container_id, enum_id) + + +def _extract_import(node: Node, graph: CodeGraph, module_id: str) -> None: + target = node_text(node).removeprefix("import").strip().strip(";").strip("\"'") + if target: + graph.dependencies.append(target) + graph.edges.append(CodeEdge(source_id=module_id, target_id=target, kind=EdgeKind.IMPORTS)) diff --git a/src/trailmark/parsers/rego/__init__.py b/src/trailmark/parsers/rego/__init__.py new file mode 100644 index 0000000..4fcba40 --- /dev/null +++ b/src/trailmark/parsers/rego/__init__.py @@ -0,0 +1,5 @@ +"""Rego language parser for Trailmark.""" + +from trailmark.parsers.rego.parser import RegoParser + +__all__ = ["RegoParser"] diff --git a/src/trailmark/parsers/rego/parser.py b/src/trailmark/parsers/rego/parser.py new file mode 100644 index 0000000..48350d9 --- /dev/null +++ b/src/trailmark/parsers/rego/parser.py @@ -0,0 +1,146 @@ +"""Rego language parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.edges import CodeEdge, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind, Parameter +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + compute_complexity, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import ( + add_call_edges, + collect_branches_and_calls, + descendants, + named_children, +) + +_EXTENSIONS = (".rego",) +_BRANCH_TYPES = frozenset({"literal"}) +_CALL_TYPES = frozenset({"expr_call"}) + + +class RegoParser: + """Parses Rego policy files into CodeGraph using tree-sitter.""" + + @property + def language(self) -> str: + return "rego" + + def __init__(self) -> None: + self._parser = Parser(get_language("rego")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="rego", root_path=file_path) + module_id = module_id_from_path(file_path) + _visit_root(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "rego", dir_path, _EXTENSIONS) + + +def _visit_root(root: Node, file_path: str, fallback_module_id: str, graph: CodeGraph) -> None: + module = _find_module(root) + package_name = _package_name(module) if module is not None else "" + module_id = package_name or fallback_module_id + add_module_node(root, file_path, module_id, graph) + if module is None: + return + for child in named_children(module): + if child.type == "import": + _extract_import(child, graph, module_id) + elif child.type == "policy": + for rule in named_children(child): + if rule.type == "rule": + _extract_rule(rule, file_path, module_id, graph) + + +def _find_module(root: Node) -> Node | None: + if root.type == "module": + return root + for child in named_children(root): + if child.type == "module": + return child + return None + + +def _package_name(module: Node | None) -> str: + if module is None: + return "" + seen_package = False + for child in named_children(module): + if child.type == "package": + seen_package = True + continue + if seen_package and child.type == "ref": + return node_text(child) + return "" + + +def _extract_rule(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + head = next((child for child in named_children(node) if child.type == "rule_head"), None) + if head is None: + return + name = _rule_name(head) + if not name: + return + rule_id = f"{module_id}:{name}" + body = next((child for child in named_children(node) if child.type == "rule_body"), None) + branches, calls = collect_branches_and_calls( + body, + file_path, + branch_types=_BRANCH_TYPES, + call_types=_CALL_TYPES, + ) + graph.nodes[rule_id] = CodeUnit( + id=rule_id, + name=name, + kind=NodeKind.FUNCTION, + location=make_location(node, file_path), + parameters=_rule_parameters(head), + cyclomatic_complexity=compute_complexity(branches), + branches=tuple(branches), + ) + add_contains_edge(graph, module_id, rule_id) + add_call_edges(graph, calls, rule_id, module_id, file_path) + + +def _rule_name(head: Node) -> str: + for child in named_children(head): + if child.type == "var": + return node_text(child) + return "" + + +def _rule_parameters(head: Node) -> tuple[Parameter, ...]: + for child in descendants(head): + if child.type == "rule_args": + params = [] + for arg in named_children(child): + text = node_text(arg) + if text: + params.append(Parameter(name=text)) + return tuple(params) + return () + + +def _extract_import(node: Node, graph: CodeGraph, module_id: str) -> None: + target = node.next_named_sibling + if target is not None and target.type == "ref": + text = node_text(target) + graph.dependencies.append(text) + graph.edges.append(CodeEdge(source_id=module_id, target_id=text, kind=EdgeKind.IMPORTS)) diff --git a/src/trailmark/parsers/sway/__init__.py b/src/trailmark/parsers/sway/__init__.py new file mode 100644 index 0000000..3c5eef1 --- /dev/null +++ b/src/trailmark/parsers/sway/__init__.py @@ -0,0 +1,5 @@ +"""Sway language parser for Trailmark.""" + +from trailmark.parsers.sway.parser import SwayParser + +__all__ = ["SwayParser"] diff --git a/src/trailmark/parsers/sway/parser.py b/src/trailmark/parsers/sway/parser.py new file mode 100644 index 0000000..3f76625 --- /dev/null +++ b/src/trailmark/parsers/sway/parser.py @@ -0,0 +1,185 @@ +"""Sway language parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.edges import CodeEdge, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + compute_complexity, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import ( + add_call_edges, + child_text, + collect_branches_and_calls, + extract_parameters, + first_child_type, + named_children, + return_type_after_params, +) + +_EXTENSIONS = (".sw",) +_BRANCH_TYPES = frozenset({"if_expression", "match_expression", "match_branch"}) +_CALL_TYPES = frozenset({"call_expression", "abi_call_expression"}) + + +class SwayParser: + """Parses Sway source files into CodeGraph using tree-sitter.""" + + @property + def language(self) -> str: + return "sway" + + def __init__(self) -> None: + self._parser = Parser(get_language("sway")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="sway", root_path=file_path) + module_id = module_id_from_path(file_path) + _visit_root(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "sway", dir_path, _EXTENSIONS) + + +def _visit_root(root: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + add_module_node(root, file_path, module_id, graph) + for child in named_children(root): + if child.type in {"use_statement", "use_item"}: + _extract_import(child, graph, module_id) + elif child.type == "abi_item": + _extract_abi(child, file_path, module_id, graph) + elif child.type == "struct_item": + _extract_struct(child, file_path, module_id, module_id, graph) + elif child.type == "impl_item": + _extract_impl(child, file_path, module_id, graph) + elif child.type == "function_item": + _extract_function(child, file_path, module_id, module_id, graph) + + +def _extract_abi(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"type_identifier", "identifier"})) + if not name: + return + abi_id = f"{module_id}:{name}" + graph.nodes[abi_id] = CodeUnit( + id=abi_id, + name=name, + kind=NodeKind.INTERFACE, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, abi_id) + body = first_child_type(node, frozenset({"declaration_list"})) + if body is None: + return + for child in named_children(body): + if child.type == "function_signature_item": + _extract_signature(child, file_path, module_id, abi_id, graph) + + +def _extract_impl(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + names = [node_text(child) for child in named_children(node) if child.type == "type_identifier"] + container_id = f"{module_id}:{names[0]}" if names else module_id + body = first_child_type(node, frozenset({"declaration_list"})) + if body is None: + return + for child in named_children(body): + if child.type == "function_item": + _extract_function(child, file_path, module_id, container_id, graph) + + +def _extract_struct( + node: Node, + file_path: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> None: + name = child_text(node, frozenset({"type_identifier", "identifier"})) + if not name: + return + struct_id = f"{module_id}:{name}" + graph.nodes[struct_id] = CodeUnit( + id=struct_id, + name=name, + kind=NodeKind.STRUCT, + location=make_location(node, file_path), + ) + add_contains_edge(graph, container_id, struct_id) + + +def _extract_signature( + node: Node, + file_path: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + method_id = f"{container_id}.{name}" + graph.nodes[method_id] = CodeUnit( + id=method_id, + name=name, + kind=NodeKind.METHOD, + location=make_location(node, file_path), + parameters=extract_parameters(node), + return_type=return_type_after_params(node), + ) + add_contains_edge(graph, container_id, method_id) + + +def _extract_function( + node: Node, + file_path: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + func_id = f"{container_id}.{name}" if container_id != module_id else f"{module_id}:{name}" + body = first_child_type(node, frozenset({"block"})) + branches, calls = collect_branches_and_calls( + body, + file_path, + branch_types=_BRANCH_TYPES, + call_types=_CALL_TYPES, + ) + graph.nodes[func_id] = CodeUnit( + id=func_id, + name=name, + kind=NodeKind.METHOD if container_id != module_id else NodeKind.FUNCTION, + location=make_location(node, file_path), + parameters=extract_parameters(node), + return_type=return_type_after_params(node), + cyclomatic_complexity=compute_complexity(branches), + branches=tuple(branches), + ) + add_contains_edge(graph, container_id, func_id) + for call in calls: + call_container = container_id if node_text(call[1]).strip().startswith("self.") else None + add_call_edges(graph, [call], func_id, module_id, file_path, call_container) + + +def _extract_import(node: Node, graph: CodeGraph, module_id: str) -> None: + target = node_text(node).strip().strip(";") + if target: + graph.dependencies.append(target) + graph.edges.append(CodeEdge(source_id=module_id, target_id=target, kind=EdgeKind.IMPORTS)) diff --git a/src/trailmark/parsers/tact/__init__.py b/src/trailmark/parsers/tact/__init__.py new file mode 100644 index 0000000..a4a76fe --- /dev/null +++ b/src/trailmark/parsers/tact/__init__.py @@ -0,0 +1,5 @@ +"""Tact language parser for Trailmark.""" + +from trailmark.parsers.tact.parser import TactParser + +__all__ = ["TactParser"] diff --git a/src/trailmark/parsers/tact/parser.py b/src/trailmark/parsers/tact/parser.py new file mode 100644 index 0000000..40dcf6c --- /dev/null +++ b/src/trailmark/parsers/tact/parser.py @@ -0,0 +1,193 @@ +"""Tact language parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.edges import CodeEdge, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + compute_complexity, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import ( + add_call_edges, + child_text, + collect_branches_and_calls, + extract_parameters, + first_child_type, + named_children, + return_type_after_params, +) + +_EXTENSIONS = (".tact",) +_BRANCH_TYPES = frozenset({"if_statement"}) +_CALL_TYPES = frozenset({"method_call_expression", "static_call_expression"}) +_FUNCTION_TYPES = frozenset( + { + "global_function", + "storage_function", + "native_function", + "asm_function", + "init_function", + "receive_function", + "bounced_function", + "external_function", + } +) + + +class TactParser: + """Parses Tact source files into CodeGraph using tree-sitter.""" + + @property + def language(self) -> str: + return "tact" + + def __init__(self) -> None: + self._parser = Parser(get_language("tact")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="tact", root_path=file_path) + module_id = module_id_from_path(file_path) + _visit_root(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "tact", dir_path, _EXTENSIONS) + + +def _visit_root(root: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + add_module_node(root, file_path, module_id, graph) + for child in named_children(root): + if child.type == "import": + _extract_import(child, graph, module_id) + elif child.type == "struct": + _extract_struct(child, file_path, module_id, module_id, graph) + elif child.type == "contract": + _extract_contract(child, file_path, module_id, graph) + elif child.type in _FUNCTION_TYPES: + _extract_function(child, file_path, module_id, module_id, graph) + + +def _extract_contract(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + contract_id = f"{module_id}:{name}" + graph.nodes[contract_id] = CodeUnit( + id=contract_id, + name=name, + kind=NodeKind.CONTRACT, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, contract_id) + body = first_child_type(node, frozenset({"contract_body"})) + if body is None: + return + for child in named_children(body): + if child.type == "struct": + _extract_struct(child, file_path, module_id, contract_id, graph) + elif child.type in _FUNCTION_TYPES: + _extract_function(child, file_path, module_id, contract_id, graph) + + +def _extract_struct( + node: Node, + file_path: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> None: + name = child_text(node, frozenset({"type_identifier", "identifier"})) + if not name: + return + struct_id = f"{module_id}:{name}" + graph.nodes[struct_id] = CodeUnit( + id=struct_id, + name=name, + kind=NodeKind.STRUCT, + location=make_location(node, file_path), + ) + add_contains_edge(graph, container_id, struct_id) + + +def _extract_function( + node: Node, + file_path: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> None: + base_name = _function_name(node) + func_id, name = _unique_function_id(base_name, module_id, container_id, graph) + body = first_child_type(node, frozenset({"function_body", "asm_function_body"})) + branches, calls = collect_branches_and_calls( + body, + file_path, + branch_types=_BRANCH_TYPES, + call_types=_CALL_TYPES, + ) + graph.nodes[func_id] = CodeUnit( + id=func_id, + name=name, + kind=NodeKind.METHOD if container_id != module_id else NodeKind.FUNCTION, + location=make_location(node, file_path), + parameters=extract_parameters(node), + return_type=return_type_after_params(node), + cyclomatic_complexity=compute_complexity(branches), + branches=tuple(branches), + attributes=(("tact_role", base_name),), + ) + add_contains_edge(graph, container_id, func_id) + add_call_edges(graph, calls, func_id, module_id, file_path, container_id) + + +def _function_name(node: Node) -> str: + if node.type == "init_function": + return "init" + if node.type == "receive_function": + return "receive" + if node.type == "bounced_function": + return "bounced" + if node.type == "external_function": + return "external" + return child_text(node, frozenset({"identifier", "func_identifier"})) or node.type + + +def _unique_function_id( + base_name: str, + module_id: str, + container_id: str, + graph: CodeGraph, +) -> tuple[str, str]: + prefix = f"{container_id}." if container_id != module_id else f"{module_id}:" + candidate = f"{prefix}{base_name}" + if candidate not in graph.nodes: + return candidate, base_name + + index = 2 + while True: + name = f"{base_name}_{index}" + candidate = f"{prefix}{name}" + if candidate not in graph.nodes: + return candidate, name + index += 1 + + +def _extract_import(node: Node, graph: CodeGraph, module_id: str) -> None: + target = node_text(node).removeprefix("import").strip().strip(";").strip("\"'") + if target: + graph.dependencies.append(target) + graph.edges.append(CodeEdge(source_id=module_id, target_id=target, kind=EdgeKind.IMPORTS)) diff --git a/src/trailmark/parsers/thrift/__init__.py b/src/trailmark/parsers/thrift/__init__.py new file mode 100644 index 0000000..3dedf7f --- /dev/null +++ b/src/trailmark/parsers/thrift/__init__.py @@ -0,0 +1,5 @@ +"""Thrift parser for Trailmark.""" + +from trailmark.parsers.thrift.parser import ThriftParser + +__all__ = ["ThriftParser"] diff --git a/src/trailmark/parsers/thrift/parser.py b/src/trailmark/parsers/thrift/parser.py new file mode 100644 index 0000000..45f9e1f --- /dev/null +++ b/src/trailmark/parsers/thrift/parser.py @@ -0,0 +1,171 @@ +"""Thrift parser using tree-sitter.""" + +from __future__ import annotations + +from pathlib import Path + +from tree_sitter import Node, Parser +from tree_sitter_language_pack import get_language + +from trailmark.models.edges import CodeEdge, EdgeKind +from trailmark.models.graph import CodeGraph +from trailmark.models.nodes import CodeUnit, NodeKind, TypeRef +from trailmark.parsers._common import ( + add_contains_edge, + add_module_node, + make_location, + module_id_from_path, + node_text, + parse_directory, +) +from trailmark.parsers._simple import ( + child_text, + extract_parameters, + first_child_type, + named_children, + return_type_after_params, +) + +_EXTENSIONS = (".thrift",) + + +class ThriftParser: + """Parses Thrift schemas into CodeGraph.""" + + @property + def language(self) -> str: + return "thrift" + + def __init__(self) -> None: + self._parser = Parser(get_language("thrift")) + + def parse_file(self, file_path: str) -> CodeGraph: + source = Path(file_path).read_bytes() + tree = self._parser.parse(source) + graph = CodeGraph(language="thrift", root_path=file_path) + module_id = _namespace_id(tree.root_node) or module_id_from_path(file_path) + _visit_root(tree.root_node, file_path, module_id, graph) + return graph + + def parse_directory(self, dir_path: str) -> CodeGraph: + return parse_directory(self.parse_file, "thrift", dir_path, _EXTENSIONS) + + +def _visit_root(root: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + add_module_node(root, file_path, module_id, graph) + for child in named_children(root): + if child.type == "include_statement": + _extract_import(child, graph, module_id) + elif child.type in {"struct_definition", "exception_definition", "union_definition"}: + _extract_struct(child, file_path, module_id, graph) + elif child.type == "service_definition": + _extract_service(child, file_path, module_id, graph) + elif child.type == "enum_definition": + _extract_enum(child, file_path, module_id, graph) + + +def _namespace_id(root: Node) -> str: + for child in named_children(root): + if child.type == "namespace_declaration": + parts = [ + node_text(part).lstrip(".") + for part in named_children(child) + if part.type == "namespace" + ] + return ".".join(parts) + return "" + + +def _extract_struct(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + struct_id = f"{module_id}:{name}" + graph.nodes[struct_id] = CodeUnit( + id=struct_id, + name=name, + kind=NodeKind.STRUCT, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, struct_id) + for child in named_children(node): + if child.type == "field": + _extract_field(child, file_path, struct_id, graph) + + +def _extract_field(node: Node, file_path: str, struct_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + field_id = f"{struct_id}.{name}" + graph.nodes[field_id] = CodeUnit( + id=field_id, + name=name, + kind=NodeKind.METHOD, + location=make_location(node, file_path), + return_type=return_type_after_params(node) or _field_type(node), + attributes=(("schema_role", "field"),), + ) + add_contains_edge(graph, struct_id, field_id) + + +def _field_type(node: Node) -> TypeRef | None: + type_node = first_child_type(node, frozenset({"type"})) + if type_node is None: + return None + return TypeRef(name=node_text(type_node)) + + +def _extract_service(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + service_id = f"{module_id}:{name}" + graph.nodes[service_id] = CodeUnit( + id=service_id, + name=name, + kind=NodeKind.INTERFACE, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, service_id) + for child in named_children(node): + if child.type == "function_definition": + _extract_function(child, file_path, service_id, graph) + + +def _extract_function(node: Node, file_path: str, service_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + func_id = f"{service_id}.{name}" + graph.nodes[func_id] = CodeUnit( + id=func_id, + name=name, + kind=NodeKind.METHOD, + location=make_location(node, file_path), + parameters=extract_parameters(node), + return_type=_field_type(node), + attributes=(("schema_role", "service_function"),), + ) + add_contains_edge(graph, service_id, func_id) + + +def _extract_enum(node: Node, file_path: str, module_id: str, graph: CodeGraph) -> None: + name = child_text(node, frozenset({"identifier"})) + if not name: + return + enum_id = f"{module_id}:{name}" + graph.nodes[enum_id] = CodeUnit( + id=enum_id, + name=name, + kind=NodeKind.ENUM, + location=make_location(node, file_path), + ) + add_contains_edge(graph, module_id, enum_id) + + +def _extract_import(node: Node, graph: CodeGraph, module_id: str) -> None: + target = node_text(node).removeprefix("include").strip().strip("\"'") + if target: + graph.dependencies.append(target) + graph.edges.append(CodeEdge(source_id=module_id, target_id=target, kind=EdgeKind.IMPORTS)) diff --git a/tests/fixtures/kat/func/taxonomy.expected.json b/tests/fixtures/kat/func/taxonomy.expected.json new file mode 100644 index 0000000..75ef2c5 --- /dev/null +++ b/tests/fixtures/kat/func/taxonomy.expected.json @@ -0,0 +1,171 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "imports", + "source": "taxonomy", + "target": "stdlib.fc" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:recv_internal" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "taxonomy:recv_internal", + "target": "taxonomy:helper" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:get_balance" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:helper" + } + ], + "language": "func", + "nodes": { + "taxonomy": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 14, + "file_path": "taxonomy.fc", + "start_col": 0, + "start_line": 1 + }, + "name": "taxonomy", + "parameters": [], + "return_type": null + }, + "taxonomy:get_balance": { + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:get_balance", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 9, + "file_path": "taxonomy.fc", + "start_col": 0, + "start_line": 7 + }, + "name": "get_balance", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "int" + } + }, + "taxonomy:helper": { + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:helper", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 13, + "file_path": "taxonomy.fc", + "start_col": 0, + "start_line": 11 + }, + "name": "helper", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "()" + } + }, + "taxonomy:recv_internal": { + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:recv_internal", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 5, + "file_path": "taxonomy.fc", + "start_col": 0, + "start_line": 3 + }, + "name": "recv_internal", + "parameters": [ + { + "default": null, + "name": "my_balance", + "type_ref": { + "generic_args": [], + "module": null, + "name": "int" + } + }, + { + "default": null, + "name": "msg_value", + "type_ref": { + "generic_args": [], + "module": null, + "name": "int" + } + }, + { + "default": null, + "name": "in_msg_full", + "type_ref": { + "generic_args": [], + "module": null, + "name": "cell" + } + }, + { + "default": null, + "name": "in_msg_body", + "type_ref": { + "generic_args": [], + "module": null, + "name": "slice" + } + } + ], + "return_type": { + "generic_args": [], + "module": null, + "name": "()" + } + } + }, + "root_path": "taxonomy.fc", + "subgraphs": {}, + "summary": { + "call_edges": 1, + "classes": 0, + "dependencies": [ + "stdlib.fc" + ], + "entrypoints": 0, + "functions": 3, + "proxies": 0, + "total_nodes": 4 + } +} diff --git a/tests/fixtures/kat/func/taxonomy.fc b/tests/fixtures/kat/func/taxonomy.fc new file mode 100644 index 0000000..c03520e --- /dev/null +++ b/tests/fixtures/kat/func/taxonomy.fc @@ -0,0 +1,13 @@ +#include "stdlib.fc"; + +() recv_internal(int my_balance, int msg_value, cell in_msg_full, slice in_msg_body) impure { + helper(); +} + +int get_balance() method_id { + return 1; +} + +() helper() impure { + return (); +} diff --git a/tests/fixtures/kat/graphql/taxonomy.expected.json b/tests/fixtures/kat/graphql/taxonomy.expected.json new file mode 100644 index 0000000..584c7a9 --- /dev/null +++ b/tests/fixtures/kat/graphql/taxonomy.expected.json @@ -0,0 +1,325 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:Query" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Query", + "target": "taxonomy:Query.user" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:Mutation" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Mutation", + "target": "taxonomy:Mutation.login" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:User" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:User", + "target": "taxonomy:User.id" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:User", + "target": "taxonomy:User.name" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:LoginResult" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:LoginResult", + "target": "taxonomy:LoginResult.ok" + }, + { + "attributes": { + "type_name": "User" + }, + "confidence": "certain", + "kind": "type_uses", + "source": "taxonomy:Query.user", + "target": "taxonomy:User" + }, + { + "attributes": { + "type_name": "LoginResult" + }, + "confidence": "certain", + "kind": "type_uses", + "source": "taxonomy:Mutation.login", + "target": "taxonomy:LoginResult" + } + ], + "language": "graphql", + "nodes": { + "taxonomy": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 22, + "file_path": "taxonomy.graphql", + "start_col": 0, + "start_line": 1 + }, + "name": "taxonomy", + "parameters": [], + "return_type": null + }, + "taxonomy:LoginResult": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:LoginResult", + "kind": "interface", + "location": { + "end_col": 1, + "end_line": 21, + "file_path": "taxonomy.graphql", + "start_col": 0, + "start_line": 19 + }, + "name": "LoginResult", + "parameters": [], + "return_type": null + }, + "taxonomy:LoginResult.ok": { + "attributes": { + "schema_role": "field" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:LoginResult.ok", + "kind": "method", + "location": { + "end_col": 14, + "end_line": 20, + "file_path": "taxonomy.graphql", + "start_col": 2, + "start_line": 20 + }, + "name": "ok", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "Boolean" + } + }, + "taxonomy:Mutation": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Mutation", + "kind": "interface", + "location": { + "end_col": 1, + "end_line": 12, + "file_path": "taxonomy.graphql", + "start_col": 0, + "start_line": 10 + }, + "name": "Mutation", + "parameters": [], + "return_type": null + }, + "taxonomy:Mutation.login": { + "attributes": { + "schema_role": "root_operation" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Mutation.login", + "kind": "method", + "location": { + "end_col": 36, + "end_line": 11, + "file_path": "taxonomy.graphql", + "start_col": 2, + "start_line": 11 + }, + "name": "login", + "parameters": [ + { + "default": null, + "name": "token", + "type_ref": { + "generic_args": [], + "module": null, + "name": "String" + } + } + ], + "return_type": { + "generic_args": [], + "module": null, + "name": "LoginResult" + } + }, + "taxonomy:Query": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Query", + "kind": "interface", + "location": { + "end_col": 1, + "end_line": 8, + "file_path": "taxonomy.graphql", + "start_col": 0, + "start_line": 6 + }, + "name": "Query", + "parameters": [], + "return_type": null + }, + "taxonomy:Query.user": { + "attributes": { + "schema_role": "root_operation" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Query.user", + "kind": "method", + "location": { + "end_col": 21, + "end_line": 7, + "file_path": "taxonomy.graphql", + "start_col": 2, + "start_line": 7 + }, + "name": "user", + "parameters": [ + { + "default": null, + "name": "id", + "type_ref": { + "generic_args": [], + "module": null, + "name": "ID" + } + } + ], + "return_type": { + "generic_args": [], + "module": null, + "name": "User" + } + }, + "taxonomy:User": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:User", + "kind": "interface", + "location": { + "end_col": 1, + "end_line": 17, + "file_path": "taxonomy.graphql", + "start_col": 0, + "start_line": 14 + }, + "name": "User", + "parameters": [], + "return_type": null + }, + "taxonomy:User.id": { + "attributes": { + "schema_role": "field" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:User.id", + "kind": "method", + "location": { + "end_col": 9, + "end_line": 15, + "file_path": "taxonomy.graphql", + "start_col": 2, + "start_line": 15 + }, + "name": "id", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "ID" + } + }, + "taxonomy:User.name": { + "attributes": { + "schema_role": "field" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:User.name", + "kind": "method", + "location": { + "end_col": 14, + "end_line": 16, + "file_path": "taxonomy.graphql", + "start_col": 2, + "start_line": 16 + }, + "name": "name", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "String" + } + } + }, + "root_path": "taxonomy.graphql", + "subgraphs": {}, + "summary": { + "call_edges": 0, + "classes": 0, + "dependencies": [], + "entrypoints": 0, + "functions": 5, + "proxies": 0, + "total_nodes": 10 + } +} diff --git a/tests/fixtures/kat/graphql/taxonomy.graphql b/tests/fixtures/kat/graphql/taxonomy.graphql new file mode 100644 index 0000000..4bd4823 --- /dev/null +++ b/tests/fixtures/kat/graphql/taxonomy.graphql @@ -0,0 +1,21 @@ +schema { + query: Query + mutation: Mutation +} + +type Query { + user(id: ID!): User +} + +type Mutation { + login(token: String!): LoginResult +} + +type User { + id: ID! + name: String +} + +type LoginResult { + ok: Boolean! +} diff --git a/tests/fixtures/kat/move/taxonomy.expected.json b/tests/fixtures/kat/move/taxonomy.expected.json new file mode 100644 index 0000000..d6dbb94 --- /dev/null +++ b/tests/fixtures/kat/move/taxonomy.expected.json @@ -0,0 +1,128 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:0x1::taxonomy" + }, + { + "confidence": "certain", + "kind": "imports", + "source": "taxonomy:0x1::taxonomy", + "target": "0x1::string" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:0x1::taxonomy", + "target": "taxonomy:0x1::taxonomy:main" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "taxonomy:0x1::taxonomy:main", + "target": "taxonomy:0x1::taxonomy:helper" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:0x1::taxonomy", + "target": "taxonomy:0x1::taxonomy:helper" + } + ], + "language": "move", + "nodes": { + "taxonomy": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 11, + "file_path": "taxonomy.move", + "start_col": 0, + "start_line": 1 + }, + "name": "taxonomy", + "parameters": [], + "return_type": null + }, + "taxonomy:0x1::taxonomy": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:0x1::taxonomy", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 11, + "file_path": "taxonomy.move", + "start_col": 0, + "start_line": 1 + }, + "name": "taxonomy", + "parameters": [], + "return_type": null + }, + "taxonomy:0x1::taxonomy:helper": { + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:0x1::taxonomy:helper", + "kind": "function", + "location": { + "end_col": 5, + "end_line": 9, + "file_path": "taxonomy.move", + "start_col": 4, + "start_line": 8 + }, + "name": "helper", + "parameters": [], + "return_type": null + }, + "taxonomy:0x1::taxonomy:main": { + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:0x1::taxonomy:main", + "kind": "function", + "location": { + "end_col": 5, + "end_line": 6, + "file_path": "taxonomy.move", + "start_col": 4, + "start_line": 4 + }, + "name": "main", + "parameters": [ + { + "default": null, + "name": "account", + "type_ref": null + } + ], + "return_type": null + } + }, + "root_path": "taxonomy.move", + "subgraphs": {}, + "summary": { + "call_edges": 1, + "classes": 0, + "dependencies": [ + "0x1::string" + ], + "entrypoints": 0, + "functions": 2, + "proxies": 0, + "total_nodes": 4 + } +} diff --git a/tests/fixtures/kat/move/taxonomy.move b/tests/fixtures/kat/move/taxonomy.move new file mode 100644 index 0000000..f67a60e --- /dev/null +++ b/tests/fixtures/kat/move/taxonomy.move @@ -0,0 +1,10 @@ +module 0x1::taxonomy { + use 0x1::string; + + public entry fun main(account: &signer) { + helper(); + } + + fun helper() { + } +} diff --git a/tests/fixtures/kat/proto/taxonomy.expected.json b/tests/fixtures/kat/proto/taxonomy.expected.json new file mode 100644 index 0000000..557e1e5 --- /dev/null +++ b/tests/fixtures/kat/proto/taxonomy.expected.json @@ -0,0 +1,237 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "imports", + "source": "example.auth", + "target": "common.proto" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:AuthService" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth:AuthService", + "target": "example.auth:AuthService.Login" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:LoginRequest" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth:LoginRequest", + "target": "example.auth:LoginRequest.token" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:LoginResponse" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth:LoginResponse", + "target": "example.auth:LoginResponse.ok" + }, + { + "attributes": { + "type_name": "LoginRequest" + }, + "confidence": "certain", + "kind": "type_uses", + "source": "example.auth:AuthService.Login", + "target": "example.auth:LoginRequest" + }, + { + "attributes": { + "type_name": "LoginResponse" + }, + "confidence": "certain", + "kind": "type_uses", + "source": "example.auth:AuthService.Login", + "target": "example.auth:LoginResponse" + } + ], + "language": "proto", + "nodes": { + "example.auth": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 18, + "file_path": "taxonomy.proto", + "start_col": 0, + "start_line": 1 + }, + "name": "example.auth", + "parameters": [], + "return_type": null + }, + "example.auth:AuthService": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:AuthService", + "kind": "interface", + "location": { + "end_col": 1, + "end_line": 9, + "file_path": "taxonomy.proto", + "start_col": 0, + "start_line": 7 + }, + "name": "AuthService", + "parameters": [], + "return_type": null + }, + "example.auth:AuthService.Login": { + "attributes": { + "schema_role": "rpc" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:AuthService.Login", + "kind": "method", + "location": { + "end_col": 51, + "end_line": 8, + "file_path": "taxonomy.proto", + "start_col": 2, + "start_line": 8 + }, + "name": "Login", + "parameters": [ + { + "default": null, + "name": "request", + "type_ref": { + "generic_args": [], + "module": null, + "name": "LoginRequest" + } + } + ], + "return_type": { + "generic_args": [], + "module": null, + "name": "LoginResponse" + } + }, + "example.auth:LoginRequest": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:LoginRequest", + "kind": "struct", + "location": { + "end_col": 1, + "end_line": 13, + "file_path": "taxonomy.proto", + "start_col": 0, + "start_line": 11 + }, + "name": "LoginRequest", + "parameters": [], + "return_type": null + }, + "example.auth:LoginRequest.token": { + "attributes": { + "schema_role": "field" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:LoginRequest.token", + "kind": "method", + "location": { + "end_col": 19, + "end_line": 12, + "file_path": "taxonomy.proto", + "start_col": 2, + "start_line": 12 + }, + "name": "token", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "string" + } + }, + "example.auth:LoginResponse": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:LoginResponse", + "kind": "struct", + "location": { + "end_col": 1, + "end_line": 17, + "file_path": "taxonomy.proto", + "start_col": 0, + "start_line": 15 + }, + "name": "LoginResponse", + "parameters": [], + "return_type": null + }, + "example.auth:LoginResponse.ok": { + "attributes": { + "schema_role": "field" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:LoginResponse.ok", + "kind": "method", + "location": { + "end_col": 14, + "end_line": 16, + "file_path": "taxonomy.proto", + "start_col": 2, + "start_line": 16 + }, + "name": "ok", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "bool" + } + } + }, + "root_path": "taxonomy.proto", + "subgraphs": {}, + "summary": { + "call_edges": 0, + "classes": 0, + "dependencies": [ + "common.proto" + ], + "entrypoints": 0, + "functions": 3, + "proxies": 0, + "total_nodes": 7 + } +} diff --git a/tests/fixtures/kat/proto/taxonomy.proto b/tests/fixtures/kat/proto/taxonomy.proto new file mode 100644 index 0000000..c58baea --- /dev/null +++ b/tests/fixtures/kat/proto/taxonomy.proto @@ -0,0 +1,17 @@ +syntax = "proto3"; + +package example.auth; + +import "common.proto"; + +service AuthService { + rpc Login (LoginRequest) returns (LoginResponse); +} + +message LoginRequest { + string token = 1; +} + +message LoginResponse { + bool ok = 1; +} diff --git a/tests/fixtures/kat/rego/taxonomy.expected.json b/tests/fixtures/kat/rego/taxonomy.expected.json new file mode 100644 index 0000000..eec88b5 --- /dev/null +++ b/tests/fixtures/kat/rego/taxonomy.expected.json @@ -0,0 +1,164 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "imports", + "source": "example.auth", + "target": "future.keywords.if" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:allow" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "example.auth:allow", + "target": "example.auth:helper" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:deny" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:helper" + } + ], + "language": "rego", + "nodes": { + "example.auth": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 16, + "file_path": "taxonomy.rego", + "start_col": 0, + "start_line": 1 + }, + "name": "example.auth", + "parameters": [], + "return_type": null + }, + "example.auth:allow": { + "branches": [ + { + "complexity_contribution": 1, + "condition": "literal", + "location": { + "end_col": 24, + "end_line": 6, + "file_path": "taxonomy.rego", + "start_col": 4, + "start_line": 6 + } + } + ], + "cyclomatic_complexity": 2, + "docstring": null, + "exception_types": [], + "id": "example.auth:allow", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 7, + "file_path": "taxonomy.rego", + "start_col": 0, + "start_line": 5 + }, + "name": "allow", + "parameters": [], + "return_type": null + }, + "example.auth:deny": { + "branches": [ + { + "complexity_contribution": 1, + "condition": "literal", + "location": { + "end_col": 15, + "end_line": 10, + "file_path": "taxonomy.rego", + "start_col": 4, + "start_line": 10 + } + } + ], + "cyclomatic_complexity": 2, + "docstring": null, + "exception_types": [], + "id": "example.auth:deny", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 11, + "file_path": "taxonomy.rego", + "start_col": 0, + "start_line": 9 + }, + "name": "deny", + "parameters": [], + "return_type": null + }, + "example.auth:helper": { + "branches": [ + { + "complexity_contribution": 1, + "condition": "literal", + "location": { + "end_col": 20, + "end_line": 14, + "file_path": "taxonomy.rego", + "start_col": 4, + "start_line": 14 + } + } + ], + "cyclomatic_complexity": 2, + "docstring": null, + "exception_types": [], + "id": "example.auth:helper", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 15, + "file_path": "taxonomy.rego", + "start_col": 0, + "start_line": 13 + }, + "name": "helper", + "parameters": [ + { + "default": null, + "name": "action", + "type_ref": null + } + ], + "return_type": null + } + }, + "root_path": "taxonomy.rego", + "subgraphs": {}, + "summary": { + "call_edges": 1, + "classes": 0, + "dependencies": [ + "future.keywords.if" + ], + "entrypoints": 0, + "functions": 3, + "proxies": 0, + "total_nodes": 4 + } +} diff --git a/tests/fixtures/kat/rego/taxonomy.rego b/tests/fixtures/kat/rego/taxonomy.rego new file mode 100644 index 0000000..af508e5 --- /dev/null +++ b/tests/fixtures/kat/rego/taxonomy.rego @@ -0,0 +1,15 @@ +package example.auth + +import future.keywords.if + +allow if { + helper(input.action) +} + +deny[msg] if { + msg := "no" +} + +helper(action) if { + action == "read" +} diff --git a/tests/fixtures/kat/sway/taxonomy.expected.json b/tests/fixtures/kat/sway/taxonomy.expected.json new file mode 100644 index 0000000..b27686b --- /dev/null +++ b/tests/fixtures/kat/sway/taxonomy.expected.json @@ -0,0 +1,180 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:Wallet" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Wallet", + "target": "taxonomy:Wallet.deposit" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:Balance" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Wallet", + "target": "taxonomy:Wallet.deposit" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "taxonomy:Wallet.deposit", + "target": "taxonomy:helper" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:helper" + } + ], + "language": "sway", + "nodes": { + "taxonomy": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 20, + "file_path": "taxonomy.sw", + "start_col": 0, + "start_line": 1 + }, + "name": "taxonomy", + "parameters": [], + "return_type": null + }, + "taxonomy:Balance": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Balance", + "kind": "struct", + "location": { + "end_col": 1, + "end_line": 9, + "file_path": "taxonomy.sw", + "start_col": 0, + "start_line": 7 + }, + "name": "Balance", + "parameters": [], + "return_type": null + }, + "taxonomy:Wallet": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet", + "kind": "interface", + "location": { + "end_col": 1, + "end_line": 5, + "file_path": "taxonomy.sw", + "start_col": 0, + "start_line": 3 + }, + "name": "Wallet", + "parameters": [], + "return_type": null + }, + "taxonomy:Wallet.deposit": { + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet.deposit", + "kind": "method", + "location": { + "end_col": 5, + "end_line": 14, + "file_path": "taxonomy.sw", + "start_col": 4, + "start_line": 12 + }, + "name": "deposit", + "parameters": [ + { + "default": null, + "name": "amount", + "type_ref": { + "generic_args": [], + "module": null, + "name": "u64" + } + } + ], + "return_type": null + }, + "taxonomy:helper": { + "branches": [ + { + "complexity_contribution": 1, + "condition": "amount > 0", + "location": { + "end_col": 39, + "end_line": 18, + "file_path": "taxonomy.sw", + "start_col": 4, + "start_line": 18 + } + } + ], + "cyclomatic_complexity": 2, + "docstring": null, + "exception_types": [], + "id": "taxonomy:helper", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 19, + "file_path": "taxonomy.sw", + "start_col": 0, + "start_line": 17 + }, + "name": "helper", + "parameters": [ + { + "default": null, + "name": "amount", + "type_ref": { + "generic_args": [], + "module": null, + "name": "u64" + } + } + ], + "return_type": { + "generic_args": [], + "module": null, + "name": "u64" + } + } + }, + "root_path": "taxonomy.sw", + "subgraphs": {}, + "summary": { + "call_edges": 1, + "classes": 0, + "dependencies": [], + "entrypoints": 0, + "functions": 2, + "proxies": 0, + "total_nodes": 5 + } +} diff --git a/tests/fixtures/kat/sway/taxonomy.sw b/tests/fixtures/kat/sway/taxonomy.sw new file mode 100644 index 0000000..f78e09b --- /dev/null +++ b/tests/fixtures/kat/sway/taxonomy.sw @@ -0,0 +1,19 @@ +contract; + +abi Wallet { + fn deposit(amount: u64); +} + +struct Balance { + amount: u64 +} + +impl Wallet for Contract { + fn deposit(amount: u64) { + helper(amount); + } +} + +pub fn helper(amount: u64) -> u64 { + if amount > 0 { amount } else { 0 } +} diff --git a/tests/fixtures/kat/tact/taxonomy.expected.json b/tests/fixtures/kat/tact/taxonomy.expected.json new file mode 100644 index 0000000..9011cf6 --- /dev/null +++ b/tests/fixtures/kat/tact/taxonomy.expected.json @@ -0,0 +1,288 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "imports", + "source": "taxonomy", + "target": "@stdlib/deploy" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:Ping" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:Wallet" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Wallet", + "target": "taxonomy:Wallet.init" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "taxonomy:Wallet.init", + "target": "taxonomy:Wallet.reply" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Wallet", + "target": "taxonomy:Wallet.receive" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "taxonomy:Wallet.receive", + "target": "taxonomy:Wallet.reply" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Wallet", + "target": "taxonomy:Wallet.external" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "taxonomy:Wallet.external", + "target": "taxonomy:Wallet.reply" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Wallet", + "target": "taxonomy:Wallet.bounced" + }, + { + "confidence": "certain", + "kind": "calls", + "source": "taxonomy:Wallet.bounced", + "target": "taxonomy:Wallet.reply" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy:Wallet", + "target": "taxonomy:Wallet.reply" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "taxonomy", + "target": "taxonomy:top" + } + ], + "language": "tact", + "nodes": { + "taxonomy": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 31, + "file_path": "taxonomy.tact", + "start_col": 0, + "start_line": 1 + }, + "name": "taxonomy", + "parameters": [], + "return_type": null + }, + "taxonomy:Ping": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Ping", + "kind": "struct", + "location": { + "end_col": 1, + "end_line": 5, + "file_path": "taxonomy.tact", + "start_col": 0, + "start_line": 3 + }, + "name": "Ping", + "parameters": [], + "return_type": null + }, + "taxonomy:Wallet": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet", + "kind": "contract", + "location": { + "end_col": 1, + "end_line": 27, + "file_path": "taxonomy.tact", + "start_col": 0, + "start_line": 7 + }, + "name": "Wallet", + "parameters": [], + "return_type": null + }, + "taxonomy:Wallet.bounced": { + "attributes": { + "tact_role": "bounced" + }, + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet.bounced", + "kind": "method", + "location": { + "end_col": 5, + "end_line": 22, + "file_path": "taxonomy.tact", + "start_col": 4, + "start_line": 20 + }, + "name": "bounced", + "parameters": [], + "return_type": null + }, + "taxonomy:Wallet.external": { + "attributes": { + "tact_role": "external" + }, + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet.external", + "kind": "method", + "location": { + "end_col": 5, + "end_line": 18, + "file_path": "taxonomy.tact", + "start_col": 4, + "start_line": 16 + }, + "name": "external", + "parameters": [], + "return_type": null + }, + "taxonomy:Wallet.init": { + "attributes": { + "tact_role": "init" + }, + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet.init", + "kind": "method", + "location": { + "end_col": 5, + "end_line": 10, + "file_path": "taxonomy.tact", + "start_col": 4, + "start_line": 8 + }, + "name": "init", + "parameters": [ + { + "default": null, + "name": "owner", + "type_ref": { + "generic_args": [], + "module": null, + "name": "Address" + } + } + ], + "return_type": null + }, + "taxonomy:Wallet.receive": { + "attributes": { + "tact_role": "receive" + }, + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet.receive", + "kind": "method", + "location": { + "end_col": 5, + "end_line": 14, + "file_path": "taxonomy.tact", + "start_col": 4, + "start_line": 12 + }, + "name": "receive", + "parameters": [], + "return_type": null + }, + "taxonomy:Wallet.reply": { + "attributes": { + "tact_role": "reply" + }, + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:Wallet.reply", + "kind": "method", + "location": { + "end_col": 5, + "end_line": 26, + "file_path": "taxonomy.tact", + "start_col": 4, + "start_line": 24 + }, + "name": "reply", + "parameters": [], + "return_type": null + }, + "taxonomy:top": { + "attributes": { + "tact_role": "top" + }, + "branches": [], + "cyclomatic_complexity": 1, + "docstring": null, + "exception_types": [], + "id": "taxonomy:top", + "kind": "function", + "location": { + "end_col": 1, + "end_line": 30, + "file_path": "taxonomy.tact", + "start_col": 0, + "start_line": 29 + }, + "name": "top", + "parameters": [], + "return_type": null + } + }, + "root_path": "taxonomy.tact", + "subgraphs": {}, + "summary": { + "call_edges": 4, + "classes": 0, + "dependencies": [ + "@stdlib/deploy" + ], + "entrypoints": 0, + "functions": 6, + "proxies": 0, + "total_nodes": 9 + } +} diff --git a/tests/fixtures/kat/tact/taxonomy.tact b/tests/fixtures/kat/tact/taxonomy.tact new file mode 100644 index 0000000..2b5616c --- /dev/null +++ b/tests/fixtures/kat/tact/taxonomy.tact @@ -0,0 +1,30 @@ +import "@stdlib/deploy"; + +struct Ping { + value: Int as uint32; +} + +contract Wallet { + init(owner: Address) { + self.reply(); + } + + receive("ping") { + self.reply(); + } + + external(msg: Ping) { + reply(); + } + + bounced(msg: bounced) { + reply(); + } + + fun reply() { + return; + } +} + +fun top() { +} diff --git a/tests/fixtures/kat/thrift/taxonomy.expected.json b/tests/fixtures/kat/thrift/taxonomy.expected.json new file mode 100644 index 0000000..9b95ac4 --- /dev/null +++ b/tests/fixtures/kat/thrift/taxonomy.expected.json @@ -0,0 +1,173 @@ +{ + "edges": [ + { + "confidence": "certain", + "kind": "imports", + "source": "example.auth", + "target": "common.thrift" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:LoginRequest" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth:LoginRequest", + "target": "example.auth:LoginRequest.token" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth", + "target": "example.auth:AuthService" + }, + { + "confidence": "certain", + "kind": "contains", + "source": "example.auth:AuthService", + "target": "example.auth:AuthService.login" + }, + { + "attributes": { + "type_name": "LoginRequest" + }, + "confidence": "certain", + "kind": "type_uses", + "source": "example.auth:AuthService.login", + "target": "example.auth:LoginRequest" + } + ], + "language": "thrift", + "nodes": { + "example.auth": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth", + "kind": "module", + "location": { + "end_col": 0, + "end_line": 12, + "file_path": "taxonomy.thrift", + "start_col": 0, + "start_line": 1 + }, + "name": "example.auth", + "parameters": [], + "return_type": null + }, + "example.auth:AuthService": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:AuthService", + "kind": "interface", + "location": { + "end_col": 1, + "end_line": 11, + "file_path": "taxonomy.thrift", + "start_col": 0, + "start_line": 9 + }, + "name": "AuthService", + "parameters": [], + "return_type": null + }, + "example.auth:AuthService.login": { + "attributes": { + "schema_role": "service_function" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:AuthService.login", + "kind": "method", + "location": { + "end_col": 33, + "end_line": 10, + "file_path": "taxonomy.thrift", + "start_col": 2, + "start_line": 10 + }, + "name": "login", + "parameters": [ + { + "default": null, + "name": "req", + "type_ref": { + "generic_args": [], + "module": null, + "name": "LoginRequest" + } + } + ], + "return_type": { + "generic_args": [], + "module": null, + "name": "bool" + } + }, + "example.auth:LoginRequest": { + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:LoginRequest", + "kind": "struct", + "location": { + "end_col": 1, + "end_line": 7, + "file_path": "taxonomy.thrift", + "start_col": 0, + "start_line": 5 + }, + "name": "LoginRequest", + "parameters": [], + "return_type": null + }, + "example.auth:LoginRequest.token": { + "attributes": { + "schema_role": "field" + }, + "branches": [], + "cyclomatic_complexity": null, + "docstring": null, + "exception_types": [], + "id": "example.auth:LoginRequest.token", + "kind": "method", + "location": { + "end_col": 17, + "end_line": 6, + "file_path": "taxonomy.thrift", + "start_col": 2, + "start_line": 6 + }, + "name": "token", + "parameters": [], + "return_type": { + "generic_args": [], + "module": null, + "name": "string" + } + } + }, + "root_path": "taxonomy.thrift", + "subgraphs": {}, + "summary": { + "call_edges": 0, + "classes": 0, + "dependencies": [ + "common.thrift" + ], + "entrypoints": 0, + "functions": 2, + "proxies": 0, + "total_nodes": 5 + } +} diff --git a/tests/fixtures/kat/thrift/taxonomy.thrift b/tests/fixtures/kat/thrift/taxonomy.thrift new file mode 100644 index 0000000..5b321f7 --- /dev/null +++ b/tests/fixtures/kat/thrift/taxonomy.thrift @@ -0,0 +1,11 @@ +include "common.thrift" + +namespace py example.auth + +struct LoginRequest { + 1: string token +} + +service AuthService { + bool login(1: LoginRequest req) +} diff --git a/tests/test_entrypoints.py b/tests/test_entrypoints.py index 973a800..f541ac3 100644 --- a/tests/test_entrypoints.py +++ b/tests/test_entrypoints.py @@ -74,6 +74,82 @@ def test_non_main_function_not_detected(self, tmp_path: Path) -> None: assert engine.attack_surface() == [] +class TestNewLanguageEntrypoints: + def test_move_public_entry_detected(self, tmp_path: Path) -> None: + (tmp_path / "m.move").write_text( + "module 0x1::m { public entry fun main(account: &signer) { } }\n" + ) + engine = QueryEngine.from_directory(str(tmp_path), language="move") + ids = {ep["node_id"] for ep in engine.attack_surface()} + assert "m:0x1::m:main" in ids + + def test_tact_receivers_detected(self, tmp_path: Path) -> None: + (tmp_path / "wallet.tact").write_text( + 'contract Wallet { init(owner: Address) { } receive("ping") { } }\n' + ) + engine = QueryEngine.from_directory(str(tmp_path), language="tact") + ids = {ep["node_id"] for ep in engine.attack_surface()} + assert "wallet:Wallet.init" in ids + assert "wallet:Wallet.receive" in ids + + def test_func_receivers_and_getters_detected(self, tmp_path: Path) -> None: + (tmp_path / "wallet.fc").write_text( + "() recv_internal() impure { }\nint get_balance() method_id { return 1; }\n" + ) + engine = QueryEngine.from_directory(str(tmp_path), language="func") + ids = {ep["node_id"] for ep in engine.attack_surface()} + assert "wallet:recv_internal" in ids + assert "wallet:get_balance" in ids + + def test_sway_abi_and_public_functions_detected(self, tmp_path: Path) -> None: + (tmp_path / "wallet.sw").write_text( + "contract;\nabi Wallet { fn deposit(amount: u64); }\n" + "pub fn helper(amount: u64) -> u64 { amount }\n" + ) + engine = QueryEngine.from_directory(str(tmp_path), language="sway") + ids = {ep["node_id"] for ep in engine.attack_surface()} + assert "wallet:Wallet.deposit" in ids + assert "wallet:helper" in ids + + def test_rego_policy_rules_detected(self, tmp_path: Path) -> None: + (tmp_path / "policy.rego").write_text( + 'package example.auth\nallow if { true }\ndeny[msg] if { msg := "no" }\n' + ) + engine = QueryEngine.from_directory(str(tmp_path), language="rego") + ids = {ep["node_id"] for ep in engine.attack_surface()} + assert "example.auth:allow" in ids + assert "example.auth:deny" in ids + + def test_schema_entrypoints_detect_operations_only(self, tmp_path: Path) -> None: + (tmp_path / "auth.proto").write_text( + 'syntax = "proto3"; package example.auth; ' + "service Auth { rpc Login (Request) returns (Response); } " + "message Request { string token = 1; } message Response { bool ok = 1; }\n" + ) + (tmp_path / "auth.thrift").write_text( + "namespace py example.auth\n" + "struct Request { 1: string token }\n" + "service Auth { bool login(1: Request req) }\n" + ) + (tmp_path / "schema.graphql").write_text( + "type Query { user(id: ID!): User }\ntype User { id: ID! }\n" + ) + + proto = QueryEngine.from_directory(str(tmp_path), language="proto") + thrift = QueryEngine.from_directory(str(tmp_path), language="thrift") + graphql = QueryEngine.from_directory(str(tmp_path), language="graphql") + + assert {ep["node_id"] for ep in proto.attack_surface()} == { + "example.auth:Auth.Login", + } + assert {ep["node_id"] for ep in thrift.attack_surface()} == { + "example.auth:Auth.login", + } + assert {ep["node_id"] for ep in graphql.attack_surface()} == { + "schema:Query.user", + } + + class TestPyprojectScripts: def test_pyproject_script_overrides_main_heuristic(self, tmp_path: Path) -> None: """A pyproject.toml script target beats the generic main heuristic.""" diff --git a/tests/test_kat_parsers.py b/tests/test_kat_parsers.py index 670efc5..086550f 100644 --- a/tests/test_kat_parsers.py +++ b/tests/test_kat_parsers.py @@ -30,19 +30,27 @@ ("c_sharp", "taxonomy"), ("dart", "taxonomy"), ("erlang", "taxonomy"), + ("func", "taxonomy"), ("go", "taxonomy"), + ("graphql", "taxonomy"), ("haskell", "taxonomy"), ("java", "Taxonomy"), ("javascript", "taxonomy"), ("kotlin", "taxonomy"), ("masm", "taxonomy"), + ("move", "taxonomy"), ("objc", "Taxonomy"), ("php", "taxonomy"), + ("proto", "taxonomy"), ("python", "taxonomy"), + ("rego", "taxonomy"), ("ruby", "taxonomy"), ("rust", "taxonomy"), ("solidity", "taxonomy"), ("swift", "taxonomy"), + ("sway", "taxonomy"), + ("tact", "taxonomy"), + ("thrift", "taxonomy"), ("typescript", "taxonomy"), ] diff --git a/tests/test_new_language_parsers.py b/tests/test_new_language_parsers.py new file mode 100644 index 0000000..3df0aa1 --- /dev/null +++ b/tests/test_new_language_parsers.py @@ -0,0 +1,61 @@ +"""Focused tests for newer language parser guarantees.""" + +from __future__ import annotations + +from pathlib import Path + +from trailmark.models.edges import EdgeKind +from trailmark.parse import parse_file + + +def test_tact_direct_calls_are_call_edges(tmp_path: Path) -> None: + source = tmp_path / "wallet.tact" + source.write_text( + 'contract Wallet { receive("ping") { self.reply(); } fun reply() { return; } }\n' + ) + + graph = parse_file(str(source), language="tact") + call_edges = { + (edge.source_id, edge.target_id) for edge in graph.edges if edge.kind == EdgeKind.CALLS + } + + assert ("wallet:Wallet.receive", "wallet:Wallet.reply") in call_edges + + +def test_tact_repeated_receivers_get_distinct_nodes(tmp_path: Path) -> None: + source = tmp_path / "wallet.tact" + source.write_text( + 'contract Wallet { receive("ping") { } receive("pong") { } fun reply() { } }\n' + ) + + graph = parse_file(str(source), language="tact") + + assert "wallet:Wallet.receive" in graph.nodes + assert "wallet:Wallet.receive_2" in graph.nodes + + +def test_schema_languages_do_not_emit_runtime_call_edges(tmp_path: Path) -> None: + samples = [ + ( + "auth.proto", + "proto", + 'syntax = "proto3"; service Auth { rpc Login (Request) returns (Response); } ' + "message Request { string token = 1; } message Response { bool ok = 1; }\n", + ), + ( + "auth.thrift", + "thrift", + "struct Request { 1: string token }\nservice Auth { bool login(1: Request req) }\n", + ), + ( + "schema.graphql", + "graphql", + "type Query { user(id: ID!): User }\ntype User { id: ID! }\n", + ), + ] + + for filename, language, text in samples: + path = tmp_path / filename + path.write_text(text) + graph = parse_file(str(path), language=language) + assert [edge for edge in graph.edges if edge.kind == EdgeKind.CALLS] == [] diff --git a/tests/test_parse_api.py b/tests/test_parse_api.py index bab5af6..d745730 100644 --- a/tests/test_parse_api.py +++ b/tests/test_parse_api.py @@ -105,6 +105,29 @@ def test_detect_languages_matches_javascript_parser_extensions(self, tmp_path: P # detect_languages should mirror that and report "javascript". assert detect_languages(str(tmp_path)) == ["javascript"] + def test_detect_languages_matches_new_parser_extensions(self, tmp_path: Path) -> None: + (tmp_path / "a.move").write_text("") + (tmp_path / "b.tact").write_text("") + (tmp_path / "c.fc").write_text("") + (tmp_path / "d.func").write_text("") + (tmp_path / "e.sw").write_text("") + (tmp_path / "f.rego").write_text("") + (tmp_path / "g.proto").write_text("") + (tmp_path / "h.thrift").write_text("") + (tmp_path / "i.graphql").write_text("") + (tmp_path / "j.gql").write_text("") + + detected = detect_languages(str(tmp_path)) + + assert "move" in detected + assert "tact" in detected + assert "func" in detected + assert "sway" in detected + assert "rego" in detected + assert "proto" in detected + assert "thrift" in detected + assert "graphql" in detected + class TestFileExtensionHelper: """Direct coverage of trailmark.parse._file_extension."""