Skip to content

CODEBUDDY.md

Latest commit

 

History

History
187 lines (138 loc) · 9.88 KB

CODEBUDDY.md

File metadata and controls

187 lines (138 loc) · 9.88 KB

CODEBUDDY.md This file provides guidance to CodeBuddy when working with code in this repository.

Project Overview

AI-Infra-Guard (A.I.G) is an AI Red Teaming Platform by Tencent Zhuque Lab. It is a Go + Python hybrid application with a Server-Agent distributed architecture. The Go backend handles the web server, task scheduling, and WebSocket communication, while three Python sub-projects provide specialized security scanning capabilities. Licensed under MIT.

Common Commands

Build & Run (Docker — primary method)

# Pull pre-built images and start
docker-compose -f docker-compose.images.yml up -d

# Build from source and start
docker-compose up -d

# Access web UI at http://localhost:8088
# Access Swagger docs at http://localhost:8088/docs/index.html

Build Go Binary Locally

# Build the CLI / web server binary
CGO_ENABLED=0 go build -ldflags="-s -w" -trimpath -buildvcs=false -o ai-infra-guard ./cmd/cli/main.go

# Build the agent binary
CGO_ENABLED=0 go build -ldflags="-s -w" -trimpath -buildvcs=false -o agent ./cmd/agent

Run Go Tests

# Run all Go tests
go test ./...

# Run tests for a specific package
go test ./common/runner/...
go test ./common/fingerprints/parser/...
go test ./pkg/vulstruct/...
go test ./internal/mcp/...

CLI Usage

# Run infrastructure scan from CLI
./ai-infra-guard scan -t <target_url>

# Start the web server
./ai-infra-guard webserver

Python Sub-projects

# AIG-PromptSecurity (jailbreak evaluation) — uses uv package manager
cd AIG-PromptSecurity
uv sync                  # install dependencies
uv run cli_run.py --base_url <url> --api_key <key> --model <model> --scenarios <...> --techniques <...>

# agent-scan
cd agent-scan
pip install -r requirements.txt
python main.py -m <model> -k <api_key> -u <base_url> --agent_provider dify --language en

# mcp-scan
cd mcp-scan
pip install -r requirements.txt
python main.py --model <model> --base_url <url> --api_key <key> --repo <folder>

YAML Rule Validation (CI)

go run cmd/yamlcheck/main.go

Architecture

Distributed Server-Agent Architecture

The system deploys as two Docker containers communicating over WebSocket:

  • webserver (cmd/cli/main.go → webserver subcommand): A Gin-based HTTP server on port 8088 that serves the SPA frontend (embedded via embed.FS), REST API (/api/v1/), SSE task progress, and manages the AgentManager/TaskManager.
  • agent (cmd/agent/main.go): A worker process that connects to the webserver via WebSocket (/api/v1/agents/ws), registers 4 task executors, and waits for dispatched tasks. The agent connects using AIG_SERVER environment variable.

Task lifecycle: Frontend → REST API → TaskManager → AgentManager dispatches via WebSocket → Agent executes → streams results back via WebSocket → SSE pushes to frontend.

Go Backend (common/, internal/, pkg/)

common/websocket/ — Core web server layer:

  • server.go: Gin router setup, route registration, static file serving, CORS
  • api.go: REST API handlers for tasks, knowledge base CRUD, file upload (chunked), models
  • task_manager.go (~1500 LOC): Central task orchestrator — creates tasks, manages SSE channels, persists state to SQLite, assigns tasks to agents
  • agent.go: AgentManager handles WebSocket connections, agent registration, task assignment, heartbeat

common/agent/ — Agent-side task execution framework:

  • agent.go: WebSocket client, task executor registration, message dispatch loop
  • types.go: Shared types for all 4 task types and WebSocket message protocol
  • tasks.go: AIInfraScanTask — orchestrates fingerprint scan + vuln matching + AI analysis
  • mcp_task.go: McpTask — invokes mcp-scan/main.py via subprocess
  • prompt_tasks.go: PromptTask — invokes AIG-PromptSecurity/cli_run.py via subprocess
  • agent_task.go: AgentTask — invokes agent-scan/main.py via subprocess

common/runner/ — CLI scan runner (for non-agent direct scans):

  • runner.go: Target parsing (URL, CIDR, local port scan), parallel fingerprint probing, vuln matching, security scoring

common/fingerprints/ — Fingerprint identification engine:

  • parser/: Custom DSL lexer/parser that compiles YAML fingerprint rules into executable Rule objects. Supports body/header/icon-hash matching with boolean expressions.
  • preload/: Concurrent HTTP probing engine that sends requests and matches responses against compiled rules, extracting component name + version.

internal/mcp/ — MCP security scanning engine:

  • scanner.go: AI-driven code audit using LLM, connects to MCP servers via Stdio/SSE/Stream
  • plugins.go: Loads security detection plugins from data/mcp/*.yaml, each defining rules and prompt templates

pkg/database/ — SQLite persistence via GORM (tasks, sessions, models, agent configs).

pkg/vulstruct/ — Vulnerability matching engine:

  • advisory.go: AdvisoryEngine loads vuln YAML files from data/vuln/, matches component+version against known CVEs using DSL-compiled version range expressions.

pkg/httpx/ — HTTP client with proxy, retry, redirect control, and encoding detection.

Python Sub-projects

All three Python sub-projects integrate with Go via the same pattern: Go spawns a subprocess (uv run <script> or python <script>), Python outputs structured JSON lines to stdout (newPlanStep, statusUpdate, toolUsed, actionLog, resultUpdate), and Go's ParseStdoutLine() parses and forwards them through WebSocket to the frontend.

AIG-PromptSecurity/ — Jailbreak / Red Team Evaluation:

  • Entry: cli_run.py, managed by uv (has pyproject.toml)
  • deepteam/red_teamer/: Core red team engine — orchestrates attack simulation, scenario generation, evaluation
  • deepteam/attacks/: Attack strategies (single-turn: Raw, Stego, StrataSword, Zalgo, Caesar, etc.; multi-turn; attack simulator)
  • deepteam/vulnerabilities/: 17+ vulnerability scenario types (Bias, Toxicity, PII, PromptLeakage, IllegalActivity, etc.)
  • deepteam/metrics/: 26+ evaluation metrics (BiasMetric, HarmMetric, SQLInjection, SSRF, Jailbreak, etc.)
  • deepteam/plugin_system/: Extensible plugin framework for custom attacks/metrics/vulnerabilities
  • Key deps: deepeval, openai, jieba, pyahocorasick

agent-scan/ — Agent Security Scanning:

  • Entry: main.py, a multi-agent framework for testing AI agent workflows on platforms like Dify/Coze
  • core/agent.py: Three-stage ScanPipeline — Info Collection → Parallel Vulnerability Detection (4 skill workers) → Vulnerability Review
  • core/agent_adapter/adapter.py: Multi-provider adapter (OpenAI/Anthropic/Google/Cohere/Ollama/HTTP/WebSocket)
  • core/report/: OWASP ASI-compliant vulnerability report generation
  • tools/: XML-schema-based tool registry with dialogue, batch, skill, task, thinking, finish tools
  • prompt/skills/: Detection skill templates (data-leakage, tool-abuse, indirect-injection, authorization-bypass, owasp-asi)
  • Key deps: mcp, openai, httpx

mcp-scan/ — MCP Server Security Scanning:

  • Entry: main.py, supports code audit mode (uploaded/cloned repos) and URL dynamic scanning
  • agent/agent.py: Multi-stage ScanPipeline similar to agent-scan
  • tools/: Shell execution, file reading, MCP tool bridging, thinking, finish tools
  • prompt/agents/: Prompt templates for code audit, vuln review, dynamic verification, malicious behavior testing
  • Key deps: mcp, openai

Data Directory (data/)

  • data/fingerprints/: ~60+ YAML fingerprint rules for AI components (Ollama, vLLM, Dify, Gradio, Jupyter, etc.). Each defines HTTP request paths and DSL match conditions.
  • data/vuln/ and data/vuln_en/: CVE vulnerability database organized by component (589+ CVEs across 43+ components). Each YAML file describes one CVE with severity, version ranges, and remediation.
  • data/eval/: 13 JSON datasets for jailbreak evaluation (JailBench, ChatGPT-Jailbreak-Prompts, CBRN, cyberattack, privacy-leakage, etc.)
  • data/mcp/: MCP security scan plugin rules (YAML with prompt templates)

Task Types

Task Type Go Handler Python Module Description
AI-Infra-Scan tasks.go N/A (pure Go) Fingerprint identification + CVE matching for exposed AI services
Mcp-Scan mcp_task.go mcp-scan/ MCP Server security scan (code audit or remote URL)
Model-Redteam-Report prompt_tasks.go AIG-PromptSecurity/ LLM jailbreak evaluation with attack strategies
Agent-Scan agent_task.go agent-scan/ AI agent workflow security testing (Dify, Coze, etc.)

Contribution: Adding Plugin Rules

  • Fingerprint rules: Add YAML files to data/fingerprints/ following existing format (DSL match expressions)
  • Vulnerability rules: Add YAML files to data/vuln/<component>/ with CVE details and version ranges
  • MCP plugins: Add YAML files to data/mcp/ with security detection rules and prompt templates
  • Evaluation datasets: Add JSON files to data/eval/ for jailbreak testing scenarios
  • Run go run cmd/yamlcheck/main.go to validate YAML format before submitting

Key Technical Details

  • Go version: 1.23.2, module: github.com/Tencent/AI-Infra-Guard
  • Web framework: Gin (gin-gonic/gin)
  • Database: SQLite via GORM, stored at DB_PATH (default /app/db/tasks.db)
  • Frontend: Embedded SPA via embed.FS (no separate frontend build step in this repo)
  • Python env: AIG-PromptSecurity uses uv (pyproject.toml); agent-scan and mcp-scan use pip (requirements.txt)
  • WebSocket protocol messages: register, register_ack, task_assign, resultUpdate, actionLog, toolUsed, newPlanStep, statusUpdate, planUpdate, error, terminate
  • Logging: tRPC-Go logging framework, configured via trpc_go.yaml
  • Config: Environment variables (APP_ENV, UPLOAD_DIR, DB_PATH, AIG_SERVER), no separate config file for app settings