fix(taint): retain sources across reassignment

src/skillspector/nodes/analyzers/behavioral_taint_tracking.py

@@ -392,7 +392,7 @@ def _find_nested_sources(
 
 def _find_tainted_names_in_args(
     node: ast.Call,
-    tainted: dict[str, _TaintedVar],
+    tainted: dict[str, dict[str, _TaintedVar]],
     check_runtime: Callable[[], None] | None = None,
 ) -> list[_TaintedVar]:
     """Find references to tainted variables in a call's arguments and keywords."""
@@ -409,35 +409,37 @@ def _find_tainted_names_in_args(
         elif isinstance(child, ast.Subscript):
             var_name = resolve_dotted_name(child.value)
         if var_name and var_name not in seen:
-            tv = tainted.get(var_name)
-            if tv:
+            sources = tainted.get(var_name)
+            if sources:
                 seen.add(var_name)
-                hits.append(tv)
+                hits.extend(sources.values())
     return hits
 
 
 def _mark_targets(
     targets: list[ast.expr],
-    tainted: dict[str, _TaintedVar],
+    tainted: dict[str, dict[str, _TaintedVar]],
     src_name: str,
     lineno: int,
-) -> list[str]:
-    """Taint each assignment target name that is not tainted yet.
+) -> list[tuple[str, str]]:
+    """Add a source to each assignment target and return newly added facts.
 
-    Add-only: an existing entry is never overwritten, so taint can only grow.
-    Returns the names newly added, for the worklist to propagate from.
+    Each target retains one location for each distinct source. Add-only taint
+    keeps the fixpoint finite while allowing reassignment to add stronger
+    sources instead of silently discarding them.
     """
-    newly_tainted: list[str] = []
+    newly_tainted: list[tuple[str, str]] = []
     names: list[str] = []
     for target in targets:
         if isinstance(target, ast.Name):
             names.append(target.id)
         elif isinstance(target, ast.Tuple):
             names.extend(elt.id for elt in target.elts if isinstance(elt, ast.Name))
     for name in names:
-        if name not in tainted:
-            tainted[name] = _TaintedVar(name, src_name, lineno)
-            newly_tainted.append(name)
+        sources = tainted.setdefault(name, {})
+        if src_name not in sources:
+            sources[src_name] = _TaintedVar(name, src_name, lineno)
+            newly_tainted.append((name, src_name))
     return newly_tainted
 
 
@@ -463,7 +465,7 @@ def _collect_tainted(
     type_map: dict[str, str],
     aliases: dict[str, str],
     check_runtime: Callable[[], None] | None = None,
-) -> dict[str, _TaintedVar]:
+) -> dict[str, dict[str, _TaintedVar]]:
     """Record ``Assign``-based taint, independent of AST visit order.
 
     Any single ordered pass over the tree — breadth-first (``ast.walk``) or
@@ -484,23 +486,18 @@ def _collect_tainted(
       assignment (``b = a``; ``payload = {"k": secret}``) once, keyed by a
       stable id, then indexes it by each name its value reads, as
       ``referenced_name -> [assignment_id, ...]``.
-    * The worklist then drains newly tainted names, firing each assignment that
-      reads a drained name AT MOST ONCE total: the first time any name it reads
-      becomes tainted, its targets are marked and enqueued; later drains of its
-      other reading names skip it. Only names not already tainted are enqueued.
-
-    Taint is add-only — an entry is never overwritten or removed — so the set
-    can only grow and is bounded by the number of assigned names. The loop
-    therefore cannot oscillate and is guaranteed to terminate. Firing an
-    assignment a second time is sound to skip: its targets are all tainted
-    after the first firing, so a later firing could only re-taint names and
-    add nothing. Each name is dequeued once and each propagating assignment
-    fires at most once, so the work is linear in the number of assignments
-    plus references — not quadratic in the longest chain, nor K×K' for a wide
-    ``a, b, ... = source`` statement read by many tainted names.
+    * The worklist drains newly discovered (name, source) facts. Each
+      assignment fires at most once for each source, so a stronger source
+      discovered through reassignment propagates without repeating work for
+      already-seen sources.
+
+    Taint is add-only and each variable has at most one entry per distinct
+    source call, so cycles terminate. Each source/assignment pair is processed
+    once; the work is bounded by the finite set of source calls and assignment
+    references.
     """
-    tainted: dict[str, _TaintedVar] = {}
-    worklist: deque[str] = deque()
+    tainted: dict[str, dict[str, _TaintedVar]] = {}
+    worklist: deque[tuple[str, str]] = deque()
     # Each propagating assignment, stored once as (targets, lineno) and keyed by
     # its index here (a stable id). ``propagators`` maps a name read by an
     # assignment's value to the ids of the assignments that read it.
@@ -535,18 +532,16 @@ def _collect_tainted(
             for ref in _referenced_names(ast_node.value):
                 propagators.setdefault(ref, []).append(assignment_id)
 
-    fired: set[int] = set()
+    fired: set[tuple[int, str]] = set()
     while worklist:
         if check_runtime is not None:
             check_runtime()
-        name = worklist.popleft()
-        src_call = tainted[name].source_call
+        name, src_call = worklist.popleft()
         for assignment_id in propagators.get(name, ()):
-            if assignment_id in fired:
-                # Already propagated once: its targets are all tainted, so
-                # firing again marks nothing new. Skip to stay linear.
+            assignment_source = (assignment_id, src_call)
+            if assignment_source in fired:
                 continue
-            fired.add(assignment_id)
+            fired.add(assignment_source)
             targets, lineno = propagating[assignment_id]
             worklist.extend(_mark_targets(targets, tainted, src_call, lineno))
 

tests/nodes/analyzers/test_behavioral_taint_tracking.py

@@ -824,13 +824,13 @@ def check_runtime() -> None:
 
 
 def _collect(code: str, check_runtime=None) -> dict:
-    """Run `_collect_tainted` directly on *code* and return name -> source_call."""
+    """Run `_collect_tainted` directly on *code* and return name -> source calls."""
     parsed = get_python_ast(None, code, "t.py")
     type_map = build_type_map(parsed.tree, parsed.import_aliases)
     tainted = behavioral_taint_tracking._collect_tainted(
         parsed.tree, type_map, parsed.import_aliases, check_runtime
     )
-    return {name: tv.source_call for name, tv in tainted.items()}
+    return {name: set(sources) for name, sources in tainted.items()}
 
 
 class TestFixpointTermination:
@@ -855,7 +855,25 @@ def test_cyclic_reassignment_terminates_and_taints_all(self) -> None:
         sources = _collect(code, _capped_check_runtime(2000))
         assert set(sources) == {"x", "y", "z"}
         # Every name traces back to one of the two credential sources.
-        assert set(sources.values()) <= {"os.getenv", "os.environ"}
+        assert all(source_set <= {"os.getenv", "os.environ"} for source_set in sources.values())
+        assert sources["x"] == {"os.getenv", "os.environ"}
+        assert sources["y"] == {"os.getenv", "os.environ"}
+        assert sources["z"] == {"os.getenv", "os.environ"}
+
+    def test_reassignment_preserves_stronger_source_through_alias(self) -> None:
+        """A later user-input source must not be hidden by an earlier credential source."""
+        for assignments in (
+            'payload = os.getenv("KEY")\npayload = input()\n',
+            'payload = input()\npayload = os.getenv("KEY")\n',
+        ):
+            code = (
+                "import os, subprocess\n"
+                + assignments
+                + "command = payload\n"
+                + "subprocess.run(command, shell=True)\n"
+            )
+            findings = _run(code)
+            assert any(f.rule_id == "TT5" for f in findings)
 
     def test_cyclic_reassignment_flows_to_sink(self) -> None:
         """End to end: the oscillating module plus a sink still reports TT3."""
@@ -892,7 +910,7 @@ def test_reverse_ordered_chain_is_linear(self) -> None:
         elapsed = time.monotonic() - start
 
         assert len(sources) == depth + 1
-        assert all(src == "os.getenv" for src in sources.values())
+        assert all(srcs == {"os.getenv"} for srcs in sources.values())
         assert elapsed < 1.0
 
     def test_cyclic_reassignment_does_not_hang_without_cap(self) -> None:
@@ -947,5 +965,5 @@ def counting_mark_targets(*args, **kwargs):
         assert calls["n"] <= n_assignments
         # All read and target names are tainted, all tracing to os.getenv.
         assert len(sources) == 2 * width
-        assert all(src == "os.getenv" for src in sources.values())
+        assert all(srcs == {"os.getenv"} for srcs in sources.values())
         assert elapsed < 1.0