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

Add files via upload

Browse Source
This commit is contained in:
公明
2025-11-15 19:40:59 +08:00
committed by GitHub
parent 401a898947
commit 7fb3987309
10 changed files with 1668 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/agent/agent.go
+109 -8
View File
@@ -14,6 +14,7 @@ import (
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/storage"
"go.uber.org/zap"
)
@@ -21,21 +22,55 @@ import (
type Agent struct {
openAIClient *http.Client
config *config.OpenAIConfig
agentConfig *config.AgentConfig
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
logger *zap.Logger
maxIterations int
resultStorage ResultStorage // 结果存储
largeResultThreshold int // 大结果阈值（字节）
mu sync.RWMutex // 添加互斥锁以支持并发更新
toolNameMapping map[string]string // 工具名称映射：OpenAI格式 -> 原始格式（用于外部MCP工具）
}
// ResultStorage 结果存储接口（直接使用 storage 包的类型）
type ResultStorage interface {
SaveResult(executionID string, toolName string, result string) error
GetResult(executionID string) (string, error)
GetResultPage(executionID string, page int, limit int) (*storage.ResultPage, error)
SearchResult(executionID string, keyword string) ([]string, error)
FilterResult(executionID string, filter string) ([]string, error)
GetResultMetadata(executionID string) (*storage.ResultMetadata, error)
DeleteResult(executionID string) error
}
// NewAgent 创建新的Agent
func NewAgent(cfg *config.OpenAIConfig, mcpServer *mcp.Server, externalMCPMgr *mcp.ExternalMCPManager, logger *zap.Logger, maxIterations int) *Agent {
func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer *mcp.Server, externalMCPMgr *mcp.ExternalMCPManager, logger *zap.Logger, maxIterations int) *Agent {
// 如果 maxIterations 为 0 或负数，使用默认值 30
if maxIterations <= 0 {
maxIterations = 30
}
// 设置大结果阈值，默认50KB
largeResultThreshold := 50 * 1024
if agentCfg != nil && agentCfg.LargeResultThreshold > 0 {
largeResultThreshold = agentCfg.LargeResultThreshold
}
// 设置结果存储目录，默认tmp
resultStorageDir := "tmp"
if agentCfg != nil && agentCfg.ResultStorageDir != "" {
resultStorageDir = agentCfg.ResultStorageDir
}
// 初始化结果存储
var resultStorage ResultStorage
if resultStorageDir != "" {
// 导入storage包（避免循环依赖，使用接口）
// 这里需要在实际使用时初始化
// 暂时设为nil，在需要时初始化
}
// 配置HTTP Transport，优化连接管理和超时设置
transport := &http.Transport{
DialContext: (&net.Dialer{
@@ -57,15 +92,25 @@ func NewAgent(cfg *config.OpenAIConfig, mcpServer *mcp.Server, externalMCPMgr *m
Timeout: 30 * time.Minute, // 从5分钟增加到30分钟
Transport: transport,
},
config: cfg,
mcpServer: mcpServer,
externalMCPMgr: externalMCPMgr,
logger: logger,
maxIterations: maxIterations,
toolNameMapping: make(map[string]string), // 初始化工具名称映射
config: cfg,
agentConfig: agentCfg,
mcpServer: mcpServer,
externalMCPMgr: externalMCPMgr,
logger: logger,
maxIterations: maxIterations,
resultStorage: resultStorage,
largeResultThreshold: largeResultThreshold,
toolNameMapping: make(map[string]string), // 初始化工具名称映射
}
}
// SetResultStorage 设置结果存储（用于避免循环依赖）
func (a *Agent) SetResultStorage(storage ResultStorage) {
a.mu.Lock()
defer a.mu.Unlock()
a.resultStorage = storage
}
// ChatMessage 聊天消息
type ChatMessage struct {
Role string `json:"role"`
@@ -1037,14 +1082,70 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
resultText.WriteString(content.Text)
resultText.WriteString("\n")
}
resultStr := resultText.String()
resultSize := len(resultStr)
// 检测大结果并保存
a.mu.RLock()
threshold := a.largeResultThreshold
storage := a.resultStorage
a.mu.RUnlock()
if resultSize > threshold && storage != nil {
// 异步保存大结果
go func() {
if err := storage.SaveResult(executionID, toolName, resultStr); err != nil {
a.logger.Warn("保存大结果失败",
zap.String("executionID", executionID),
zap.String("toolName", toolName),
zap.Error(err),
)
} else {
a.logger.Info("大结果已保存",
zap.String("executionID", executionID),
zap.String("toolName", toolName),
zap.Int("size", resultSize),
)
}
}()
// 返回最小化通知
lines := strings.Split(resultStr, "\n")
notification := a.formatMinimalNotification(executionID, toolName, resultSize, len(lines))
return &ToolExecutionResult{
Result: notification,
ExecutionID: executionID,
IsError: result != nil && result.IsError,
}, nil
}
return &ToolExecutionResult{
Result: resultText.String(),
Result: resultStr,
ExecutionID: executionID,
IsError: result != nil && result.IsError,
}, nil
}
// formatMinimalNotification 格式化最小化通知
func (a *Agent) formatMinimalNotification(executionID string, toolName string, size int, lineCount int) string {
var sb strings.Builder
sb.WriteString(fmt.Sprintf("工具执行完成。结果已保存（ID: %s）。\n\n", executionID))
sb.WriteString("结果信息：\n")
sb.WriteString(fmt.Sprintf(" - 工具: %s\n", toolName))
sb.WriteString(fmt.Sprintf(" - 大小: %d 字节 (%.2f KB)\n", size, float64(size)/1024))
sb.WriteString(fmt.Sprintf(" - 行数: %d 行\n", lineCount))
sb.WriteString("\n")
sb.WriteString("使用以下工具查询完整结果：\n")
sb.WriteString(fmt.Sprintf(" - 查询第一页: query_execution_result(execution_id=\"%s\", page=1, limit=100)\n", executionID))
sb.WriteString(fmt.Sprintf(" - 搜索关键词: query_execution_result(execution_id=\"%s\", search=\"关键词\")\n", executionID))
sb.WriteString(fmt.Sprintf(" - 过滤条件: query_execution_result(execution_id=\"%s\", filter=\"error\")\n", executionID))
return sb.String()
}
// UpdateConfig 更新OpenAI配置
func (a *Agent) UpdateConfig(cfg *config.OpenAIConfig) {
a.mu.Lock()
internal/agent/agent_test.go
+284
View File
@@ -0,0 +1,284 @@
package agent
import (
"os"
"path/filepath"
"strings"
"testing"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/storage"
"go.uber.org/zap"
)
// setupTestAgent 创建测试用的Agent
func setupTestAgent(t *testing.T) (*Agent, *storage.FileResultStorage) {
logger := zap.NewNop()
mcpServer := mcp.NewServer(logger)
openAICfg := &config.OpenAIConfig{
APIKey: "test-key",
BaseURL: "https://api.test.com/v1",
Model: "test-model",
}
agentCfg := &config.AgentConfig{
MaxIterations: 10,
LargeResultThreshold: 100, // 设置较小的阈值便于测试
ResultStorageDir: "",
}
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 10)
// 创建测试存储
tmpDir := filepath.Join(os.TempDir(), "test_agent_storage_"+time.Now().Format("20060102_150405"))
testStorage, err := storage.NewFileResultStorage(tmpDir, logger)
if err != nil {
t.Fatalf("创建测试存储失败: %v", err)
}
agent.SetResultStorage(testStorage)
return agent, testStorage
}
func TestAgent_FormatMinimalNotification(t *testing.T) {
agent, testStorage := setupTestAgent(t)
_ = testStorage // 避免未使用变量警告
executionID := "test_exec_001"
toolName := "nmap_scan"
size := 50000
lineCount := 1000
notification := agent.formatMinimalNotification(executionID, toolName, size, lineCount)
// 验证通知包含必要信息
if !strings.Contains(notification, executionID) {
t.Errorf("通知中应该包含执行ID: %s", executionID)
}
if !strings.Contains(notification, toolName) {
t.Errorf("通知中应该包含工具名称: %s", toolName)
}
if !strings.Contains(notification, "50000") {
t.Errorf("通知中应该包含大小信息")
}
if !strings.Contains(notification, "1000") {
t.Errorf("通知中应该包含行数信息")
}
if !strings.Contains(notification, "query_execution_result") {
t.Errorf("通知中应该包含查询工具的使用说明")
}
}
func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
agent, _ := setupTestAgent(t)
// 创建模拟的MCP工具结果（大结果）
largeResult := &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: strings.Repeat("This is a test line with some content.\n", 1000), // 约50KB
},
},
IsError: false,
}
// 模拟MCP服务器返回大结果
// 由于我们需要模拟CallTool的行为，这里需要创建一个mock或者使用实际的MCP服务器
// 为了简化测试，我们直接测试结果处理逻辑
// 设置阈值
agent.mu.Lock()
agent.largeResultThreshold = 1000 // 设置较小的阈值
agent.mu.Unlock()
// 创建执行ID
executionID := "test_exec_large_001"
toolName := "test_tool"
// 格式化结果
var resultText strings.Builder
for _, content := range largeResult.Content {
resultText.WriteString(content.Text)
resultText.WriteString("\n")
}
resultStr := resultText.String()
resultSize := len(resultStr)
// 检测大结果并保存
agent.mu.RLock()
threshold := agent.largeResultThreshold
storage := agent.resultStorage
agent.mu.RUnlock()
if resultSize > threshold && storage != nil {
// 保存大结果
err := storage.SaveResult(executionID, toolName, resultStr)
if err != nil {
t.Fatalf("保存大结果失败: %v", err)
}
// 生成通知
lines := strings.Split(resultStr, "\n")
notification := agent.formatMinimalNotification(executionID, toolName, resultSize, len(lines))
// 验证通知格式
if !strings.Contains(notification, executionID) {
t.Errorf("通知中应该包含执行ID")
}
// 验证结果已保存
savedResult, err := storage.GetResult(executionID)
if err != nil {
t.Fatalf("获取保存的结果失败: %v", err)
}
if savedResult != resultStr {
t.Errorf("保存的结果与原始结果不匹配")
}
} else {
t.Fatal("大结果应该被检测到并保存")
}
}
func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
agent, _ := setupTestAgent(t)
// 创建小结果
smallResult := &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: "Small result content",
},
},
IsError: false,
}
// 设置较大的阈值
agent.mu.Lock()
agent.largeResultThreshold = 100000 // 100KB
agent.mu.Unlock()
// 格式化结果
var resultText strings.Builder
for _, content := range smallResult.Content {
resultText.WriteString(content.Text)
resultText.WriteString("\n")
}
resultStr := resultText.String()
resultSize := len(resultStr)
// 检测大结果
agent.mu.RLock()
threshold := agent.largeResultThreshold
storage := agent.resultStorage
agent.mu.RUnlock()
if resultSize > threshold && storage != nil {
t.Fatal("小结果不应该被保存")
}
// 小结果应该直接返回
if resultSize <= threshold {
// 这是预期的行为
if resultStr == "" {
t.Fatal("小结果应该直接返回，不应该为空")
}
}
}
func TestAgent_SetResultStorage(t *testing.T) {
agent, _ := setupTestAgent(t)
// 创建新的存储
tmpDir := filepath.Join(os.TempDir(), "test_new_storage_"+time.Now().Format("20060102_150405"))
newStorage, err := storage.NewFileResultStorage(tmpDir, zap.NewNop())
if err != nil {
t.Fatalf("创建新存储失败: %v", err)
}
// 设置新存储
agent.SetResultStorage(newStorage)
// 验证存储已更新
agent.mu.RLock()
currentStorage := agent.resultStorage
agent.mu.RUnlock()
if currentStorage != newStorage {
t.Fatal("存储未正确更新")
}
// 清理
os.RemoveAll(tmpDir)
}
func TestAgent_NewAgent_DefaultValues(t *testing.T) {
logger := zap.NewNop()
mcpServer := mcp.NewServer(logger)
openAICfg := &config.OpenAIConfig{
APIKey: "test-key",
BaseURL: "https://api.test.com/v1",
Model: "test-model",
}
// 测试默认配置
agent := NewAgent(openAICfg, nil, mcpServer, nil, logger, 0)
if agent.maxIterations != 30 {
t.Errorf("默认迭代次数不匹配。期望: 30, 实际: %d", agent.maxIterations)
}
agent.mu.RLock()
threshold := agent.largeResultThreshold
agent.mu.RUnlock()
if threshold != 50*1024 {
t.Errorf("默认阈值不匹配。期望: %d, 实际: %d", 50*1024, threshold)
}
}
func TestAgent_NewAgent_CustomConfig(t *testing.T) {
logger := zap.NewNop()
mcpServer := mcp.NewServer(logger)
openAICfg := &config.OpenAIConfig{
APIKey: "test-key",
BaseURL: "https://api.test.com/v1",
Model: "test-model",
}
agentCfg := &config.AgentConfig{
MaxIterations: 20,
LargeResultThreshold: 100 * 1024, // 100KB
ResultStorageDir: "custom_tmp",
}
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 15)
if agent.maxIterations != 15 {
t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations)
}
agent.mu.RLock()
threshold := agent.largeResultThreshold
agent.mu.RUnlock()
if threshold != 100*1024 {
t.Errorf("阈值不匹配。期望: %d, 实际: %d", 100*1024, threshold)
}
}