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.")
            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_test2(): # Using Test 2 to test sys redirection
            nonlocal result
            original_stdin = sys.stdin
            original_stdout = sys.stdout
            original_stderr = sys.stderr
            try:
                sys.stdin = websocket_io
                sys.stdout = websocket_io
                sys.stderr = websocket_io
                print(f"DEBUG: Executor thread '{threading.current_thread().name}' started. IO redirected.")
                print("DEBUG: This should now go via websocket_io.write!")
                
                # Test print to new sys.stdout
                print("Hello from redirected stdout!")
                
                result['success'] = True
            except Exception as e:
                error_msg = f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
                print(f"ERROR: Exception in _execute_in_thread_test2 during sys redirection test: {error_msg}")
                websocket_io.error_buffer.write(error_msg)
                result['success'] = False
            finally:
                sys.stdin = original_stdin
                sys.stdout = original_stdout
                sys.stderr = original_stderr
                print(f"DEBUG: Executor thread '{threading.current_thread().name}' finished. IO restored.")


        # 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_test2 in executor with timeout {self.timeout}s.")
            await asyncio.wait_for(
                loop.run_in_executor(None, _execute_in_thread_test2), # Call _execute_in_thread_test2
                timeout=self.timeout
            )
            print("DEBUG: _execute_in_thread_test2 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
        except Exception as e:
            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']:
            result['output'] = websocket_io.output_buffer.getvalue()
        if websocket_io.error_buffer.getvalue():
            result['error'] = result['error'] + "\n" + websocket_io.error_buffer.getvalue() if result['error'] else websocket_io.error_buffer.getvalue()
            result['success'] = False
        
        # 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.websocket.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
        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."""
    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]}...")
                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]}...")
                await websocket_io.input_queue.put(user_input + "\n")
            elif data["type"] == "ping":
                await websocket_io.websocket.send_text(json.dumps({"type": "pong"}))
                print(f"DEBUG: Sent pong to {websocket_io.websocket.client}")
            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
        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}")

if __name__ == "__main__":
    import uvicorn
    print("Starting Interactive Python Code Executor API (FastAPI)...")
    print(f"Python version: {sys.version}")
    print(f"Working directory: {os.getcwd()}")
    
    server_name = os.getenv('FASTAPI_SERVER_NAME', '0.0.0.0')
    server_port = int(os.getenv('FASTAPI_SERVER_PORT', 7860))
    
    print(f"Starting server on {server_name}:{server_port}")
    
    try:
        uvicorn.run(
            "app:app",
            host=server_name,
            port=server_port,
            reload=False,
            log_level="info"
        )
    except Exception as e:
        print(f"Failed to start server: {e}")
        sys.exit(1)