Blame for app.py

import sys
import io
import traceback
import ast
import time
import json
import os
from contextlib import redirect_stdout, redirect_stderr
from typing import Dict, Any, Optional
import threading
import asyncio
import signal
from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from pydantic import BaseModel
from starlette.responses import HTMLResponse

# Set up proper signal handling for Docker
def signal_handler(signum, frame):
    print(f"Received signal {signum}, shutting down gracefully...")
    sys.exit(0)

signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGINT, signal_handler)

# --- Custom IO Redirection for Interactive Execution ---

class WebSocketInputOutput:
    """
    Redirects stdin/stdout/stderr to/from a WebSocket connection.
    Handles interactive input by waiting for messages from the client.
    """
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
        self.input_queue = asyncio.Queue()
        self.output_buffer = io.StringIO()
        self.error_buffer = io.StringIO()
        self.closed = False
        print("DEBUG: WebSocketInputOutput initialized.")

    async def write(self, s: str):
        """Writes to stdout/stderr and sends to WebSocket."""
        self.output_buffer.write(s)
        try:
            # Send output in chunks or lines to prevent large messages
            await self.websocket.send_text(json.dumps({"type": "output", "content": s}))
            print(f"DEBUG: Sent output chunk: {s.strip()[:100]}...")
        except WebSocketDisconnect:
            self.closed = True
            print("ERROR: WebSocket disconnected during write.")
        except Exception as e:
            print(f"ERROR: Error sending output over WebSocket: {e}")
            self.closed = True

    def flush(self):
        """Flushes the buffer (no-op for now, as we send immediately)."""
        pass

    async def readline(self) -> str:
        """Reads a line from stdin, waiting for input from WebSocket."""
        try:
            print("DEBUG: Sending input_request to client.")
            await self.websocket.send_text(json.dumps({"type": "input_request"}))
            line = await self.input_queue.get() # Wait for input from client
            print(f"DEBUG: Received input from queue: {line.strip()[:50]}...")
            return line
        except WebSocketDisconnect:
            self.closed = True
            print("ERROR: WebSocket disconnected during readline.")
            raise EOFError("Input stream closed due to WebSocket disconnect")
        except Exception as e:
            print(f"ERROR: Error requesting input over WebSocket: {e}")
            self.closed = True
            raise EOFError(f"Input stream error: {e}")

    # For compatibility with sys.stdin, which expects a file-like object
    def _get_input_line(self):
        """Synchronous wrapper for readline for use with exec."""
        try:
            print("DEBUG: Synchronous input wrapper called.")
            # This is called from the executor thread, so asyncio.run is appropriate here
            return asyncio.run(self.readline())
        except Exception as e:
            print(f"ERROR: Synchronous input wrapper error: {e}")
            raise

    def read(self, n=-1):
        """Reads n characters. For simplicity, we'll treat it as readline."""
        print(f"DEBUG: Read called with n={n}. Delegating to readline.")
        return self._get_input_line() # Or implement more complex buffering

    def __getattr__(self, name):
        """Delegate other attributes if needed, or raise AttributeError."""
        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")


# --- Python Syntax Validation FSM (unchanged) ---

class PythonSyntaxFSM:
    """Finite State Machine for Python syntax validation"""
    
    def __init__(self):
        self.reset()
        
    def reset(self):
        self.state = 'normal'
        self.bracket_stack = []
        self.in_string = False
        self.string_delimiter = None
        self.line_number = 0
        self.errors = []
        self.warnings = []
        
    def validate_code(self, code: str) -> Dict[str, Any]:
        """Validate Python code using FSM approach"""
        self.reset()
        
        if not code or not code.strip():
            return {
                'valid': False,
                'errors': [{'line': 0, 'message': 'Empty code input', 'type': 'input_error'}],
                'warnings': [],
                'total_issues': 1
            }
        
        lines = code.split('\n')
        
        for i, line in enumerate(lines):
            self.line_number = i + 1
            self.validate_line(line)
            
        # Check for unclosed brackets at end
        if self.bracket_stack:
            self.errors.append({
                'line': self.line_number,
                'type': 'syntax_error',
                'message': f'Unclosed bracket: {self.bracket_stack[-1]}',
                'severity': 'error'
            })
            
        return {
            'valid': len(self.errors) == 0,
            'errors': self.errors,
            'warnings': self.warnings,
            'total_issues': len(self.errors) + len(self.warnings)
        }
    
    def validate_line(self, line: str):
        """Validate a single line using FSM logic"""
        if not line.strip():
            return
            
        # Check indentation
        leading_spaces = len(line) - len(line.lstrip())
        if leading_spaces % 4 != 0 and line.strip():
            self.warnings.append({
                'line': self.line_number,
                'type': 'style_warning',
                'message': 'Inconsistent indentation (PEP 8 recommends 4 spaces)',
                'severity': 'warning'
            })
        
        # Process character by character for brackets and strings
        i = 0
        while i < len(line):
            char = line[i]
            
            # Handle string states
            if self.in_string:
                if char == self.string_delimiter and (i == 0 or line[i-1] != '\\'):
                    self.in_string = False
                    self.string_delimiter = None
                i += 1
                continue
                
            # Handle comment detection
            if char == '#':
                break   # Rest of line is comment
                
            # Handle string detection
            if char in ['"', "'"]:
                # Check for triple quotes
                if i + 2 < len(line) and line[i:i+3] == char * 3:
                    self.string_delimiter = char * 3
                    i += 3
                else:
                    self.string_delimiter = char
                    i += 1
                self.in_string = True
                continue
                
            # Handle brackets
            if char in '([{':
                self.bracket_stack.append(char)
            elif char in ')]}':
                if not self.bracket_stack:
                    self.errors.append({
                        'line': self.line_number,
                        'type': 'syntax_error',
                        'message': f'Unmatched closing bracket: {char}',
                        'severity': 'error'
                    })
                else:
                    expected = {'(': ')', '[': ']', '{': '}'}
                    last_open = self.bracket_stack[-1]
                    if expected[last_open] != char:
                        self.errors.append({
                            'line': self.line_number,
                            'type': 'syntax_error',
                            'message': f'Mismatched brackets: expected {expected[last_open]}, got {char}',
                            'severity': 'error'
                        })
                    else:
                        self.bracket_stack.pop()
                        
            i += 1
        
        # Check Python syntax patterns
        stripped = line.strip()
        control_keywords = ['def ', 'class ', 'if ', 'elif ', 'for ', 'while ', 'try:', 'except']
        
        for keyword in control_keywords:
            if stripped.startswith(keyword):
                if not stripped.endswith(':') and keyword != 'except':
                    self.errors.append({
                        'line': self.line_number,
                        'type': 'syntax_error',
                        'message': f'{keyword.strip()} statement must end with colon',
                        'severity': 'error'
                    })
                break

# --- Secure Python Executor (modified for WebSockets) ---

class SecurePythonExecutor:
    """Secure Python code executor with sandboxing and timeout, now interactive."""
    
    def __init__(self, timeout: int = 10):
        self.timeout = timeout
        self.restricted_modules = {
            'os', 'sys', 'subprocess', 'socket', 'urllib', 'requests',
            'shutil', 'glob', 'pickle', 'marshal', 'shelve', 'dbm',
            'sqlite3', 'threading', 'multiprocessing', 'ctypes', 'importlib',
            'builtins', '__builtin__', 'imp', 'zipimport'
        }
        
        # Create safe builtins
        # IMPORTANT: 'input' is now included and will be redirected!
        self.safe_builtins = {
            'print': print,
            'len': len,
            'range': range,
            'str': str,
            'int': int,
            'float': float,
            'bool': bool,
            'list': list,
            'dict': dict,
            'tuple': tuple,
            'set': set,
            'frozenset': frozenset,
            'abs': abs,
            'max': max,
            'min': min,
            'sum': sum,
            'sorted': sorted,
            'reversed': reversed,
            'enumerate': enumerate,
            'zip': zip,
            'map': map,
            'filter': filter,
            'all': all,
            'any': any,
            'type': type,
            'isinstance': isinstance,
            'hasattr': hasattr,
            'getattr': getattr,
            'setattr': setattr,
            'delattr': delattr,
            'round': round,
            'pow': pow,
            'divmod': divmod,
            'chr': chr,
            'ord': ord,
            'hex': hex,
            'oct': oct,
            'bin': bin,
            'format': format,
            'repr': repr,
            'ascii': ascii,
            'iter': iter,
            'next': next,
            'slice': slice,
            'callable': callable,
            'id': id,
            'hash': hash,
            'vars': vars,
            'dir': dir,
            'help': help,
            'Exception': Exception,
            'ValueError': ValueError,
            'TypeError': TypeError,
            'IndexError': IndexError,
            'KeyError': KeyError,
            'AttributeError': AttributeError,
            'NameError': NameError,
            'ZeroDivisionError': ZeroDivisionError,
            'input': input # Now included for interactive use
        }
        
    async def execute_code(self, code: str, websocket_io: WebSocketInputOutput) -> Dict[str, Any]:
        """Execute Python code in a secure environment with WebSocket IO."""
        start_time = time.time()
        print(f"DEBUG: execute_code called for client: {websocket_io.websocket.client}")
        
        # Input validation
        if not code or not code.strip():
            error_msg = 'No code provided'
            await websocket_io.write(f"Error: {error_msg}\n")
            print(f"DEBUG: No code provided for execution. Client: {websocket_io.websocket.client}")
            return {
                'success': False,
                'output': '',
                'error': error_msg,
                'execution_time': 0,
                'validation': {
                    'valid': False,
                    'errors': [{'message': error_msg, 'line': 0}],
                    'warnings': [],
                    'total_issues': 1
                }
            }
        
        # Pre-execution validation
        fsm = PythonSyntaxFSM()
        validation_result = fsm.validate_code(code)
        
        # Check for restricted imports
        restricted_check = self.check_restricted_imports(code)
        if not restricted_check['allowed']:
            error_msg = f"Security violation: Restricted module '{restricted_check['module']}' is not allowed"
            await websocket_io.write(f"Error: {error_msg}\n")
            print(f"DEBUG: Restricted import detected: {restricted_check['module']}. Client: {websocket_io.websocket.client}")
            return {
                'success': False,
                'output': '',
                'error': error_msg,
                'execution_time': 0,
                'validation': validation_result
            }
        
        # Execute code with timeout
        result = await self.run_with_timeout(code, websocket_io) # Pass websocket_io
        execution_time = time.time() - start_time
        print(f"DEBUG: Code execution finished. Success: {result['success']}, Client: {websocket_io.websocket.client}")
        
        return {
            'success': result['success'],
            'output': result['output'],
            'error': result['error'],
            'execution_time': round(execution_time, 3),
            'validation': validation_result
        }
    
    def check_restricted_imports(self, code: str) -> Dict[str, Any]:
        """Check for restricted module imports"""
        try:
            tree = ast.parse(code)
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for alias in node.names:
                        if alias.name.split('.')[0] in self.restricted_modules:
                            return {'allowed': False, 'module': alias.name}
                elif isinstance(node, ast.ImportFrom):
                    if node.module and node.module.split('.')[0] in self.restricted_modules:
                        return {'allowed': False, 'module': node.module}
            return {'allowed': True, 'module': None}
        except Exception as e:
            print(f"WARNING: AST parsing failed for restricted import check: {e}")
            return {'allowed': True, 'module': None}   # If parsing fails, let execution handle it
    
    async def run_with_timeout(self, code: str, websocket_io: WebSocketInputOutput) -> Dict[str, Any]:
        """Run code with timeout protection, interacting via WebSocketIO."""
        result = {'success': False, 'output': '', 'error': ''}
        print(f"DEBUG: run_with_timeout called. Client: {websocket_io.websocket.client}")
        
        # Define the synchronous target function to be run in the executor thread
        def _execute_in_thread_test1():
            # This function runs in a separate thread managed by the event loop's executor
            nonlocal result
            
            # --- START TEST 1 MODIFICATIONS ---
            # DO NOT redirect sys.stdin/stdout/stderr here yet
            print(f"DEBUG: Executor thread '{threading.current_thread().name}' started. (No IO redirected yet)") # UPDATED PRINT
            result['success'] = True
            # Do not call exec yet. Just write directly to output_buffer for testing.
            websocket_io.output_buffer.write("DEBUG: Test message from executor thread (no exec).\n") # Write to buffer directly
            # Do not call asyncio.run(websocket_io.write) here. The main loop will check buffer.
            print("DEBUG: _execute_in_thread finished with direct buffer write.") # ADDED THIS LINE
            # --- END TEST 1 MODIFICATIONS ---

        # Use the event loop's default ThreadPoolExecutor to run the synchronous code
        loop = asyncio.get_running_loop()
        try:
            print(f"DEBUG: Scheduling _execute_in_thread_test1 in executor with timeout {self.timeout}s.")
            # Run the synchronous _execute_in_thread_test1 in a separate thread
            await asyncio.wait_for(
                loop.run_in_executor(None, _execute_in_thread_test1), # None uses the default ThreadPoolExecutor
                timeout=self.timeout
            )
            print("DEBUG: _execute_in_thread_test1 completed via executor.")
        except asyncio.TimeoutError:
            print(f"ERROR: Code execution timed out after {self.timeout} seconds in run_in_executor.")
            result['error'] = f"Code execution timed out after {self.timeout} seconds"
            result['success'] = False
            try:
                await websocket_io.write(f"Error: {result['error']}\n")
            except Exception:
                pass # Ignore if websocket is already closed
        except Exception as e:
            # This catches exceptions from _execute_in_thread_test1 that propagate up
            # (e.g., if the thread itself crashes in an unrecoverable way before our try/except)
            print(f"ERROR: Unexpected exception from executor: {e}")
            result['error'] = f"Unexpected executor error: {type(e).__name__}: {str(e)}"
            result['success'] = False
            try:
                await websocket_io.write(f"Error: {result['error']}\n")
            except Exception:
                pass

        # After execution (or timeout), collect output
        if result['success']: # Only if execution was successful in _execute_in_thread_test1
            result['output'] = websocket_io.output_buffer.getvalue()
        if websocket_io.error_buffer.getvalue(): # Error buffer always checked
            result['error'] = result['error'] + "\n" + websocket_io.error_buffer.getvalue() if result['error'] else websocket_io.error_buffer.getvalue()
            result['success'] = False # Ensure success is false if there's any error output
        
        # Send a final message to the client indicating execution is complete
        try:
            print(f"DEBUG: Attempting to send execution_complete message (success={result['success']}).")
            await websocket_io.send_text(json.dumps({"type": "execution_complete", "result": result}))
            print("DEBUG: execution_complete message sent.")
        except WebSocketDisconnect:
            print("DEBUG: WebSocketDisconnect when sending execution_complete.")
            pass # Ignore if websocket is already closed
        except Exception as e:
            print(f"ERROR: Error sending execution_complete message: {e}")

        print(f"DEBUG: run_with_timeout returning result for client: {websocket_io.websocket.client}")
        return result


# Initialize the executor
executor = SecurePythonExecutor(timeout=10)

# Initialize FastAPI app
app = FastAPI(
    title="Interactive Python Code Executor API",
    description="A secure WebSocket API for executing Python code interactively.",
    version="1.0.0",
)

# Serve the HTML frontend
@app.get("/", response_class=HTMLResponse)
async def get_index():
    """Serves the interactive frontend."""
    # This assumes index.html is in the same directory as app.py
    # Explicitly copy index.html to /app/index.html to avoid any ambiguity with '.'
    with open("index.html", "r") as f:
        return HTMLResponse(content=f.read())

# WebSocket endpoint for interactive execution
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    websocket_io = WebSocketInputOutput(websocket)
    print(f"WebSocket connection established: {websocket.client}")

    try:
        while True:
            message = await websocket.receive_text()
            data = json.loads(message)
            print(f"DEBUG: Received message from {websocket.client}: type={data['type']}")

            if data["type"] == "code":
                code = data["content"]
                print(f"Received code from {websocket_io.websocket.client}: {code[:50]}...") # Use websocket_io.websocket.client here
                # Execute code asynchronously in the background
                asyncio.create_task(executor.execute_code(code, websocket_io))
            elif data["type"] == "input":
                user_input = data["content"]
                print(f"Received input from {websocket_io.websocket.client}: {user_input.strip()[:50]}...") # Use websocket_io.websocket.client here
                await websocket_io.input_queue.put(user_input + "\n") # Add newline for readline
            elif data["type"] == "ping":
                # Respond to pings to keep connection alive
                await websocket.send_text(json.dumps({"type": "pong"}))
                print(f"DEBUG: Sent pong to {websocket_io.websocket.client}") # Use websocket_io.websocket.client here
            else:
                await websocket_io.write(f"Unknown message type: {data['type']}\n")

    except WebSocketDisconnect:
        print(f"WebSocket disconnected: {websocket.client}")
    except json.JSONDecodeError:
        print(f"Received invalid JSON from {websocket.client}: {message}")
    except Exception as e:
        print(f"WebSocket error for {websocket.client}: {e}")
    finally:
        websocket_io.closed = True
        # Clean up any pending input requests if the client disconnects
        while not websocket_io.input_queue.empty():
            try:
                websocket_io.input_queue.get_nowait()
            except asyncio.QueueEmpty:
                pass
        print(f"WebSocket connection closed for {websocket_io.websocket.client}")