Update app.py

app.py

@@ -1,3 +1,477 @@
+#!/usr/bin/env python3
+"""
+LibreChat Custom Code Interpreter - HuggingFace Space
+A secure code execution server using Gradio + MCP with SSE support
+"""
+
+import json
+import subprocess
+import tempfile
+import os
+import sys
+import uuid
+import time
+import base64
+import io
+import re
+from pathlib import Path
+from typing import Dict, Any, Optional, List
+import gradio as gr
+import matplotlib
+matplotlib.use('Agg')  # Use non-interactive backend
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import base64
+from PIL import Image
+import io
+ 
+class SecureCodeExecutor:
+    def __init__(self):
+        self.sessions = {}
+        self.max_execution_time = 30
+        self.max_output_length = 10000
+        self.allowed_languages = ["python", "javascript", "bash"]
+        
+        # Security: List of blocked commands/imports
+        self.blocked_imports = [
+            'subprocess', 'os', 'sys', 'shutil', 'glob', 'pickle', 
+            'marshal', 'imp', 'importlib', '__import__'
+        ]
+        self.blocked_bash_commands = [
+            'rm', 'sudo', 'chmod', 'chown', 'dd', 'mkfs', 'fdisk',
+            'curl', 'wget', 'ssh', 'scp', 'nc', 'netcat'
+        ]
+ 
+    def create_session(self) -> str:
+        """Create a new execution session"""
+        session_id = str(uuid.uuid4())[:8]  # Shorter ID for HF
+        self.sessions[session_id] = {
+            'created_at': time.time(),
+            'variables': {},
+            'history': [],
+            'files': {}
+        }
+        return session_id
+ 
+    def cleanup_old_sessions(self):
+        """Remove sessions older than 1 hour"""
+        current_time = time.time()
+        old_sessions = [
+            sid for sid, session in self.sessions.items()
+            if current_time - session['created_at'] > 3600
+        ]
+        for sid in old_sessions:
+            del self.sessions[sid]
+ 
+    def is_code_safe(self, code: str, language: str) -> tuple[bool, str]:
+        """Check if code is safe to execute using whole-word matching."""
+        if language == "python":
+            # Use regex to find whole-word matches for blocked imports
+            for blocked in self.blocked_imports:
+                # \b matches a word boundary. This prevents 'imp' from matching 'import'
+                pattern = r'\b' + re.escape(blocked) + r'\b'
+                if re.search(pattern, code):
+                    return False, f"Blocked import/function: {blocked}"
+            
+            # Check for dangerous patterns that are not whole words
+            dangerous_patterns = ['exec(', 'eval(', 'open(', 'file(', '__']
+            for pattern in dangerous_patterns:
+                if pattern in code:
+                    return False, f"Dangerous pattern detected: {pattern}"
+        
+        elif language == "bash":
+            # Use regex to find whole-word matches for blocked commands
+            for blocked in self.blocked_bash_commands:
+                pattern = r'\b' + re.escape(blocked) + r'\b'
+                if re.search(pattern, code.lower()):
+                    return False, f"Blocked command: {blocked}"
+        
+        return True, ""
+ 
+    def execute_python_code(self, code: str, session_id: Optional[str] = None) -> Dict[str, Any]:
+        """Execute Python code with visualization support"""
+        # Security check
+        is_safe, reason = self.is_code_safe(code, "python")
+        if not is_safe:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": f"Security violation: {reason}",
+                "execution_time": time.time()
+            }
+ 
+        # Prepare execution environment
+        setup_code = '''
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import json
+import math
+import random
+import base64
+import io
+from datetime import datetime, timedelta
+ 
+# Custom print function to capture output
+_output_buffer = []
+_original_print = print
+def print(*args, **kwargs):
+    _output_buffer.append(' '.join(str(arg) for arg in args))
+ 
+# Function to save plots as base64
+def save_current_plot():
+    if plt.get_fignums():  # Check if there are any figures
+        buffer = io.BytesIO()
+        plt.savefig(buffer, format='png', bbox_inches='tight', dpi=100)
+        buffer.seek(0)
+        plot_data = buffer.getvalue()
+        buffer.close()
+        return base64.b64encode(plot_data).decode()
+    return None
+'''
+ 
+        # Combine setup and user code
+        full_code = setup_code + "\n" + code
+        
+        # Finalization code to capture plot and print the buffer to real stdout
+        finalization_code = '''
+# Capture the plot if one exists and print it directly to stdout for the parent process
+_plot_data = save_current_plot()
+if _plot_data:
+    _original_print('PLOT_DATA:' + _plot_data)
+
+# Print everything from the custom print buffer to stdout
+if _output_buffer:
+    _original_print('\\n'.join(_output_buffer))
+'''
+        full_code += "\n" + finalization_code
+ 
+        try:
+            # Create temporary file
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+                f.write(full_code)
+                temp_file = f.name
+ 
+            # Execute with timeout
+            result = subprocess.run(
+                [sys.executable, temp_file],
+                capture_output=True,
+                text=True,
+                timeout=self.max_execution_time,
+                cwd=tempfile.gettempdir()
+            )
+ 
+            # Process output
+            stdout = result.stdout
+            stderr = result.stderr
+            plot_data = None
+ 
+            # Extract plot data if present
+            if 'PLOT_DATA:' in stdout:
+                lines = stdout.split('\n')
+                clean_lines = []
+                for line in lines:
+                    if line.startswith('PLOT_DATA:'):
+                        plot_data = line.replace('PLOT_DATA:', '')
+                    else:
+                        clean_lines.append(line)
+                stdout = '\n'.join(clean_lines)
+ 
+            # Limit output length
+            if len(stdout) > self.max_output_length:
+                stdout = stdout[:self.max_output_length] + "\n... (output truncated)"
+ 
+            execution_result = {
+                "success": result.returncode == 0,
+                "stdout": stdout.strip(),
+                "stderr": stderr.strip() if stderr else "",
+                "execution_time": time.time(),
+                "return_code": result.returncode
+            }
+ 
+            if plot_data:
+                execution_result["plot"] = plot_data
+ 
+            return execution_result
+ 
+        except subprocess.TimeoutExpired:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": "Execution timed out (30s limit)",
+                "execution_time": time.time()
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": str(e),
+                "execution_time": time.time()
+            }
+        finally:
+            if 'temp_file' in locals():
+                try:
+                    os.unlink(temp_file)
+                except:
+                    pass
+ 
+    def execute_javascript_code(self, code: str, session_id: Optional[str] = None) -> Dict[str, Any]:
+        """Execute JavaScript code using Node.js"""
+        # Security check
+        is_safe, reason = self.is_code_safe(code, "javascript")
+        if not is_safe:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": f"Security violation: {reason}",
+                "execution_time": time.time()
+            }
+ 
+        try:
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.js', delete=False) as f:
+                f.write(code)
+                temp_file = f.name
+ 
+            result = subprocess.run(
+                ['node', temp_file],
+                capture_output=True,
+                text=True,
+                timeout=self.max_execution_time
+            )
+ 
+            stdout = result.stdout
+            if len(stdout) > self.max_output_length:
+                stdout = stdout[:self.max_output_length] + "\n... (output truncated)"
+ 
+            return {
+                "success": result.returncode == 0,
+                "stdout": stdout.strip(),
+                "stderr": result.stderr.strip() if result.stderr else "",
+                "execution_time": time.time(),
+                "return_code": result.returncode
+            }
+ 
+        except subprocess.TimeoutExpired:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": "Execution timed out (30s limit)",
+                "execution_time": time.time()
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": str(e),
+                "execution_time": time.time()
+            }
+        finally:
+            if 'temp_file' in locals():
+                try:
+                    os.unlink(temp_file)
+                except:
+                    pass
+ 
+    def execute_bash_command(self, command: str, session_id: Optional[str] = None) -> Dict[str, Any]:
+        """Execute bash commands with security restrictions"""
+        # Security check
+        is_safe, reason = self.is_code_safe(command, "bash")
+        if not is_safe:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": f"Security violation: {reason}",
+                "execution_time": time.time()
+            }
+ 
+        try:
+            result = subprocess.run(
+                command,
+                shell=True,
+                capture_output=True,
+                text=True,
+                timeout=self.max_execution_time,
+                cwd=tempfile.gettempdir()
+            )
+ 
+            stdout = result.stdout
+            if len(stdout) > self.max_output_length:
+                stdout = stdout[:self.max_output_length] + "\n... (output truncated)"
+ 
+            return {
+                "success": result.returncode == 0,
+                "stdout": stdout.strip(),
+                "stderr": result.stderr.strip() if result.stderr else "",
+                "execution_time": time.time(),
+                "return_code": result.returncode
+            }
+ 
+        except subprocess.TimeoutExpired:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": "Command timed out (30s limit)",
+                "execution_time": time.time()
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": str(e),
+                "execution_time": time.time()
+            }
+ 
+    def execute_code(self, code: str, language: str = "python", session_id: Optional[str] = None) -> str:
+        """Main execution function - returns JSON for MCP compatibility"""
+        # Cleanup old sessions periodically
+        if len(self.sessions) > 10:
+            self.cleanup_old_sessions()
+ 
+        if language not in self.allowed_languages:
+            return json.dumps({
+                "success": False,
+                "error": f"Language '{language}' not supported. Allowed: {', '.join(self.allowed_languages)}"
+            })
+ 
+        # Create session if needed
+        if session_id and session_id not in self.sessions:
+            session_id = self.create_session()
+        elif not session_id:
+            session_id = self.create_session()
+ 
+        # Execute based on language
+        if language == "python":
+            result = self.execute_python_code(code, session_id)
+        elif language == "javascript":
+            result = self.execute_javascript_code(code, session_id)
+        elif language == "bash":
+            result = self.execute_bash_command(code, session_id)
+        else:
+            result = {
+                "success": False,
+                "error": f"Execution handler for {language} not implemented"
+            }
+ 
+        # Store in session history
+        if session_id in self.sessions:
+            self.sessions[session_id]['history'].append({
+                'code': code,
+                'language': language,
+                'result': result,
+                'timestamp': time.time()
+            })
+ 
+        result['session_id'] = session_id
+        return json.dumps(result, indent=2)
+ 
+# Global executor instance
+executor = SecureCodeExecutor()
+ 
+# MCP Functions
+def execute_python_code(code: str, session_id: str = None) -> str:
+    """
+    Execute Python code safely with visualization support.
+    
+    Args:
+        code (str): Python code to execute
+        session_id (str, optional): Session ID for persistent context
+    
+    Returns:
+        str: JSON string with execution results
+    """
+    return executor.execute_code(code, "python", session_id)
+ 
+def execute_javascript_code(code: str, session_id: str = None) -> str:
+    """
+    Execute JavaScript code using Node.js.
+    
+    Args:
+        code (str): JavaScript code to execute
+        session_id (str, optional): Session ID for persistent context
+    
+    Returns:
+        str: JSON string with execution results
+    """
+    return executor.execute_code(code, "javascript", session_id)
+ 
+def execute_bash_command(command: str, session_id: str = None) -> str:
+    """
+    Execute bash commands with security restrictions.
+    
+    Args:
+        command (str): Bash command to execute
+        session_id (str, optional): Session ID for persistent context
+    
+    Returns:
+        str: JSON string with execution results
+    """
+    return executor.execute_code(command, "bash", session_id)
+ 
+def create_execution_session() -> str:
+    """
+    Create a new execution session for maintaining state.
+    
+    Returns:
+        str: JSON string containing new session ID
+    """
+    session_id = executor.create_session()
+    return json.dumps({"session_id": session_id, "created_at": time.time()})
+ 
+def list_execution_sessions() -> str:
+    """
+    List all active execution sessions.
+    
+    Returns:
+        str: JSON string containing session information
+    """
+    return json.dumps({
+        "sessions": list(executor.sessions.keys()),
+        "count": len(executor.sessions),
+        "timestamp": time.time()
+    })
+ 
+def get_execution_history(session_id: str) -> str:
+    """
+    Get execution history for a specific session.
+    
+    Args:
+        session_id (str): Session ID to get history for
+    
+    Returns:
+        str: JSON string containing execution history
+    """
+    if session_id not in executor.sessions:
+        return json.dumps({"error": "Session not found"})
+    
+    return json.dumps({
+        "session_id": session_id,
+        "history": executor.sessions[session_id]['history'],
+        "created_at": executor.sessions[session_id]['created_at']
+    })
+ 
+def get_system_info() -> str:
+    """
+    Get system information and available packages.
+    
+    Returns:
+        str: JSON string containing system information
+    """
+    return json.dumps({
+        "python_version": sys.version,
+        "available_packages": [
+            "numpy", "pandas", "matplotlib", "json", "math", 
+            "random", "datetime", "base64", "io"
+        ],
+        "execution_limits": {
+            "max_time": executor.max_execution_time,
+            "max_output": executor.max_output_length
+        },
+        "supported_languages": executor.allowed_languages
+    })
+ 
 # Replace the gradio_execute_code function with this improved version
 def gradio_execute_code(code: str, language: str, session_id: str = ""):
     """Gradio interface for code execution with proper image handling"""
@@ -128,4 +602,13 @@ console.log(JSON.stringify(result, null, 2));
         fn=gradio_execute_code,
         inputs=[code_input, language_dropdown, session_input],
         outputs=[output_display, plot_display]  # Now returns both text and image
+    )
+ 
+if __name__ == "__main__":
+    # Launch with MCP server enabled
+    demo.launch(
+        mcp_server=True,
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860
     )