File size: 8,489 Bytes
0024fe7 53dfbe8 bd3a6a3 0024fe7 53dfbe8 bd3a6a3 0024fe7 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 | """
Code validation and security checking utilities
"""
import ast
from typing import Tuple, List, Optional
from config import ALLOWED_IMPORTS, BLOCKED_IMPORTS
def validate_python_syntax(code: str) -> Tuple[bool, str]:
"""
Validate Python code syntax using AST parsing
Args:
code: Python code string to validate
Returns:
Tuple of (is_valid, error_message)
- If valid: (True, "OK")
- If invalid: (False, "SyntaxError: ...")
"""
try:
ast.parse(code)
return True, "OK"
except SyntaxError as e:
error_msg = f"SyntaxError on line {e.lineno}: {e.msg}"
if e.text:
error_msg += f"\n {e.text.strip()}"
return False, error_msg
except Exception as e:
return False, f"Parse error: {str(e)}"
def check_for_dangerous_imports(code: str) -> Tuple[bool, str]:
"""
Check for potentially dangerous imports that could pose security risks
Args:
code: Python code string to analyze
Returns:
Tuple of (is_safe, error_message)
- If safe: (True, "OK")
- If unsafe: (False, "Blocked import: ...")
"""
try:
tree = ast.parse(code)
except SyntaxError:
# Syntax errors are handled by validate_python_syntax
return True, "OK"
for node in ast.walk(tree):
# Check regular imports: import os, import sys
if isinstance(node, ast.Import):
for alias in node.names:
module_name = alias.name.split('.')[0]
if module_name in BLOCKED_IMPORTS:
return False, f"Blocked import: '{alias.name}' - This module is restricted for security reasons"
# Optionally enforce whitelist (commented out for flexibility)
# if module_name not in ALLOWED_IMPORTS:
# return False, f"Unauthorized import: '{alias.name}' - Only {ALLOWED_IMPORTS} are allowed"
# Check from imports: from os import system
elif isinstance(node, ast.ImportFrom):
if node.module:
module_name = node.module.split('.')[0]
if module_name in BLOCKED_IMPORTS:
return False, f"Blocked import: 'from {node.module}' - This module is restricted for security reasons"
# Check for dangerous built-in functions
elif isinstance(node, ast.Name):
if node.id in BLOCKED_IMPORTS:
return False, f"Blocked function: '{node.id}' - This function is restricted for security reasons"
return True, "OK"
def extract_scene_classes(code: str) -> List[str]:
"""
Extract all Scene class names from Manim code using AST
Args:
code: Python code string containing Manim scenes
Returns:
List of Scene class names found in the code
"""
try:
tree = ast.parse(code)
except SyntaxError:
return []
scene_classes = []
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef):
# Check if class inherits from Scene or any Scene-like class
for base in node.bases:
# Handle simple inheritance: class MyScene(Scene)
if isinstance(base, ast.Name):
if 'Scene' in base.id:
scene_classes.append(node.name)
break
# Handle module.Class inheritance: class MyScene(manim.Scene)
elif isinstance(base, ast.Attribute):
if 'Scene' in base.attr:
scene_classes.append(node.name)
break
return scene_classes
def validate_scene_exists(code: str, scene_name: str) -> Tuple[bool, str]:
"""
Verify that a specific scene class exists in the code
Args:
code: Python code string
scene_name: Name of the Scene class to verify
Returns:
Tuple of (exists, error_message)
- If exists: (True, "OK")
- If not found: (False, "Scene 'X' not found...")
"""
scene_classes = extract_scene_classes(code)
if scene_name in scene_classes:
return True, "OK"
else:
available = ", ".join(scene_classes) if scene_classes else "none"
return False, f"Scene '{scene_name}' not found in code. Available scenes: {available}"
def auto_detect_scene(code: str) -> Optional[str]:
"""
Automatically detect the scene to render when user doesn't specify one
Args:
code: Python code string
Returns:
Scene class name if exactly one is found, None otherwise
"""
scene_classes = extract_scene_classes(code)
if len(scene_classes) == 1:
return scene_classes[0]
else:
return None
def check_for_animation_commands(code: str) -> Tuple[bool, str]:
"""
Check if code contains animation commands (self.play() or self.wait())
Scenes without these will not generate video frames
Args:
code: Python code to check
Returns:
Tuple of (has_animations, error_message)
"""
has_play = 'self.play(' in code
has_wait = 'self.wait(' in code
if not (has_play or has_wait):
return False, (
"ERROR: Your scene must include 'self.play()' or 'self.wait()' to generate video frames. "
"Using only 'self.add()' will not create an animation. "
"Example: self.play(Create(circle)) or self.wait(1)"
)
return True, "OK"
def check_for_latex_usage(code: str) -> Tuple[bool, str]:
"""
Check if code uses LaTeX objects which may cause rendering issues
Args:
code: Python code to check
Returns:
Tuple of (has_latex, warning_message)
"""
latex_classes = ['MathTex', 'Tex', 'TexTemplate', 'Text']
found_latex = []
for latex_class in latex_classes:
if latex_class in code:
found_latex.append(latex_class)
if found_latex:
warning = (
f"Warning: Your code uses {', '.join(found_latex)} which requires LaTeX. "
f"This may cause timeouts. Consider using geometric shapes instead (Circle, Square, Line, etc.)."
)
return True, warning
return False, "OK"
def validate_code_full(code: str, scene_name: Optional[str] = None) -> Tuple[bool, str, Optional[str]]:
"""
Perform full code validation including syntax, security, and scene checks
Args:
code: Python code to validate
scene_name: Optional specific scene name to validate
Returns:
Tuple of (is_valid, error_message, detected_scene_name)
- is_valid: True if all checks pass
- error_message: Error description if validation fails, "OK" otherwise
- detected_scene_name: Auto-detected scene name (if scene_name was None)
"""
# Step 1: Check syntax
is_valid, error_msg = validate_python_syntax(code)
if not is_valid:
return False, error_msg, None
# Step 2: Check for dangerous imports
is_safe, error_msg = check_for_dangerous_imports(code)
if not is_safe:
return False, error_msg, None
# Step 2.5: Check for animation commands (REQUIRED)
has_animations, anim_error = check_for_animation_commands(code)
if not has_animations:
return False, anim_error, None
# Step 2.6: Check for LaTeX usage (warning only)
has_latex, latex_warning = check_for_latex_usage(code)
if has_latex:
# Log warning but don't fail
import logging
logger = logging.getLogger("manim_studio.validators")
logger.warning(latex_warning)
# Step 3: Check scenes
scene_classes = extract_scene_classes(code)
if not scene_classes:
return False, "No Scene classes found in code. Please define at least one class inheriting from manim.Scene", None
# Step 4: Validate or auto-detect scene
if scene_name:
# User specified a scene - verify it exists
is_valid, error_msg = validate_scene_exists(code, scene_name)
if not is_valid:
return False, error_msg, None
return True, "OK", scene_name
else:
# Auto-detect scene
if len(scene_classes) == 1:
return True, "OK", scene_classes[0]
else:
# Multiple scenes found, user must specify
available = ", ".join(scene_classes)
return False, f"Multiple Scene classes found: {available}. Please specify scene_name in your request.", None
|