feat: add safe-mode AST checks and execution timeout for MultiOutputScript

README.md

@@ -1,6 +1,6 @@
 # ComfyUI-ScriptFlow
 
-**Version:** 1.1.0
+**Version:** 1.1.1
 **License:** GPL-3.0
 
 A general-purpose calculation node that executes user-written Python code and returns multiple text and numeric outputs based on the inputs.
@@ -117,10 +117,14 @@ exec(code, globals_dict, locals_dict)
 - file access (`open`, etc.)
 - OS operations (`os`, `sys`)
 - re-running `eval`/`exec`
+- blocked syntax in safe mode: `Import`, `ImportFrom`, `Global`, `Nonlocal`, `ClassDef`, `Try`, `With`
+- blocked calls in safe mode: `__import__`, `eval`, `exec`, `compile`, `open`, `input`, `globals`, `locals`, `vars`, `getattr`, `setattr`, `delattr`
 
 ### Execution Notes
 - Using `random` or `datetime` makes outputs non-deterministic.
 - Type mismatches raise an error and stop execution.
+- Safe mode validates AST before execution and stops scripts that exceed the timeout (`1.5s` by default).
+- Use only trusted scripts/workflows. Do not run untrusted code from unknown sources.
 
 ## Examples
 
@@ -157,6 +161,13 @@ You should have received a copy of the GNU General Public License along with thi
 ---
 
 ## Release Notes
+### 1.1.1
+- Added safe-mode AST validation before script execution.
+- Blocked unsafe syntax (`Import`, `ImportFrom`, `Global`, `Nonlocal`, `ClassDef`, `Try`, `With`, `AsyncWith`).
+- Blocked unsafe calls (`__import__`, `eval`, `exec`, `compile`, `open`, `input`, `globals`, `locals`, `vars`, `getattr`, `setattr`, `delattr`).
+- Added execution timeout guard (`1.5s` default) to stop runaway scripts.
+- Updated security notes for trusted-workflow usage.
+
 ### 1.1.0
 Added `centi` node (`utils`).
 - Minimal connector-only utility node.

__init__.py

@@ -8,11 +8,14 @@
 
 from __future__ import annotations
 
-__version__ = "1.1.0"
+__version__ = "1.1.1"
 
+import ast
 import datetime
 import math
 import random
+import sys
+import time
 from typing import Any, Dict
 
 
@@ -43,6 +46,34 @@ _ALLOWED_BUILTINS = {
     "tuple": tuple,
 }
 
+_DEFAULT_EXEC_TIMEOUT_SECONDS = 1.5
+
+_BLOCKED_AST_NODES = (
+    ast.Import,
+    ast.ImportFrom,
+    ast.Global,
+    ast.Nonlocal,
+    ast.ClassDef,
+    ast.Try,
+    ast.With,
+    ast.AsyncWith,
+)
+
+_BLOCKED_CALL_NAMES = {
+    "__import__",
+    "eval",
+    "exec",
+    "compile",
+    "open",
+    "input",
+    "globals",
+    "locals",
+    "vars",
+    "getattr",
+    "setattr",
+    "delattr",
+}
+
 
 class AnyType(str):
     def __ne__(self, __value: object) -> bool:  # allow any connection type
@@ -61,6 +92,48 @@ def _build_globals() -> Dict[str, Any]:
     }
 
 
+def _validate_ast(code: str) -> None:
+    try:
+        tree = ast.parse(code, mode="exec")
+    except SyntaxError as exc:
+        raise ValueError(f"Syntax error in code: {exc}") from exc
+
+    for node in ast.walk(tree):
+        if isinstance(node, _BLOCKED_AST_NODES):
+            raise ValueError(
+                f"Unsupported syntax in safe mode: {type(node).__name__}"
+            )
+
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            if node.func.id in _BLOCKED_CALL_NAMES:
+                raise ValueError(
+                    f"Blocked function call in safe mode: {node.func.id}"
+                )
+
+
+def _exec_with_timeout(
+    code: str,
+    globals_dict: Dict[str, Any],
+    locals_dict: Dict[str, Any],
+    timeout_seconds: float = _DEFAULT_EXEC_TIMEOUT_SECONDS,
+) -> None:
+    deadline = time.perf_counter() + max(0.1, timeout_seconds)
+
+    def _trace(frame, event, arg):
+        if event == "line" and time.perf_counter() > deadline:
+            raise TimeoutError(
+                f"Script execution exceeded timeout ({timeout_seconds:.2f}s)"
+            )
+        return _trace
+
+    previous_trace = sys.gettrace()
+    sys.settrace(_trace)
+    try:
+        exec(code, globals_dict, locals_dict)
+    finally:
+        sys.settrace(previous_trace)
+
+
 def _validate_outputs(locals_dict: Dict[str, Any]) -> None:
     text_keys = ("ot1", "ot2", "ot3")
     value_keys = ("ov1", "ov2", "ov3")
@@ -180,8 +253,10 @@ class MultiOutputScript:
             "ov3": None,
         }
 
+        _validate_ast(code)
+
         globals_dict = _build_globals()
-        exec(code, globals_dict, locals_dict)
+        _exec_with_timeout(code, globals_dict, locals_dict)
 
         _validate_outputs(locals_dict)
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-ScriptFlow"
 description = "Python script node for ComfyUI supporting math, calc, logic, and text editing with multi output. Handle conditional workflows, auto-select resolution, parse LLM outputs, and generate filenames — all in one flexible node."
-version = "1.1.0"
+version = "1.1.1"
 license = {file = "LICENSE"} 
 # classifiers = [
 #     # For OS-independent nodes (works on all operating systems)