core/engine/src/primitives/_safe_eval.py - main - alexefan136/flowstack

core/engine/src/primitives/_safe_eval.py
405 строк14 KB
Alexander Efanov
upd fix
31 июл 2026, 19:17
31 июл 2026, 19:17d146d86
"""
Безопасный вычислитель выражений для FlowStack Engine.

Заменяет опасный ``eval()`` в conditional edges flow-графа.
Использует AST-парсинг с whitelist'ом безопасных узлов —
произвольный код выполнить невозможно.

Поддерживает:
- Сравнения: ==, !=, <, <=, >, >=, in, not in, is, is not
- Логику: and, or, not
- Арифметику: +, -, *, /, //, %
- Переменные из context
- Константы: числа, строки, bool, None
- Коллекции: списки, кортежи
- Индексацию: data["key"], items[0]
- Атрибуты (не-dunder): result.status
- Whitelisted функции: len, min, max, abs, str, int, float, bool, round, sum, sorted, list

НЕ поддерживает (намеренно, для безопасности):
- Импорт модулей
- Lambda, comprehensions, определения функций/классов
- Dunder-атрибуты (__class__, __bases__, ...) — блокирует sandbox escape
- Произвольные вызовы функций и методов
- Возведение в степень (защита от DoS через огромные числа)

Использование:
    from src.primitives._safe_eval import safe_eval_bool, is_safe_expression

    # В conditional edge:
    if safe_eval_bool('status == "success" and count > 10', context):
        ...

    # Валидация при сохранении flow:
    if not is_safe_expression(edge.condition):
        raise ValueError("Unsafe condition")
"""

from __future__ import annotations

import ast
import operator
from typing import Any


class SafeEvalError(Exception):
    """Ошибка безопасного вычисления выражения."""


# ============================================================================
# Whitelisted operators
# ============================================================================

_BINARY_OPS: dict[type[ast.operator], Any] = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.FloorDiv: operator.floordiv,
    ast.Mod: operator.mod,
}

_UNARY_OPS: dict[type[ast.unaryop], Any] = {
    ast.UAdd: operator.pos,
    ast.USub: operator.neg,
    ast.Not: operator.not_,
}

_COMPARE_OPS: dict[type[ast.cmpop], Any] = {
    ast.Eq: operator.eq,
    ast.NotEq: operator.ne,
    ast.Lt: operator.lt,
    ast.LtE: operator.le,
    ast.Gt: operator.gt,
    ast.GtE: operator.ge,
    ast.In: lambda a, b: a in b,
    ast.NotIn: lambda a, b: a not in b,
    ast.Is: operator.is_,
    ast.IsNot: operator.is_not,
}

# Whitelisted встроенные функции (безопасные)
_SAFE_FUNCTIONS: dict[str, Any] = {
    "len": len,
    "min": min,
    "max": max,
    "abs": abs,
    "round": round,
    "sum": sum,
    "str": str,
    "int": int,
    "float": float,
    "bool": bool,
    "list": list,
    "sorted": sorted,
}

# Константы, доступные в выражениях
_SAFE_CONSTANTS: dict[str, Any] = {
    "True": True,
    "False": False,
    "None": None,
}

# Максимальная глубина вложенности (защита от stack overflow)
_MAX_DEPTH = 20


# ============================================================================
# Public API
# ============================================================================


def safe_eval(expression: str, context: dict[str, Any] | None = None) -> Any:
    """
    Безопасно вычислить выражение.

    Args:
        expression: строка с выражением (например, ``count > 10 and ok``)
        context: переменные, доступные в выражении

    Returns:
        Результат вычисления

    Raises:
        SafeEvalError: при синтаксической ошибке или запрещённой конструкции
    """
    ctx = context or {}

    try:
        tree = ast.parse(expression, mode="eval")
    except SyntaxError as e:
        raise SafeEvalError(f"Syntax error in expression '{expression}': {e}") from e

    try:
        return _eval_node(tree.body, ctx, depth=0)
    except SafeEvalError:
        raise
    except Exception as e:
        raise SafeEvalError(f"Error evaluating expression '{expression}': {e}") from e


def safe_eval_bool(expression: str, context: dict[str, Any] | None = None) -> bool:
    """
    Безопасно вычислить условие и вернуть bool.

    Используется для conditional edges в flow-графе.
    При любой ошибке возвращает False (fail-safe).

    Args:
        expression: строка с условием
        context: переменные

    Returns:
        bool результат (False при ошибке)
    """
    try:
        return bool(safe_eval(expression, context))
    except SafeEvalError:
        return False


def is_safe_expression(expression: str) -> bool:
    """
    Проверить, является ли выражение безопасным (без вычисления).

    Полезно для валидации при сохранении flow (отклонить опасные условия
    до выполнения).

    Returns:
        True если выражение содержит только безопасные конструкции.
    """
    try:
        tree = ast.parse(expression, mode="eval")
        _validate_node(tree.body, depth=0)
        return True
    except (SyntaxError, SafeEvalError):
        return False


# ============================================================================
# AST Evaluation (whitelist)
# ============================================================================


def _eval_node(node: ast.AST, context: dict[str, Any], depth: int) -> Any:
    """Рекурсивно вычислить AST узел (только whitelisted конструкции)."""
    if depth > _MAX_DEPTH:
        raise SafeEvalError("Expression is too deeply nested")

    next_depth = depth + 1

    # Константы (числа, строки, bool, None)
    if isinstance(node, ast.Constant):
        return node.value

    # Переменные (из context)
    if isinstance(node, ast.Name):
        if node.id in _SAFE_CONSTANTS:
            return _SAFE_CONSTANTS[node.id]
        if node.id in context:
            return context[node.id]
        raise SafeEvalError(f"Unknown variable: '{node.id}'")

    # Логические операторы (and, or) — с short-circuit
    if isinstance(node, ast.BoolOp):
        return _eval_boolop(node, context, next_depth)

    # Унарные операторы (not, -, +)
    if isinstance(node, ast.UnaryOp):
        op_type = type(node.op)
        if op_type not in _UNARY_OPS:
            raise SafeEvalError(f"Unsupported unary operator: {op_type.__name__}")
        return _UNARY_OPS[op_type](_eval_node(node.operand, context, next_depth))

    # Бинарные операторы (+, -, *, /, ...)
    if isinstance(node, ast.BinOp):
        op_type = type(node.op)
        if op_type not in _BINARY_OPS:
            raise SafeEvalError(f"Unsupported binary operator: {op_type.__name__}")
        left = _eval_node(node.left, context, next_depth)
        right = _eval_node(node.right, context, next_depth)
        return _BINARY_OPS[op_type](left, right)

    # Сравнения (==, !=, <, >, in, ...) —支持 цепочки (a < b < c)
    if isinstance(node, ast.Compare):
        return _eval_compare(node, context, next_depth)

    # Тернарный оператор (x if cond else y)
    if isinstance(node, ast.IfExp):
        return (
            _eval_node(node.body, context, next_depth)
            if _eval_node(node.test, context, next_depth)
            else _eval_node(node.orelse, context, next_depth)
        )

    # Коллекции (списки, кортежи)
    if isinstance(node, (ast.List, ast.Tuple)):
        return [_eval_node(el, context, next_depth) for el in node.elts]

    # Индексация (data["key"], items[0])
    if isinstance(node, ast.Subscript):
        value = _eval_node(node.value, context, next_depth)
        slice_value = _eval_node(node.slice, context, next_depth)
        try:
            return value[slice_value]
        except (KeyError, IndexError, TypeError) as e:
            raise SafeEvalError(f"Subscript error: {e}") from e

    # Атрибуты (result.status) — только non-dunder
    if isinstance(node, ast.Attribute):
        if node.attr.startswith("_"):
            raise SafeEvalError(f"Access to private/dunder attribute '{node.attr}' is forbidden")
        value = _eval_node(node.value, context, next_depth)
        try:
            return getattr(value, node.attr)
        except AttributeError as e:
            raise SafeEvalError(f"Attribute error: {e}") from e

    # Вызовы функций — только whitelisted
    if isinstance(node, ast.Call):
        return _eval_call(node, context, next_depth)

    raise SafeEvalError(f"Unsupported expression element: {type(node).__name__}")


def _eval_boolop(node: ast.BoolOp, context: dict[str, Any], depth: int) -> Any:
    """Вычислить and/or с short-circuit evaluation."""
    if isinstance(node.op, ast.And):
        result: Any = True
        for value in node.values:
            result = _eval_node(value, context, depth)
            if not result:
                return result
        return result

    if isinstance(node.op, ast.Or):
        result = False
        for value in node.values:
            result = _eval_node(value, context, depth)
            if result:
                return result
        return result

    raise SafeEvalError(f"Unsupported boolean operator: {type(node.op).__name__}")


def _eval_compare(node: ast.Compare, context: dict[str, Any], depth: int) -> bool:
    """Вычислить цепочку сравнений (a < b < c)."""
    left = _eval_node(node.left, context, depth)
    for op, comparator in zip(node.ops, node.comparators, strict=True):
        op_type = type(op)
        if op_type not in _COMPARE_OPS:
            raise SafeEvalError(f"Unsupported comparison operator: {op_type.__name__}")
        right = _eval_node(comparator, context, depth)
        if not _COMPARE_OPS[op_type](left, right):
            return False
        left = right
    return True


def _eval_call(node: ast.Call, context: dict[str, Any], depth: int) -> Any:
    """Вычислить вызов whitelisted функции (только по имени, без методов)."""
    if not isinstance(node.func, ast.Name):
        raise SafeEvalError("Only whitelisted function calls are allowed (no methods)")

    func_name = node.func.id
    if func_name not in _SAFE_FUNCTIONS:
        raise SafeEvalError(f"Function '{func_name}' is not allowed")

    if node.keywords:
        raise SafeEvalError("Keyword arguments are not supported")

    args = [_eval_node(arg, context, depth) for arg in node.args]
    try:
        return _SAFE_FUNCTIONS[func_name](*args)
    except Exception as e:
        raise SafeEvalError(f"Error calling '{func_name}': {e}") from e


# ============================================================================
# Validation (без вычисления)
# ============================================================================

_ALLOWED_NODE_TYPES: tuple[type[ast.AST], ...] = (
    ast.Constant,
    ast.Name,
    ast.Load,
    ast.BoolOp,
    ast.And,
    ast.Or,
    ast.UnaryOp,
    ast.UAdd,
    ast.USub,
    ast.Not,
    ast.BinOp,
    ast.Add,
    ast.Sub,
    ast.Mult,
    ast.Div,
    ast.FloorDiv,
    ast.Mod,
    ast.Compare,
    ast.Eq,
    ast.NotEq,
    ast.Lt,
    ast.LtE,
    ast.Gt,
    ast.GtE,
    ast.In,
    ast.NotIn,
    ast.Is,
    ast.IsNot,
    ast.IfExp,
    ast.List,
    ast.Tuple,
    ast.Subscript,
)


def _validate_node(node: ast.AST, depth: int) -> None:
    """
    Проверить, что выражение содержит только безопасные конструкции.

    Бросает SafeEvalError при обнаружении запрещённой конструкции.
    """
    if depth > _MAX_DEPTH:
        raise SafeEvalError("Expression is too deeply nested")

    next_depth = depth + 1

    # Атрибуты — только non-dunder
    if isinstance(node, ast.Attribute):
        if node.attr.startswith("_"):
            raise SafeEvalError(f"Access to private/dunder attribute '{node.attr}' is forbidden")
        _validate_node(node.value, next_depth)
        return

    # Вызовы — только whitelisted функции
    if isinstance(node, ast.Call):
        if not isinstance(node.func, ast.Name) or node.func.id not in _SAFE_FUNCTIONS:
            raise SafeEvalError("Only whitelisted function calls are allowed")
        if node.keywords:
            raise SafeEvalError("Keyword arguments are not supported")
        for arg in node.args:
            _validate_node(arg, next_depth)
        return

    if not isinstance(node, _ALLOWED_NODE_TYPES):
        raise SafeEvalError(f"Unsupported expression element: {type(node).__name__}")

    # Рекурсивная проверка дочерних узлов
    for child in ast.iter_child_nodes(node):
        _validate_node(child, next_depth)


# ============================================================================
# Exports
# ============================================================================

__all__ = [
    "SafeEvalError",
    "is_safe_expression",
    "safe_eval",
    "safe_eval_bool",
]