main.py

#!/usr/bin/env python3
"""
Dynamic and Adaptive Python Environment in a Bubblewrap Sandbox

Overview:
    This project demonstrates a minimal, safe, and self-contained Python execution
    environment using bubblewrap (bwrap) and uv. The goal is to provide a lightweight
    alternative to Docker for running agent code—in this case, dynamically generated
    Python code along with its dependencies—within an isolated sandbox.

Intended Future Iteration:
    In future iterations, an LLM will determine which packages are required for a
    given task (e.g., installing torch or other libraries as needed) so that the
    environment grows only as necessary. This dynamic, agent-driven approach will
    allow the sandbox to evolve over time, installing new dependencies on demand,
    while remaining secure and isolated.

Key Features:
    1. **Dynamic Code Generation:** The script generates Python code and a corresponding
       requirements.txt file on the fly.
    2. **Bubblewrap Isolation:** Bubblewrap leverages Linux namespaces to create an isolated,
       container-like sandbox that restricts host access. Essential directories (e.g., /usr,
       /bin, /lib) are mounted read-only, while /etc/resolv.conf is bind-mounted to enable DNS
       resolution and /tmp is mounted as a tmpfs.
    3. **Persistent, Adaptive Environment:** A persistent environment directory is created
       based on a hash of the required packages. Dependencies are installed into a virtual
       environment using uv. In the future, an LLM could adjust these requirements dynamically,
       ensuring that only necessary packages are installed.
    4. **Self-Contained Execution:** The generated Python code is executed within the sandbox,
       providing Docker-like isolation without the need for a long-running daemon.

Usage:
    1. Ensure that:
       - 'bubblewrap' is installed (e.g., 'sudo apt-get install bubblewrap').
       - 'uv' is installed and available in one of the bind-mounted directories.
       - 'python' is accessible via the bind mounts.
    2. Run this script: `python bubblewrap_uv_example.py`.
    3. The script will generate a sample (e.g., matrix addition using numpy and pandas),
       create or update the persistent environment, and execute the code inside the sandbox.

Security Note:
    This approach minimizes the file system exposure of the host system. For even tighter
    security, additional namespaces (such as the network namespace via '--unshare-net') can be
    disabled as needed.
"""


import os
import sys
import hashlib
import subprocess
import textwrap
from pathlib import Path
from loguru import logger

# Constant bubblewrap options that remain the same for every sandbox invocation.
BWRAP_BASE_OPTS = [
    "bwrap",
    "--unshare-user",
    "--unshare-pid",
    "--unshare-uts",
    "--unshare-ipc",
    "--proc", "/proc",
    "--dev", "/dev",
    "--ro-bind", "/etc/ssl", "/etc/ssl",
    "--ro-bind", "/etc/ca-certificates", "/etc/ca-certificates",
    "--ro-bind", "/etc/resolv.conf", "/etc/resolv.conf",  # Allow DNS resolution
    "--ro-bind", "/usr", "/usr",
    "--ro-bind", "/bin", "/bin",
    "--ro-bind", "/lib", "/lib",
    "--ro-bind", "/lib64", "/lib64",
    "--tmpfs", "/tmp"  # Provide writable /tmp for temporary files
]

def create_sandbox_command(env_dir: str, command: list[str]) -> list[str]:
    """
    Build the complete bubblewrap command by combining the base options with environment-specific
    bind-mounts and the desired command.
    """
    # Bind the persistent environment directory and change into it.
    env_opts = [
        "--bind", env_dir, env_dir,
        "--chdir", env_dir
    ]
    return BWRAP_BASE_OPTS + env_opts + command

def get_or_create_persistent_env(base_dir: str, requirements: list[str], force: bool = True) -> str:
    """
    Create or reuse a persistent environment based on the hash of the requirements.
    Writes a requirements.txt file, creates a new virtual environment with uv,
    activates it, and installs dependencies.
    """
    # Generate a hash based on sorted requirements for reproducibility.
    env_hash = hashlib.md5(str(sorted(requirements)).encode()).hexdigest()
    env_dir = os.path.join(base_dir, f"env_{env_hash}")
    
    if not os.path.exists(env_dir) or force:
        logger.info("[Agent] {} persistent environment at: {}",
                    "Recreating" if force and os.path.exists(env_dir) else "Creating",
                    env_dir)
        
        if os.path.exists(env_dir) and force:
            import shutil
            shutil.rmtree(env_dir)
        os.makedirs(env_dir)
        
        # Write the requirements.txt file.
        req_file = os.path.join(env_dir, "requirements.txt")
        with open(req_file, "w") as f:
            f.write("\n".join(requirements) + "\n")
        
        # Construct the shell command to set up the virtual environment.
        # It unsets any stale VIRTUAL_ENV, creates the venv using uv, activates it, and installs dependencies.
        setup_shell_command = (
            "unset VIRTUAL_ENV; "
            "uv venv venv && "
            ". venv/bin/activate && "
            "uv pip install -r requirements.txt --verbose"
        )
        setup_cmd = create_sandbox_command(env_dir, ["/bin/sh", "-c", setup_shell_command])
        
        logger.info("[Sandbox] Setting up environment:\n{}", " ".join(setup_cmd))
        try:
            result = subprocess.run(setup_cmd, capture_output=True, text=True, check=True)
            logger.info("[Sandbox] Setup output:\n{}", result.stdout)
        except subprocess.CalledProcessError as e:
            logger.error("Failed to set up environment: {}", e)
            logger.error("stdout:\n{}", e.stdout)
            logger.error("stderr:\n{}", e.stderr)
            raise
    
    return env_dir

def run_in_sandbox(env_dir: str, script_path: str) -> str:
    """
    Run the Python script inside the sandbox using the persistent virtual environment.
    """
    # Determine the Python interpreter from the created virtual environment.
    venv_python = os.path.join(env_dir, "venv/bin/python")
    run_cmd = create_sandbox_command(env_dir, [venv_python, script_path])
    
    logger.info("[Sandbox] Running script:\n{}", " ".join(run_cmd))
    try:
        result = subprocess.run(run_cmd, capture_output=True, text=True, check=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        logger.error("Script failed with code: {}", e.returncode)
        logger.error("stdout:\n{}", e.stdout)
        logger.error("stderr:\n{}", e.stderr)
        raise

def main():
    base_dir = os.path.expanduser("~/.bubblewrap_envs")
    os.makedirs(base_dir, exist_ok=True)
    
    # Example: Matrix addition using numpy and pandas.
    packages = ["numpy", "pandas"]
    code = textwrap.dedent("""\
    import numpy as np
    import pandas as pd

    def main():
        A = np.array([[1, 2], [3, 4]])
        B = np.array([[5, 6], [7, 8]])
        result = A + B
        df = pd.DataFrame(result, columns=['C1', 'C2'])
        print("Matrix Addition Result:")
        print(df)

    if __name__ == "__main__":
        main()
    """)
    
    # Create (or recreate) the persistent environment.
    env_dir = get_or_create_persistent_env(base_dir, packages, force=True)
    
    # Write the generated Python code to a script file in the environment.
    script_path = os.path.join(env_dir, "tool.py")
    with open(script_path, "w") as f:
        f.write(code)
    
    # Run the script inside the sandbox.
    try:
        output = run_in_sandbox(env_dir, script_path)
        logger.info("==== Execution Output ====\n{}\n========================", output)
    except subprocess.CalledProcessError:
        logger.error("Script execution failed")
        sys.exit(1)

if __name__ == "__main__":
    main()