From 3aac79be8fc1d18b53d66a566adddbbdd2b38ad5 Mon Sep 17 00:00:00 2001
From: Gabriele Venturi <lele.venturi@gmail.com>
Date: Sat, 29 Jul 2023 01:29:40 +0200
Subject: [PATCH] fix: bypass the security check with prompt injection (#399)
 (#409)

---
 pandasai/__init__.py   | 20 +++++++++++++++++++-
 tests/test_pandasai.py |  9 +++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/pandasai/__init__.py b/pandasai/__init__.py
index bb325fb6a..a7748e6ac 100644
--- a/pandasai/__init__.py
+++ b/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)
 
diff --git a/tests/test_pandasai.py b/tests/test_pandasai.py
index 8c1f09cba..c4fa2bba0 100644
--- a/tests/test_pandasai.py
+++ b/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