Pull request from 'feature_main_DKRN' to 'main'

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

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

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
 

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):

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")