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

Add files via upload

Browse Source
This commit is contained in:
公明
2025-12-30 22:02:13 +08:00
committed by GitHub
parent 98713236b7
commit d48238f6a0
7 changed files with 200 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/database/monitor.go
+33 -6
View File
@@ -3,9 +3,11 @@ package database
import (
"database/sql"
"encoding/json"
"strings"
"time"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
@@ -70,13 +72,25 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
}
// CountToolExecutions 统计工具执行记录总数
func (db *DB) CountToolExecutions(status string) (int, error) {
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
query := `SELECT COUNT(*) FROM tool_executions`
args := []interface{}{}
conditions := []string{}
if status != "" {
query += ` WHERE status = ?`
conditions = append(conditions, "status = ?")
args = append(args, status)
}
if toolName != "" {
// 支持部分匹配（模糊搜索），不区分大小写
conditions = append(conditions, "LOWER(tool_name) LIKE ?")
args = append(args, "%"+strings.ToLower(toolName)+"%")
}
if len(conditions) > 0 {
query += ` WHERE ` + conditions[0]
for i := 1; i < len(conditions); i++ {
query += ` AND ` + conditions[i]
}
}
var count int
err := db.QueryRow(query, args...).Scan(&count)
if err != nil {
@@ -87,30 +101,43 @@ func (db *DB) CountToolExecutions(status string) (int, error) {
// LoadToolExecutions 加载所有工具执行记录（支持分页）
func (db *DB) LoadToolExecutions() ([]*mcp.ToolExecution, error) {
return db.LoadToolExecutionsWithPagination(0, 1000, "")
return db.LoadToolExecutionsWithPagination(0, 1000, "", "")
}
// LoadToolExecutionsWithPagination 分页加载工具执行记录
// limit: 最大返回记录数，0 表示使用默认值 1000
// offset: 跳过的记录数，用于分页
// status: 状态筛选，空字符串表示不过滤
func (db *DB) LoadToolExecutionsWithPagination(offset, limit int, status string) ([]*mcp.ToolExecution, error) {
// toolName: 工具名称筛选，空字符串表示不过滤
func (db *DB) LoadToolExecutionsWithPagination(offset, limit int, status, toolName string) ([]*mcp.ToolExecution, error) {
if limit <= 0 {
limit = 1000 // 默认限制
}
if limit > 10000 {
limit = 10000 // 最大限制，防止一次性加载过多数据
}
query := `
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms
FROM tool_executions
`
args := []interface{}{}
conditions := []string{}
if status != "" {
query += ` WHERE status = ?`
conditions = append(conditions, "status = ?")
args = append(args, status)
}
if toolName != "" {
// 支持部分匹配（模糊搜索），不区分大小写
conditions = append(conditions, "LOWER(tool_name) LIKE ?")
args = append(args, "%"+strings.ToLower(toolName)+"%")
}
if len(conditions) > 0 {
query += ` WHERE ` + conditions[0]
for i := 1; i < len(conditions); i++ {
query += ` AND ` + conditions[i]
}
}
query += ` ORDER BY start_time DESC LIMIT ? OFFSET ?`
args = append(args, limit, offset)
internal/handler/monitor.go
+21 -12
View File
@@ -3,6 +3,7 @@ package handler
import (
"net/http"
"strconv"
"strings"
"time"
"cyberstrike-ai/internal/database"
@@ -66,8 +67,10 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
// 解析状态筛选参数
status := c.Query("status")
// 解析工具筛选参数
toolName := c.Query("tool")
executions, total := h.loadExecutionsWithPagination(page, pageSize, status)
executions, total := h.loadExecutionsWithPagination(page, pageSize, status, toolName)
stats := h.loadStats()
totalPages := (total + pageSize - 1) / pageSize
@@ -87,18 +90,21 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
}
func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution {
executions, _ := h.loadExecutionsWithPagination(1, 1000, "")
executions, _ := h.loadExecutionsWithPagination(1, 1000, "", "")
return executions
}
func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status string) ([]*mcp.ToolExecution, int) {
func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
if h.db == nil {
allExecutions := h.mcpServer.GetAllExecutions()
// 如果指定了状态筛选，先进行筛选
if status != "" {
// 如果指定了状态筛选或工具筛选，先进行筛选
if status != "" || toolName != "" {
filtered := make([]*mcp.ToolExecution, 0)
for _, exec := range allExecutions {
if exec.Status == status {
matchStatus := status == "" || exec.Status == status
// 支持部分匹配（模糊搜索）
matchTool := toolName == "" || strings.Contains(strings.ToLower(exec.ToolName), strings.ToLower(toolName))
if matchStatus && matchTool {
filtered = append(filtered, exec)
}
}
@@ -117,15 +123,18 @@ func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status
}
offset := (page - 1) * pageSize
executions, err := h.db.LoadToolExecutionsWithPagination(offset, pageSize, status)
executions, err := h.db.LoadToolExecutionsWithPagination(offset, pageSize, status, toolName)
if err != nil {
h.logger.Warn("从数据库加载执行记录失败，回退到内存数据", zap.Error(err))
allExecutions := h.mcpServer.GetAllExecutions()
// 如果指定了状态筛选，先进行筛选
if status != "" {
// 如果指定了状态筛选或工具筛选，先进行筛选
if status != "" || toolName != "" {
filtered := make([]*mcp.ToolExecution, 0)
for _, exec := range allExecutions {
if exec.Status == status {
matchStatus := status == "" || exec.Status == status
// 支持部分匹配（模糊搜索）
matchTool := toolName == "" || strings.Contains(strings.ToLower(exec.ToolName), strings.ToLower(toolName))
if matchStatus && matchTool {
filtered = append(filtered, exec)
}
}
@@ -143,8 +152,8 @@ func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status
return allExecutions[offset:end], total
}
// 获取总数（考虑状态筛选）
total, err := h.db.CountToolExecutions(status)
// 获取总数（考虑状态筛选和工具筛选）
total, err := h.db.CountToolExecutions(status, toolName)
if err != nil {
h.logger.Warn("获取执行记录总数失败", zap.Error(err))
// 回退：使用已加载的记录数估算
internal/knowledge/manager.go
+17 -1
View File
@@ -639,7 +639,12 @@ func (m *Manager) UpdateItem(id, category, title, content string) (*KnowledgeIte
// 删除旧目录（如果为空）
oldDir := filepath.Dir(item.FilePath)
if entries, err := os.ReadDir(oldDir); err == nil && len(entries) == 0 {
os.Remove(oldDir)
// 只有当目录不是知识库根目录时才删除（避免删除根目录）
if oldDir != m.basePath {
if err := os.Remove(oldDir); err != nil {
m.logger.Warn("删除空目录失败", zap.String("dir", oldDir), zap.Error(err))
}
}
}
}
@@ -686,6 +691,17 @@ func (m *Manager) DeleteItem(id string) error {
return fmt.Errorf("删除知识项失败: %w", err)
}
// 删除空目录（如果为空）
dir := filepath.Dir(filePath)
if entries, err := os.ReadDir(dir); err == nil && len(entries) == 0 {
// 只有当目录不是知识库根目录时才删除（避免删除根目录）
if dir != m.basePath {
if err := os.Remove(dir); err != nil {
m.logger.Warn("删除空目录失败", zap.String("dir", dir), zap.Error(err))
}
}
}
return nil
}