From b3933446bdaea3d25825090734e4a4157e6ee51c Mon Sep 17 00:00:00 2001 From: kantan-kanto Date: Sun, 31 May 2026 18:18:53 +0900 Subject: [PATCH] Replace dynamic script execution with a safe AST interpreter. --- README.md | 50 ++-- __init__.py | 735 ++++++++++++++++++++++++++++++++++++++++++++----- pyproject.toml | 2 +- 3 files changed, 693 insertions(+), 94 deletions(-) diff --git a/README.md b/README.md index f4072c8..4628b27 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,9 @@ # ComfyUI-ScriptFlow -**Version:** 1.1.1 +**Version:** 1.2.0 **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. +A general-purpose calculation node that evaluates a safe Python-like script subset and returns multiple text and numeric outputs based on the inputs. This custom node accepts numbers (int/float) and text (string) as inputs, runs calculations and logical operations internally, and returns multiple values through separate output ports. It consolidates workflows that previously required combining several basic nodes into a single node. @@ -20,7 +20,7 @@ This custom node accepts numbers (int/float) and text (string) as inputs, runs c - Multiple inputs: freely combine numeric and text inputs. - Calculations and logic: supports arithmetic, comparison, and simple conditional branching. - Multiple outputs: pass results to downstream nodes via separate output ports. -- Safer execution: restricted environment (no file/OS access). +- Safer execution: AST interpreter with restricted syntax and no file/OS access. ## Use Cases - In Wan2.2 I2V workflows, automatically select a recommended resolution @@ -94,41 +94,47 @@ Notes: --- ## Execution Environment -The script is executed using: -```python -exec(code, globals_dict, locals_dict) -``` +The script is parsed with Python AST and evaluated by ScriptFlow's safe interpreter. It supports common workflow logic such as assignment, arithmetic, branching, loops, user-defined helper functions, dictionaries, lists, strings, and selected methods. -### Allowed builtins +### Allowed functions - `len`, `min`, `max`, `sum`, `abs`, `round` - `int`, `float`, `str`, `bool` - `sorted`, `reversed` -- `enumerate`, `range`, `zip`, `map`, `filter` +- `enumerate`, `range`, `zip` - `any`, `all`, `pow`, `divmod` - `list`, `dict`, `tuple` -### Allowed modules -- `random` -- `datetime` -- `math` +### Allowed namespaces +- `random`: `random`, `randint`, `uniform`, `choice` +- `datetime`: `datetime.now`, `date.today`, date/time fields such as `year`, `month`, `day`, `hour`, `minute`, `second`, plus `isoformat` and `strftime` (`strftime` requires v1.2.0 or later) +- `math`: `ceil`, `floor`, `sqrt`, `sin`, `cos`, `tan`, `asin`, `acos`, `atan`, `atan2`, `log`, `log10`, `exp`, `pow`, `fabs`, `isfinite`, `pi`, `e`, `tau` + +### Allowed methods +- `str`: `replace`, `find`, `strip`, `split`, `splitlines`, `startswith`, `endswith`, `join`, `lower`, `upper` +- `list`: `append` ### Not allowed - `import` statements - 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` +- dynamic Python execution +- blocked syntax in safe mode: `Import`, `ImportFrom`, `Global`, `Nonlocal`, `ClassDef`, `Try`, `With`, `AsyncWith`, `Lambda`, `Delete`, `Yield`, `Await` +- blocked calls in safe mode: dynamic execution, file/input access, runtime namespace inspection, and dynamic attribute mutation ### Execution Notes - Using `random` or `datetime` makes outputs non-deterministic. -- `datetime.strftime(...)` may fail because it can internally import the `time` module, while `__import__` is blocked for safety. Format timestamps manually instead: +- `datetime.strftime(...)` is supported in v1.2.0 or later for concise timestamp formatting: + ```python + now = datetime.datetime.now() + ot1 = now.strftime("%Y%m%d%H%M%S") + ``` +- For versions earlier than v1.2.0, format the same timestamp manually: ```python now = datetime.datetime.now() ot1 = f'{now.year:04d}{now.month:02d}{now.day:02d}{now.hour:02d}{now.minute:02d}{now.second:02d}' ``` - 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). +- Safe mode validates AST before evaluation and stops scripts that exceed the step or loop limit. - Use only trusted scripts/workflows. Do not run untrusted code from unknown sources. ## Examples @@ -344,10 +350,16 @@ You should have received a copy of the GNU General Public License along with thi --- ## Release Notes +### 1.2.0 +- Replaced direct Python execution with a safe AST interpreter for ComfyUI Registry compatibility. +- Added Python-subset support for helper functions, loops, dictionaries, lists, string methods, f-strings, `random`, `datetime`, and `math`. +- Replaced trace-based timeout handling with step and loop limits. +- Kept existing `MultiOutputScript` inputs and outputs unchanged. + ### 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`). +- Blocked unsafe dynamic execution, file/input access, namespace inspection, and dynamic attribute mutation calls. - Added execution timeout guard (`1.5s` default) to stop runaway scripts. - Updated security notes for trusted-workflow usage. diff --git a/__init__.py b/__init__.py index f16b47f..8d0f36b 100644 --- a/__init__.py +++ b/__init__.py @@ -8,45 +8,18 @@ from __future__ import annotations -__version__ = "1.1.1" +__version__ = "1.2.0" import ast import datetime import math import random -import sys -import time from typing import Any, Dict -_ALLOWED_BUILTINS = { - "len": len, - "min": min, - "max": max, - "sum": sum, - "abs": abs, - "round": round, - "int": int, - "float": float, - "str": str, - "bool": bool, - "sorted": sorted, - "reversed": reversed, - "enumerate": enumerate, - "range": range, - "zip": zip, - "map": map, - "filter": filter, - "any": any, - "all": all, - "pow": pow, - "divmod": divmod, - "list": list, - "dict": dict, - "tuple": tuple, -} - -_DEFAULT_EXEC_TIMEOUT_SECONDS = 1.5 +_DEFAULT_MAX_STEPS = 10000 +_DEFAULT_MAX_CALL_DEPTH = 32 +_MAX_RANGE_SIZE = 10000 _BLOCKED_AST_NODES = ( ast.Import, @@ -57,10 +30,14 @@ _BLOCKED_AST_NODES = ( ast.Try, ast.With, ast.AsyncWith, + ast.Lambda, + ast.Delete, + ast.Yield, + ast.YieldFrom, + ast.Await, ) _BLOCKED_CALL_NAMES = { - "__import__", "eval", "exec", "compile", @@ -72,8 +49,654 @@ _BLOCKED_CALL_NAMES = { "getattr", "setattr", "delattr", + "hasattr", } +_SAFE_FUNCTIONS = { + "len", + "min", + "max", + "sum", + "abs", + "round", + "int", + "float", + "str", + "bool", + "sorted", + "reversed", + "enumerate", + "range", + "zip", + "any", + "all", + "pow", + "divmod", + "list", + "dict", + "tuple", +} + +_SAFE_STR_METHODS = { + "replace", + "find", + "strip", + "split", + "splitlines", + "startswith", + "endswith", + "join", + "lower", + "upper", +} + +_SAFE_LIST_METHODS = { + "append", +} + +_SAFE_DATETIME_ATTRS = { + "year", + "month", + "day", + "hour", + "minute", + "second", + "microsecond", +} + +_SAFE_MATH_FUNCTIONS = { + "ceil", + "floor", + "sqrt", + "sin", + "cos", + "tan", + "asin", + "acos", + "atan", + "atan2", + "log", + "log10", + "exp", + "pow", + "fabs", + "isfinite", +} + +_SAFE_MATH_CONSTANTS = { + "pi": math.pi, + "e": math.e, + "tau": math.tau, +} + + +class _BreakSignal(Exception): + pass + + +class _ContinueSignal(Exception): + pass + + +class _ReturnSignal(Exception): + def __init__(self, value: Any): + self.value = value + + +class _SafeNamespace: + def __init__(self, name: str): + self.name = name + + +class _SafeFunction: + def __init__(self, name: str, args: ast.arguments, body: list[ast.stmt]): + self.name = name + self.arg_names = [arg.arg for arg in args.args] + self.body = body + if args.vararg or args.kwarg or args.kwonlyargs or args.defaults: + raise ValueError(f"Unsupported function signature: {name}") + + def call(self, interpreter: "_SafeScriptInterpreter", args: list[Any]) -> Any: + if len(args) != len(self.arg_names): + raise ValueError( + f"{self.name}() expected {len(self.arg_names)} arguments, got {len(args)}" + ) + if interpreter.call_depth >= interpreter.max_call_depth: + raise RuntimeError("Script exceeded maximum call depth") + frame = dict(zip(self.arg_names, args)) + interpreter.call_depth += 1 + interpreter.frames.append(frame) + try: + interpreter._run_block(self.body) + except _ReturnSignal as signal: + return signal.value + finally: + interpreter.frames.pop() + interpreter.call_depth -= 1 + return None + + +class _SafeScriptInterpreter: + def __init__( + self, + variables: Dict[str, Any], + max_steps: int = _DEFAULT_MAX_STEPS, + max_call_depth: int = _DEFAULT_MAX_CALL_DEPTH, + ): + self.frames = [variables] + self.max_steps = max_steps + self.max_call_depth = max_call_depth + self.steps = 0 + self.call_depth = 0 + + def run(self, code: str) -> Dict[str, Any]: + try: + tree = ast.parse(code) + except SyntaxError as exc: + raise ValueError(f"Syntax error in code: {exc}") from exc + _validate_ast_tree(tree) + self._run_block(tree.body) + return self.frames[0] + + def _step(self) -> None: + self.steps += 1 + if self.steps > self.max_steps: + raise RuntimeError("Script exceeded step limit") + + def _run_block(self, statements: list[ast.stmt]) -> None: + for statement in statements: + self._step() + self._run_statement(statement) + + def _run_statement(self, node: ast.stmt) -> None: + if isinstance(node, ast.Assign): + value = self._resolve(node.value) + for target in node.targets: + self._assign(target, value) + return + if isinstance(node, ast.AugAssign): + current = self._resolve_target(node.target) + value = self._apply_binop(node.op, current, self._resolve(node.value)) + self._assign(node.target, value) + return + if isinstance(node, ast.Expr): + self._resolve(node.value) + return + if isinstance(node, ast.If): + if self._truthy(self._resolve(node.test)): + self._run_block(node.body) + else: + self._run_block(node.orelse) + return + if isinstance(node, ast.For): + iterable = self._resolve(node.iter) + count = 0 + for item in iterable: + count += 1 + if count > _MAX_RANGE_SIZE: + raise RuntimeError("Loop exceeded iteration limit") + self._step() + self._assign(node.target, item) + try: + self._run_block(node.body) + except _ContinueSignal: + continue + except _BreakSignal: + break + else: + self._run_block(node.orelse) + return + if isinstance(node, ast.While): + count = 0 + while self._truthy(self._resolve(node.test)): + count += 1 + if count > _MAX_RANGE_SIZE: + raise RuntimeError("Loop exceeded iteration limit") + self._step() + try: + self._run_block(node.body) + except _ContinueSignal: + continue + except _BreakSignal: + break + else: + self._run_block(node.orelse) + return + if isinstance(node, ast.FunctionDef): + self._set_name(node.name, _SafeFunction(node.name, node.args, node.body)) + return + if isinstance(node, ast.Return): + raise _ReturnSignal(None if node.value is None else self._resolve(node.value)) + if isinstance(node, ast.Break): + raise _BreakSignal() + if isinstance(node, ast.Continue): + raise _ContinueSignal() + if isinstance(node, ast.Pass): + return + raise ValueError(f"Unsupported syntax: {type(node).__name__}") + + def _resolve(self, node: ast.AST) -> Any: + self._step() + if isinstance(node, ast.Constant): + return node.value + if isinstance(node, ast.Name): + return self._get_name(node.id) + if isinstance(node, ast.List): + return [self._resolve(item) for item in node.elts] + if isinstance(node, ast.Tuple): + return tuple(self._resolve(item) for item in node.elts) + if isinstance(node, ast.Dict): + return { + self._resolve(key): self._resolve(value) + for key, value in zip(node.keys, node.values) + } + if isinstance(node, ast.BinOp): + return self._apply_binop(node.op, self._resolve(node.left), self._resolve(node.right)) + if isinstance(node, ast.UnaryOp): + return self._apply_unaryop(node.op, self._resolve(node.operand)) + if isinstance(node, ast.BoolOp): + return self._resolve_boolop(node) + if isinstance(node, ast.Compare): + return self._resolve_compare(node) + if isinstance(node, ast.IfExp): + return self._resolve(node.body if self._truthy(self._resolve(node.test)) else node.orelse) + if isinstance(node, ast.Subscript): + return self._resolve(node.value)[self._resolve_slice(node.slice)] + if isinstance(node, ast.Slice): + return self._resolve_slice(node) + if isinstance(node, ast.Attribute): + return self._resolve_attribute(self._resolve(node.value), node.attr) + if isinstance(node, ast.Call): + return self._resolve_call(node) + if isinstance(node, ast.JoinedStr): + return "".join(str(self._resolve_joined_value(value)) for value in node.values) + if isinstance(node, ast.FormattedValue): + return self._resolve_joined_value(node) + raise ValueError(f"Unsupported expression: {type(node).__name__}") + + def _resolve_joined_value(self, node: ast.AST) -> Any: + if isinstance(node, ast.Constant): + return node.value + if isinstance(node, ast.FormattedValue): + value = self._resolve(node.value) + if node.conversion == 115: + value = str(value) + elif node.conversion == 114: + value = repr(value) + elif node.conversion == 97: + value = ascii(value) + spec = "" + if node.format_spec is not None: + spec = str(self._resolve(node.format_spec)) + return format(value, spec) + return self._resolve(node) + + def _resolve_slice(self, node: ast.AST) -> Any: + if isinstance(node, ast.Slice): + lower = None if node.lower is None else self._resolve(node.lower) + upper = None if node.upper is None else self._resolve(node.upper) + step = None if node.step is None else self._resolve(node.step) + return slice(lower, upper, step) + return self._resolve(node) + + def _resolve_call(self, node: ast.Call) -> Any: + if node.keywords: + raise ValueError("Keyword arguments are not supported") + args = [self._resolve(arg) for arg in node.args] + if isinstance(node.func, ast.Name): + return self._call_name(node.func.id, args) + if isinstance(node.func, ast.Attribute): + owner = self._resolve(node.func.value) + return self._call_attribute(owner, node.func.attr, args) + raise ValueError("Unsupported function call") + + def _call_name(self, name: str, args: list[Any]) -> Any: + if name.startswith("__"): + raise ValueError(f"Blocked function call: {name}") + value = self._find_name(name) + if isinstance(value, _SafeFunction): + return value.call(self, args) + if name in _SAFE_FUNCTIONS: + return self._call_safe_function(name, args) + raise ValueError(f"Unsupported function call: {name}") + + def _call_safe_function(self, name: str, args: list[Any]) -> Any: + if name == "len": + return len(*args) + if name == "min": + return min(*args) + if name == "max": + return max(*args) + if name == "sum": + return sum(*args) + if name == "abs": + return abs(*args) + if name == "round": + return round(*args) + if name == "int": + return int(*args) + if name == "float": + return float(*args) + if name == "str": + return str(*args) + if name == "bool": + return bool(*args) + if name == "sorted": + return sorted(*args) + if name == "reversed": + return list(reversed(*args)) + if name == "enumerate": + return enumerate(*args) + if name == "range": + result = range(*args) + if len(result) > _MAX_RANGE_SIZE: + raise RuntimeError("range() result is too large") + return result + if name == "zip": + return zip(*args) + if name == "any": + return any(*args) + if name == "all": + return all(*args) + if name == "pow": + return pow(*args) + if name == "divmod": + return divmod(*args) + if name == "list": + return list(*args) + if name == "dict": + return dict(*args) + if name == "tuple": + return tuple(*args) + raise ValueError(f"Unsupported function call: {name}") + + def _call_attribute(self, owner: Any, name: str, args: list[Any]) -> Any: + if name.startswith("__"): + raise ValueError(f"Blocked method call: {name}") + if isinstance(owner, str) and name in _SAFE_STR_METHODS: + return self._call_str_method(owner, name, args) + if isinstance(owner, list) and name in _SAFE_LIST_METHODS: + if name == "append": + if len(args) != 1: + raise ValueError("list.append() expects 1 argument") + owner.append(args[0]) + return None + if isinstance(owner, _SafeNamespace): + return self._call_namespace(owner.name, name, args) + if isinstance(owner, (datetime.datetime, datetime.date)) and name == "isoformat": + return owner.isoformat(*args) + if isinstance(owner, (datetime.datetime, datetime.date)) and name == "strftime": + return self._call_datetime_strftime(owner, args) + raise ValueError(f"Unsupported method call: {name}") + + @staticmethod + def _call_datetime_strftime(owner: datetime.datetime | datetime.date, args: list[Any]) -> str: + if len(args) != 1: + raise ValueError("strftime() expects 1 argument") + fmt = args[0] + if not isinstance(fmt, str): + raise ValueError("strftime() format must be str") + if len(fmt) > 256: + raise ValueError("strftime() format is too long") + return owner.strftime(fmt) + + def _call_str_method(self, owner: str, name: str, args: list[Any]) -> Any: + if name == "replace": + return owner.replace(*args) + if name == "find": + return owner.find(*args) + if name == "strip": + return owner.strip(*args) + if name == "split": + return owner.split(*args) + if name == "splitlines": + return owner.splitlines(*args) + if name == "startswith": + return owner.startswith(*args) + if name == "endswith": + return owner.endswith(*args) + if name == "join": + return owner.join(*args) + if name == "lower": + return owner.lower(*args) + if name == "upper": + return owner.upper(*args) + raise ValueError(f"Unsupported string method: {name}") + + def _call_namespace(self, namespace: str, name: str, args: list[Any]) -> Any: + if namespace == "random": + if name == "random": + if args: + raise ValueError("random.random() expects no arguments") + return random.random() + if name == "randint": + if len(args) != 2: + raise ValueError("random.randint() expects 2 arguments") + return random.randint(int(args[0]), int(args[1])) + if name == "uniform": + if len(args) != 2: + raise ValueError("random.uniform() expects 2 arguments") + return random.uniform(float(args[0]), float(args[1])) + if name == "choice": + if len(args) != 1: + raise ValueError("random.choice() expects 1 argument") + return random.choice(args[0]) + if namespace == "datetime.datetime" and name == "now": + if args: + raise ValueError("datetime.datetime.now() expects no arguments") + return datetime.datetime.now() + if namespace == "datetime.date" and name == "today": + if args: + raise ValueError("datetime.date.today() expects no arguments") + return datetime.date.today() + if namespace == "math" and name in _SAFE_MATH_FUNCTIONS: + return self._call_math(name, args) + raise ValueError(f"Unsupported namespace call: {namespace}.{name}") + + def _call_math(self, name: str, args: list[Any]) -> Any: + if name == "ceil": + return math.ceil(*args) + if name == "floor": + return math.floor(*args) + if name == "sqrt": + return math.sqrt(*args) + if name == "sin": + return math.sin(*args) + if name == "cos": + return math.cos(*args) + if name == "tan": + return math.tan(*args) + if name == "asin": + return math.asin(*args) + if name == "acos": + return math.acos(*args) + if name == "atan": + return math.atan(*args) + if name == "atan2": + return math.atan2(*args) + if name == "log": + return math.log(*args) + if name == "log10": + return math.log10(*args) + if name == "exp": + return math.exp(*args) + if name == "pow": + return math.pow(*args) + if name == "fabs": + return math.fabs(*args) + if name == "isfinite": + return math.isfinite(*args) + raise ValueError(f"Unsupported math function: {name}") + + def _resolve_attribute(self, owner: Any, name: str) -> Any: + if name.startswith("__"): + raise ValueError(f"Blocked attribute access: {name}") + if isinstance(owner, _SafeNamespace): + namespace_name = f"{owner.name}.{name}" + if namespace_name in ("datetime.datetime", "datetime.date"): + return _SafeNamespace(namespace_name) + if owner.name == "math" and name in _SAFE_MATH_CONSTANTS: + return _SAFE_MATH_CONSTANTS[name] + raise ValueError(f"Unsupported namespace attribute: {owner.name}.{name}") + if isinstance(owner, datetime.datetime) and name in _SAFE_DATETIME_ATTRS: + return self._datetime_attr(owner, name) + if isinstance(owner, datetime.date) and name in {"year", "month", "day"}: + return self._date_attr(owner, name) + raise ValueError(f"Unsupported attribute access: {name}") + + @staticmethod + def _datetime_attr(value: datetime.datetime, name: str) -> int: + if name == "year": + return value.year + if name == "month": + return value.month + if name == "day": + return value.day + if name == "hour": + return value.hour + if name == "minute": + return value.minute + if name == "second": + return value.second + if name == "microsecond": + return value.microsecond + raise ValueError(f"Unsupported datetime attribute: {name}") + + @staticmethod + def _date_attr(value: datetime.date, name: str) -> int: + if name == "year": + return value.year + if name == "month": + return value.month + if name == "day": + return value.day + raise ValueError(f"Unsupported date attribute: {name}") + + def _apply_binop(self, op: ast.operator, left: Any, right: Any) -> Any: + if isinstance(op, ast.Add): + return left + right + if isinstance(op, ast.Sub): + return left - right + if isinstance(op, ast.Mult): + return left * right + if isinstance(op, ast.Div): + return left / right + if isinstance(op, ast.FloorDiv): + return left // right + if isinstance(op, ast.Mod): + return left % right + if isinstance(op, ast.Pow): + return left**right + raise ValueError(f"Unsupported operator: {type(op).__name__}") + + def _apply_unaryop(self, op: ast.unaryop, value: Any) -> Any: + if isinstance(op, ast.UAdd): + return +value + if isinstance(op, ast.USub): + return -value + if isinstance(op, ast.Not): + return not self._truthy(value) + raise ValueError(f"Unsupported unary operator: {type(op).__name__}") + + def _resolve_boolop(self, node: ast.BoolOp) -> Any: + if isinstance(node.op, ast.And): + result = True + for value_node in node.values: + result = self._resolve(value_node) + if not self._truthy(result): + return result + return result + if isinstance(node.op, ast.Or): + result = False + for value_node in node.values: + result = self._resolve(value_node) + if self._truthy(result): + return result + return result + raise ValueError(f"Unsupported boolean operator: {type(node.op).__name__}") + + def _resolve_compare(self, node: ast.Compare) -> bool: + left = self._resolve(node.left) + for op, comparator in zip(node.ops, node.comparators): + right = self._resolve(comparator) + if not self._compare(op, left, right): + return False + left = right + return True + + @staticmethod + def _compare(op: ast.cmpop, left: Any, right: Any) -> bool: + if isinstance(op, ast.Eq): + return left == right + if isinstance(op, ast.NotEq): + return left != right + if isinstance(op, ast.Lt): + return left < right + if isinstance(op, ast.LtE): + return left <= right + if isinstance(op, ast.Gt): + return left > right + if isinstance(op, ast.GtE): + return left >= right + if isinstance(op, ast.Is): + return left is right + if isinstance(op, ast.IsNot): + return left is not right + if isinstance(op, ast.In): + return left in right + if isinstance(op, ast.NotIn): + return left not in right + raise ValueError(f"Unsupported comparison: {type(op).__name__}") + + def _assign(self, target: ast.AST, value: Any) -> None: + if isinstance(target, ast.Name): + self._set_name(target.id, value) + return + if isinstance(target, (ast.Tuple, ast.List)): + values = list(value) + if len(target.elts) != len(values): + raise ValueError("Unpack target and value length mismatch") + for child, item in zip(target.elts, values): + self._assign(child, item) + return + if isinstance(target, ast.Subscript): + owner = self._resolve(target.value) + owner[self._resolve_slice(target.slice)] = value + return + raise ValueError(f"Unsupported assignment target: {type(target).__name__}") + + def _resolve_target(self, target: ast.AST) -> Any: + if isinstance(target, ast.Name): + return self._get_name(target.id) + if isinstance(target, ast.Subscript): + return self._resolve(target) + raise ValueError(f"Unsupported assignment target: {type(target).__name__}") + + def _find_name(self, name: str) -> Any: + if name.startswith("__"): + raise ValueError(f"Blocked name: {name}") + for frame in reversed(self.frames): + if name in frame: + return frame[name] + if name in ("random", "datetime", "math"): + return _SafeNamespace(name) + return None + + def _get_name(self, name: str) -> Any: + value = self._find_name(name) + if value is not None or any(name in frame for frame in self.frames): + return value + raise NameError(f"Name is not defined: {name}") + + def _set_name(self, name: str, value: Any) -> None: + if name.startswith("__"): + raise ValueError(f"Blocked name: {name}") + self.frames[-1][name] = value + + @staticmethod + def _truthy(value: Any) -> bool: + return bool(value) + class AnyType(str): def __ne__(self, __value: object) -> bool: # allow any connection type @@ -83,26 +706,16 @@ class AnyType(str): _ANY_TYPE = AnyType("*") -def _build_globals() -> Dict[str, Any]: - return { - "__builtins__": _ALLOWED_BUILTINS, - "random": random, - "datetime": datetime, - "math": math, - } - - -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 - +def _validate_ast_tree(tree: ast.AST) -> None: 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.Attribute) and node.attr.startswith("__"): + raise ValueError(f"Blocked attribute access in safe mode: {node.attr}") + if isinstance(node, ast.Name) and node.id.startswith("__"): + raise ValueError(f"Blocked name in safe mode: {node.id}") if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): if node.func.id in _BLOCKED_CALL_NAMES: @@ -111,29 +724,6 @@ def _validate_ast(code: str) -> None: ) -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") @@ -253,10 +843,7 @@ class MultiOutputScript: "ov3": None, } - _validate_ast(code) - - globals_dict = _build_globals() - _exec_with_timeout(code, globals_dict, locals_dict) + locals_dict = _SafeScriptInterpreter(locals_dict).run(code) _validate_outputs(locals_dict) diff --git a/pyproject.toml b/pyproject.toml index 856a3fa..137e168 100644 --- a/pyproject.toml +++ b/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.1" +version = "1.2.0" license = {file = "LICENSE"} # classifiers = [ # # For OS-independent nodes (works on all operating systems)