diff --git a/CHANGELOG.md b/CHANGELOG.md index 725a781..6730724 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,22 @@ below). ## [Unreleased] +### BREAKING + +- **Threat conditions are now validated against a fixed set of allowed names + (#362).** A condition referencing a name outside the supported element, + enum, and safe-builtin set is rejected when the `Threat` is constructed, + rather than raising `NameError` later during evaluation. Custom threat + files carrying typos or unsupported names will now fail fast at load time; + correct the name or drop the condition. Names bound by list and generator + comprehensions inside a condition remain valid. + +### Fixed + +- `LLM` is now available to threat condition evaluation (#362). Conditions + referencing `LLM` previously raised `NameError`. No shipped threat uses + `LLM` in a condition, so this affected only user-authored threats. + ## [1.4.0] - 2026-05-21 ### BREAKING diff --git a/pytm/threat.py b/pytm/threat.py index f5f13c7..813b6d0 100644 --- a/pytm/threat.py +++ b/pytm/threat.py @@ -83,6 +83,7 @@ class _ConditionValidator(ast.NodeVisitor): def __init__(self, allowed_names: set[str]) -> None: super().__init__() self.allowed_names = allowed_names | {"target", "True", "False", "None"} + self.local_names: list[set[str]] = [] def visit(self, node: ast.AST) -> Any: # type: ignore[override] if not isinstance(node, self._ALLOWED_NODES): @@ -124,11 +125,40 @@ def visit_Name(self, node: ast.Name) -> Any: # noqa: D401 isinstance(node.ctx, ast.Load) and node.id not in self.allowed_names and node.id not in self.SAFE_CALL_NAMES + and not any(node.id in names for names in self.local_names) ): - # Allow names introduced by comprehensions; they will fail at runtime if undefined. - return + raise ValueError(f"Unknown name in threat condition: {node.id}") return None + def visit_ListComp(self, node: ast.ListComp) -> Any: # noqa: D401 + return self._visit_comprehension_expression(node) + + def visit_GeneratorExp(self, node: ast.GeneratorExp) -> Any: # noqa: D401 + return self._visit_comprehension_expression(node) + + def _visit_comprehension_expression( + self, node: ast.ListComp | ast.GeneratorExp + ) -> None: + local_names = set() + for generator in node.generators: + local_names.update(self._target_names(generator.target)) + + self.local_names.append(local_names) + try: + self.generic_visit(node) + finally: + self.local_names.pop() + + def _target_names(self, node: ast.AST) -> set[str]: + if isinstance(node, ast.Name): + return {node.id} + if isinstance(node, (ast.Tuple, ast.List)): + names = set() + for element in node.elts: + names.update(self._target_names(element)) + return names + return set() + @staticmethod def _attribute_chain(node: ast.Attribute) -> List[str]: chain: List[str] = [node.attr] @@ -280,6 +310,7 @@ def _build_eval_globals(cls) -> dict[str, Any]: "Element": pytm.Element, "ExternalEntity": pytm.ExternalEntity, "Lambda": pytm.Lambda, + "LLM": pytm.LLM, "Process": pytm.Process, "Server": pytm.Server, "SetOfProcesses": pytm.SetOfProcesses, diff --git a/tests/test_pydantic_models.py b/tests/test_pydantic_models.py index 0942b9e..75adf76 100644 --- a/tests/test_pydantic_models.py +++ b/tests/test_pydantic_models.py @@ -391,7 +391,7 @@ class TestConditionValidator: def _validate(self, condition: str): import ast tree = ast.parse(condition, mode="eval") - validator = _ConditionValidator(allowed_names=set()) + validator = _ConditionValidator(allowed_names={"Boundary", "Server"}) validator.visit(tree) def test_simple_comparison_is_valid(self): @@ -403,6 +403,14 @@ def test_allowed_method_call_is_valid(self): def test_builtin_any_is_valid(self): self._validate("any(f.isEncrypted for f in target.inputs)") + def test_unknown_name_raises(self): + with pytest.raises(ValueError, match="Unknown name"): + self._validate("missing_name") + + def test_threat_rejects_unknown_condition_name(self): + with pytest.raises(ValueError, match="Unknown name"): + Threat(SID="T1", condition="missing_name") + def test_dunder_attribute_raises(self): with pytest.raises(ValueError, match="dunder"): self._validate("target.__class__")