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

Add files via upload

Browse Source
This commit is contained in:
公明
2025-11-15 20:15:55 +08:00
committed by GitHub
parent 0455549c18
commit ac2c62f882
8 changed files with 274 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/storage/result_storage.go
+67 -40
View File
@@ -5,6 +5,7 @@ import (
"fmt"
"os"
"path/filepath"
"regexp"
"strings"
"sync"
"time"
@@ -16,22 +17,27 @@ import (
type ResultStorage interface {
// SaveResult 保存工具执行结果
SaveResult(executionID string, toolName string, result string) error
// GetResult 获取完整结果
GetResult(executionID string) (string, error)
// GetResultPage 分页获取结果
GetResultPage(executionID string, page int, limit int) (*ResultPage, error)
// SearchResult 搜索结果
SearchResult(executionID string, keyword string) ([]string, error)
// useRegex: 如果为 true，将 keyword 作为正则表达式使用；如果为 false，使用简单的字符串包含匹配
SearchResult(executionID string, keyword string, useRegex bool) ([]string, error)
// FilterResult 过滤结果
FilterResult(executionID string, filter string) ([]string, error)
// useRegex: 如果为 true，将 filter 作为正则表达式使用；如果为 false，使用简单的字符串包含匹配
FilterResult(executionID string, filter string, useRegex bool) ([]string, error)
// GetResultMetadata 获取结果元信息
GetResultMetadata(executionID string) (*ResultMetadata, error)
// GetResultPath 获取结果文件路径
GetResultPath(executionID string) string
// DeleteResult 删除结果
DeleteResult(executionID string) error
}
@@ -67,7 +73,7 @@ func NewFileResultStorage(baseDir string, logger *zap.Logger) (*FileResultStorag
if err := os.MkdirAll(baseDir, 0755); err != nil {
return nil, fmt.Errorf("创建存储目录失败: %w", err)
}
return &FileResultStorage{
baseDir: baseDir,
logger: logger,
@@ -88,13 +94,13 @@ func (s *FileResultStorage) getMetadataPath(executionID string) string {
func (s *FileResultStorage) SaveResult(executionID string, toolName string, result string) error {
s.mu.Lock()
defer s.mu.Unlock()
// 保存结果文件
resultPath := s.getResultPath(executionID)
if err := os.WriteFile(resultPath, []byte(result), 0644); err != nil {
return fmt.Errorf("保存结果文件失败: %w", err)
}
// 计算统计信息
lines := strings.Split(result, "\n")
metadata := &ResultMetadata{
@@ -104,25 +110,25 @@ func (s *FileResultStorage) SaveResult(executionID string, toolName string, resu
TotalLines: len(lines),
CreatedAt: time.Now(),
}
// 保存元数据
metadataPath := s.getMetadataPath(executionID)
metadataJSON, err := json.Marshal(metadata)
if err != nil {
return fmt.Errorf("序列化元数据失败: %w", err)
}
if err := os.WriteFile(metadataPath, metadataJSON, 0644); err != nil {
return fmt.Errorf("保存元数据文件失败: %w", err)
}
s.logger.Info("保存工具执行结果",
zap.String("executionID", executionID),
zap.String("toolName", toolName),
zap.Int("size", len(result)),
zap.Int("lines", len(lines)),
)
return nil
}
@@ -130,7 +136,7 @@ func (s *FileResultStorage) SaveResult(executionID string, toolName string, resu
func (s *FileResultStorage) GetResult(executionID string) (string, error) {
s.mu.RLock()
defer s.mu.RUnlock()
resultPath := s.getResultPath(executionID)
data, err := os.ReadFile(resultPath)
if err != nil {
@@ -139,7 +145,7 @@ func (s *FileResultStorage) GetResult(executionID string) (string, error) {
}
return "", fmt.Errorf("读取结果文件失败: %w", err)
}
return string(data), nil
}
@@ -147,7 +153,7 @@ func (s *FileResultStorage) GetResult(executionID string) (string, error) {
func (s *FileResultStorage) GetResultMetadata(executionID string) (*ResultMetadata, error) {
s.mu.RLock()
defer s.mu.RUnlock()
metadataPath := s.getMetadataPath(executionID)
data, err := os.ReadFile(metadataPath)
if err != nil {
@@ -156,12 +162,12 @@ func (s *FileResultStorage) GetResultMetadata(executionID string) (*ResultMetada
}
return nil, fmt.Errorf("读取元数据文件失败: %w", err)
}
var metadata ResultMetadata
if err := json.Unmarshal(data, &metadata); err != nil {
return nil, fmt.Errorf("解析元数据失败: %w", err)
}
return &metadata, nil
}
@@ -169,17 +175,17 @@ func (s *FileResultStorage) GetResultMetadata(executionID string) (*ResultMetada
func (s *FileResultStorage) GetResultPage(executionID string, page int, limit int) (*ResultPage, error) {
s.mu.RLock()
defer s.mu.RUnlock()
// 获取完整结果
result, err := s.GetResult(executionID)
if err != nil {
return nil, err
}
// 分割为行
lines := strings.Split(result, "\n")
totalLines := len(lines)
// 计算分页
totalPages := (totalLines + limit - 1) / limit
if page < 1 {
@@ -188,14 +194,14 @@ func (s *FileResultStorage) GetResultPage(executionID string, page int, limit in
if page > totalPages && totalPages > 0 {
page = totalPages
}
// 计算起始和结束索引
start := (page - 1) * limit
end := start + limit
if end > totalLines {
end = totalLines
}
// 提取指定页的行
var pageLines []string
if start < totalLines {
@@ -203,7 +209,7 @@ func (s *FileResultStorage) GetResultPage(executionID string, page int, limit in
} else {
pageLines = []string{}
}
return &ResultPage{
Lines: pageLines,
Page: page,
@@ -214,57 +220,78 @@ func (s *FileResultStorage) GetResultPage(executionID string, page int, limit in
}
// SearchResult 搜索结果
func (s *FileResultStorage) SearchResult(executionID string, keyword string) ([]string, error) {
func (s *FileResultStorage) SearchResult(executionID string, keyword string, useRegex bool) ([]string, error) {
s.mu.RLock()
defer s.mu.RUnlock()
// 获取完整结果
result, err := s.GetResult(executionID)
if err != nil {
return nil, err
}
// 如果使用正则表达式，先编译正则
var regex *regexp.Regexp
if useRegex {
compiledRegex, err := regexp.Compile(keyword)
if err != nil {
return nil, fmt.Errorf("无效的正则表达式: %w", err)
}
regex = compiledRegex
}
// 分割为行并搜索
lines := strings.Split(result, "\n")
var matchedLines []string
for _, line := range lines {
if strings.Contains(line, keyword) {
var matched bool
if useRegex {
matched = regex.MatchString(line)
} else {
matched = strings.Contains(line, keyword)
}
if matched {
matchedLines = append(matchedLines, line)
}
}
return matchedLines, nil
}
// FilterResult 过滤结果
func (s *FileResultStorage) FilterResult(executionID string, filter string) ([]string, error) {
func (s *FileResultStorage) FilterResult(executionID string, filter string, useRegex bool) ([]string, error) {
// 过滤和搜索逻辑相同，都是查找包含关键词的行
return s.SearchResult(executionID, filter)
return s.SearchResult(executionID, filter, useRegex)
}
// GetResultPath 获取结果文件路径
func (s *FileResultStorage) GetResultPath(executionID string) string {
return s.getResultPath(executionID)
}
// DeleteResult 删除结果
func (s *FileResultStorage) DeleteResult(executionID string) error {
s.mu.Lock()
defer s.mu.Unlock()
resultPath := s.getResultPath(executionID)
metadataPath := s.getMetadataPath(executionID)
// 删除结果文件
if err := os.Remove(resultPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("删除结果文件失败: %w", err)
}
// 删除元数据文件
if err := os.Remove(metadataPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("删除元数据文件失败: %w", err)
}
s.logger.Info("删除工具执行结果",
zap.String("executionID", executionID),
)
return nil
}
internal/storage/result_storage_test.go
+15 -5
View File
@@ -256,8 +256,8 @@ func TestFileResultStorage_SearchResult(t *testing.T) {
t.Fatalf("保存结果失败: %v", err)
}
// 搜索包含"error"的行
matchedLines, err := storage.SearchResult(executionID, "error")
// 搜索包含"error"的行（简单字符串匹配）
matchedLines, err := storage.SearchResult(executionID, "error", false)
if err != nil {
t.Fatalf("搜索失败: %v", err)
}
@@ -274,7 +274,7 @@ func TestFileResultStorage_SearchResult(t *testing.T) {
}
// 测试搜索不存在的关键词
noMatch, err := storage.SearchResult(executionID, "nonexistent")
noMatch, err := storage.SearchResult(executionID, "nonexistent", false)
if err != nil {
t.Fatalf("搜索失败: %v", err)
}
@@ -282,6 +282,16 @@ func TestFileResultStorage_SearchResult(t *testing.T) {
if len(noMatch) != 0 {
t.Errorf("搜索不存在的关键词应该返回空结果。实际: %d行", len(noMatch))
}
// 测试正则表达式搜索
regexMatched, err := storage.SearchResult(executionID, "error.*again", true)
if err != nil {
t.Fatalf("正则搜索失败: %v", err)
}
if len(regexMatched) != 1 {
t.Errorf("正则搜索结果数量不匹配。期望: 1, 实际: %d", len(regexMatched))
}
}
func TestFileResultStorage_FilterResult(t *testing.T) {
@@ -298,8 +308,8 @@ func TestFileResultStorage_FilterResult(t *testing.T) {
t.Fatalf("保存结果失败: %v", err)
}
// 过滤包含"warning"的行
filteredLines, err := storage.FilterResult(executionID, "warning")
// 过滤包含"warning"的行（简单字符串匹配）
filteredLines, err := storage.FilterResult(executionID, "warning", false)
if err != nil {
t.Fatalf("过滤失败: %v", err)
}