fix: bypass the security check with prompt injection (#399) (#409)

pandasai/__init__.py

@@ -586,6 +586,24 @@ def _is_df_overwrite(self, node: ast.stmt) -> bool:
             and re.match(r"df\d{0,2}$", node.targets[0].id)
         )
 
+    def _is_jailbreak(self, node: ast.stmt) -> bool:
+        """
+        Remove jailbreaks from the code to prevent malicious code execution.
+
+        Args:
+            node (object): ast.stmt
+
+        Returns (bool):
+        """
+
+        DANGEROUS_BUILTINS = ["__subclasses__", "__builtins__", "__import__"]
+
+        for child in ast.walk(node):
+            if isinstance(child, ast.Name) and child.id in DANGEROUS_BUILTINS:
+                return True
+
+        return False
+
     def _clean_code(self, code: str) -> str:
         """
         A method to clean the code to prevent malicious code execution
@@ -608,7 +626,7 @@ def _clean_code(self, code: str) -> str:
             if isinstance(node, (ast.Import, ast.ImportFrom)):
                 self._check_imports(node)
                 continue
-            if self._is_df_overwrite(node):
+            if self._is_df_overwrite(node) or self._is_jailbreak(node):
                 continue
             new_body.append(node)
 

tests/test_pandasai.py

@@ -327,6 +327,15 @@ def test_clean_code_remove_builtins(self, pandasai):
         assert pandasai.run_code(builtins_code, pd.DataFrame()) == {1, 2, 3}
         assert pandasai.last_code_executed == "print(set([1, 2, 3]))"
 
+    def test_clean_code_removes_jailbreak_code(self, pandasai):
+        malicious_code = """
+__builtins__['str'].__class__.__mro__[-1].__subclasses__()[140].__init__.__globals__['system']('ls')
+print(df)
+"""
+        pandasai._llm._output = malicious_code
+        pandasai.run_code(malicious_code, pd.DataFrame())
+        assert pandasai.last_code_executed == "print(df)"
+
     def test_clean_code_remove_environment_defaults(self, pandasai):
         pandas_code = """
 import pandas as pd