diff --git a/src/crewai_tools/printer.py b/src/crewai_tools/printer.py new file mode 100644 index 000000000..c67005ddd --- /dev/null +++ b/src/crewai_tools/printer.py @@ -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)) diff --git a/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py b/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py index 2a0f9ffe6..95559f2a7 100644 --- a/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py +++ b/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py @@ -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)}" \ No newline at end of file + return f"An error occurred: {str(e)}" diff --git a/tests/tools/test_code_interpreter_tool.py b/tests/tools/test_code_interpreter_tool.py index e281fffaf..e46c8bde4 100644 --- a/tests/tools/test_code_interpreter_tool.py +++ b/tests/tools/test_code_interpreter_tool.py @@ -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