From b366dc0287cca7e4cc03dfcba0775097f79a6660 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, 23 Jun 2026 15:35:12 +0800 Subject: [PATCH] Add files via upload --- internal/config/config.go | 22 ++++ internal/database/monitor.go | 70 +++++++++++ internal/database/monitor_retention_test.go | 122 ++++++++++++++++++++ 3 files changed, 214 insertions(+) create mode 100644 internal/database/monitor_retention_test.go diff --git a/internal/config/config.go b/internal/config/config.go index 59f986bb..46b6ec0f 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -27,6 +27,7 @@ type Config struct { Database DatabaseConfig `yaml:"database"` Auth AuthConfig `yaml:"auth"` Audit AuditConfig `yaml:"audit,omitempty" json:"audit,omitempty"` + Monitor MonitorConfig `yaml:"monitor,omitempty" json:"monitor,omitempty"` ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"` Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"` C2 C2Config `yaml:"c2,omitempty" json:"c2,omitempty"` // 内置 C2 总开关；未配置时默认启用 @@ -623,6 +624,23 @@ type AuthConfig struct { GeneratedPasswordPersistErr string `yaml:"-" json:"-"` } +// MonitorConfig MCP 状态监控（tool_executions）保留策略。 +type MonitorConfig struct { + // RetentionDays 执行记录保留天数；省略时默认 90；0 表示不自动清理。 + RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"` +} + +// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90. +func (m MonitorConfig) RetentionDaysEffective() int { + if m.RetentionDays == nil { + return 90 + } + if *m.RetentionDays < 0 { + return 0 + } + return *m.RetentionDays +} + // AuditConfig platform operation audit log settings (not chat/tool execution bodies). type AuditConfig struct { // Enabled nil or true enables persistence; explicit false disables. @@ -1274,6 +1292,10 @@ func Default() *Config { Enabled: &on, } }(), + Monitor: func() MonitorConfig { + days := 90 + return MonitorConfig{RetentionDays: &days} + }(), Robots: RobotsConfig{ Session: RobotSessionConfig{ StrictUserIdentity: &strictRobotIdentity, diff --git a/internal/database/monitor.go b/internal/database/monitor.go index b215674e..75ad29ae 100644 --- a/internal/database/monitor.go +++ b/internal/database/monitor.go @@ -410,6 +410,76 @@ func (db *DB) GetToolExecutionsByIds(ids []string) ([]*mcp.ToolExecution, error) return executions, nil } +type toolExecutionStatDelta struct { + totalCalls int + successCalls int + failedCalls int +} + +// PurgeToolExecutionsBefore deletes executions older than cutoff and adjusts tool_stats. +func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) { + query := ` + SELECT tool_name, status, COUNT(*) AS cnt + FROM tool_executions + WHERE ` + sqliteEpochGE("start_time", "<") + ` + GROUP BY tool_name, status + ` + rows, err := db.Query(query, formatSQLiteUTC(cutoff)) + if err != nil { + return 0, err + } + defer rows.Close() + + deltas := make(map[string]*toolExecutionStatDelta) + for rows.Next() { + var toolName, status string + var count int + if err := rows.Scan(&toolName, &status, &count); err != nil { + db.logger.Warn("读取待清理执行记录统计失败", zap.Error(err)) + continue + } + toolName = strings.TrimSpace(toolName) + if toolName == "" || count <= 0 { + continue + } + delta := deltas[toolName] + if delta == nil { + delta = &toolExecutionStatDelta{} + deltas[toolName] = delta + } + delta.totalCalls += count + switch status { + case "failed", "cancelled": + delta.failedCalls += count + case "completed": + delta.successCalls += count + } + } + if err := rows.Err(); err != nil { + return 0, err + } + + res, err := db.Exec(`DELETE FROM tool_executions WHERE `+sqliteEpochGE("start_time", "<"), formatSQLiteUTC(cutoff)) + if err != nil { + return 0, err + } + deleted, err := res.RowsAffected() + if err != nil { + return 0, err + } + + for toolName, delta := range deltas { + if err := db.DecreaseToolStats(toolName, delta.totalCalls, delta.successCalls, delta.failedCalls); err != nil { + db.logger.Warn("清理过期执行记录后更新统计失败", + zap.Error(err), + zap.String("toolName", toolName), + ) + } + } + + return deleted, nil +} + // SaveToolStats 保存工具统计信息 func (db *DB) SaveToolStats(toolName string, stats *mcp.ToolStats) error { var lastCallTime sql.NullTime diff --git a/internal/database/monitor_retention_test.go b/internal/database/monitor_retention_test.go new file mode 100644 index 00000000..20de7cad --- /dev/null +++ b/internal/database/monitor_retention_test.go @@ -0,0 +1,122 @@ +package database + +import ( + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/mcp" + + "go.uber.org/zap" +) + +func TestPurgeToolExecutionsBefore(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "monitor.db") + db, err := NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer db.Close() + + oldStart := time.Now().AddDate(0, 0, -100) + newStart := time.Now().AddDate(0, 0, -1) + + oldExec := &mcp.ToolExecution{ + ID: "old-completed", + ToolName: "nmap::scan", + Arguments: map[string]interface{}{"target": "127.0.0.1"}, + Status: "completed", + StartTime: oldStart, + } + oldFailed := &mcp.ToolExecution{ + ID: "old-failed", + ToolName: "nmap::scan", + Arguments: map[string]interface{}{"target": "127.0.0.1"}, + Status: "failed", + Error: "timeout", + StartTime: oldStart, + } + newExec := &mcp.ToolExecution{ + ID: "new-completed", + ToolName: "nmap::scan", + Arguments: map[string]interface{}{"target": "127.0.0.1"}, + Status: "completed", + StartTime: newStart, + } + for _, exec := range []*mcp.ToolExecution{oldExec, oldFailed, newExec} { + if err := db.SaveToolExecution(exec); err != nil { + t.Fatalf("SaveToolExecution(%s): %v", exec.ID, err) + } + } + if err := db.UpdateToolStats("nmap::scan", 3, 2, 1, &newStart); err != nil { + t.Fatalf("UpdateToolStats: %v", err) + } + + cutoff := time.Now().AddDate(0, 0, -90) + deleted, err := db.PurgeToolExecutionsBefore(cutoff) + if err != nil { + t.Fatalf("PurgeToolExecutionsBefore: %v", err) + } + if deleted != 2 { + t.Fatalf("deleted = %d, want 2", deleted) + } + + if _, err := db.GetToolExecution("old-completed"); err == nil { + t.Fatal("old-completed should be deleted") + } + if _, err := db.GetToolExecution("old-failed"); err == nil { + t.Fatal("old-failed should be deleted") + } + if _, err := db.GetToolExecution("new-completed"); err != nil { + t.Fatalf("new-completed should remain: %v", err) + } + + stats, err := db.LoadToolStats() + if err != nil { + t.Fatalf("LoadToolStats: %v", err) + } + stat := stats["nmap::scan"] + if stat == nil { + t.Fatal("expected stats for nmap::scan") + } + if stat.TotalCalls != 1 || stat.SuccessCalls != 1 || stat.FailedCalls != 0 { + t.Fatalf("stats after purge = %+v, want total=1 success=1 failed=0", stat) + } + + total, err := db.CountToolExecutions("", "") + if err != nil { + t.Fatalf("CountToolExecutions: %v", err) + } + if total != 1 { + t.Fatalf("remaining executions = %d, want 1", total) + } +} + +func TestPurgeToolExecutionsBefore_zeroRetentionSkipsViaService(t *testing.T) { + // RetentionDaysEffective: 0 means no purge at service layer; DB method still works when called directly. + dbPath := filepath.Join(t.TempDir(), "monitor.db") + db, err := NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer db.Close() + + exec := &mcp.ToolExecution{ + ID: "ancient", + ToolName: "curl::get", + Arguments: map[string]interface{}{}, + Status: "completed", + StartTime: time.Now().AddDate(-1, 0, 0), + } + if err := db.SaveToolExecution(exec); err != nil { + t.Fatalf("SaveToolExecution: %v", err) + } + + deleted, err := db.PurgeToolExecutionsBefore(time.Now()) + if err != nil { + t.Fatalf("PurgeToolExecutionsBefore: %v", err) + } + if deleted != 1 { + t.Fatalf("deleted = %d, want 1", deleted) + } +}