| """ |
| 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: |
| |
| return True, "OK" |
|
|
| for node in ast.walk(tree): |
| |
| 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" |
| |
| |
| |
|
|
| |
| 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" |
|
|
| |
| 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): |
| |
| for base in node.bases: |
| |
| if isinstance(base, ast.Name): |
| if 'Scene' in base.id: |
| scene_classes.append(node.name) |
| break |
| |
| 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) |
| """ |
| |
| is_valid, error_msg = validate_python_syntax(code) |
| if not is_valid: |
| return False, error_msg, None |
|
|
| |
| is_safe, error_msg = check_for_dangerous_imports(code) |
| if not is_safe: |
| return False, error_msg, None |
|
|
| |
| has_animations, anim_error = check_for_animation_commands(code) |
| if not has_animations: |
| return False, anim_error, None |
|
|
| |
| has_latex, latex_warning = check_for_latex_usage(code) |
| if has_latex: |
| |
| import logging |
| logger = logging.getLogger("manim_studio.validators") |
| logger.warning(latex_warning) |
|
|
| |
| 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 |
|
|
| |
| if scene_name: |
| |
| 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: |
| |
| if len(scene_classes) == 1: |
| return True, "OK", scene_classes[0] |
| else: |
| |
| available = ", ".join(scene_classes) |
| return False, f"Multiple Scene classes found: {available}. Please specify scene_name in your request.", None |
|
|
