Update app.py

app.py

@@ -3,7 +3,7 @@
 LibreChat Custom Code Interpreter - HuggingFace Space
 A secure code execution server using Gradio + MCP with SSE support
 """
- 
+
 import json
 import subprocess
 import tempfile
@@ -13,6 +13,7 @@ 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
@@ -61,23 +62,26 @@ class SecureCodeExecutor:
             del self.sessions[sid]
  
     def is_code_safe(self, code: str, language: str) -> tuple[bool, str]:
-        """Check if code is safe to execute"""
+        """Check if code is safe to execute using whole-word matching."""
         if language == "python":
-            # Check for blocked imports
+            # Use regex to find whole-word matches for blocked imports
             for blocked in self.blocked_imports:
-                if blocked in code:
+                # \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
+            # 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":
-            # Check for blocked commands
+            # Use regex to find whole-word matches for blocked commands
             for blocked in self.blocked_bash_commands:
-                if blocked in code.lower():
+                pattern = r'\b' + re.escape(blocked) + r'\b'
+                if re.search(pattern, code.lower()):
                     return False, f"Blocked command: {blocked}"
         
         return True, ""
@@ -127,11 +131,20 @@ def save_current_plot():
 '''
  
         # Combine setup and user code
-        full_code = setup_code + "\n" + code + "\n"
+        full_code = setup_code + "\n" + code
         
-        # Add plot capture if plotting commands detected
-        if any(cmd in code for cmd in ['plt.', 'plot(', 'scatter(', 'bar(', 'hist(']):
-            full_code += "\n_plot_data = save_current_plot()\nif _plot_data: _output_buffer.append('PLOT_DATA:' + _plot_data)\n"
+        # 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