"""Powershell extractor. Moved verbatim from graphify/extract.py.""" from __future__ import annotations import re from pathlib import Path from typing import Any from graphify.extractors.base import _file_stem, _make_id, _read_text def extract_powershell(path: Path) -> dict: """Extract functions, classes, methods, and using statements from a .ps1 file.""" try: import tree_sitter_powershell as tsps from tree_sitter import Language, Parser except ImportError: return {"nodes": [], "edges": [], "error": "tree_sitter_powershell not installed"} try: language = Language(tsps.language()) parser = Parser(language) source = path.read_bytes() tree = parser.parse(source) root = tree.root_node except Exception as e: return {"nodes": [], "edges": [], "error": str(e)} stem = _file_stem(path) str_path = str(path) nodes: list[dict] = [] edges: list[dict] = [] seen_ids: set[str] = set() function_bodies: list[tuple[str, Any]] = [] def add_node(nid: str, label: str, line: int) -> None: if nid not in seen_ids: seen_ids.add(nid) nodes.append({"id": nid, "label": label, "file_type": "code", "source_file": str_path, "source_location": f"L{line}"}) def add_edge(src: str, tgt: str, relation: str, line: int, confidence: str = "EXTRACTED", weight: float = 1.0, context: str | None = None) -> None: edge = {"source": src, "target": tgt, "relation": relation, "confidence": confidence, "source_file": str_path, "source_location": f"L{line}", "weight": weight} if context: edge["context"] = context edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) _PS_SKIP = frozenset({ "using", "return", "if", "else", "elseif", "foreach", "for", "while", "do", "switch", "try", "catch", "finally", "throw", "break", "continue", "exit", "param", "begin", "process", "end", # Import commands — handled as import edges, not function calls "import-module", }) def _find_script_block_body(node): for child in node.children: if child.type == "script_block": for sc in child.children: if sc.type == "script_block_body": return sc return child return None def ensure_named_node(name: str, line: int) -> str: nid = _make_id(stem, name) if nid in seen_ids: return nid nid = _make_id(name) if nid not in seen_ids: # The name isn't defined in this file, so this is a cross-file reference # (e.g. a `Thing` type annotation imported from another module). Emit a # SOURCELESS stub — like the inheritance-base path below — so the # corpus-level rewire can collapse it onto the real definition. A sourced # stub here makes _disambiguate_colliding_node_ids bake the referencing # file's path (with extension) into the id and blocks the rewire, which is # the phantom-duplicate-node bug (#1402). seen_ids.add(nid) nodes.append({ "id": nid, "label": name, "file_type": "code", "source_file": "", "source_location": "", "origin_file": str_path, }) return nid def _ps_type_name(type_literal_node) -> str | None: """Drill into a type_literal node and return the inner type_identifier text.""" if type_literal_node is None: return None for spec in type_literal_node.children: if spec.type != "type_spec": continue for tname in spec.children: if tname.type != "type_name": continue for tid in tname.children: if tid.type == "type_identifier": return _read_text(tid, source) return None def walk(node, parent_class_nid: str | None = None) -> None: t = node.type if t == "function_statement": name_node = next((c for c in node.children if c.type == "function_name"), None) if name_node: func_name = _read_text(name_node, source) line = node.start_point[0] + 1 func_nid = _make_id(stem, func_name) add_node(func_nid, f"{func_name}()", line) add_edge(file_nid, func_nid, "contains", line) body = _find_script_block_body(node) if body: function_bodies.append((func_nid, body)) # Also walk the body during the main pass so that # Import-Module / dot-source inside functions emit # file-level imports_from edges (#1331). walk(body, parent_class_nid) return if t == "class_statement": name_node = next((c for c in node.children if c.type == "simple_name"), None) if name_node: class_name = _read_text(name_node, source) line = node.start_point[0] + 1 class_nid = _make_id(stem, class_name) add_node(class_nid, class_name, line) add_edge(file_nid, class_nid, "contains", line) # Base type(s) after ':'. PowerShell has no syntactic base vs # interface split, so (matching the C# convention) treat the # first base as the superclass (inherits) and the rest as # interfaces (implements). Bases are the simple_name children # after the ':' token. colon_seen = False base_index = 0 for child in node.children: if child.type == ":": colon_seen = True elif colon_seen and child.type == "simple_name": base_nid = ensure_named_node(_read_text(child, source), line) if base_nid != class_nid: rel = "inherits" if base_index == 0 else "implements" add_edge(class_nid, base_nid, rel, line) base_index += 1 for child in node.children: walk(child, parent_class_nid=class_nid) return if t == "class_property_definition" and parent_class_nid: type_literal = next((c for c in node.children if c.type == "type_literal"), None) type_name = _ps_type_name(type_literal) if type_name: line = node.start_point[0] + 1 target_nid = ensure_named_node(type_name, line) if target_nid != parent_class_nid: add_edge(parent_class_nid, target_nid, "references", line, context="field") return if t == "class_method_definition": name_node = next((c for c in node.children if c.type == "simple_name"), None) if name_node: method_name = _read_text(name_node, source) line = node.start_point[0] + 1 if parent_class_nid: method_nid = _make_id(parent_class_nid, method_name) add_node(method_nid, f".{method_name}()", line) add_edge(parent_class_nid, method_nid, "method", line) else: method_nid = _make_id(stem, method_name) add_node(method_nid, f"{method_name}()", line) add_edge(file_nid, method_nid, "contains", line) # Return type: type_literal sibling of simple_name return_type_literal = next( (c for c in node.children if c.type == "type_literal"), None) return_type_name = _ps_type_name(return_type_literal) if return_type_name: target_nid = ensure_named_node(return_type_name, line) if target_nid != method_nid: add_edge(method_nid, target_nid, "references", line, context="return_type") # Parameter types: class_method_parameter_list param_list = next( (c for c in node.children if c.type == "class_method_parameter_list"), None) if param_list is not None: for p in param_list.children: if p.type != "class_method_parameter": continue ptype_literal = next( (c for c in p.children if c.type == "type_literal"), None) ptype_name = _ps_type_name(ptype_literal) if not ptype_name: continue p_line = p.start_point[0] + 1 target_nid = ensure_named_node(ptype_name, p_line) if target_nid != method_nid: add_edge(method_nid, target_nid, "references", p_line, context="parameter_type") body = _find_script_block_body(node) if body: function_bodies.append((method_nid, body)) return if t == "command": # Dot-sourcing: `. ./Shared.psm1` # Uses command_invokation_operator '.' + command_name_expr (not command_name) invoke_op = next( (c for c in node.children if c.type == "command_invokation_operator"), None ) if invoke_op is not None and _read_text(invoke_op, source).strip() == ".": name_expr = next( (c for c in node.children if c.type == "command_name_expr"), None ) if name_expr is not None: name_node = next( (c for c in name_expr.children if c.type == "command_name"), None ) if name_node: raw_path = _read_text(name_node, source) # Strip relative path prefix (./ or .\ or just the dot) module_stem = re.sub(r'^[./\\]+', '', raw_path) # Drop extension to get bare module name module_stem = re.sub(r'\.[^.]+$', '', module_stem).replace('\\', '/') module_name = module_stem.split('/')[-1] if module_name: add_edge(file_nid, _make_id(module_name), "imports_from", node.start_point[0] + 1) return cmd_name_node = next((c for c in node.children if c.type == "command_name"), None) if cmd_name_node: cmd_text = _read_text(cmd_name_node, source).lower() if cmd_text == "using": tokens = [] for child in node.children: if child.type == "command_elements": for el in child.children: if el.type == "generic_token": tokens.append(_read_text(el, source)) module_tokens = [t for t in tokens if t.lower() not in ("namespace", "module", "assembly")] if module_tokens: module_name = module_tokens[-1].split(".")[-1] add_edge(file_nid, _make_id(module_name), "imports_from", node.start_point[0] + 1) elif cmd_text == "import-module": # Collect generic_token args; skip command_parameter flags like -Name # The module name is the first generic_token (or the one after -Name) module_name: str | None = None expect_name = False for child in node.children: if child.type != "command_elements": continue for el in child.children: if el.type == "command_parameter": param_text = _read_text(el, source).lstrip("-").lower() expect_name = param_text in ("name", "n") elif el.type == "generic_token": token = _read_text(el, source) if module_name is None or expect_name: module_name = token expect_name = False if module_name: # Strip extension; keep only the stem for the node ID bare = re.sub(r'\.[^.]+$', '', module_name).split('/')[-1].split('\\')[-1] if bare: add_edge(file_nid, _make_id(bare), "imports_from", node.start_point[0] + 1) return for child in node.children: walk(child, parent_class_nid) walk(root) label_to_nid = {n["label"].strip("()").lstrip(".").lower(): n["id"] for n in nodes} seen_call_pairs: set[tuple[str, str]] = set() raw_calls: list[dict] = [] def walk_calls(node, caller_nid: str) -> None: if node.type in ("function_statement", "class_statement"): return if node.type == "command": cmd_name_node = next((c for c in node.children if c.type == "command_name"), None) if cmd_name_node: cmd_text = _read_text(cmd_name_node, source) if cmd_text.lower() not in _PS_SKIP: tgt_nid = label_to_nid.get(cmd_text.lower()) if tgt_nid and tgt_nid != caller_nid: pair = (caller_nid, tgt_nid) if pair not in seen_call_pairs: seen_call_pairs.add(pair) add_edge(caller_nid, tgt_nid, "calls", node.start_point[0] + 1, confidence="EXTRACTED", weight=1.0) elif cmd_text: raw_calls.append({ "caller_nid": caller_nid, "callee": cmd_text, "is_member_call": False, "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", }) for child in node.children: walk_calls(child, caller_nid) for caller_nid, body_node in function_bodies: walk_calls(body_node, caller_nid) clean_edges = [e for e in edges if e["source"] in seen_ids and (e["target"] in seen_ids or e["relation"] in ("imports_from", "imports"))] return {"nodes": nodes, "edges": clean_edges, "raw_calls": raw_calls} _PSD1_IMPORT_KEYS = frozenset({"RootModule", "NestedModules", "RequiredModules"}) def _psd1_collect_string_literals(node, source: bytes) -> list[str]: """Recursively collect all string_literal text values under *node*.""" results: list[str] = [] def _walk(n) -> None: if n.type == "string_literal": raw = source[n.start_byte:n.end_byte].decode(errors="replace") # Strip surrounding quote chars (' or ") results.append(raw.strip("'\"")) return for child in n.children: _walk(child) _walk(node) return results def _psd1_module_name(raw: str) -> str: """Derive a bare module name from a raw string value. e.g. 'MyModule.psm1' → 'MyModule', './sub/Util.psm1' → 'Util', 'PSReadLine' → 'PSReadLine' """ # Strip path prefix and extension name = raw.replace("\\", "/").split("/")[-1] name = re.sub(r"\.[^.]+$", "", name) # remove last extension return name.strip() def extract_powershell_manifest(path: Path) -> dict: """Extract module dependency edges from a PowerShell .psd1 manifest file. .psd1 files are PowerShell data hashtables, not scripts. tree-sitter-powershell parses them correctly (they are syntactically valid PS). We walk the AST looking for RootModule, NestedModules, and RequiredModules keys and emit imports_from edges for every referenced module. RequiredModules supports two forms: - Simple string: 'PSReadLine' - Module specification: @{ ModuleName = 'Pester'; ModuleVersion = '5.0' } For the hashtable form we only follow the ModuleName key. """ try: import tree_sitter_powershell as tsps from tree_sitter import Language, Parser except ImportError: return {"nodes": [], "edges": [], "error": "tree_sitter_powershell not installed"} try: language = Language(tsps.language()) parser = Parser(language) source = path.read_bytes() tree = parser.parse(source) root = tree.root_node except Exception as e: return {"nodes": [], "edges": [], "error": str(e)} str_path = str(path) nodes: list[dict] = [] edges: list[dict] = [] seen_ids: set[str] = set() def add_node(nid: str, label: str, line: int) -> None: if nid not in seen_ids: seen_ids.add(nid) nodes.append({"id": nid, "label": label, "file_type": "code", "source_file": str_path, "source_location": f"L{line}"}) def add_import_edge(src: str, module_raw: str, line: int) -> None: name = _psd1_module_name(module_raw) if not name: return tgt_nid = _make_id(name) edges.append({ "source": src, "target": tgt_nid, "relation": "imports_from", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{line}", "weight": 1.0, "context": "import", }) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) def walk_manifest(node) -> None: """Walk the AST and emit edges for import-relevant hash_entry nodes.""" if node.type != "hash_entry": for child in node.children: walk_manifest(child) return # Identify the key key_node = next((c for c in node.children if c.type == "key_expression"), None) if key_node is None: return key_text = source[key_node.start_byte:key_node.end_byte].decode(errors="replace").strip() if key_text not in _PSD1_IMPORT_KEYS: # Still recurse in case there are nested hashes (e.g. ModuleVersion entries # contain sub-hashes, but we only care about top-level keys for imports) return line = node.start_point[0] + 1 value_node = next((c for c in node.children if c.type == "pipeline"), None) if value_node is None: return if key_text == "RootModule": # Value is a single string strings = _psd1_collect_string_literals(value_node, source) for s in strings: add_import_edge(file_nid, s, line) elif key_text == "NestedModules": # Value is a string or @('a', 'b', ...) array — collect all string literals strings = _psd1_collect_string_literals(value_node, source) for s in strings: add_import_edge(file_nid, s, line) elif key_text == "RequiredModules": # Two forms: # 1) 'SimpleModule' — direct string literals in the array # 2) @{ ModuleName = 'Foo'; ModuleVersion = '2.0' } — use ModuleName only # # Strategy: walk the value for hash_entry nodes whose key is 'ModuleName'; # collect their string values. For the remaining string_literal nodes that # are NOT inside a hash_entry subtree, treat them as simple module names. module_name_strings: list[str] = [] inside_hash_entries: set[int] = set() # byte offsets of handled strings def find_modulename_entries(n) -> None: if n.type == "hash_entry": sub_key = next((c for c in n.children if c.type == "key_expression"), None) if sub_key is not None: sk_text = source[sub_key.start_byte:sub_key.end_byte].decode(errors="replace").strip() # Collect strings inside *all* sub-keys so we can exclude them for c in n.children: if c.type == "pipeline": for s_node in _collect_string_nodes(c): inside_hash_entries.add(s_node.start_byte) if sk_text == "ModuleName": for c in n.children: if c.type == "pipeline": for s in _psd1_collect_string_literals(c, source): module_name_strings.append(s) return # don't recurse further into this hash_entry for child in n.children: find_modulename_entries(child) def _collect_string_nodes(n): """Return all string_literal nodes in subtree.""" if n.type == "string_literal": yield n return for child in n.children: yield from _collect_string_nodes(child) find_modulename_entries(value_node) # Now gather direct string literals not inside hash entries direct_strings: list[str] = [] for s_node in _collect_string_nodes(value_node): if s_node.start_byte not in inside_hash_entries: raw = source[s_node.start_byte:s_node.end_byte].decode(errors="replace") direct_strings.append(raw.strip("'\"")) for s in direct_strings + module_name_strings: add_import_edge(file_nid, s, line) walk_manifest(root) return {"nodes": nodes, "edges": edges, "raw_calls": []}