mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-19 18:38:00 +02:00
Add files via upload
This commit is contained in:
公明
@@ -0,0 +1,297 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// ResultStorage 结果存储接口
|
||||
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 搜索结果
|
||||
// useRegex: 如果为 true，将 keyword 作为正则表达式使用；如果为 false，使用简单的字符串包含匹配
|
||||
SearchResult(executionID string, keyword string, useRegex bool) ([]string, error)
|
||||
|
||||
// FilterResult 过滤结果
|
||||
// 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
|
||||
}
|
||||
|
||||
// ResultPage 分页结果
|
||||
type ResultPage struct {
|
||||
Lines []string `json:"lines"`
|
||||
Page int `json:"page"`
|
||||
Limit int `json:"limit"`
|
||||
TotalLines int `json:"total_lines"`
|
||||
TotalPages int `json:"total_pages"`
|
||||
}
|
||||
|
||||
// ResultMetadata 结果元信息
|
||||
type ResultMetadata struct {
|
||||
ExecutionID string `json:"execution_id"`
|
||||
ToolName string `json:"tool_name"`
|
||||
TotalSize int `json:"total_size"`
|
||||
TotalLines int `json:"total_lines"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// FileResultStorage 基于文件的结果存储实现
|
||||
type FileResultStorage struct {
|
||||
baseDir string
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewFileResultStorage 创建新的文件结果存储
|
||||
func NewFileResultStorage(baseDir string, logger *zap.Logger) (*FileResultStorage, error) {
|
||||
// 确保目录存在
|
||||
if err := os.MkdirAll(baseDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("创建存储目录失败: %w", err)
|
||||
}
|
||||
|
||||
return &FileResultStorage{
|
||||
baseDir: baseDir,
|
||||
logger: logger,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// getResultPath 获取结果文件路径
|
||||
func (s *FileResultStorage) getResultPath(executionID string) string {
|
||||
return filepath.Join(s.baseDir, executionID+".txt")
|
||||
}
|
||||
|
||||
// getMetadataPath 获取元数据文件路径
|
||||
func (s *FileResultStorage) getMetadataPath(executionID string) string {
|
||||
return filepath.Join(s.baseDir, executionID+".meta.json")
|
||||
}
|
||||
|
||||
// SaveResult 保存工具执行结果
|
||||
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{
|
||||
ExecutionID: executionID,
|
||||
ToolName: toolName,
|
||||
TotalSize: len(result),
|
||||
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
|
||||
}
|
||||
|
||||
// GetResult 获取完整结果
|
||||
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 {
|
||||
if os.IsNotExist(err) {
|
||||
return "", fmt.Errorf("结果不存在: %s", executionID)
|
||||
}
|
||||
return "", fmt.Errorf("读取结果文件失败: %w", err)
|
||||
}
|
||||
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
// GetResultMetadata 获取结果元信息
|
||||
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 {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("结果不存在: %s", executionID)
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
// GetResultPage 分页获取结果
|
||||
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 {
|
||||
page = 1
|
||||
}
|
||||
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 {
|
||||
pageLines = lines[start:end]
|
||||
} else {
|
||||
pageLines = []string{}
|
||||
}
|
||||
|
||||
return &ResultPage{
|
||||
Lines: pageLines,
|
||||
Page: page,
|
||||
Limit: limit,
|
||||
TotalLines: totalLines,
|
||||
TotalPages: totalPages,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// SearchResult 搜索结果
|
||||
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 {
|
||||
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, useRegex bool) ([]string, error) {
|
||||
// 过滤和搜索逻辑相同，都是查找包含关键词的行
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user