mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-04 03:27:54 +02:00
公明
461 lines
17 KiB
Go
461 lines
17 KiB
Go
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.Preconditions != "" {
|
||
b.WriteString("\n--- 前置条件 ---\n")
|
||
b.WriteString(v.Preconditions)
|
||
b.WriteString("\n")
|
||
}
|
||
if v.ReproSteps != "" {
|
||
b.WriteString("\n--- 复现步骤 ---\n")
|
||
b.WriteString(v.ReproSteps)
|
||
b.WriteString("\n")
|
||
}
|
||
if v.Evidence != "" {
|
||
b.WriteString("\n--- 证据 / POC ---\n")
|
||
b.WriteString(v.Evidence)
|
||
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")
|
||
}
|
||
if v.RetestNotes != "" {
|
||
b.WriteString("\n--- 复测方式 ---\n")
|
||
b.WriteString(v.RetestNotes)
|
||
b.WriteString("\n")
|
||
}
|
||
return b.String()
|
||
}
|
||
|
||
func missingVulnerabilityReproFields(args map[string]interface{}) []string {
|
||
required := []struct {
|
||
key string
|
||
label string
|
||
}{
|
||
{"target", "target（受影响的 URL/IP/服务/接口）"},
|
||
{"vulnerability_type", "vulnerability_type（漏洞类型）"},
|
||
{"description", "description（漏洞摘要与触发点）"},
|
||
{"reproduction_steps", "reproduction_steps（可逐步执行的复现步骤）"},
|
||
{"evidence", "evidence（POC、原始请求/响应、命令输出或截图/日志证据）"},
|
||
{"impact", "impact（确认后的实际影响）"},
|
||
{"recommendation", "recommendation（修复建议）"},
|
||
}
|
||
missing := make([]string, 0)
|
||
for _, item := range required {
|
||
if strings.TrimSpace(strArg(args, item.key)) == "" {
|
||
missing = append(missing, item.label)
|
||
}
|
||
}
|
||
return missing
|
||
}
|
||
|
||
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": "漏洞标题（必需）。建议格式：<资产/接口> 存在 <漏洞类型>，例如“/api/login 存在 SQL 注入”。",
|
||
},
|
||
"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:端口、服务名、接口路径和参数名。",
|
||
},
|
||
"preconditions": map[string]interface{}{
|
||
"type": "string",
|
||
"description": "前置条件：登录状态、权限、账号、Header/Cookie、特定数据、网络位置、环境/版本等；无前置条件写“无”。",
|
||
},
|
||
"reproduction_steps": map[string]interface{}{
|
||
"type": "string",
|
||
"description": "复现步骤（必需）：按 1/2/3 编号，写清入口、参数、payload、执行命令、观察点。应让未参与对话的人照做即可复现。",
|
||
},
|
||
"evidence": map[string]interface{}{
|
||
"type": "string",
|
||
"description": "证据 / POC（必需）：原始 HTTP 请求/响应、curl/工具命令、截图文字说明、日志、DNSLog/回连记录、数据库结果、文件路径、时间戳等。优先放最小可验证证据。",
|
||
},
|
||
"impact": map[string]interface{}{
|
||
"type": "string",
|
||
"description": "漏洞影响说明（必需）：结合已验证事实说明可造成什么后果，避免泛泛而谈。",
|
||
},
|
||
"recommendation": map[string]interface{}{
|
||
"type": "string",
|
||
"description": "修复建议（必需）：给出针对该触发点/参数/组件的具体修复和复测建议。",
|
||
},
|
||
"retest_notes": map[string]interface{}{
|
||
"type": "string",
|
||
"description": "复测方式：修复后如何验证漏洞已关闭，包括应返回的状态码、错误信息或访问控制结果。",
|
||
},
|
||
},
|
||
"required": []string{"title", "description", "severity", "vulnerability_type", "target", "reproduction_steps", "evidence", "impact", "recommendation"},
|
||
},
|
||
}
|
||
|
||
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
|
||
}
|
||
if missing := missingVulnerabilityReproFields(args); len(missing) > 0 {
|
||
return textResult("错误: 漏洞记录缺少复现所需信息，请补充后再记录：\n- "+strings.Join(missing, "\n- ")+"\n\n最佳实践：漏洞管理中的单条记录应独立包含目标、前置条件、复现步骤、证据/POC、影响和修复/复测方式。", 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"),
|
||
Preconditions: strArg(args, "preconditions"),
|
||
ReproSteps: strArg(args, "reproduction_steps"),
|
||
Evidence: strArg(args, "evidence"),
|
||
Impact: strArg(args, "impact"),
|
||
Recommendation: strArg(args, "recommendation"),
|
||
RetestNotes: strArg(args, "retest_notes"),
|
||
}
|
||
|
||
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、ignored",
|
||
"enum": []string{"open", "confirmed", "fixed", "false_positive", "ignored"},
|
||
},
|
||
"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
|
||
})
|
||
}
|