From dc5791f85c2e5b42b9899167d388756830d47f1b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:39:05 +0800 Subject: [PATCH] Add files via upload --- internal/handler/config.go | 52 +++++++++++++++++++++------- internal/handler/hitl_audit_agent.go | 27 ++++++++------- 2 files changed, 54 insertions(+), 25 deletions(-) diff --git a/internal/handler/config.go b/internal/handler/config.go index 284c5e22..7bba7e0b 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -246,7 +246,7 @@ type GetConfigResponse struct { Knowledge config.KnowledgeConfig `json:"knowledge"` Robots config.RobotsConfig `json:"robots,omitempty"` MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"` - C2 config.C2Public `json:"c2"` + C2 config.C2Public `json:"c2"` } // ToolConfigInfo 工具配置信息 @@ -320,7 +320,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) { } multiPub := config.MultiAgentPublic{ Enabled: h.config.MultiAgent.Enabled, - RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent), + RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent), BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent, SubAgentCount: subAgentCount, Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration), @@ -673,16 +673,17 @@ func (h *ConfigHandler) GetTools(c *gin.Context) { // UpdateConfigRequest 更新配置请求 type UpdateConfigRequest struct { - OpenAI *config.OpenAIConfig `json:"openai,omitempty"` - Vision *config.VisionConfig `json:"vision,omitempty"` - FOFA *config.FofaConfig `json:"fofa,omitempty"` - MCP *config.MCPConfig `json:"mcp,omitempty"` - Tools []ToolEnableStatus `json:"tools,omitempty"` - Agent *AgentConfigUpdate `json:"agent,omitempty"` - Knowledge *config.KnowledgeConfig `json:"knowledge,omitempty"` - Robots *config.RobotsConfig `json:"robots,omitempty"` - MultiAgent *config.MultiAgentAPIUpdate `json:"multi_agent,omitempty"` - C2 *config.C2APIUpdate `json:"c2,omitempty"` + OpenAI *config.OpenAIConfig `json:"openai,omitempty"` + Vision *config.VisionConfig `json:"vision,omitempty"` + FOFA *config.FofaConfig `json:"fofa,omitempty"` + MCP *config.MCPConfig `json:"mcp,omitempty"` + Tools []ToolEnableStatus `json:"tools,omitempty"` + Agent *AgentConfigUpdate `json:"agent,omitempty"` + Hitl *config.HitlConfig `json:"hitl,omitempty"` + Knowledge *config.KnowledgeConfig `json:"knowledge,omitempty"` + Robots *config.RobotsConfig `json:"robots,omitempty"` + MultiAgent *config.MultiAgentAPIUpdate `json:"multi_agent,omitempty"` + C2 *config.C2APIUpdate `json:"c2,omitempty"` } // AgentConfigUpdate 用于 PATCH /api/config 的 agent 段：仅 JSON 中出现的字段（指针非 nil）覆盖内存配置。 @@ -775,6 +776,25 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) { } } + if req.Hitl != nil { + h.config.Hitl.AuditModel = req.Hitl.AuditModel + h.config.Hitl.ToolWhitelist = mergeHitlToolWhitelistSlice(nil, req.Hitl.ToolWhitelist) + h.config.Hitl.DefaultReviewer = req.Hitl.EffectiveDefaultReviewer() + h.config.Hitl.AuditAgentPrompt = strings.TrimSpace(req.Hitl.AuditAgentPrompt) + h.config.Hitl.AuditAgentPromptReviewEdit = strings.TrimSpace(req.Hitl.AuditAgentPromptReviewEdit) + if req.Hitl.RetentionDays != nil { + v := *req.Hitl.RetentionDays + if v < 0 { + v = 0 + } + h.config.Hitl.RetentionDays = &v + } + h.logger.Info("更新HITL配置", + zap.String("default_reviewer", h.config.Hitl.DefaultReviewer), + zap.Int("tool_whitelist", len(h.config.Hitl.ToolWhitelist)), + ) + } + // 更新Knowledge配置 if req.Knowledge != nil { // 保存旧的嵌入模型配置（用于检测变更） @@ -1471,7 +1491,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) { Result: "success", Message: "配置已应用", Detail: map[string]interface{}{ - "tools_count": len(h.config.Security.Tools), + "tools_count": len(h.config.Security.Tools), "knowledge_enabled": h.config.Knowledge.Enabled, "c2_enabled": h.config.C2.EnabledEffective(), }, @@ -1804,9 +1824,15 @@ func (h *ConfigHandler) MergeHitlToolWhitelistIntoConfig(add []string) error { func updateHitlConfig(doc *yaml.Node, cfg config.HitlConfig) { root := doc.Content[0] hitlNode := ensureMap(root, "hitl") + auditModelNode := ensureMap(hitlNode, "audit_model") + setStringInMap(auditModelNode, "provider", cfg.AuditModel.Provider) + setStringInMap(auditModelNode, "base_url", cfg.AuditModel.BaseURL) + setStringInMap(auditModelNode, "api_key", cfg.AuditModel.APIKey) + setStringInMap(auditModelNode, "model", cfg.AuditModel.Model) // flow 样式 [a, b, c] 单行展示，工具多时比块序列省行数 setFlowStringSliceInMap(hitlNode, "tool_whitelist", cfg.ToolWhitelist) setStringInMap(hitlNode, "default_reviewer", cfg.EffectiveDefaultReviewer()) + setIntInMap(hitlNode, "retention_days", cfg.RetentionDaysEffective()) setStringInMap(hitlNode, "audit_agent_prompt", cfg.AuditAgentPrompt) setStringInMap(hitlNode, "audit_agent_prompt_review_edit", cfg.AuditAgentPromptReviewEdit) } diff --git a/internal/handler/hitl_audit_agent.go b/internal/handler/hitl_audit_agent.go index c06bd312..8bbd37b9 100644 --- a/internal/handler/hitl_audit_agent.go +++ b/internal/handler/hitl_audit_agent.go @@ -10,6 +10,7 @@ import ( "time" "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/openai" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -26,7 +27,8 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName if h.config != nil { prompt = h.config.Hitl.EffectiveAuditAgentPromptForMode(mode) } - if h.auditLLM == nil { + llmCfg := h.auditLLMConfig() + if strings.TrimSpace(llmCfg.APIKey) == "" || strings.TrimSpace(llmCfg.Model) == "" { return hitlDecision{Decision: "reject", Comment: "audit agent: LLM 未配置"} } if ctx == nil { @@ -37,7 +39,7 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName userContent := buildAuditAgentReviewInput(mode, toolName, payload) requestBody := map[string]interface{}{ - "model": h.auditLLMModel(), + "model": strings.TrimSpace(llmCfg.Model), "messages": []map[string]interface{}{ {"role": "system", "content": prompt}, {"role": "user", "content": userContent}, @@ -56,7 +58,8 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName } `json:"message"` } `json:"choices"` } - if err := h.auditLLM.ChatCompletion(callCtx, requestBody, &apiResponse); err != nil { + client := openai.NewClient(&llmCfg, nil, h.logger) + if err := client.ChatCompletion(callCtx, requestBody, &apiResponse); err != nil { h.logger.Warn("审计 Agent LLM 调用失败", zap.Error(err), zap.String("tool", toolName)) return hitlDecision{ Decision: "reject", @@ -99,11 +102,11 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName return dec } -func (h *AgentHandler) auditLLMModel() string { - if h.config != nil && strings.TrimSpace(h.config.OpenAI.Model) != "" { - return strings.TrimSpace(h.config.OpenAI.Model) +func (h *AgentHandler) auditLLMConfig() config.OpenAIConfig { + if h != nil && h.config != nil { + return h.config.Hitl.AuditModelEffective(h.config.OpenAI) } - return "" + return config.OpenAIConfig{} } func buildAuditAgentReviewInput(hitlMode, toolName string, payload map[string]interface{}) string { @@ -338,11 +341,11 @@ func (h *AgentHandler) UpdateHITLAuditStrategy(c *gin.Context) { h.config.Hitl.AuditAgentPromptReviewEdit = reviewEditPrompt } c.JSON(http.StatusOK, gin.H{ - "ok": true, - "auditAgentPrompt": config.HitlConfig{AuditAgentPrompt: approvalPrompt}.EffectiveAuditAgentPromptForMode("approval"), - "auditAgentPromptCustom": approvalPrompt != "", - "auditAgentPromptReviewEdit": config.HitlConfig{AuditAgentPromptReviewEdit: reviewEditPrompt}.EffectiveAuditAgentPromptForMode("review_edit"), - "auditAgentPromptReviewEditCustom": reviewEditPrompt != "", + "ok": true, + "auditAgentPrompt": config.HitlConfig{AuditAgentPrompt: approvalPrompt}.EffectiveAuditAgentPromptForMode("approval"), + "auditAgentPromptCustom": approvalPrompt != "", + "auditAgentPromptReviewEdit": config.HitlConfig{AuditAgentPromptReviewEdit: reviewEditPrompt}.EffectiveAuditAgentPromptForMode("review_edit"), + "auditAgentPromptReviewEditCustom": reviewEditPrompt != "", }) }