Issue with custom_whitelisted_dependencies Not Working as Expected in PandasAI

[BennisonDevadoss]

System Info

OS version: Ubuntu 22.04.4 LTS
Python version: 3.11
PandasAI version: 2.4.2

🐛 Describe the bug

I encountered an issue while using the custom_whitelisted_dependencies parameter in PandasAI. Despite including certain libraries in the whitelist, the MaliciousQueryError is raised when those libraries are used.

Here’s the relevant portion of my code:

from pandasai import Agent
from pandasai.llm import LocalLLM

llm = LocalLLM(model="deepseek-coder:33b", api_base="http://localhost:11434/v1/")

config = {
    "llm": llm,
    "enable_cache": False,
    "custom_whitelisted_dependencies": [
        "os",
        "io",
        "sys",
        "chr",
        "glob",
        "b64decoder",
        "collections",
    ],
}

agent = Agent(
    dfs=[df], config=config
)

When I run the agent, I receive the following error:

Traceback (most recent call last):
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/generate_chat_pipeline.py", line 335, in run
    ).run(input)
      ^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/pipeline.py", line 137, in run
    raise e
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/pipeline.py", line 101, in run
    step_output = logic.execute(
                  ^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/code_execution.py", line 125, in execute
    code_to_run = self._retry_run_code(
                  ^^^^^^^^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/code_execution.py", line 349, in _retry_run_code
    return self.on_retry(code, e)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/generate_chat_pipeline.py", line 162, in on_code_retry
    return self.code_exec_error_pipeline.run(correction_input)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/error_correction_pipeline/error_correction_pipeline.py", line 48, in run
    return self.pipeline.run(input)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/pipeline.py", line 137, in run
    raise e
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/pipeline.py", line 101, in run
    step_output = logic.execute(
                  ^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/code_cleaning.py", line 95, in execute
    code_to_run = self.get_code_to_run(input, code_context)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/code_cleaning.py", line 127, in get_code_to_run
    ] and self._is_malicious_code(code):
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/bennison/miniconda3/envs/hdfc-ai/lib/python3.11/site-packages/pandasai/pipelines/chat/code_cleaning.py", line 212, in _is_malicious_code
    raise MaliciousQueryError(
pandasai.exceptions.MaliciousQueryError: Restricted library import detected: io

Steps to Reproduce

1. Configure a PandasAI Agent with the custom_whitelisted_dependencies parameter.
2. Include libraries such as os, io, etc., in the whitelist.
3. Attempt to execute a query that uses one of the whitelisted libraries, such as io.

Expected Behavior

The libraries listed in custom_whitelisted_dependencies should be allowed during code execution without triggering the MaliciousQueryError.

Actual Behavior

The MaliciousQueryError is raised for a library (e.g., io) even though it is included in custom_whitelisted_dependencies.

Thank you for your attention to this issue! Looking forward to any guidance or fixes.

[dosubot]

The issue you're encountering with the MaliciousQueryError being raised despite including libraries like os and io in the custom_whitelisted_dependencies is due to the way PandasAI handles restricted libraries. Certain libraries, such as os, io, sys, and others, are considered inherently risky and are explicitly restricted to prevent potential security vulnerabilities. These libraries are not allowed to be whitelisted through the custom_whitelisted_dependencies parameter ^[1].

The MaliciousQueryError is triggered when the code attempts to import these restricted libraries, regardless of their presence in the custom whitelist. This is a security measure to prevent unauthorized or unsafe code execution ^[1].

To resolve this, you should avoid including inherently risky libraries like os, io, and sys in the custom_whitelisted_dependencies. Instead, consider alternative approaches that do not require these libraries, or ensure that any necessary operations are performed in a secure and controlled environment outside of PandasAI.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Other}

[gdcsinaptik]

@BennisonDevadoss please try to set the "security" param to "none" in the config.
This will disable the keywords-based checks against prompt injection.
Let me know if it works: https://docs.pandas-ai.com/library#config