Add files via upload

This commit is contained in:
公明
2026-07-07 11:39:05 +08:00
committed by GitHub
parent 45e5c1bf39
commit dc5791f85c
2 changed files with 54 additions and 25 deletions
+39 -13
View File
@@ -246,7 +246,7 @@ type GetConfigResponse struct {
Knowledge config.KnowledgeConfig `json:"knowledge"` Knowledge config.KnowledgeConfig `json:"knowledge"`
Robots config.RobotsConfig `json:"robots,omitempty"` Robots config.RobotsConfig `json:"robots,omitempty"`
MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"` MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"`
C2 config.C2Public `json:"c2"` C2 config.C2Public `json:"c2"`
} }
// ToolConfigInfo 工具配置信息 // ToolConfigInfo 工具配置信息
@@ -320,7 +320,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
} }
multiPub := config.MultiAgentPublic{ multiPub := config.MultiAgentPublic{
Enabled: h.config.MultiAgent.Enabled, Enabled: h.config.MultiAgent.Enabled,
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent), RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent, BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
SubAgentCount: subAgentCount, SubAgentCount: subAgentCount,
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration), Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
@@ -673,16 +673,17 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
// UpdateConfigRequest 更新配置请求 // UpdateConfigRequest 更新配置请求
type UpdateConfigRequest struct { type UpdateConfigRequest struct {
OpenAI *config.OpenAIConfig `json:"openai,omitempty"` OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
Vision *config.VisionConfig `json:"vision,omitempty"` Vision *config.VisionConfig `json:"vision,omitempty"`
FOFA *config.FofaConfig `json:"fofa,omitempty"` FOFA *config.FofaConfig `json:"fofa,omitempty"`
MCP *config.MCPConfig `json:"mcp,omitempty"` MCP *config.MCPConfig `json:"mcp,omitempty"`
Tools []ToolEnableStatus `json:"tools,omitempty"` Tools []ToolEnableStatus `json:"tools,omitempty"`
Agent *AgentConfigUpdate `json:"agent,omitempty"` Agent *AgentConfigUpdate `json:"agent,omitempty"`
Knowledge *config.KnowledgeConfig `json:"knowledge,omitempty"` Hitl *config.HitlConfig `json:"hitl,omitempty"`
Robots *config.RobotsConfig `json:"robots,omitempty"` Knowledge *config.KnowledgeConfig `json:"knowledge,omitempty"`
MultiAgent *config.MultiAgentAPIUpdate `json:"multi_agent,omitempty"` Robots *config.RobotsConfig `json:"robots,omitempty"`
C2 *config.C2APIUpdate `json:"c2,omitempty"` MultiAgent *config.MultiAgentAPIUpdate `json:"multi_agent,omitempty"`
C2 *config.C2APIUpdate `json:"c2,omitempty"`
} }
// AgentConfigUpdate 用于 PATCH /api/config 的 agent 段:仅 JSON 中出现的字段(指针非 nil)覆盖内存配置。 // 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配置 // 更新Knowledge配置
if req.Knowledge != nil { if req.Knowledge != nil {
// 保存旧的嵌入模型配置(用于检测变更) // 保存旧的嵌入模型配置(用于检测变更)
@@ -1471,7 +1491,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
Result: "success", Result: "success",
Message: "配置已应用", Message: "配置已应用",
Detail: map[string]interface{}{ Detail: map[string]interface{}{
"tools_count": len(h.config.Security.Tools), "tools_count": len(h.config.Security.Tools),
"knowledge_enabled": h.config.Knowledge.Enabled, "knowledge_enabled": h.config.Knowledge.Enabled,
"c2_enabled": h.config.C2.EnabledEffective(), "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) { func updateHitlConfig(doc *yaml.Node, cfg config.HitlConfig) {
root := doc.Content[0] root := doc.Content[0]
hitlNode := ensureMap(root, "hitl") 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] 单行展示,工具多时比块序列省行数 // flow 样式 [a, b, c] 单行展示,工具多时比块序列省行数
setFlowStringSliceInMap(hitlNode, "tool_whitelist", cfg.ToolWhitelist) setFlowStringSliceInMap(hitlNode, "tool_whitelist", cfg.ToolWhitelist)
setStringInMap(hitlNode, "default_reviewer", cfg.EffectiveDefaultReviewer()) setStringInMap(hitlNode, "default_reviewer", cfg.EffectiveDefaultReviewer())
setIntInMap(hitlNode, "retention_days", cfg.RetentionDaysEffective())
setStringInMap(hitlNode, "audit_agent_prompt", cfg.AuditAgentPrompt) setStringInMap(hitlNode, "audit_agent_prompt", cfg.AuditAgentPrompt)
setStringInMap(hitlNode, "audit_agent_prompt_review_edit", cfg.AuditAgentPromptReviewEdit) setStringInMap(hitlNode, "audit_agent_prompt_review_edit", cfg.AuditAgentPromptReviewEdit)
} }
+15 -12
View File
@@ -10,6 +10,7 @@ import (
"time" "time"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/openai"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"go.uber.org/zap" "go.uber.org/zap"
@@ -26,7 +27,8 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName
if h.config != nil { if h.config != nil {
prompt = h.config.Hitl.EffectiveAuditAgentPromptForMode(mode) 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 未配置"} return hitlDecision{Decision: "reject", Comment: "audit agent: LLM 未配置"}
} }
if ctx == nil { if ctx == nil {
@@ -37,7 +39,7 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName
userContent := buildAuditAgentReviewInput(mode, toolName, payload) userContent := buildAuditAgentReviewInput(mode, toolName, payload)
requestBody := map[string]interface{}{ requestBody := map[string]interface{}{
"model": h.auditLLMModel(), "model": strings.TrimSpace(llmCfg.Model),
"messages": []map[string]interface{}{ "messages": []map[string]interface{}{
{"role": "system", "content": prompt}, {"role": "system", "content": prompt},
{"role": "user", "content": userContent}, {"role": "user", "content": userContent},
@@ -56,7 +58,8 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName
} `json:"message"` } `json:"message"`
} `json:"choices"` } `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)) h.logger.Warn("审计 Agent LLM 调用失败", zap.Error(err), zap.String("tool", toolName))
return hitlDecision{ return hitlDecision{
Decision: "reject", Decision: "reject",
@@ -99,11 +102,11 @@ func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName
return dec return dec
} }
func (h *AgentHandler) auditLLMModel() string { func (h *AgentHandler) auditLLMConfig() config.OpenAIConfig {
if h.config != nil && strings.TrimSpace(h.config.OpenAI.Model) != "" { if h != nil && h.config != nil {
return strings.TrimSpace(h.config.OpenAI.Model) return h.config.Hitl.AuditModelEffective(h.config.OpenAI)
} }
return "" return config.OpenAIConfig{}
} }
func buildAuditAgentReviewInput(hitlMode, toolName string, payload map[string]interface{}) string { 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 h.config.Hitl.AuditAgentPromptReviewEdit = reviewEditPrompt
} }
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"ok": true, "ok": true,
"auditAgentPrompt": config.HitlConfig{AuditAgentPrompt: approvalPrompt}.EffectiveAuditAgentPromptForMode("approval"), "auditAgentPrompt": config.HitlConfig{AuditAgentPrompt: approvalPrompt}.EffectiveAuditAgentPromptForMode("approval"),
"auditAgentPromptCustom": approvalPrompt != "", "auditAgentPromptCustom": approvalPrompt != "",
"auditAgentPromptReviewEdit": config.HitlConfig{AuditAgentPromptReviewEdit: reviewEditPrompt}.EffectiveAuditAgentPromptForMode("review_edit"), "auditAgentPromptReviewEdit": config.HitlConfig{AuditAgentPromptReviewEdit: reviewEditPrompt}.EffectiveAuditAgentPromptForMode("review_edit"),
"auditAgentPromptReviewEditCustom": reviewEditPrompt != "", "auditAgentPromptReviewEditCustom": reviewEditPrompt != "",
}) })
} }