The source code of Make Agent Defeat Agent: Automatic Detection of Taint-Style Vulnerabilities in LLM-based Agents
@inproceedings{liu2025make,
title={Make Agent Defeat Agent: Automatic Detection of $\{$Taint-Style$\}$ Vulnerabilities in $\{$LLM-based$\}$ Agents},
author={Liu, Fengyu and Zhang, Yuan and Luo, Jiaqi and Dai, Jiarun and Chen, Tian and Yuan, Letian and Yu, Zhengmin and Shi, Youkun and Li, Ke and Zhou, Chengyuan and others},
booktitle={34th USENIX Security Symposium (USENIX Security 25)},
pages={3767--3786},
year={2025}
}
Before Static Analysis, you should install CodeQL-CLI, and Miniconda.
Then, run the following command to create a Python 3.10.12 environment:
conda create -n py31012 python=3.10.12Then execute the following command to enter the virtual environment:
conda activate py31012Execute the following command within the virtual environment
pip install -r requirements.txtgit clone your analysis target. For example:
git clone https://github.com/microsoft/TaskWeaver.gitcd into it.
cd TaskWeaverCreate a CodeQL database:
codeql database create /path/to/database/TaskWeaver --language=python --source-root=.Modify file auto_analyze.py, change DB_HOME and arg_dbname to your previously created database:
DB_HOME = '/path/to/database'
arg_dbname = "TaskWeaver"Finally, run it:
python3 auto_analyze.pyFinally, you'll get 4 files under directory output/TaskWeaver:
enter_hook.json
oracle.json
TaskWeaver-dsc.json
TaskWeaver-if.json
You should start you target application manually.
Move trace/cetracer.py , enter_hook.json, oracle.json, TaskWeaver-if.json to where the application is running.
For example:
docker cp trace/cetracer.py taskweaver:/app/playground/UI/cetracer.py
docker cp output/TaskWeaver/enter_hook.json taskweaver:/app/enter_hook.json
docker cp output/TaskWeaver/oracle.json taskweaver:/app/oracle.json
docker cp output/TaskWeaver/TaskWeaver-if.json taskweaver:/app/TaskWeaver-if.jsonNext, add the following code to the agent thread. (For multi-threaded applications, this code should be added to the thread where the agent is running.)
import cetracer
cetracer.start_ce_trace(conf="/app/if.json", enter_input_conf="/app/enter_hook.json", oracle_rule_conf = "/app/oracle.json", log="/tmp/if.log", match_log = "/tmp/hook.log", call_stack_log = "/tmp/callstack.log", oracle_name = "/tmp/oracle.log")For example, add codes to /app/playground/UIapp.py
from taskweaver.app.app import TaskWeaverApp
......
import cetracer
cetracer.start_ce_trace(conf="/app/if.json", enter_input_conf="/app/enter_hook.json", oracle_rule_conf = "/app/oracle.json", log="/tmp/if.log", match_log = "/tmp/hook.log", call_stack_log = "/tmp/callstack.log", oracle_name = "/tmp/oracle.log")
......
if __name__ == "__main__":
......Finally, restart the target application and wait for 40 seconds. (Because, to avoid introducing excessive overhead during application startup, our instrumentation remains inactive for the first 40 seconds after the target application starts.)
- SMT-solver installed (Z3)
Environment for py-conbyte
- Exit the current virtual environment.
conda deactivate
- Creating a Python 3.7.3 Virtual Environment with Miniconda:
conda create -n py373 python=3.7.3
- Enter the virtual environment and obtain the Python address:
conda activate py373
- Install pipenv:
pip install pipenv
- Enter the py conbyte directory and install the required virtual environment
cd py-conbyte pipenv shell - Install required packages for this environment.
conda activate py373 pipenv install
- Leave this virtual environment.
exit conda deactivate - Enter the virtual environment for agentfuzz.
conda activate py31012
Fill in OPENAI_API_BASE and OPENAI_API_KEY in ./config/__init__.py.
Create a script under ./poc to tell us how to send message to target agent.
For example, ./poc/TaskWeaver/poc.py:
with sync_playwright() as playwright:
browser = playwright.chromium.launch(headless=True)
context = browser.new_context()
page = context.new_page()
page.route(
"**/*",
lambda route: route.abort()
if route.request.resource_type in ["image", "stylesheet"] else route.continue_()
)
page.goto("http://localhost:8000/")
page.get_by_placeholder("Type your message here...").click()
page.get_by_placeholder("Type your message here...").fill(payload)
page.keyboard.press("Enter")
time.sleep(60)
context.close()
browser.close()It just open the browser, type in the message and click button to send to the agent.
Open ./poc/poc_factory.py and add an entry to the dict factory.
The application name and agent name can be any name you prefer. For example, "Taskweaver" and "CodeInterpreter".
The call_chain of MetaData is the target call chain in oracle.json.
call_chainofMetaDatacan be any content when runningbatchmain.py, because we will traverse all call chains. Only if you are runningmain.py, you should fill incall_chainofMetaData.
The container_name of MetaData means which docker the application is running.
For example:
"Taskweaver": MetaData(
"DEADBEEF",
"taskweaver",
"output/TaskWeaver/oracle.json",
"output/TaskWeaver/TaskWeaver-if.json",
"output/TaskWeaver/enter_hook.json",
"output/TaskWeaver/TaskWeaver-dsc.json",
poc.TaskWeaver.CodeInterpreter.poc.connect_with_auth
)Then, modify the commands and callchain file path in batchmain.py. For example:
with open('./output/TaskWeaver/oracle.json') as f:
......
subprocess.run([PYTHON_EXECUTABLE, "main.py", "-app", "Taskweaver"], env={"CALLCHAIN": c}, timeout=1200)This enables our tool to traverse all call chains, send the prompt to the target agent via poc.TaskWeaver.CodeInterpreter.poc.connect_with_auth, and detect potential vulnerabilities.
Just run:
python3 batchmain.pyAnd results will be shown in command line and ./log.
If you see the following output, it means we have triggered the callchain and it may be a vulnerability.
**********exploration successful**********
True