CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-08-01 00:27:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-21 20:49:10 +08:00
committed by GitHub
parent 4bbbe3bcf8
commit c925c2e74b
5 changed files with 80 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/app/app.go
+12
View File
@@ -148,10 +148,13 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
mcpServer.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
mcpServer.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
// 创建安全工具执行器
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds)
executor.SetToolOutputMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
// 注册工具
executor.RegisterTools(mcpServer)
@@ -165,6 +168,15 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 创建外部MCP管理器（使用与内部MCP服务器相同的存储）
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
externalMCPMgr.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
externalMCPMgr.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
MaxConcurrentPerServer: cfg.Agent.ExternalMCPMaxConcurrentPerServer,
MaxConcurrentTotal: cfg.Agent.ExternalMCPMaxConcurrentTotal,
CircuitFailureThreshold: cfg.Agent.ExternalMCPCircuitFailureThreshold,
CircuitCooldown: time.Duration(cfg.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
})
mcp.RegisterExecutionControlTools(mcpServer, externalMCPMgr)
if cfg.ExternalMCP.Servers != nil {
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
// 启动所有启用的外部MCP客户端
internal/app/mcp_authorization.go
+14
View File
@@ -34,6 +34,16 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
}
return nil
}
toolExecutionResource := func(permission string) error {
if err := require(permission); err != nil {
return err
}
id := mcpAuthorizationString(args, "execution_id")
if id == "" || db == nil || !db.UserCanAccessToolExecution(principal.UserID, principal.ScopeFor(permission), id) {
return fmt.Errorf("no access to tool execution %s", id)
}
return nil
}
switch toolName {
case builtin.ToolWebshellExec, builtin.ToolWebshellFileWrite:
@@ -108,6 +118,10 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
return require("knowledge:read")
case builtin.ToolAnalyzeImage:
return require("agent:execute")
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution:
return toolExecutionResource("monitor:read")
case builtin.ToolCancelToolExecution:
return toolExecutionResource("monitor:write")
case builtin.ToolBatchTaskList:
return require("tasks:read")
case builtin.ToolBatchTaskGet:
internal/app/mcp_authorization_test.go
+41 -1
View File
@@ -9,6 +9,7 @@ import (
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/security"
@@ -62,7 +63,7 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
authorize := mcpToolAuthorizer(db)
args := map[string]interface{}{
"action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x",
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x",
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x", "execution_id": "x",
}
for _, toolName := range builtin.GetAllBuiltinTools() {
err := authorize(ctx, toolName, args)
@@ -72,6 +73,45 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
}
}
func TestMCPExecutionControlAuthorizationUsesExecutionScope(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-exec-authz.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
user, err := db.CreateRBACUser("exec-user", "Exec User", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "exec-owned",
ToolName: "lab::slow",
Status: "running",
StartTime: time.Now(),
OwnerUserID: user.ID,
}); err != nil {
t.Fatal(err)
}
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "exec-hidden",
ToolName: "lab::slow",
Status: "running",
StartTime: time.Now(),
}); err != nil {
t.Fatal(err)
}
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"monitor:read": true, "monitor:write": true})
ctx := authctx.WithPrincipal(context.Background(), principal)
authorize := mcpToolAuthorizer(db)
if err := authorize(ctx, builtin.ToolWaitToolExecution, map[string]interface{}{"execution_id": "exec-owned"}); err != nil {
t.Fatalf("owned execution denied: %v", err)
}
if err := authorize(ctx, builtin.ToolCancelToolExecution, map[string]interface{}{"execution_id": "exec-hidden"}); err == nil {
t.Fatal("foreign execution was allowed")
}
}
func TestMCPAssetToolAuthorizationUsesAssetPermissionsAndScope(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop())
if err != nil {
internal/database/monitor.go
+9 -9
View File
@@ -359,7 +359,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
var lastCall sql.NullString
err := db.QueryRow(`SELECT COUNT(*),
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
&lastCall, &result.Summary.ToolCount,
@@ -373,7 +373,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
}
rows, err := db.Query(`SELECT tool_name, COUNT(*),
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), MAX(start_time)`+
SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
if err != nil {
return nil, err
@@ -565,7 +565,7 @@ func (db *DB) UserCanAccessToolExecution(userID, scope, executionID string) bool
return conversation != "" && db.UserCanAccessResource(userID, scope, "conversation", conversation)
}
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 cancelled（如进程重启后无对应执行协程）。
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 orphaned（如进程重启后无对应执行协程）。
func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg string) (int64, error) {
errMsg = strings.TrimSpace(errMsg)
if errMsg == "" {
@@ -573,7 +573,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri
}
query := `
UPDATE tool_executions
SET status = 'cancelled',
SET status = 'orphaned',
error = ?,
end_time = ?,
duration_ms = MAX(0, CAST((julianday(?) - julianday(start_time)) * 86400000 AS INTEGER))
@@ -586,7 +586,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri
return res.RowsAffected()
}
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 cancelled。
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 orphaned。
// activeIDs 为当前进程内仍登记 cancel 的 executionId；不在集合内且已超时的视为孤儿记录。
func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.Duration, activeIDs map[string]struct{}, errMsg string) (int64, error) {
errMsg = strings.TrimSpace(errMsg)
@@ -639,7 +639,7 @@ func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.
}
res, err := db.Exec(`
UPDATE tool_executions
SET status = 'cancelled', error = ?, end_time = ?, duration_ms = ?
SET status = 'orphaned', error = ?, end_time = ?, duration_ms = ?
WHERE id = ? AND status = 'running'
`, errMsg, endTime, durationMs, row.id)
if err != nil {
@@ -815,7 +815,7 @@ func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) {
}
delta.totalCalls += count
switch status {
case "failed", "cancelled":
case "failed", "cancelled", "hard_timeout", "orphaned":
delta.failedCalls += count
case "completed":
delta.successCalls += count
@@ -971,7 +971,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT date(start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
@@ -981,7 +981,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
internal/database/monitor_reconcile_test.go
+4 -4
View File
@@ -43,8 +43,8 @@ func TestCancelOrphanedRunningToolExecutions(t *testing.T) {
if err != nil {
t.Fatalf("GetToolExecution: %v", err)
}
if got.Status != "cancelled" {
t.Fatalf("expected cancelled, got %s", got.Status)
if got.Status != "orphaned" {
t.Fatalf("expected orphaned, got %s", got.Status)
}
if got.EndTime == nil {
t.Fatal("expected end_time to be set")
@@ -88,8 +88,8 @@ func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) {
if err != nil {
t.Fatalf("GetToolExecution stale: %v", err)
}
if stale.Status != "cancelled" {
t.Fatalf("stale expected cancelled, got %s", stale.Status)
if stale.Status != "orphaned" {
t.Fatalf("stale expected orphaned, got %s", stale.Status)
}
activeExec, err := db.GetToolExecution("active")