mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
Add secure Python Sandbox and Enhanced Logging in CodeInterpreterTool (#281)
* feat: add a safety sandbox to run Python code This sandbox blocks a bunch of dangerous imports and built-in functions * feat: add more logs and warning about code execution * test: add tests to cover sandbox code execution * docs: add Google-style docstrings and type hints to printer and code_interpreter * chore: renaming globals and locals paramenters --------- Co-authored-by: Greyson Lalonde <greyson.r.lalonde@gmail.com>
This commit is contained in:
Lucas Gomide
131
src/crewai_tools/printer.py
Normal file
131
src/crewai_tools/printer.py
Normal file
@@ -0,0 +1,131 @@
|
||||
"""Utility for colored console output."""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class Printer:
|
||||
"""Handles colored console output formatting."""
|
||||
|
||||
@staticmethod
|
||||
def print(content: str, color: Optional[str] = None) -> None:
|
||||
"""Prints content with optional color formatting.
|
||||
|
||||
Args:
|
||||
content: The string to be printed.
|
||||
color: Optional color name to format the output. If provided,
|
||||
must match one of the _print_* methods available in this class.
|
||||
If not provided or if the color is not supported, prints without
|
||||
formatting.
|
||||
"""
|
||||
if hasattr(Printer, f"_print_{color}"):
|
||||
getattr(Printer, f"_print_{color}")(content)
|
||||
else:
|
||||
print(content)
|
||||
|
||||
@staticmethod
|
||||
def _print_bold_purple(content: str) -> None:
|
||||
"""Prints content in bold purple color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in bold purple.
|
||||
"""
|
||||
print("\033[1m\033[95m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_bold_green(content: str) -> None:
|
||||
"""Prints content in bold green color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in bold green.
|
||||
"""
|
||||
print("\033[1m\033[92m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_purple(content: str) -> None:
|
||||
"""Prints content in purple color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in purple.
|
||||
"""
|
||||
print("\033[95m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_red(content: str) -> None:
|
||||
"""Prints content in red color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in red.
|
||||
"""
|
||||
print("\033[91m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_bold_blue(content: str) -> None:
|
||||
"""Prints content in bold blue color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in bold blue.
|
||||
"""
|
||||
print("\033[1m\033[94m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_yellow(content: str) -> None:
|
||||
"""Prints content in yellow color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in yellow.
|
||||
"""
|
||||
print("\033[93m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_bold_yellow(content: str) -> None:
|
||||
"""Prints content in bold yellow color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in bold yellow.
|
||||
"""
|
||||
print("\033[1m\033[93m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_cyan(content: str) -> None:
|
||||
"""Prints content in cyan color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in cyan.
|
||||
"""
|
||||
print("\033[96m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_bold_cyan(content: str) -> None:
|
||||
"""Prints content in bold cyan color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in bold cyan.
|
||||
"""
|
||||
print("\033[1m\033[96m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_magenta(content: str) -> None:
|
||||
"""Prints content in magenta color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in magenta.
|
||||
"""
|
||||
print("\033[35m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_bold_magenta(content: str) -> None:
|
||||
"""Prints content in bold magenta color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in bold magenta.
|
||||
"""
|
||||
print("\033[1m\033[35m {}\033[00m".format(content))
|
||||
|
||||
@staticmethod
|
||||
def _print_green(content: str) -> None:
|
||||
"""Prints content in green color.
|
||||
|
||||
Args:
|
||||
content: The string to be printed in green.
|
||||
"""
|
||||
print("\033[32m {}\033[00m".format(content))
|
||||
@@ -1,18 +1,31 @@
|
||||
"""Code Interpreter Tool for executing Python code in isolated environments.
|
||||
|
||||
This module provides a tool for executing Python code either in a Docker container for
|
||||
safe isolation or directly in a restricted sandbox. It includes mechanisms for blocking
|
||||
potentially unsafe operations and importing restricted modules.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
from typing import List, Optional, Type
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
|
||||
from crewai.tools import BaseTool
|
||||
from docker import from_env as docker_from_env
|
||||
from docker import DockerClient
|
||||
from docker.models.containers import Container
|
||||
from docker import from_env as docker_from_env
|
||||
from docker.errors import ImageNotFound, NotFound
|
||||
from docker.models.containers import Container
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from crewai_tools.printer import Printer
|
||||
|
||||
|
||||
class CodeInterpreterSchema(BaseModel):
|
||||
"""Input for CodeInterpreterTool."""
|
||||
"""Schema for defining inputs to the CodeInterpreterTool.
|
||||
|
||||
This schema defines the required parameters for code execution,
|
||||
including the code to run and any libraries that need to be installed.
|
||||
"""
|
||||
|
||||
code: str = Field(
|
||||
...,
|
||||
@@ -25,7 +38,102 @@ class CodeInterpreterSchema(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class SandboxPython:
|
||||
"""A restricted Python execution environment for running code safely.
|
||||
|
||||
This class provides methods to safely execute Python code by restricting access to
|
||||
potentially dangerous modules and built-in functions. It creates a sandboxed
|
||||
environment where harmful operations are blocked.
|
||||
"""
|
||||
|
||||
BLOCKED_MODULES = {
|
||||
"os",
|
||||
"sys",
|
||||
"subprocess",
|
||||
"shutil",
|
||||
"importlib",
|
||||
"inspect",
|
||||
"tempfile",
|
||||
"sysconfig",
|
||||
"builtins",
|
||||
}
|
||||
|
||||
UNSAFE_BUILTINS = {
|
||||
"exec",
|
||||
"eval",
|
||||
"open",
|
||||
"compile",
|
||||
"input",
|
||||
"globals",
|
||||
"locals",
|
||||
"vars",
|
||||
"help",
|
||||
"dir",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def restricted_import(
|
||||
name: str,
|
||||
custom_globals: Optional[Dict[str, Any]] = None,
|
||||
custom_locals: Optional[Dict[str, Any]] = None,
|
||||
fromlist: Optional[List[str]] = None,
|
||||
level: int = 0,
|
||||
) -> ModuleType:
|
||||
"""A restricted import function that blocks importing of unsafe modules.
|
||||
|
||||
Args:
|
||||
name: The name of the module to import.
|
||||
custom_globals: Global namespace to use.
|
||||
custom_locals: Local namespace to use.
|
||||
fromlist: List of items to import from the module.
|
||||
level: The level value passed to __import__.
|
||||
|
||||
Returns:
|
||||
The imported module if allowed.
|
||||
|
||||
Raises:
|
||||
ImportError: If the module is in the blocked modules list.
|
||||
"""
|
||||
if name in SandboxPython.BLOCKED_MODULES:
|
||||
raise ImportError(f"Importing '{name}' is not allowed.")
|
||||
return __import__(name, custom_globals, custom_locals, fromlist or (), level)
|
||||
|
||||
@staticmethod
|
||||
def safe_builtins() -> Dict[str, Any]:
|
||||
"""Creates a dictionary of built-in functions with unsafe ones removed.
|
||||
|
||||
Returns:
|
||||
A dictionary of safe built-in functions and objects.
|
||||
"""
|
||||
import builtins
|
||||
|
||||
safe_builtins = {
|
||||
k: v
|
||||
for k, v in builtins.__dict__.items()
|
||||
if k not in SandboxPython.UNSAFE_BUILTINS
|
||||
}
|
||||
safe_builtins["__import__"] = SandboxPython.restricted_import
|
||||
return safe_builtins
|
||||
|
||||
@staticmethod
|
||||
def exec(code: str, locals: Dict[str, Any]) -> None:
|
||||
"""Executes Python code in a restricted environment.
|
||||
|
||||
Args:
|
||||
code: The Python code to execute as a string.
|
||||
locals: A dictionary that will be used for local variable storage.
|
||||
"""
|
||||
exec(code, {"__builtins__": SandboxPython.safe_builtins()}, locals)
|
||||
|
||||
|
||||
class CodeInterpreterTool(BaseTool):
|
||||
"""A tool for executing Python code in isolated environments.
|
||||
|
||||
This tool provides functionality to run Python code either in a Docker container
|
||||
for safe isolation or directly in a restricted sandbox. It can handle installing
|
||||
Python packages and executing arbitrary Python code.
|
||||
"""
|
||||
|
||||
name: str = "Code Interpreter"
|
||||
description: str = "Interprets Python3 code strings with a final print statement."
|
||||
args_schema: Type[BaseModel] = CodeInterpreterSchema
|
||||
@@ -36,18 +144,28 @@ class CodeInterpreterTool(BaseTool):
|
||||
unsafe_mode: bool = False
|
||||
|
||||
@staticmethod
|
||||
def _get_installed_package_path():
|
||||
def _get_installed_package_path() -> str:
|
||||
"""Gets the installation path of the crewai_tools package.
|
||||
|
||||
Returns:
|
||||
The directory path where the package is installed.
|
||||
"""
|
||||
spec = importlib.util.find_spec("crewai_tools")
|
||||
return os.path.dirname(spec.origin)
|
||||
|
||||
def _verify_docker_image(self) -> None:
|
||||
"""
|
||||
Verify if the Docker image is available. Optionally use a user-provided Dockerfile.
|
||||
"""Verifies if the Docker image is available or builds it if necessary.
|
||||
|
||||
Checks if the required Docker image exists. If not, builds it using either a
|
||||
user-provided Dockerfile or the default one included with the package.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the Dockerfile cannot be found.
|
||||
"""
|
||||
|
||||
client = (
|
||||
docker_from_env()
|
||||
if self.user_docker_base_url == None
|
||||
if self.user_docker_base_url is None
|
||||
else DockerClient(base_url=self.user_docker_base_url)
|
||||
)
|
||||
|
||||
@@ -74,22 +192,41 @@ class CodeInterpreterTool(BaseTool):
|
||||
)
|
||||
|
||||
def _run(self, **kwargs) -> str:
|
||||
"""Runs the code interpreter tool with the provided arguments.
|
||||
|
||||
Args:
|
||||
**kwargs: Keyword arguments that should include 'code' and 'libraries_used'.
|
||||
|
||||
Returns:
|
||||
The output of the executed code as a string.
|
||||
"""
|
||||
code = kwargs.get("code", self.code)
|
||||
libraries_used = kwargs.get("libraries_used", [])
|
||||
|
||||
if self.unsafe_mode:
|
||||
return self.run_code_unsafe(code, libraries_used)
|
||||
else:
|
||||
return self.run_code_in_docker(code, libraries_used)
|
||||
return self.run_code_safety(code, libraries_used)
|
||||
|
||||
def _install_libraries(self, container: Container, libraries: List[str]) -> None:
|
||||
"""
|
||||
Install missing libraries in the Docker container
|
||||
"""Installs required Python libraries in the Docker container.
|
||||
|
||||
Args:
|
||||
container: The Docker container where libraries will be installed.
|
||||
libraries: A list of library names to install using pip.
|
||||
"""
|
||||
for library in libraries:
|
||||
container.exec_run(["pip", "install", library])
|
||||
|
||||
def _init_docker_container(self) -> Container:
|
||||
"""Initializes and returns a Docker container for code execution.
|
||||
|
||||
Stops and removes any existing container with the same name before creating
|
||||
a new one. Maps the current working directory to /workspace in the container.
|
||||
|
||||
Returns:
|
||||
A Docker container object ready for code execution.
|
||||
"""
|
||||
container_name = "code-interpreter"
|
||||
client = docker_from_env()
|
||||
current_path = os.getcwd()
|
||||
@@ -111,7 +248,68 @@ class CodeInterpreterTool(BaseTool):
|
||||
volumes={current_path: {"bind": "/workspace", "mode": "rw"}}, # type: ignore
|
||||
)
|
||||
|
||||
def _check_docker_available(self) -> bool:
|
||||
"""Checks if Docker is available and running on the system.
|
||||
|
||||
Attempts to run the 'docker info' command to verify Docker availability.
|
||||
Prints appropriate messages if Docker is not installed or not running.
|
||||
|
||||
Returns:
|
||||
True if Docker is available and running, False otherwise.
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker", "info"],
|
||||
check=True,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
timeout=1,
|
||||
)
|
||||
return True
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
||||
Printer.print(
|
||||
"Docker is installed but not running or inaccessible.",
|
||||
color="bold_purple",
|
||||
)
|
||||
return False
|
||||
except FileNotFoundError:
|
||||
Printer.print("Docker is not installed", color="bold_purple")
|
||||
return False
|
||||
|
||||
def run_code_safety(self, code: str, libraries_used: List[str]) -> str:
|
||||
"""Runs code in the safest available environment.
|
||||
|
||||
Attempts to run code in Docker if available, falls back to a restricted
|
||||
sandbox if Docker is not available.
|
||||
|
||||
Args:
|
||||
code: The Python code to execute as a string.
|
||||
libraries_used: A list of Python library names to install before execution.
|
||||
|
||||
Returns:
|
||||
The output of the executed code as a string.
|
||||
"""
|
||||
if self._check_docker_available():
|
||||
return self.run_code_in_docker(code, libraries_used)
|
||||
else:
|
||||
return self.run_code_in_restricted_sandbox(code)
|
||||
|
||||
def run_code_in_docker(self, code: str, libraries_used: List[str]) -> str:
|
||||
"""Runs Python code in a Docker container for safe isolation.
|
||||
|
||||
Creates a Docker container, installs the required libraries, executes the code,
|
||||
and then cleans up by stopping and removing the container.
|
||||
|
||||
Args:
|
||||
code: The Python code to execute as a string.
|
||||
libraries_used: A list of Python library names to install before execution.
|
||||
|
||||
Returns:
|
||||
The output of the executed code as a string, or an error message if execution failed.
|
||||
"""
|
||||
Printer.print("Running code in Docker environment", color="bold_blue")
|
||||
self._verify_docker_image()
|
||||
container = self._init_docker_container()
|
||||
self._install_libraries(container, libraries_used)
|
||||
@@ -125,10 +323,43 @@ class CodeInterpreterTool(BaseTool):
|
||||
return f"Something went wrong while running the code: \n{exec_result.output.decode('utf-8')}"
|
||||
return exec_result.output.decode("utf-8")
|
||||
|
||||
def run_code_in_restricted_sandbox(self, code: str) -> str:
|
||||
"""Runs Python code in a restricted sandbox environment.
|
||||
|
||||
Executes the code with restricted access to potentially dangerous modules and
|
||||
built-in functions for basic safety when Docker is not available.
|
||||
|
||||
Args:
|
||||
code: The Python code to execute as a string.
|
||||
|
||||
Returns:
|
||||
The value of the 'result' variable from the executed code,
|
||||
or an error message if execution failed.
|
||||
"""
|
||||
Printer.print("Running code in restricted sandbox", color="yellow")
|
||||
exec_locals = {}
|
||||
try:
|
||||
SandboxPython.exec(code=code, locals=exec_locals)
|
||||
return exec_locals.get("result", "No result variable found.")
|
||||
except Exception as e:
|
||||
return f"An error occurred: {str(e)}"
|
||||
|
||||
def run_code_unsafe(self, code: str, libraries_used: List[str]) -> str:
|
||||
"""Runs code directly on the host machine without any safety restrictions.
|
||||
|
||||
WARNING: This mode is unsafe and should only be used in trusted environments
|
||||
with code from trusted sources.
|
||||
|
||||
Args:
|
||||
code: The Python code to execute as a string.
|
||||
libraries_used: A list of Python library names to install before execution.
|
||||
|
||||
Returns:
|
||||
The value of the 'result' variable from the executed code,
|
||||
or an error message if execution failed.
|
||||
"""
|
||||
Run the code directly on the host machine (unsafe mode).
|
||||
"""
|
||||
|
||||
Printer.print("WARNING: Running code in unsafe mode", color="bold_magenta")
|
||||
# Install libraries on the host machine
|
||||
for library in libraries_used:
|
||||
os.system(f"pip install {library}")
|
||||
@@ -139,4 +370,4 @@ class CodeInterpreterTool(BaseTool):
|
||||
exec(code, {}, exec_locals)
|
||||
return exec_locals.get("result", "No result variable found.")
|
||||
except Exception as e:
|
||||
return f"An error occurred: {str(e)}"
|
||||
return f"An error occurred: {str(e)}"
|
||||
|
||||
@@ -1,57 +1,175 @@
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from crewai_tools.tools.code_interpreter_tool.code_interpreter_tool import (
|
||||
CodeInterpreterTool,
|
||||
SandboxPython,
|
||||
)
|
||||
|
||||
|
||||
class TestCodeInterpreterTool(unittest.TestCase):
|
||||
@patch(
|
||||
"crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env"
|
||||
@pytest.fixture
|
||||
def printer_mock():
|
||||
with patch("crewai_tools.printer.Printer.print") as mock:
|
||||
yield mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def docker_unavailable_mock():
|
||||
with patch(
|
||||
"crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.CodeInterpreterTool._check_docker_available",
|
||||
return_value=False,
|
||||
) as mock:
|
||||
yield mock
|
||||
|
||||
|
||||
@patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
|
||||
def test_run_code_in_docker(docker_mock, printer_mock):
|
||||
tool = CodeInterpreterTool()
|
||||
code = "print('Hello, World!')"
|
||||
libraries_used = ["numpy", "pandas"]
|
||||
expected_output = "Hello, World!\n"
|
||||
|
||||
docker_mock().containers.run().exec_run().exit_code = 0
|
||||
docker_mock().containers.run().exec_run().output = expected_output.encode()
|
||||
|
||||
result = tool.run_code_in_docker(code, libraries_used)
|
||||
assert result == expected_output
|
||||
printer_mock.assert_called_with(
|
||||
"Running code in Docker environment", color="bold_blue"
|
||||
)
|
||||
def test_run_code_in_docker(self, docker_mock):
|
||||
tool = CodeInterpreterTool()
|
||||
code = "print('Hello, World!')"
|
||||
libraries_used = ["numpy", "pandas"]
|
||||
expected_output = "Hello, World!\n"
|
||||
|
||||
docker_mock().containers.run().exec_run().exit_code = 0
|
||||
docker_mock().containers.run().exec_run().output = expected_output.encode()
|
||||
result = tool.run_code_in_docker(code, libraries_used)
|
||||
|
||||
self.assertEqual(result, expected_output)
|
||||
@patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
|
||||
def test_run_code_in_docker_with_error(docker_mock, printer_mock):
|
||||
tool = CodeInterpreterTool()
|
||||
code = "print(1/0)"
|
||||
libraries_used = ["numpy", "pandas"]
|
||||
expected_output = "Something went wrong while running the code: \nZeroDivisionError: division by zero\n"
|
||||
|
||||
@patch(
|
||||
"crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env"
|
||||
docker_mock().containers.run().exec_run().exit_code = 1
|
||||
docker_mock().containers.run().exec_run().output = (
|
||||
b"ZeroDivisionError: division by zero\n"
|
||||
)
|
||||
def test_run_code_in_docker_with_error(self, docker_mock):
|
||||
tool = CodeInterpreterTool()
|
||||
code = "print(1/0)"
|
||||
libraries_used = ["numpy", "pandas"]
|
||||
expected_output = "Something went wrong while running the code: \nZeroDivisionError: division by zero\n"
|
||||
|
||||
docker_mock().containers.run().exec_run().exit_code = 1
|
||||
docker_mock().containers.run().exec_run().output = (
|
||||
b"ZeroDivisionError: division by zero\n"
|
||||
)
|
||||
result = tool.run_code_in_docker(code, libraries_used)
|
||||
|
||||
self.assertEqual(result, expected_output)
|
||||
|
||||
@patch(
|
||||
"crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env"
|
||||
result = tool.run_code_in_docker(code, libraries_used)
|
||||
assert result == expected_output
|
||||
printer_mock.assert_called_with(
|
||||
"Running code in Docker environment", color="bold_blue"
|
||||
)
|
||||
def test_run_code_in_docker_with_script(self, docker_mock):
|
||||
tool = CodeInterpreterTool()
|
||||
code = """print("This is line 1")
|
||||
|
||||
|
||||
@patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
|
||||
def test_run_code_in_docker_with_script(docker_mock, printer_mock):
|
||||
tool = CodeInterpreterTool()
|
||||
code = """print("This is line 1")
|
||||
print("This is line 2")"""
|
||||
libraries_used = [] # No additional libraries needed for this test
|
||||
expected_output = "This is line 1\nThis is line 2\n"
|
||||
libraries_used = []
|
||||
expected_output = "This is line 1\nThis is line 2\n"
|
||||
|
||||
# Mock Docker responses
|
||||
docker_mock().containers.run().exec_run().exit_code = 0
|
||||
docker_mock().containers.run().exec_run().output = expected_output.encode()
|
||||
docker_mock().containers.run().exec_run().exit_code = 0
|
||||
docker_mock().containers.run().exec_run().output = expected_output.encode()
|
||||
|
||||
result = tool.run_code_in_docker(code, libraries_used)
|
||||
self.assertEqual(result, expected_output)
|
||||
result = tool.run_code_in_docker(code, libraries_used)
|
||||
assert result == expected_output
|
||||
printer_mock.assert_called_with(
|
||||
"Running code in Docker environment", color="bold_blue"
|
||||
)
|
||||
|
||||
|
||||
def test_restricted_sandbox_basic_code_execution(printer_mock, docker_unavailable_mock):
|
||||
"""Test basic code execution."""
|
||||
tool = CodeInterpreterTool()
|
||||
code = """
|
||||
result = 2 + 2
|
||||
print(result)
|
||||
"""
|
||||
result = tool.run(code=code, libraries_used=[])
|
||||
printer_mock.assert_called_with(
|
||||
"Running code in restricted sandbox", color="yellow"
|
||||
)
|
||||
assert result == 4
|
||||
|
||||
|
||||
def test_restricted_sandbox_running_with_blocked_modules(
|
||||
printer_mock, docker_unavailable_mock
|
||||
):
|
||||
"""Test that restricted modules cannot be imported."""
|
||||
tool = CodeInterpreterTool()
|
||||
restricted_modules = SandboxPython.BLOCKED_MODULES
|
||||
|
||||
for module in restricted_modules:
|
||||
code = f"""
|
||||
import {module}
|
||||
result = "Import succeeded"
|
||||
"""
|
||||
result = tool.run(code=code, libraries_used=[])
|
||||
printer_mock.assert_called_with(
|
||||
"Running code in restricted sandbox", color="yellow"
|
||||
)
|
||||
|
||||
assert f"An error occurred: Importing '{module}' is not allowed" in result
|
||||
|
||||
|
||||
def test_restricted_sandbox_running_with_blocked_builtins(
|
||||
printer_mock, docker_unavailable_mock
|
||||
):
|
||||
"""Test that restricted builtins are not available."""
|
||||
tool = CodeInterpreterTool()
|
||||
restricted_builtins = SandboxPython.UNSAFE_BUILTINS
|
||||
|
||||
for builtin in restricted_builtins:
|
||||
code = f"""
|
||||
{builtin}("test")
|
||||
result = "Builtin available"
|
||||
"""
|
||||
result = tool.run(code=code, libraries_used=[])
|
||||
printer_mock.assert_called_with(
|
||||
"Running code in restricted sandbox", color="yellow"
|
||||
)
|
||||
assert f"An error occurred: name '{builtin}' is not defined" in result
|
||||
|
||||
|
||||
def test_restricted_sandbox_running_with_no_result_variable(
|
||||
printer_mock, docker_unavailable_mock
|
||||
):
|
||||
"""Test behavior when no result variable is set."""
|
||||
tool = CodeInterpreterTool()
|
||||
code = """
|
||||
x = 10
|
||||
"""
|
||||
result = tool.run(code=code, libraries_used=[])
|
||||
printer_mock.assert_called_with(
|
||||
"Running code in restricted sandbox", color="yellow"
|
||||
)
|
||||
assert result == "No result variable found."
|
||||
|
||||
|
||||
def test_unsafe_mode_running_with_no_result_variable(
|
||||
printer_mock, docker_unavailable_mock
|
||||
):
|
||||
"""Test behavior when no result variable is set."""
|
||||
tool = CodeInterpreterTool(unsafe_mode=True)
|
||||
code = """
|
||||
x = 10
|
||||
"""
|
||||
result = tool.run(code=code, libraries_used=[])
|
||||
printer_mock.assert_called_with(
|
||||
"WARNING: Running code in unsafe mode", color="bold_magenta"
|
||||
)
|
||||
assert result == "No result variable found."
|
||||
|
||||
|
||||
def test_unsafe_mode_running_unsafe_code(printer_mock, docker_unavailable_mock):
|
||||
"""Test behavior when no result variable is set."""
|
||||
tool = CodeInterpreterTool(unsafe_mode=True)
|
||||
code = """
|
||||
import os
|
||||
os.system("ls -la")
|
||||
result = eval("5/1")
|
||||
"""
|
||||
result = tool.run(code=code, libraries_used=[])
|
||||
printer_mock.assert_called_with(
|
||||
"WARNING: Running code in unsafe mode", color="bold_magenta"
|
||||
)
|
||||
assert 5.0 == result
|
||||
|
||||
Reference in New Issue
Block a user