src/validation.py

"""
Input validation and sanitization module for Prompt2Frame.

This module provides comprehensive validation for user prompts and generated code
to prevent security vulnerabilities and ensure safe execution.
"""

import re
from typing import List, Tuple
import logging

logger = logging.getLogger(__name__)


class PromptValidator:
    """Validates and sanitizes user prompts."""
    
    # Dangerous patterns that might be used for code injection
    DANGEROUS_PATTERNS = [
        # File system operations
        r'\bopen\s*\(',
        r'\bfile\s*\(',
        r'\bread\s*\(',
        r'\bwrite\s*\(',
        r'\bos\.',
        r'\bpath\.',
        r'__file__',
        r'__path__',
        
        # Network operations
        r'\burllib\b',
        r'\brequests\b',
        r'\bsocket\b',
        r'\bhttp\.',
        r'\bftp\b',
        
        # System/subprocess
        r'\bsubprocess\b',
        r'\bsystem\(',
        r'\bexec\(',
        r'\beval\(',
        r'\bcompile\(',
        r'__import__',
        
        # Database
        r'\bsql\b',
        r'\binsert\b.*\binto\b',
        r'\bselect\b.*\bfrom\b',
        r'\bdrop\b.*\btable\b',
        
        # Code execution
        r'globals\(',
        r'locals\(',
        r'vars\(',
        r'dir\(',
    ]
    
    @classmethod
    def validate_prompt(cls, prompt: str) -> Tuple[bool, str]:
        """
        Validate a user prompt for security and content.
        
        Args:
            prompt: The user's prompt
            
        Returns:
            Tuple of (is_valid, error_message)
        """
        # Check length
        if not prompt or not prompt.strip():
            return False, "Prompt cannot be empty"
        
        if len(prompt) < 3:
            return False, "Prompt is too short (minimum 3 characters)"
        
        if len(prompt) > 500:
            return False, "Prompt is too long (maximum 500 characters)"
        
        # Check for dangerous patterns
        prompt_lower = prompt.lower()
        for pattern in cls.DANGEROUS_PATTERNS:
            if re.search(pattern, prompt_lower, re.IGNORECASE):
                logger.warning(f"Dangerous pattern detected in prompt: {pattern}")
                return False, (
                    "Your prompt contains potentially unsafe content. "
                    "Please rephrase to describe visual animations only."
                )
        
        # Check for excessive punctuation (spam indicator)
        special_char_ratio = sum(c in '!@#$%^&*()_+=' for c in prompt) / len(prompt)
        if special_char_ratio > 0.3:
            return False, "Prompt contains too many special characters"
        
        # Check for repeated characters (spam indicator)
        if re.search(r'(.)\1{10,}', prompt):
            return False, "Prompt contains excessive repeated characters"
        
        return True, ""
    
    @classmethod
    def sanitize_prompt(cls, prompt: str) -> str:
        """
        Sanitize a prompt by removing potentially harmful content.
        
        Args:
            prompt: The raw prompt
            
        Returns:
            Sanitized prompt
        """
        # Remove null bytes
        prompt = prompt.replace('\x00', '')
        
        # Remove control characters except newlines and tabs
        prompt = ''.join(char for char in prompt if char.isprintable() or char in '\n\t')
        
        # Normalize whitespace
        prompt = ' '.join(prompt.split())
        
        # Trim to max length
        prompt = prompt[:500]
        
        return prompt.strip()


class CodeSecurityValidator:
    """Enhanced security validation for generated code."""
    
    # Comprehensive list of dangerous operations
    DANGEROUS_OPERATIONS = [
        # File I/O
        'open(', 'file(', 'with open',
        
        # OS operations  
        'os.', 'sys.', 'subprocess.', 'shutil.', 'pathlib.',
        
        # Network
        'urllib', 'requests', 'socket', 'http.',
        
        # Code execution
        'exec', 'eval', 'compile', '__import__',
        'globals()', 'locals()', 'vars()',
        
        # Dangerous imports
        'import os', 'import sys', 'import subprocess',
        'import shutil', 'import requests', 'import urllib',
        'import socket', 'import pickle',
        
        # Shell commands
        'popen(', 'shell=',
    ]
    
    # Allowed imports only
    ALLOWED_IMPORTS = [
        'from manim import *',
        'import random',
        'import numpy as np',
        'import math',  # Safe math operations
    ]
    
    @classmethod
    def validate_code_safety(cls, code: str) -> Tuple[bool, str]:
        """
        Validate that generated code is safe to execute.
        
        Args:
            code: The generated Python code
            
        Returns:
            Tuple of (is_safe, error_message)
        """
        # Check for dangerous operations
        code_lower = code.lower()
        for operation in cls.DANGEROUS_OPERATIONS:
            if operation in code_lower:
                logger.error(f"Dangerous operation detected: {operation}")
                return False, f"Code contains forbidden operation: {operation}"
        
        # Validate imports
        import_lines = [line.strip() for line in code.split('\n') 
                       if line.strip().startswith(('import ', 'from '))]
        
        for imp_line in import_lines:
            # Check if it's an allowed import
            is_allowed = any(allowed in imp_line for allowed in cls.ALLOWED_IMPORTS)
            if not is_allowed:
                logger.error(f"Unauthorized import: {imp_line}")
                return False, f"Unauthorized import statement: {imp_line}"
        
        # Check for attempts to access internals
        if '__' in code and any(dangerous in code for dangerous in ['__file__', '__path__', '__dict__', '__class__']):
            return False, "Code attempts to access Python internals"
        
        # Check code length (prevent DoS)
        if len(code) > 10000:
            return False, "Generated code is too large"
        
        # Count class definitions (should be exactly 1 Scene subclass)
        class_count = code.count('class ')
        if class_count < 1:
            return False, "No class definition found"
        if class_count > 3:
            return False, "Too many class definitions"
        
        return True, ""
    
    @classmethod
    def analyze_code_complexity(cls, code: str) -> dict:
        """
        Analyze code complexity to prevent resource exhaustion.
        
        Args:
            code: The generated code
            
        Returns:
            Dict with complexity metrics
        """
        return {
            'line_count': len(code.split('\n')),
            'object_count': code.count('Circle(') + code.count('Square(') + 
                          code.count('Rectangle(') + code.count('Triangle('),
            'animation_count': code.count('self.play('),
            'loop_count': code.count('for ') + code.count('while '),
            'function_count': code.count('def '),
        }
    
    @classmethod
    def validate_code_complexity(cls, code: str) -> Tuple[bool, str]:
        """
        Validate that code complexity is within acceptable limits.
        
        Args:
            code: The generated code
            
        Returns:
            Tuple of (is_valid, error_message)
        """
        metrics = cls.analyze_code_complexity(code)
        
        # Check limits (Relaxed for complex animations)
        if metrics['object_count'] > 100:
            return False, "Too many objects (limit: 100)"
        
        if metrics['animation_count'] > 100:
            return False, "Too many animations (limit: 100)"
        
        if metrics['loop_count'] > 50:
            return False, "Too many loops (limit: 50)"
        
        if metrics['function_count'] > 50:
            return False, "Too many function definitions (limit: 50)"
        
        return True, ""