CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-09-17 15:12:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: add configurable tool call blocking and monitoring

Browse Source
This commit is contained in:
Ed1s0nZ
2026-09-08 09:44:32 +08:00
parent c70da22de7
commit 6ad9ea2d13
54 changed files with 4635 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/app/app.go
+11
View File
@@ -33,6 +33,7 @@ import (
"cyberstrike-ai/internal/robot"
"cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/skillpackage"
"cyberstrike-ai/internal/toolguard"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
@@ -76,6 +77,10 @@ type App struct {
// New 创建新应用
func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error) {
toolGuard, err := toolguard.NewManager(cfg.EffectiveToolGuard())
if err != nil {
return nil, fmt.Errorf("初始化调用拦截规则: %w", err)
}
if err := multiagent.InitADK(); err != nil {
return nil, fmt.Errorf("初始化 Eino ADK: %w", err)
}
@@ -147,6 +152,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 创建MCP服务器（带数据库持久化）
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
mcpServer.SetToolGuard(toolGuard)
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
mcpServer.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
mcpServer.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
@@ -170,6 +176,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 创建外部MCP管理器（使用与内部MCP服务器相同的存储）
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
externalMCPMgr.SetToolGuard(toolGuard)
externalMCPMgr.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
externalMCPMgr.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
externalMCPMgr.ConfigureToolResultSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
@@ -412,6 +419,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger)
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger)
configHandler.SetDB(db)
configHandler.SetToolGuard(toolGuard)
configHandler.SetAudit(auditSvc)
agentHandler.SetHitlToolWhitelistSaver(configHandler)
agentHandler.SetHitlAuditStrategySaver(configHandler)
@@ -1054,6 +1062,9 @@ func setupRoutes(
// 配置管理
protected.GET("/config", configHandler.GetConfig)
protected.GET("/tool-guard", configHandler.GetToolGuard)
protected.PUT("/tool-guard", configHandler.UpdateToolGuard)
protected.POST("/tool-guard/test", configHandler.TestToolGuard)
protected.GET("/config/tools", configHandler.GetTools)
protected.GET("/config/tools/:name/schema", configHandler.GetToolSchema)
protected.PUT("/config", configHandler.UpdateConfig)