Spaces:
Sleeping
Sleeping
File size: 3,557 Bytes
9bed109 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | from __future__ import annotations
import ast
from dataclasses import dataclass
FORBIDDEN_CALL_NAMES = frozenset({"exec", "eval", "__import__", "compile"})
ALLOWED_IMPORT_ROOTS = frozenset(
{
"__future__",
"manim",
"primitives",
"typing",
"math",
"numpy",
},
)
class SandboxValidationError(ValueError):
"""Raised when generated Manim code fails static safety checks."""
@dataclass(frozen=True)
class SandboxLimits:
max_bytes: int
def _import_root(name: str) -> str:
return name.split(".", 1)[0]
def _validate_imports(tree: ast.Module) -> None:
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
root = _import_root(alias.name)
if root not in ALLOWED_IMPORT_ROOTS:
msg = f"Disallowed import root: {root!r}"
raise SandboxValidationError(msg)
if isinstance(node, ast.ImportFrom):
if node.level and node.level > 0:
msg = "Relative imports are not allowed in generated code"
raise SandboxValidationError(msg)
if node.module is None:
msg = "Wildcard or ambiguous imports are not allowed"
raise SandboxValidationError(msg)
root = _import_root(node.module)
if root not in ALLOWED_IMPORT_ROOTS:
msg = f"Disallowed import-from module root: {root!r}"
raise SandboxValidationError(msg)
if isinstance(node, ast.Call):
func = node.func
if isinstance(func, ast.Name) and func.id in FORBIDDEN_CALL_NAMES:
msg = f"Disallowed call: {func.id}()"
raise SandboxValidationError(msg)
def _validate_generated_scene_class(tree: ast.Module) -> None:
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name == "GeneratedScene":
return
msg = "Generated code must define class GeneratedScene"
raise SandboxValidationError(msg)
def validate_manim_code(source: str, *, limits: SandboxLimits) -> None:
"""Static checks: size, syntax, import policy, required class name."""
data = source.encode("utf-8")
if len(data) > limits.max_bytes:
msg = f"Code exceeds max size ({limits.max_bytes} bytes)"
raise SandboxValidationError(msg)
try:
tree = ast.parse(source)
except SyntaxError as exc:
msg = f"Invalid Python syntax: {exc.msg} at line {exc.lineno}, col {exc.offset}"
raise SandboxValidationError(msg) from exc
if not isinstance(tree, ast.Module):
msg = "Expected a module"
raise SandboxValidationError(msg)
_validate_imports(tree)
_validate_generated_scene_class(tree)
def static_check_split(source: str, *, limits: SandboxLimits) -> tuple[bool, bool, str | None]:
"""Return (syntax_ok, policy_ok, error_code).
``syntax_ok`` means Python parses. ``policy_ok`` adds import rules + GeneratedScene class.
"""
data = source.encode("utf-8")
if len(data) > limits.max_bytes:
return False, False, "size_exceeded"
try:
tree = ast.parse(source)
except SyntaxError:
return False, False, "syntax_error"
if not isinstance(tree, ast.Module):
return False, False, "syntax_error"
try:
_validate_imports(tree)
_validate_generated_scene_class(tree)
except SandboxValidationError:
return True, False, "policy_error"
return True, True, None
|