Skip to content

chrispangg/secure-code-execution-fastapi

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

13 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Secure Code Execution FastAPI

A production-ready FastAPI service for executing Python code in isolated Jupyter kernel sessions with comprehensive security controls and resource monitoring. This project combines the best practices from HuggingFace's smolagent secure execution tutorial and sandboxed-jupyter-code-exec to create a robust, multi-layered security architecture.

A personal note: There are more robust solutions out there, such as microvm/Firecracker, but I wanted a solution that I can deploy to a container service such as Azure Container Apps without any baggage. This aims to push the boundary of what is possible in executing code in a containerized, sandboxed environment.

Note that this solution uses Azure OpenAI API for the agent in the Streamlit frontend, you will need to set up your own Azure OpenAI API key and endpoint. Otherwise, feel free to swap it out for your own OpenAI's API key and setup.

Features

Security

  • Multi-layered Protection

    • AST-based pre-execution code validation
    • Runtime kernel restrictions with custom builtins
    • Module import whitelist/blacklist enforcement
    • Filesystem sandboxing (session-isolated directories)
    • Blocked dangerous functions (eval, exec, compile, etc.)
  • Resource Monitoring

    • CPU, memory, and disk usage tracking per session
    • Configurable resource quotas and limits
    • Automatic cleanup of inactive sessions
    • Output size limitations to prevent buffer overflow
  • Security Logging

    • Comprehensive security event tracking
    • Violation detection and reporting
    • Configurable log levels and destinations

Functionality

  • Session Management

    • Per-user isolated Jupyter kernel sessions
    • Stateful code execution (variables persist across requests)
    • Automatic session cleanup after inactivity timeout
    • Session reset without losing uploaded files
  • File Operations

    • Upload files to user sessions
    • Download files from user sessions
    • List files in session directory
    • Session-isolated file access
  • Data Analysis Support

    • Pre-installed data science libraries (pandas, numpy, matplotlib, scipy, scikit-learn)
    • Excel file support (openpyxl, xlrd)
    • Visualization libraries (seaborn, plotly)
    • Efficient data formats (pyarrow)

Streamlit Frontend

  • Interactive Chat UI

    • OpenAI Agents integration for conversational analysis
    • Real-time streaming of agent responses
    • Tool execution visualization
    • Chat history persistence across sessions
  • Session Management UI

    • Start/stop sessions with visual feedback
    • File upload interface
    • Session status indicators
    • Multiple chat sessions support

Architecture

Security Layers

  1. Pre-execution Validation (app/core/code_validator.py:39-60)

    • Parses code into AST before execution
    • Checks for dangerous imports, functions, and attributes
    • Validates file operations stay within session directory
    • Estimates complexity to prevent infinite loops
  2. Runtime Restrictions (app/core/secure_kernel.py:195-255)

    • Injects restricted builtins into kernel namespace
    • Overrides open(), __import__(), and dangerous functions
    • Enforces filesystem sandboxing at runtime
    • Tracks operation counts and output size
  3. Resource Monitoring (app/core/resource_monitor.py:109-127)

    • Monitors CPU, memory, and disk usage
    • Enforces configurable quotas
    • Tracks usage history for analytics
    • Automatic process termination on quota violation
  4. Session Isolation (app/core/jupyter_controller.py:69-109)

    • Each user gets dedicated Jupyter kernel
    • Isolated working directory
    • Separate environment variables
    • Independent resource tracking

Installation

Prerequisites

  • Python 3.13+
  • uv package manager
  • Docker (for containerized deployment)

How To Run

Option 1: Docker Compose (Recommended)

  1. Build and run with Docker Compose:
docker-compose up --build
  1. Access the API:

Option 2: Local Development

  1. Install dependencies using uv:
uv sync --frozen
uv sync --frozen --group user-runtime
uv sync --frozen --group dev
  1. Install Jupyter kernel:
uv run python -m ipykernel install --user
  1. Start the FastAPI backend:
uv run fastapi dev app/main.py
  1. (Optional) Start the Streamlit frontend:
# Set up Azure OpenAI credentials in .env file
cp .env.example .env
# Edit .env with your credentials

# Run Streamlit
uv run streamlit run frontend/streamlit_app.py

Configuration

Environment Variables

Configure the service via environment variables (see docker-compose.yml:16-36):

Session Management:

  • SESSIONS_FOLDER: Path to store user sessions (default: ./jupyter_sessions)
  • SESSION_INACTIVITY_TIMEOUT: Session timeout in seconds (default: 3600)

Security:

  • ENABLE_CODE_VALIDATION: Enable AST-based validation (default: true)
  • ENABLE_RESOURCE_MONITORING: Enable resource tracking (default: true)
  • ENABLE_KERNEL_RESTRICTIONS: Enable runtime restrictions (default: true)
  • ALLOWED_MODULES: Comma-separated list of allowed imports

Resource Limits:

  • MAX_MEMORY_MB: Maximum memory per session (default: 512)
  • MAX_DISK_MB: Maximum disk space per session (default: 100)
  • MAX_OUTPUT_SIZE: Maximum output characters (default: 50000)
  • MAX_OPERATIONS: Maximum operation count (default: 10000000)

Logging:

  • ENABLE_SECURITY_LOGGING: Enable security logging (default: true)
  • SECURITY_LOG_FILE: Path to security log file

Security Configuration

Customize security policies in app/core/security_config.py:

  • Dangerous Modules (app/core/security_config.py:8-30): Blocked imports (os, subprocess, socket, etc.)
  • Allowed Modules (app/core/security_config.py:71-105): Whitelisted imports
  • Dangerous Functions (app/core/security_config.py:33-51): Blocked builtins
  • Dangerous Attributes (app/core/security_config.py:54-68): Blocked attribute access

API Reference

Session Endpoints

Start Session

POST /start_session
Content-Type: application/x-www-form-urlencoded

user_id=user123

End Session

POST /end_session
Content-Type: application/x-www-form-urlencoded

user_id=user123

Reset Session

POST /reset
Content-Type: application/x-www-form-urlencoded

user_id=user123

Execution Endpoints

Execute Code

POST /execute
Content-Type: application/json

{
  "user_id": "user123",
  "code": "import pandas as pd\ndf = pd.DataFrame({'a': [1, 2, 3]})\nprint(df)"
}

File Endpoints

Upload Files

POST /upload_files
Content-Type: multipart/form-data

user_id=user123
files=@data.csv
files=@analysis.xlsx

List Files

GET /list_files/{user_id}

Download File

GET /download_file/{user_id}/{filename}

Usage Examples

Python Client

import requests

BASE_URL = "http://localhost:8000"
USER_ID = "alice"

# Start session
response = requests.post(f"{BASE_URL}/start_session", data={"user_id": USER_ID})
print(response.json())

# Execute code
code = """
import pandas as pd
import numpy as np

# Create sample data
data = pd.DataFrame({
    'x': np.random.randn(100),
    'y': np.random.randn(100)
})

# Calculate correlation
correlation = data['x'].corr(data['y'])
print(f"Correlation: {correlation:.3f}")
"""

response = requests.post(
    f"{BASE_URL}/execute",
    json={"user_id": USER_ID, "code": code}
)
print(response.json()["output"])

# End session
response = requests.post(f"{BASE_URL}/end_session", data={"user_id": USER_ID})

Streamlit Frontend

  1. Configure environment variables in .env:
BACKEND_URL=http://localhost:8000
AZURE_OPENAI_API_KEY=your_key_here
AZURE_OPENAI_API_VERSION=2024-08-01-preview
AZURE_OPENAI_API_BASE=https://your-resource.openai.azure.com
AZURE_OPENAI_MODEL=gpt-4o-mini
  1. Run the Streamlit app:
uv run streamlit run frontend/streamlit_app.py
  1. Use the chat interface to:
    • Start/stop Jupyter sessions
    • Upload data files
    • Ask questions about data analysis
    • Execute code through natural language
    • View tool execution and results in real-time

Security Considerations

Defense in Depth

This project implements multiple security layers inspired by smolagent's approach:

  1. Input Validation: AST-based code analysis before execution
  2. Runtime Restrictions: Custom builtins and import controls
  3. Filesystem Isolation: Chroot-like session directories
  4. Resource Limits: Quotas for CPU, memory, disk, and operations
  5. Containerization: Docker isolation (optional but recommended)

Known Limitations

As noted in the smolagent documentation:

"No local Python sandbox can ever be completely secure."

While this implementation provides robust protections against common attacks and accidental damage, determined attackers may:

  • Exploit vulnerabilities in allowed packages
  • Find edge cases in the AST validator
  • Perform timing attacks or resource exhaustion within limits
  • Abuse allowed functionality for unintended purposes

Recommended Deployment Practices

  • Network Isolation: Deploy behind firewall with strict ingress rules
  • Authentication: Add proper authentication/authorization (not included)
  • Rate Limiting: Implement API rate limiting per user
  • Monitoring: Set up alerting for security violations and resource anomalies
  • Updates: Keep dependencies updated for security patches
  • Least Privilege: Run containers with minimal permissions
  • Audit Logging: Enable comprehensive logging for forensics

For Production Use

Add these components before production deployment:

  1. Authentication and authorization (JWT, OAuth, API keys)
  2. Rate limiting and DDoS protection
  3. Input sanitization and request validation
  4. HTTPS/TLS encryption
  5. Database for persistent session management
  6. Distributed tracing and monitoring
  7. Backup and disaster recovery
  8. Regular security audits and penetration testing

Testing

Run tests with pytest:

uv run pytest tests/

Tests cover:

  • Security validation (tests/test_security.py:1-100)
  • Code execution and session management
  • File operations
  • Resource monitoring
  • Error handling

Troubleshooting

Common Issues

Session startup timeout

  • Increase KERNEL_STARTUP_TIMEOUT in config.py:10
  • Check if Jupyter kernel is properly installed

Resource quota exceeded

  • Adjust MAX_MEMORY_MB and MAX_DISK_MB limits
  • Clean up old sessions manually if needed

Import errors for allowed modules

  • Ensure packages are installed in user-runtime group
  • Verify module name in ALLOWED_MODULES configuration

Docker build fails

  • Ensure uv is properly installed in container
  • Check network connectivity for package downloads

Inspiration and Credits

This project combines approaches from:

Support

For issues and questions:

  • Open an issue on GitHub
  • Check existing issues and discussions
  • Review security documentation in docs/

Disclaimer

This software is provided as-is for educational and development purposes. While comprehensive security measures are implemented, no system is completely secure. Use at your own risk and implement additional security controls for production deployments.

About

Secure Code Execution for your LLM Agent using FastAPI in a container

Resources

Security policy

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published