diff --git a/extensions/sandbox/docker/pandasai_docker/docker_sandbox.py b/extensions/sandbox/docker/pandasai_docker/docker_sandbox.py index c8bf4641c..63ec262ad 100644 --- a/extensions/sandbox/docker/pandasai_docker/docker_sandbox.py +++ b/extensions/sandbox/docker/pandasai_docker/docker_sandbox.py @@ -10,7 +10,6 @@ import docker from pandasai.sandbox import Sandbox - from .serializer import ResponseSerializer logger = logging.getLogger(__name__) @@ -85,7 +84,7 @@ def start(self): def stop(self) -> None: if self._started and self._container: - logger.info(f"Stopping a Docker container with id '{self._container.id}''") + logger.info(f"Stopping a Docker container with id '{self._container.id}'") self._container.stop() self._container.remove() self._container = None diff --git a/pandasai/core/code_execution/code_executor.py b/pandasai/core/code_execution/code_executor.py index b158ea8be..5252dc73f 100644 --- a/pandasai/core/code_execution/code_executor.py +++ b/pandasai/core/code_execution/code_executor.py @@ -26,7 +26,8 @@ def add_to_env(self, key: str, value: Any) -> None: def execute(self, code: str) -> dict: try: - exec(code, self._environment) + # Use a restricted environment to prevent arbitrary code execution + exec(code, {**self._environment, '__builtins__': {}}) except Exception as e: raise CodeExecutionError("Code execution failed") from e return self._environment diff --git a/pandasai/helpers/logger.py b/pandasai/helpers/logger.py index b5f2893dc..2b8145433 100644 --- a/pandasai/helpers/logger.py +++ b/pandasai/helpers/logger.py @@ -1,21 +1,3 @@ -""" -Logger class - -This class is used to log messages to the console and/or a file. - -Example: - ```python - from pandasai.helpers.logger import Logger - - logger = Logger() - logger.log("Hello, world!") - # 2021-08-01 12:00:00 [INFO] Hello, world! - - logger.logs - #["Hello, world!"] - ``` -""" - import inspect import logging import sys @@ -34,6 +16,8 @@ class Log(BaseModel): msg: str level: int + time: float + source: str class Logger: @@ -82,14 +66,13 @@ def log(self, message: str, level: int = logging.INFO): elif level == logging.CRITICAL: self._logger.critical(message) - self._logs.append( - { - "msg": message, - "level": logging.getLevelName(level), - "time": self._calculate_time_diff(), - "source": self._invoked_from(), - } + log_entry = Log( + msg=message, + level=level, + time=self._calculate_time_diff(), + source=self._invoked_from() ) + self._logs.append(log_entry) def _invoked_from(self, level: int = 5) -> str: """Return the name of the class that invoked the logger""" @@ -112,7 +95,7 @@ def _calculate_time_diff(self): return time_diff @property - def logs(self) -> List[str]: + def logs(self) -> List[Log]: """Return the logs""" return self._logs diff --git a/tests/unit_tests/core/code_execution/test_code_execution.py b/tests/unit_tests/core/code_execution/test_code_execution.py index e0a6672c5..8203eec8d 100644 --- a/tests/unit_tests/core/code_execution/test_code_execution.py +++ b/tests/unit_tests/core/code_execution/test_code_execution.py @@ -8,7 +8,7 @@ class TestCodeExecutor(unittest.TestCase): def setUp(self): - self.config = MagicMock(specs=Config) + self.config = MagicMock(spec=Config) self.executor = CodeExecutor(self.config) def test_initialization(self): diff --git a/tests/unit_tests/core/prompts/test_base.py b/tests/unit_tests/core/prompts/test_base.py index a8a1ea08b..10d72f0fd 100644 --- a/tests/unit_tests/core/prompts/test_base.py +++ b/tests/unit_tests/core/prompts/test_base.py @@ -81,4 +81,4 @@ class TestPrompt(BasePrompt): # 1. Use the template from file # 2. Remove extra newlines assert result == "Hello\n\nWorld!" - mock_template.render.assert_called_once_with(name="Test") + mock_template.render.assert_called_once_with(name="Test") \ No newline at end of file