adl-toolkit 0.0.0-stage → 0.1.1

package/adl_drift/lang/python/cfg.py

@@ -0,0 +1,537 @@
+"""Python CFG construction (P2.18, ported unchanged onto the P4.6 seam).
+
+`PythonCFGBuilder` is the exact pre-P4.6 `graph.cfg.CFGBuilder` dispatch
+logic (statement-kind CST walk, literal/taint-source detection on RHS text),
+now a `graph.cfg.CFGBuilderBase` subclass reading its taint-source list from
+`self.config.taint_sources` (a `CFGConfig`) instead of a hardcoded module
+constant. No behavioral change for Python — `default_python_cfg_config()`
+below packages the original `TAINT_SOURCES` list into that config, so every
+existing call site (which never passed a config) sees byte-identical output.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+from tree_sitter import Node
+
+from adl_drift.graph.cfg import (
+    CFGBlock,
+    CFGBuilderBase,
+    CFGDef,
+    _is_literal,
+    _node_pos,
+    _text,
+)
+from adl_drift.lang.base import CFGConfig
+
+__all__ = ["TAINT_SOURCES", "PythonCFGBuilder", "default_python_cfg_config"]
+
+# ─── taint source prefixes (P2.18/P2.19 original table) ──────────────────────
+
+TAINT_SOURCES: Tuple[str, ...] = (
+    # HTTP request objects (Django / Flask / ASGI / Starlette …)
+    "request.GET", "request.POST", "request.body", "request.json",
+    "request.form", "request.args", "request.data", "request.values",
+    "request.FILES", "request.COOKIES",
+    "request.get_json", "request.get_data",
+    # WSGI / ASGI environ
+    "environ.get", "scope.get",
+    # OS env-vars
+    "os.environ.get", "os.getenv",
+    # Other user inputs
+    "sys.argv", "input(",
+    # DB results (can contain user data that should not be trusted)
+    "cursor.fetchone", "cursor.fetchall", "cursor.fetchmany",
+    # Path / file args coming in from outside
+    "open(",
+)
+
+# ─── statement types that terminate a block (no normal fall-through) ─────────
+
+_TERMINATOR_TYPES: frozenset = frozenset({
+    "return_statement", "raise_statement",
+    "break_statement", "continue_statement",
+})
+
+# ─── statement types with no control-flow split (fall straight through) ──────
+# Note: expression_statement is NOT here — it's handled separately in
+# _process_stmt because tree-sitter Python wraps assignment / augmented_assignment
+# / call nodes inside expression_statement.
+
+_SIMPLE_STMT_TYPES: frozenset = frozenset({
+    "delete_statement",
+    "import_statement", "import_from_statement",
+    "assert_statement",
+    "global_statement", "nonlocal_statement",
+    "pass_statement",
+    "print_statement",  # Python 2 guard
+})
+
+
+# ─── helper functions (Python-CST-coupled) ────────────────────────────────────
+
+def _real_children(node: Node) -> List[Node]:
+    """Non-extra, non-punctuation children."""
+    return [c for c in node.children
+            if not c.is_extra and c.type not in (",", ";", ":", "\n", "(", ")")]
+
+
+def _lhs_names(node: Node, source: bytes) -> List[str]:
+    """Extract identifier names from an assignment LHS (handles tuples, lists)."""
+    if node.type == "identifier":
+        return [_text(node, source)]
+    names: List[str] = []
+    for child in node.children:
+        if child.type not in (",", "(", ")", "[", "]", "*"):
+            names.extend(_lhs_names(child, source))
+    return names
+
+
+def _uses_of(node: Node) -> List:
+    """Convenience wrapper: collect uses from *node* and return as a new list."""
+    from adl_drift.graph.cfg import CFGUse
+    out: List = []
+    _collect_uses(node, out)
+    return out
+
+
+def _collect_uses(node: Node, uses: List, _visited: Optional[set] = None) -> None:
+    """Walk *node* and append identifier reads to *uses*.
+
+    Skips:
+      - LHS of assignments (handled separately)
+      - Keyword argument names (``kwarg=val``)
+      - Attribute name parts (``obj.attr`` → only ``obj`` is a read)
+      - Type annotation nodes
+      - Function / class definition names
+    """
+    from adl_drift.graph.cfg import CFGUse
+
+    if _visited is None:
+        _visited = set()
+    nid = id(node)
+    if nid in _visited:
+        return
+    _visited.add(nid)
+
+    t = node.type
+
+    if t == "assignment":
+        rhs = node.child_by_field_name("right")
+        if rhs:
+            _collect_uses(rhs, uses, _visited)
+        return
+
+    if t == "augmented_assignment":
+        lhs = node.child_by_field_name("left")
+        rhs = node.child_by_field_name("right")
+        if lhs and lhs.type == "identifier":
+            uses.append(CFGUse(_text(lhs, b""), _node_pos(lhs)))
+        if rhs:
+            _collect_uses(rhs, uses, _visited)
+        return
+
+    if t == "attribute":
+        # obj.attr — only obj is a read; attr is a field name
+        obj = node.child_by_field_name("object")
+        if obj:
+            _collect_uses(obj, uses, _visited)
+        return
+
+    if t == "keyword_argument":
+        val = node.child_by_field_name("value")
+        if val:
+            _collect_uses(val, uses, _visited)
+        return
+
+    if t in ("function_definition", "class_definition", "lambda"):
+        # Don't descend into nested defs (separate function analysis)
+        name_node = node.child_by_field_name("name")
+        if name_node and name_node.type == "identifier":
+            uses.append(CFGUse(_text(name_node, b""), _node_pos(name_node)))
+        return
+
+    if t == "type" or t == "type_annotation":
+        return
+
+    if t == "identifier":
+        name = node.text.decode("utf-8", errors="replace") if node.text else ""
+        if name:
+            uses.append(CFGUse(name, _node_pos(node)))
+        return
+
+    for child in node.children:
+        _collect_uses(child, uses, _visited)
+
+
+def _collect_param_names(params_node: Node, source: bytes) -> List[str]:
+    """Extract all parameter identifier names from a ``parameters`` CST node."""
+    names: List[str] = []
+    for child in params_node.children:
+        ct = child.type
+        if ct == "identifier":
+            names.append(_text(child, source))
+        elif ct in ("typed_parameter", "default_parameter", "typed_default_parameter",
+                    "list_splat_pattern", "dictionary_splat_pattern"):
+            inner = child.child_by_field_name("name") or (
+                child.children[0] if child.children else None
+            )
+            if inner and inner.type == "identifier":
+                names.append(_text(inner, source))
+    return names
+
+
+# ─── CFG builder ──────────────────────────────────────────────────────────────
+
+class PythonCFGBuilder(CFGBuilderBase):
+    """Converts a ``function_definition`` CST node into a basic-block CFG."""
+
+    def build(
+        self, func_node: Node, source: bytes
+    ) -> Tuple[int, int, Dict[int, CFGBlock], List[CFGDef]]:
+        self._reset()
+
+        entry = self._new_block()
+        exit_b = self._new_block()
+
+        # Parameters are defs in the entry block
+        params = func_node.child_by_field_name("parameters")
+        if params:
+            for name in _collect_param_names(params, source):
+                d = self._new_def(name, _node_pos(params), False, None)
+                entry.events.append(d)
+
+        body = func_node.child_by_field_name("body")
+        if body:
+            tail = self._process_stmts(_real_children(body), entry, exit_b, source)
+            if tail is not None:
+                self._link(tail, exit_b)
+        else:
+            self._link(entry, exit_b)
+
+        return entry.id, exit_b.id, self._blocks, self._all_defs
+
+    # ── internal helpers ──────────────────────────────────────────────────────
+
+    def _taint_of(self, node: Node, source: bytes) -> Optional[str]:
+        """Return a taint-source label if node's text starts with a known source."""
+        text = _text(node, source).strip()
+        for src in self.config.taint_sources:
+            if text.startswith(src):
+                return src
+        return None
+
+    def _process_stmts(
+        self,
+        stmts: List[Node],
+        cur: CFGBlock,
+        exit_b: CFGBlock,
+        source: bytes,
+    ) -> Optional[CFGBlock]:
+        for stmt in stmts:
+            if cur is None:
+                return None
+            cur = self._process_stmt(stmt, cur, exit_b, source)
+        return cur
+
+    def _process_stmt(
+        self,
+        stmt: Node,
+        cur: CFGBlock,
+        exit_b: CFGBlock,
+        source: bytes,
+    ) -> Optional[CFGBlock]:
+        t = stmt.type
+
+        # ── tree-sitter Python wraps ALL simple stmts in expression_statement ──
+        # Unwrap and re-dispatch so assignment/augmented_assignment defs are
+        # captured, and calls are treated as expressions (uses only).
+        if t == "expression_statement":
+            real_children = [c for c in stmt.children
+                             if not c.is_extra and c.type not in ("\n", "comment")]
+            if len(real_children) == 1:
+                return self._process_stmt(real_children[0], cur, exit_b, source)
+            # Multiple children (e.g. type-annotated assignment) — collect uses
+            cur.events.extend(_uses_of(stmt))
+            return cur
+
+        # ── simple assignment ─────────────────────────────────────────────────
+        if t == "assignment":
+            rhs = stmt.child_by_field_name("right")
+            lhs = stmt.child_by_field_name("left")
+            is_lit = _is_literal(rhs) if rhs else False
+            taint = self._taint_of(rhs, source) if rhs else None
+            if rhs:
+                cur.events.extend(_uses_of(rhs))
+            if lhs:
+                for name in _lhs_names(lhs, source):
+                    if name:
+                        cur.events.append(
+                            self._new_def(name, _node_pos(lhs), is_lit, taint)
+                        )
+            return cur
+
+        # ── augmented assignment ──────────────────────────────────────────────
+        if t == "augmented_assignment":
+            lhs = stmt.child_by_field_name("left")
+            rhs = stmt.child_by_field_name("right")
+            if rhs:
+                cur.events.extend(_uses_of(rhs))
+            if lhs and lhs.type == "identifier":
+                name = _text(lhs, source)
+                if name:
+                    from adl_drift.graph.cfg import CFGUse
+                    cur.events.append(CFGUse(name, _node_pos(lhs)))
+                    cur.events.append(self._new_def(name, _node_pos(lhs), False, None))
+            return cur
+
+        # ── named expression (walrus) ─────────────────────────────────────────
+        if t == "named_expression":
+            name_node = stmt.child_by_field_name("name")
+            val = stmt.child_by_field_name("value")
+            if val:
+                cur.events.extend(_uses_of(val))
+            if name_node and name_node.type == "identifier":
+                name = _text(name_node, source)
+                is_lit = _is_literal(val) if val else False
+                taint = self._taint_of(val, source) if val else None
+                cur.events.append(self._new_def(name, _node_pos(name_node), is_lit, taint))
+            return cur
+
+        # ── terminators ───────────────────────────────────────────────────────
+        if t in _TERMINATOR_TYPES:
+            cur.events.extend(_uses_of(stmt))
+            self._link(cur, exit_b)
+            return None
+
+        # ── simple pass-through ───────────────────────────────────────────────
+        if t in _SIMPLE_STMT_TYPES:
+            cur.events.extend(_uses_of(stmt))
+            return cur
+
+        # ── if / elif / else ──────────────────────────────────────────────────
+        if t == "if_statement":
+            return self._process_if(stmt, cur, exit_b, source)
+
+        # ── for / while loops ─────────────────────────────────────────────────
+        if t in ("for_statement", "while_statement"):
+            return self._process_loop(stmt, cur, exit_b, source)
+
+        # ── try / except / finally ────────────────────────────────────────────
+        if t == "try_statement":
+            return self._process_try(stmt, cur, exit_b, source)
+
+        # ── with statement ────────────────────────────────────────────────────
+        if t == "with_statement":
+            return self._process_with(stmt, cur, exit_b, source)
+
+        # ── decorated / nested definitions ────────────────────────────────────
+        if t in ("function_definition", "async_function_definition",
+                 "class_definition", "decorated_definition"):
+            inner = stmt
+            if t == "decorated_definition":
+                inner = stmt.child_by_field_name("definition") or stmt
+            name_node = inner.child_by_field_name("name")
+            if name_node and name_node.type == "identifier":
+                name = _text(name_node, source)
+                cur.events.append(self._new_def(name, _node_pos(name_node), False, None))
+            return cur
+
+        # ── async wrapper ─────────────────────────────────────────────────────
+        if t == "async_statement":
+            for child in stmt.children:
+                if child.type not in ("async",):
+                    cur = self._process_stmt(child, cur, exit_b, source)
+                    if cur is None:
+                        return None
+            return cur
+
+        # ── fallback: conservative — collect all identifiers as uses ──────────
+        cur.events.extend(_uses_of(stmt))
+        return cur
+
+    # ── compound statement helpers ────────────────────────────────────────────
+
+    def _process_if(
+        self, node: Node, cur: CFGBlock, exit_b: CFGBlock, source: bytes
+    ) -> Optional[CFGBlock]:
+        cond = node.child_by_field_name("condition")
+        if cond:
+            cur.events.extend(_uses_of(cond))
+
+        merge = self._new_block()
+
+        # then-branch
+        body = node.child_by_field_name("consequence")
+        if body:
+            then_b = self._new_block()
+            self._link(cur, then_b)
+            tail = self._process_stmts(_real_children(body), then_b, exit_b, source)
+            if tail is not None:
+                self._link(tail, merge)
+        else:
+            self._link(cur, merge)
+
+        # else / elif  (alternative field)
+        alt = node.child_by_field_name("alternative")
+        if alt is None:
+            self._link(cur, merge)  # no else → fall through
+        else:
+            if alt.type == "else_clause":
+                alt_b = self._new_block()
+                self._link(cur, alt_b)
+                body_of_else = alt.child_by_field_name("body")
+                children = _real_children(body_of_else) if body_of_else else _real_children(alt)
+                tail = self._process_stmts(children, alt_b, exit_b, source)
+                if tail is not None:
+                    self._link(tail, merge)
+            elif alt.type == "elif_clause":
+                # Treat as nested if
+                alt_b = self._new_block()
+                self._link(cur, alt_b)
+                tail = self._process_stmt(alt, alt_b, exit_b, source)
+                if tail is not None:
+                    self._link(tail, merge)
+            else:
+                self._link(cur, merge)
+
+        return merge
+
+    def _process_loop(
+        self, node: Node, cur: CFGBlock, exit_b: CFGBlock, source: bytes
+    ) -> Optional[CFGBlock]:
+        header = self._new_block()
+        body_b = self._new_block()
+        after = self._new_block()
+
+        self._link(cur, header)
+
+        # Iterator expression (for) or condition (while)
+        iter_node = (node.child_by_field_name("right") or
+                     node.child_by_field_name("condition"))
+        if iter_node:
+            header.events.extend(_uses_of(iter_node))
+
+        # Loop variable defs (for-loop only)
+        left = node.child_by_field_name("left")
+        if left:
+            for name in _lhs_names(left, source):
+                if name:
+                    header.defs.append(self._new_def(name, _node_pos(left), False, None))
+
+        # header → body (iterations exist) or after (zero iterations)
+        self._link(header, body_b)
+        self._link(header, after)
+
+        body = node.child_by_field_name("body")
+        if body:
+            tail = self._process_stmts(_real_children(body), body_b, exit_b, source)
+            if tail is not None:
+                self._link(tail, header)  # back-edge
+
+        # else-clause (loop not broken early)
+        alt = node.child_by_field_name("alternative")
+        if alt:
+            else_b = self._new_block()
+            self._link(after, else_b)
+            alt_body = alt.child_by_field_name("body")
+            children = _real_children(alt_body) if alt_body else _real_children(alt)
+            tail = self._process_stmts(children, else_b, exit_b, source)
+            if tail:
+                after = tail  # merge point shifts to end of else
+
+        return after
+
+    def _process_try(
+        self, node: Node, cur: CFGBlock, exit_b: CFGBlock, source: bytes
+    ) -> Optional[CFGBlock]:
+        merge = self._new_block()
+
+        # try body
+        body = node.child_by_field_name("body")
+        if body:
+            try_b = self._new_block()
+            self._link(cur, try_b)
+            tail = self._process_stmts(_real_children(body), try_b, exit_b, source)
+            if tail:
+                self._link(tail, merge)
+        else:
+            self._link(cur, merge)
+
+        # except handlers — conservative: all are reachable from try entry
+        for child in node.children:
+            if child.type in ("except_clause", "except_group_clause"):
+                exc_b = self._new_block()
+                self._link(cur, exc_b)
+                # "as <name>" exception variable def
+                saw_as = False
+                for gc in child.children:
+                    if saw_as and gc.type == "identifier":
+                        exc_b.defs.append(
+                            self._new_def(_text(gc, source), _node_pos(gc), False, None)
+                        )
+                        saw_as = False
+                    elif gc.type == "as":
+                        saw_as = True
+                exc_body = child.child_by_field_name("body")
+                if exc_body:
+                    tail = self._process_stmts(
+                        _real_children(exc_body), exc_b, exit_b, source
+                    )
+                    if tail:
+                        self._link(tail, merge)
+                else:
+                    self._link(exc_b, merge)
+
+            elif child.type == "finally_clause":
+                fin_body = child.child_by_field_name("body")
+                if fin_body:
+                    fin_b = self._new_block()
+                    self._link(merge, fin_b)
+                    tail = self._process_stmts(
+                        _real_children(fin_body), fin_b, exit_b, source
+                    )
+                    if tail:
+                        merge = tail
+
+        return merge
+
+    def _process_with(
+        self, node: Node, cur: CFGBlock, exit_b: CFGBlock, source: bytes
+    ) -> Optional[CFGBlock]:
+        # Process each with_item: value uses + alias defs
+        for child in node.children:
+            if child.type == "with_item":
+                val = child.child_by_field_name("value")
+                alias = child.child_by_field_name("alias")
+                if val:
+                    cur.events.extend(_uses_of(val))
+                    taint = self._taint_of(val, source)
+                    if alias and alias.type == "identifier":
+                        name = _text(alias, source)
+                        cur.events.append(
+                            self._new_def(name, _node_pos(alias), False, taint)
+                        )
+
+        body = node.child_by_field_name("body")
+        if body:
+            cur = self._process_stmts(_real_children(body), cur, exit_b, source)
+        return cur
+
+
+def default_python_cfg_config() -> CFGConfig:
+    """The real Python `CFGConfig` — returned by `PythonBackend.cfg_config()`
+    and defaulted to by `graph/cfg.py::analyze_function`/`analyze_file` when
+    no config is passed (every pre-P4.6 call site)."""
+    from adl_drift.lang.python.taint import (
+        TAINT_SINKS, SINK_METHODS, default_python_taint_strategy,
+    )
+
+    return CFGConfig(
+        taint_sources=TAINT_SOURCES,
+        taint_sinks=dict(TAINT_SINKS),
+        sink_methods=dict(SINK_METHODS),
+        cfg_builder_factory=lambda config: PythonCFGBuilder(config),
+        taint_strategy_factory=lambda config: default_python_taint_strategy(config),
+    )