Agent-Scan

English | 中文

Command usage: USAGE.md · 中文

A dynamic black-box security testing tool for running AI Agent platforms, developed by Tencent Zhuque Lab.

✨ Features

• 🤖 Dynamic Black-Box Testing: Connects to running AI Agent platforms via API, sends crafted test prompts, and analyzes responses
• 🎯 Multi-Platform Support: Supports Dify, Coze, OpenAI-compatible apps, Anthropic, Google Gemini, HTTP endpoints, WebSocket, and more
• 🔍 Three-Stage Pipeline: Info Collection → Parallel Vulnerability Detection → Vulnerability Review
• 🛡️ 10 Detection Skills: Authorization bypass, data leakage, indirect prompt injection, tool abuse, web exfiltration, agentic supply chain, unexpected code execution, inter-agent communication security, cascading failure, human-agent trust exploit
• 📊 Structured JSON Report: CLI mode outputs a comprehensive JSON security report with OWASP ASI mapping
• 🔌 AIG Platform Integration: Interactive scanning via the Web UI with real-time pipeline results
• 🔐 Built-in Injection Defense: The scanning agent has system-level indirect prompt injection defense
• 📝 Detailed Logging: Full execution process logged with loguru
• 🐛 Debug Mode: Integrated Laminar tracing for easy debugging

🚀 Quick Start

Install from PyPI

pip install aig-agent-scan
aig-agent-scan --agent_provider providers.yaml -k "$LLM_API_KEY" -m deepseek-v3.2-exp

The flags below are the same when the command is aig-agent-scan instead of python main.py.

1. Clone the Project

git clone <your-repo>
cd agent-scan

2. Install Dependencies

pip install -r requirements.txt

Or install as a package:

pip install -e .

3. Configure Environment

Copy the example env file and fill in your API key:

cp env.example .env

Edit .env:

# Required
OPENROUTER_API_KEY=your-api-key-here
DEFAULT_MODEL=deepseek/deepseek-v3.2-exp
DEFAULT_BASE_URL=https://openrouter.ai/api/v1

4. Run a Scan

CLI Mode (Standalone)

# Basic scan with a provider YAML
python main.py --agent_provider providers.yaml -k sk-xxx -m deepseek-v3.2-exp

# Save result to file
python main.py --agent_provider providers.yaml -k sk-xxx -m deepseek-v3.2-exp -o result.json

# Specify language
python main.py --agent_provider providers.yaml -k sk-xxx -m deepseek-v3.2-exp --language en

# Run specific detection skills only
python main.py --agent_provider providers.yaml -k sk-xxx -m deepseek-v3.2-exp --skills "data-leakage-detection,tool-abuse-detection"

AIG Platform Integration

For interactive scanning via the AIG Web UI:

1. Start the AIG backend and Agent process:

# Terminal 1: Start the Go web server
go build -o ai-infra-guard ./cmd/cli/main.go
./ai-infra-guard webserver --server 127.0.0.1:8088

# Terminal 2: Start the Agent worker
go build -o agent ./cmd/agent
AIG_SERVER=127.0.0.1:8088 ./agent

2. Open the AIG Web UI in your browser, select "Agent Scan", configure the target agent's provider settings, and click "Start Scan".

The Go backend automatically invokes agent-scan with --aig-mode, so you do not need to pass this flag manually.

📖 Parameter Reference

Parameter	Required	Description
--agent_provider	Yes	Path to the provider YAML file describing the target agent's API
-k, --api_key	No	LLM API key (falls back to LLM_API_KEY / OPENAI_API_KEY / OPENROUTER_API_KEY env var)
-m, --model	No	LLM model name (default: deepseek/deepseek-v3.2-exp)
-u, --base_url	No	API base URL (default: https://openrouter.ai/api/v1)
--repo	No	Project directory path for optional source code audit
-p, --prompt	No	Custom scan prompt
--language	No	Output language: zh (default) or en
--skills	No	Comma-separated skill names to run (default: core 5 skills — data-leakage, tool-abuse, indirect-injection, authorization-bypass, web-exfiltration); 10 skills available total
--debug	No	Enable debug mode
--aig-mode	No	Enable AIG integration mode (structured JSON logs for Go backend)
-o, --output	No	Save result as JSON file

🔧 Provider Configuration

The --agent_provider parameter points to a YAML file that describes how to connect to the target agent. Example:

# HTTP endpoint provider
targets:
  - id: "http"
    config:
      url: "http://127.0.0.1:18081"
      endpoint: "/chat"
      method: "POST"
      headers:
        Content-Type: "application/json"
      body:
        message: "{{prompt}}"
      transform_response: "reply"

Supported provider types: http, dify, coze, openai, anthropic, google, cohere, huggingface, replicate, ollama, localai, litellm, openrouter, orcarouter, api_route, atlascloud, cheaperinference, runcrate, websocket, and more.

API Route

API Route is an optional OpenAI-compatible provider using https://global.api-route.com/v1. Set your own key in the environment:

export API_ROUTE_API_KEY="your-api-route-key"

To use it as the target provider, save the following in a separate target YAML file and pass that file to --agent_provider:

targets:
  - id: "api_route:gpt-6.1-sol"

The bare provider ID api_route uses the default model gpt-6.1-sol. To choose another model, use api_route:<exact-model-id> without adding an openai/ or anthropic/ prefix. Verify the model ID with an authenticated GET /v1/models request first: availability depends on your key's group and permissions. The preset lists only the default model, not the full gateway catalog.

This configures the target being scanned. The scanning LLM's -k, -m, and -u options remain separate and unchanged. API Route requests use Authorization: Bearer <API_ROUTE_API_KEY> through the existing adapter; no extra SDK is required.

Atlas Cloud

Atlas Cloud is an optional OpenAI-compatible provider using https://api.atlascloud.ai/v1. Set your own key in the environment:

export ATLASCLOUD_API_KEY="your-atlas-cloud-key"

To use it as the target provider, save the following in a separate target YAML file and pass that file to --agent_provider:

targets:
  - id: "atlascloud:deepseek-ai/deepseek-v4-flash"

The bare provider ID atlascloud uses the default model deepseek-ai/deepseek-v4-flash. To choose another model, use atlascloud:<exact-model-id> with the vendor prefix as listed by GET https://api.atlascloud.ai/v1/models (for example moonshotai/kimi-k3). That listing is public, so it does not prove a key is valid. The preset lists only a few models, not the full gateway catalog.

This configures the target being scanned. The scanning LLM's -k, -m, and -u options remain separate and unchanged. Atlas Cloud requests use Authorization: Bearer <ATLASCLOUD_API_KEY> through the existing adapter; no extra SDK is required.

Cheaper Inference

Cheaper Inference is an optional OpenAI-compatible LLM gateway using https://api.cheaperinference.com/v1. Set your own key in the environment:

export CHEAPER_INFERENCE_API_KEY="your-cheaper-inference-key"

To use it as the target provider, save the following in a separate target YAML file and pass that file to --agent_provider:

targets:
  - id: "cheaperinference:gpt-5.4-mini"

The bare provider ID cheaperinference uses the default model gpt-5.4-mini. To choose another model, use cheaperinference:<exact-model-id> as listed by GET https://api.cheaperinference.com/v1/models (for example claude-sonnet-5). The preset lists only a few models, not the full gateway catalog.

This configures the target being scanned. The scanning LLM's -k, -m, and -u options remain separate and unchanged. Cheaper Inference requests use Authorization: Bearer <CHEAPER_INFERENCE_API_KEY> through the existing adapter; no extra SDK is required.

Runcrate

Runcrate is an optional OpenAI-compatible provider using https://api.runcrate.ai/v1. Open models are free with your GitHub account. Set your own key in the environment:

export RUNCRATE_API_KEY="your-runcrate-key"

To use it as the target provider, save the following in a separate target YAML file and pass that file to --agent_provider:

targets:
  - id: "runcrate:deepseek-ai/DeepSeek-V4.1-Flash"

The bare provider ID runcrate uses the default model deepseek-ai/DeepSeek-V4.1-Flash. To choose another model, use runcrate:<exact-model-id> with the vendor prefix as listed by GET https://api.runcrate.ai/v1/models (for example openai/gpt-oss-120b). That listing is public, so it does not prove a key is valid. The preset lists only a few models, not the full catalog.

This configures the target being scanned. The scanning LLM's -k, -m, and -u options remain separate and unchanged. Runcrate requests use Authorization: Bearer <RUNCRATE_API_KEY> through the existing adapter; no extra SDK is required.

⚠️ transform_response — Critical for Custom HTTP Providers

The transform_response field tells agent-scan which JSON key to extract the agent's reply from. This is critical and must not be left empty for custom HTTP providers:

Provider Type	Response Format	transform_response Value
openai / dify / coze etc.	Standard API format	Leave empty (auto-detected)
http (custom endpoint)	{"reply": "..."}	"reply"
http (custom endpoint)	{"response": "..."}	"response"
http (custom endpoint)	{"data": {"text": "..."}}	"data.text"

If left empty for a custom HTTP provider, all responses will be parsed as None and the scan will produce zero results.

When using the AIG Web UI, this field corresponds to the "Response Parser" (响应解析器) input in the Agent configuration form — make sure to fill it with the correct key (e.g., reply).

🧪 Test Cases

The testcase/ directory contains ready-to-use vulnerable target agents for testing agent-scan's detection capabilities.

Case 2: Memory Heist Agent

A target agent vulnerable to the "Memory Heist" attack chain. It has a web_fetch tool (no URL allowlist), user PII embedded in the system prompt, and follows instructions from fetched web pages without any indirect prompt injection defense.

Requires: DEEPSEEK_API_KEY environment variable (or edit the default key in the script).

# 1. Start the target agent (port 18081)
cd testcase/case2
python memory_heist_agent.py

# 2. Run agent-scan against it (CLI mode)
cd ../..
python main.py \
    --agent_provider testcase/case2/provider.yaml \
    -k <your_llm_api_key> \
    -m <model_name> \
    -u <api_base_url>

# 3. Or via AIG Web UI
#    - Start AIG backend + Agent worker (see "AIG Platform Integration" above)
#    - In Web UI, create Agent Scan with:
#      URL: http://127.0.0.1:18081  Endpoint: /chat  Method: POST
#      Body: {"message": "{{prompt}}"}  Response Parser: reply

Expected detections: Tool abuse (SSRF via web_fetch), indirect prompt injection, web data exfiltration.

Case 3: Vulnerable Customer-Service Agent

A deliberately vulnerable HTTP agent simulating an AI customer-service bot with 8 security weaknesses across multiple OWASP ASI categories. See case3/README.md for full details.

# 1. Start the target agent (port 18091)
cd testcase/case3
python main.py

# 2. Run agent-scan against it (CLI mode)
cd ../..
python main.py \
    --agent_provider testcase/case3/provider.yaml \
    -k <your_llm_api_key> \
    -m <model_name> \
    -u <api_base_url>

Expected detections: System prompt leakage, API key leakage, database credential leak, authorization bypass, indirect prompt injection, tool abuse, web exfiltration.

Via AIG Web UI

When using test cases with the AIG Web UI, note the following required fields:

Field	Case 2	Case 3
Agent URL	http://127.0.0.1:18081	http://127.0.0.1:18091
Endpoint	/chat	/chat
Method	POST	POST
Request Body	{"message": "{{prompt}}"}	{"message": "{{prompt}}"}
Response Parser	reply (required)	reply (required)

⚠️ The "Response Parser" field must be set to reply for both test cases. Leaving it empty will cause all responses to be parsed as None, resulting in zero detections.

📊 Output Format

CLI Mode

Outputs a structured JSON report containing:

• Agent metadata (name, type, model)
• Detection results with severity levels
• OWASP ASI vulnerability mapping
• Scan statistics (duration, dialogue count)
• Source language analysis

AIG Mode

Emits structured JSON log lines to stdout, consumed by the Go backend via ParseStdoutLine():

• newPlanStep — pipeline stage started
• statusUpdate — stage progress
• toolUsed — tool invocation tracking
• actionLog — tool action details
• resultUpdate — final scan result
• error — error messages

🏗️ Project Structure

agent-scan/
├── agent_scan/              # Python package
│   ├── __init__.py
│   ├── __main__.py          # python -m agent_scan entry
│   ├── main.py              # CLI entry point (dual-mode)
│   ├── core/                # Core scanning pipeline
│   │   ├── agent.py         # Top-level scan orchestrator
│   │   ├── base_agent.py    # ReAct agent loop
│   │   ├── agent_adapter/   # Provider adapters
│   │   └── report/          # Report generation
│   ├── tools/               # Agent tools (dialogue, scan, thinking, etc.)
│   ├── utils/               # Utilities (LLM, config, logging, etc.)
│   ├── prompt/              # Prompt templates and skill definitions
│   │   ├── system/          # System prompts
│   │   └── skills/          # Detection skill definitions
│   └── config/              # Provider configuration schemas
├── main.py                  # Dev-mode thin shell (for uv run)
├── pyproject.toml           # Package configuration
├── requirements.txt         # Dependencies
├── env.example              # Environment variable template
└── testcase/                # Test cases

🔍 How It Works

Stage 1: Info Collection
  └─ Recon agent explores the target agent's capabilities and endpoints

Stage 2: Parallel Vulnerability Detection (10 skills available, 5 run by default)
  ├─ Skill Worker A: data-leakage-detection             [default]
  ├─ Skill Worker B: tool-abuse-detection              [default]
  ├─ Skill Worker C: indirect-injection-detection      [default]
  ├─ Skill Worker D: authorization-bypass-detection    [default]
  ├─ Skill Worker E: web-exfiltration-detection        [default]
  ├─ Skill Worker F: agentic-supply-chain-detection        (opt-in via --skills)
  ├─ Skill Worker G: unexpected-code-execution-detection  (opt-in via --skills)
  ├─ Skill Worker H: inter-agent-comm-security-detection (opt-in via --skills)
  ├─ Skill Worker I: cascading-failure-detection          (opt-in via --skills)
  └─ Skill Worker J: human-agent-trust-exploit-detection (opt-in via --skills)
  (max 4 concurrent dialogue calls)

Stage 3: Vulnerability Review
  └─ Reviewer consolidates findings, maps to OWASP ASI, assigns severity

🛠️ Development

# Install in editable mode
pip install -e .

# Run tests
pytest pytests/

# Lint
ruff check agent_scan/

📄 License

Apache License 2.0 — see LICENSE for details.

🙏 Attribution

Developed by Tencent Zhuque Lab. Any integration or derivative work must explicitly attribute Tencent Zhuque Lab in its documentation and user interface.