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

Add files via upload

Browse Source
This commit is contained in:
公明
2025-11-14 01:34:16 +08:00
committed by GitHub
parent d61c85e73e
commit 1b14070cee
9 changed files with 1315 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/app/app.go
+3 -4
View File
@@ -60,8 +60,8 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
return nil, fmt.Errorf("初始化数据库失败: %w", err)
}
// 创建MCP服务器
mcpServer := mcp.NewServer(log.Logger)
// 创建MCP服务器（带数据库持久化）
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
// 创建安全工具执行器
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
@@ -91,7 +91,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
// 创建处理器
agentHandler := handler.NewAgentHandler(agent, db, log.Logger)
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, log.Logger)
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, db, log.Logger)
conversationHandler := handler.NewConversationHandler(db, log.Logger)
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, log.Logger)
@@ -188,7 +188,6 @@ func setupRoutes(
protected.GET("/monitor", monitorHandler.Monitor)
protected.GET("/monitor/execution/:id", monitorHandler.GetExecution)
protected.GET("/monitor/stats", monitorHandler.GetStats)
protected.GET("/monitor/vulnerabilities", monitorHandler.GetVulnerabilities)
// 配置管理
protected.GET("/config", configHandler.GetConfig)
internal/database/database.go
+37
View File
@@ -75,12 +75,41 @@ func (db *DB) initTables() error {
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
);`
// 创建工具执行记录表
createToolExecutionsTable := `
CREATE TABLE IF NOT EXISTS tool_executions (
id TEXT PRIMARY KEY,
tool_name TEXT NOT NULL,
arguments TEXT NOT NULL,
status TEXT NOT NULL,
result TEXT,
error TEXT,
start_time DATETIME NOT NULL,
end_time DATETIME,
duration_ms INTEGER,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);`
// 创建工具统计表
createToolStatsTable := `
CREATE TABLE IF NOT EXISTS tool_stats (
tool_name TEXT PRIMARY KEY,
total_calls INTEGER NOT NULL DEFAULT 0,
success_calls INTEGER NOT NULL DEFAULT 0,
failed_calls INTEGER NOT NULL DEFAULT 0,
last_call_time DATETIME,
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);`
// 创建索引
createIndexes := `
CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id);
CREATE INDEX IF NOT EXISTS idx_conversations_updated_at ON conversations(updated_at);
CREATE INDEX IF NOT EXISTS idx_process_details_message_id ON process_details(message_id);
CREATE INDEX IF NOT EXISTS idx_process_details_conversation_id ON process_details(conversation_id);
CREATE INDEX IF NOT EXISTS idx_tool_executions_tool_name ON tool_executions(tool_name);
CREATE INDEX IF NOT EXISTS idx_tool_executions_start_time ON tool_executions(start_time);
CREATE INDEX IF NOT EXISTS idx_tool_executions_status ON tool_executions(status);
`
if _, err := db.Exec(createConversationsTable); err != nil {
@@ -95,6 +124,14 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建process_details表失败: %w", err)
}
if _, err := db.Exec(createToolExecutionsTable); err != nil {
return fmt.Errorf("创建tool_executions表失败: %w", err)
}
if _, err := db.Exec(createToolStatsTable); err != nil {
return fmt.Errorf("创建tool_stats表失败: %w", err)
}
if _, err := db.Exec(createIndexes); err != nil {
return fmt.Errorf("创建索引失败: %w", err)
}
internal/database/monitor.go
+309
View File
@@ -0,0 +1,309 @@
package database
import (
"database/sql"
"encoding/json"
"time"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
// SaveToolExecution 保存工具执行记录
func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
argsJSON, err := json.Marshal(exec.Arguments)
if err != nil {
db.logger.Warn("序列化执行参数失败", zap.Error(err))
argsJSON = []byte("{}")
}
var resultJSON sql.NullString
if exec.Result != nil {
resultBytes, err := json.Marshal(exec.Result)
if err != nil {
db.logger.Warn("序列化执行结果失败", zap.Error(err))
} else {
resultJSON = sql.NullString{String: string(resultBytes), Valid: true}
}
}
var errorText sql.NullString
if exec.Error != "" {
errorText = sql.NullString{String: exec.Error, Valid: true}
}
var endTime sql.NullTime
if exec.EndTime != nil {
endTime = sql.NullTime{Time: *exec.EndTime, Valid: true}
}
var durationMs sql.NullInt64
if exec.Duration > 0 {
durationMs = sql.NullInt64{Int64: exec.Duration.Milliseconds(), Valid: true}
}
query := `
INSERT OR REPLACE INTO tool_executions
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`
_, err = db.Exec(query,
exec.ID,
exec.ToolName,
string(argsJSON),
exec.Status,
resultJSON,
errorText,
exec.StartTime,
endTime,
durationMs,
time.Now(),
)
if err != nil {
db.logger.Error("保存工具执行记录失败", zap.Error(err), zap.String("executionId", exec.ID))
return err
}
return nil
}
// LoadToolExecutions 加载所有工具执行记录
func (db *DB) LoadToolExecutions() ([]*mcp.ToolExecution, error) {
query := `
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms
FROM tool_executions
ORDER BY start_time DESC
LIMIT 1000
`
rows, err := db.Query(query)
if err != nil {
return nil, err
}
defer rows.Close()
var executions []*mcp.ToolExecution
for rows.Next() {
var exec mcp.ToolExecution
var argsJSON string
var resultJSON sql.NullString
var errorText sql.NullString
var endTime sql.NullTime
var durationMs sql.NullInt64
err := rows.Scan(
&exec.ID,
&exec.ToolName,
&argsJSON,
&exec.Status,
&resultJSON,
&errorText,
&exec.StartTime,
&endTime,
&durationMs,
)
if err != nil {
db.logger.Warn("加载执行记录失败", zap.Error(err))
continue
}
// 解析参数
if err := json.Unmarshal([]byte(argsJSON), &exec.Arguments); err != nil {
db.logger.Warn("解析执行参数失败", zap.Error(err))
exec.Arguments = make(map[string]interface{})
}
// 解析结果
if resultJSON.Valid && resultJSON.String != "" {
var result mcp.ToolResult
if err := json.Unmarshal([]byte(resultJSON.String), &result); err != nil {
db.logger.Warn("解析执行结果失败", zap.Error(err))
} else {
exec.Result = &result
}
}
// 设置错误
if errorText.Valid {
exec.Error = errorText.String
}
// 设置结束时间
if endTime.Valid {
exec.EndTime = &endTime.Time
}
// 设置持续时间
if durationMs.Valid {
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
}
executions = append(executions, &exec)
}
return executions, nil
}
// GetToolExecution 根据ID获取单条工具执行记录
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
query := `
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms
FROM tool_executions
WHERE id = ?
`
row := db.QueryRow(query, id)
var exec mcp.ToolExecution
var argsJSON string
var resultJSON sql.NullString
var errorText sql.NullString
var endTime sql.NullTime
var durationMs sql.NullInt64
err := row.Scan(
&exec.ID,
&exec.ToolName,
&argsJSON,
&exec.Status,
&resultJSON,
&errorText,
&exec.StartTime,
&endTime,
&durationMs,
)
if err != nil {
return nil, err
}
if err := json.Unmarshal([]byte(argsJSON), &exec.Arguments); err != nil {
db.logger.Warn("解析执行参数失败", zap.Error(err))
exec.Arguments = make(map[string]interface{})
}
if resultJSON.Valid && resultJSON.String != "" {
var result mcp.ToolResult
if err := json.Unmarshal([]byte(resultJSON.String), &result); err != nil {
db.logger.Warn("解析执行结果失败", zap.Error(err))
} else {
exec.Result = &result
}
}
if errorText.Valid {
exec.Error = errorText.String
}
if endTime.Valid {
exec.EndTime = &endTime.Time
}
if durationMs.Valid {
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
}
return &exec, nil
}
// SaveToolStats 保存工具统计信息
func (db *DB) SaveToolStats(toolName string, stats *mcp.ToolStats) error {
var lastCallTime sql.NullTime
if stats.LastCallTime != nil {
lastCallTime = sql.NullTime{Time: *stats.LastCallTime, Valid: true}
}
query := `
INSERT OR REPLACE INTO tool_stats
(tool_name, total_calls, success_calls, failed_calls, last_call_time, updated_at)
VALUES (?, ?, ?, ?, ?, ?)
`
_, err := db.Exec(query,
toolName,
stats.TotalCalls,
stats.SuccessCalls,
stats.FailedCalls,
lastCallTime,
time.Now(),
)
if err != nil {
db.logger.Error("保存工具统计信息失败", zap.Error(err), zap.String("toolName", toolName))
return err
}
return nil
}
// LoadToolStats 加载所有工具统计信息
func (db *DB) LoadToolStats() (map[string]*mcp.ToolStats, error) {
query := `
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
FROM tool_stats
`
rows, err := db.Query(query)
if err != nil {
return nil, err
}
defer rows.Close()
stats := make(map[string]*mcp.ToolStats)
for rows.Next() {
var stat mcp.ToolStats
var lastCallTime sql.NullTime
err := rows.Scan(
&stat.ToolName,
&stat.TotalCalls,
&stat.SuccessCalls,
&stat.FailedCalls,
&lastCallTime,
)
if err != nil {
db.logger.Warn("加载统计信息失败", zap.Error(err))
continue
}
if lastCallTime.Valid {
stat.LastCallTime = &lastCallTime.Time
}
stats[stat.ToolName] = &stat
}
return stats, nil
}
// UpdateToolStats 更新工具统计信息（累加模式）
func (db *DB) UpdateToolStats(toolName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error {
var lastCallTimeSQL sql.NullTime
if lastCallTime != nil {
lastCallTimeSQL = sql.NullTime{Time: *lastCallTime, Valid: true}
}
query := `
INSERT INTO tool_stats (tool_name, total_calls, success_calls, failed_calls, last_call_time, updated_at)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT(tool_name) DO UPDATE SET
total_calls = total_calls + ?,
success_calls = success_calls + ?,
failed_calls = failed_calls + ?,
last_call_time = COALESCE(?, last_call_time),
updated_at = ?
`
_, err := db.Exec(query,
toolName, totalCalls, successCalls, failedCalls, lastCallTimeSQL, time.Now(),
totalCalls, successCalls, failedCalls, lastCallTimeSQL, time.Now(),
)
if err != nil {
db.logger.Error("更新工具统计信息失败", zap.Error(err), zap.String("toolName", toolName))
return err
}
return nil
}
internal/handler/monitor.go
+41 -37
View File
@@ -4,6 +4,7 @@ import (
"net/http"
"time"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
@@ -14,61 +15,70 @@ import (
type MonitorHandler struct {
mcpServer *mcp.Server
executor *security.Executor
db *database.DB
logger *zap.Logger
vulns []security.Vulnerability
}
// NewMonitorHandler 创建新的监控处理器
func NewMonitorHandler(mcpServer *mcp.Server, executor *security.Executor, logger *zap.Logger) *MonitorHandler {
func NewMonitorHandler(mcpServer *mcp.Server, executor *security.Executor, db *database.DB, logger *zap.Logger) *MonitorHandler {
return &MonitorHandler{
mcpServer: mcpServer,
executor: executor,
db: db,
logger: logger,
vulns: []security.Vulnerability{},
}
}
// MonitorResponse 监控响应
type MonitorResponse struct {
Executions []*mcp.ToolExecution `json:"executions"`
Stats map[string]*mcp.ToolStats `json:"stats"`
Vulnerabilities []security.Vulnerability `json:"vulnerabilities"`
Report map[string]interface{} `json:"report"`
Timestamp time.Time `json:"timestamp"`
Executions []*mcp.ToolExecution `json:"executions"`
Stats map[string]*mcp.ToolStats `json:"stats"`
Timestamp time.Time `json:"timestamp"`
}
// Monitor 获取监控信息
func (h *MonitorHandler) Monitor(c *gin.Context) {
// 获取所有执行记录
executions := h.mcpServer.GetAllExecutions()
// 分析执行结果，提取漏洞
for _, exec := range executions {
if exec.Status == "completed" && exec.Result != nil {
vulns := h.executor.AnalyzeResults(exec.ToolName, exec.Result)
h.vulns = append(h.vulns, vulns...)
}
}
// 获取统计信息
stats := h.mcpServer.GetStats()
// 生成报告
report := h.executor.GetVulnerabilityReport(h.vulns)
executions := h.loadExecutions()
stats := h.loadStats()
c.JSON(http.StatusOK, MonitorResponse{
Executions: executions,
Stats: stats,
Vulnerabilities: h.vulns,
Report: report,
Timestamp: time.Now(),
Executions: executions,
Stats: stats,
Timestamp: time.Now(),
})
}
func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution {
if h.db == nil {
return h.mcpServer.GetAllExecutions()
}
executions, err := h.db.LoadToolExecutions()
if err != nil {
h.logger.Warn("从数据库加载执行记录失败，回退到内存数据", zap.Error(err))
return h.mcpServer.GetAllExecutions()
}
return executions
}
func (h *MonitorHandler) loadStats() map[string]*mcp.ToolStats {
if h.db == nil {
return h.mcpServer.GetStats()
}
stats, err := h.db.LoadToolStats()
if err != nil {
h.logger.Warn("从数据库加载统计信息失败，回退到内存数据", zap.Error(err))
return h.mcpServer.GetStats()
}
return stats
}
// GetExecution 获取特定执行记录
func (h *MonitorHandler) GetExecution(c *gin.Context) {
id := c.Param("id")
exec, exists := h.mcpServer.GetExecution(id)
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "执行记录未找到"})
@@ -80,13 +90,7 @@ func (h *MonitorHandler) GetExecution(c *gin.Context) {
// GetStats 获取统计信息
func (h *MonitorHandler) GetStats(c *gin.Context) {
stats := h.mcpServer.GetStats()
stats := h.loadStats()
c.JSON(http.StatusOK, stats)
}
// GetVulnerabilities 获取漏洞列表
func (h *MonitorHandler) GetVulnerabilities(c *gin.Context) {
report := h.executor.GetVulnerabilityReport(h.vulns)
c.JSON(http.StatusOK, report)
}
internal/mcp/server.go
+322 -183
View File
@@ -15,14 +15,25 @@ import (
"go.uber.org/zap"
)
// MonitorStorage 监控数据存储接口
type MonitorStorage interface {
SaveToolExecution(exec *ToolExecution) error
LoadToolExecutions() ([]*ToolExecution, error)
GetToolExecution(id string) (*ToolExecution, error)
SaveToolStats(toolName string, stats *ToolStats) error
LoadToolStats() (map[string]*ToolStats, error)
UpdateToolStats(toolName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error
}
// Server MCP服务器
type Server struct {
tools map[string]ToolHandler
toolDefs map[string]Tool // 工具定义
toolDefs map[string]Tool // 工具定义
executions map[string]*ToolExecution
stats map[string]*ToolStats
prompts map[string]*Prompt // 提示词模板
prompts map[string]*Prompt // 提示词模板
resources map[string]*Resource // 资源
storage MonitorStorage // 可选的持久化存储
mu sync.RWMutex
logger *zap.Logger
}
@@ -32,6 +43,11 @@ type ToolHandler func(ctx context.Context, args map[string]interface{}) (*ToolRe
// NewServer 创建新的MCP服务器
func NewServer(logger *zap.Logger) *Server {
return NewServerWithStorage(logger, nil)
}
// NewServerWithStorage 创建新的MCP服务器（带持久化存储）
func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
s := &Server{
tools: make(map[string]ToolHandler),
toolDefs: make(map[string]Tool),
@@ -39,13 +55,14 @@ func NewServer(logger *zap.Logger) *Server {
stats: make(map[string]*ToolStats),
prompts: make(map[string]*Prompt),
resources: make(map[string]*Resource),
storage: storage,
logger: logger,
}
// 初始化默认提示词和资源
s.initDefaultPrompts()
s.initDefaultResources()
return s
}
@@ -55,7 +72,7 @@ func (s *Server) RegisterTool(tool Tool, handler ToolHandler) {
defer s.mu.Unlock()
s.tools[tool.Name] = handler
s.toolDefs[tool.Name] = tool
// 自动为工具创建资源文档
resourceURI := fmt.Sprintf("tool://%s", tool.Name)
s.resources[resourceURI] = &Resource{
@@ -70,11 +87,11 @@ func (s *Server) RegisterTool(tool Tool, handler ToolHandler) {
func (s *Server) ClearTools() {
s.mu.Lock()
defer s.mu.Unlock()
// 清空工具和工具定义
s.tools = make(map[string]ToolHandler)
s.toolDefs = make(map[string]Tool)
// 清空工具相关的资源（保留其他资源）
newResources := make(map[string]*Resource)
for uri, resource := range s.resources {
@@ -107,7 +124,7 @@ func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) {
// 处理消息
response := s.handleMessage(&msg)
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(response)
}
@@ -116,7 +133,7 @@ func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) {
func (s *Server) handleMessage(msg *Message) *Message {
// 检查是否是通知（notification）- 通知没有id字段，不需要响应
isNotification := msg.ID.Value() == nil || msg.ID.String() == ""
// 如果不是通知且ID为空，生成新的UUID
if !isNotification && msg.ID.String() == "" {
msg.ID = MessageID{value: uuid.New().String()}
@@ -239,7 +256,6 @@ func (s *Server) handleCallTool(msg *Message) *Message {
}
}
// 创建执行记录
executionID := uuid.New().String()
execution := &ToolExecution{
ID: executionID,
@@ -253,10 +269,12 @@ func (s *Server) handleCallTool(msg *Message) *Message {
s.executions[executionID] = execution
s.mu.Unlock()
// 更新统计
s.updateStats(req.Name, false)
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
// 执行工具
s.mu.RLock()
handler, exists := s.tools[req.Name]
s.mu.RUnlock()
@@ -266,6 +284,19 @@ func (s *Server) handleCallTool(msg *Message) *Message {
execution.Error = "Tool not found"
now := time.Now()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
s.mu.Lock()
delete(s.executions, executionID)
s.mu.Unlock()
}
s.updateStats(req.Name, true)
return &Message{
ID: msg.ID,
Type: MessageTypeError,
@@ -274,7 +305,6 @@ func (s *Server) handleCallTool(msg *Message) *Message {
}
}
// 同步执行所有工具，确保错误能正确返回
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
defer cancel()
@@ -284,24 +314,63 @@ func (s *Server) handleCallTool(msg *Message) *Message {
)
result, err := handler(ctx, req.Arguments)
s.mu.Lock()
now := time.Now()
var failed bool
var finalResult *ToolResult
s.mu.Lock()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
if err != nil {
execution.Status = "failed"
execution.Error = err.Error()
s.updateStats(req.Name, true)
failed = true
} else if result != nil && result.IsError {
execution.Status = "failed"
if len(result.Content) > 0 {
execution.Error = result.Content[0].Text
} else {
execution.Error = "工具执行返回错误结果"
}
execution.Result = result
failed = true
} else {
execution.Status = "completed"
if result == nil {
result = &ToolResult{
Content: []Content{
{Type: "text", Text: "工具执行完成，但未返回结果"},
},
}
}
execution.Result = result
failed = false
}
finalResult = execution.Result
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
s.updateStats(req.Name, failed)
if s.storage != nil {
s.mu.Lock()
delete(s.executions, executionID)
s.mu.Unlock()
}
if err != nil {
s.logger.Error("工具执行失败",
zap.String("toolName", req.Name),
zap.Error(err),
)
// 返回错误结果
errorResult, _ := json.Marshal(CallToolResponse{
Content: []Content{
{Type: "text", Text: fmt.Sprintf("工具执行失败: %v", err)},
@@ -315,40 +384,42 @@ func (s *Server) handleCallTool(msg *Message) *Message {
Result: errorResult,
}
}
// 检查result是否为错误
if result != nil && result.IsError {
execution.Status = "failed"
if len(result.Content) > 0 {
execution.Error = result.Content[0].Text
if finalResult != nil && finalResult.IsError {
s.logger.Warn("工具执行返回错误结果",
zap.String("toolName", req.Name),
)
errorResult, _ := json.Marshal(CallToolResponse{
Content: finalResult.Content,
IsError: true,
})
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: errorResult,
}
s.updateStats(req.Name, true)
} else {
execution.Status = "completed"
execution.Result = result
s.updateStats(req.Name, false)
}
s.mu.Unlock()
// 返回执行结果
if result == nil {
result = &ToolResult{
if finalResult == nil {
finalResult = &ToolResult{
Content: []Content{
{Type: "text", Text: "工具执行完成，但未返回结果"},
},
}
}
resultJSON, _ := json.Marshal(CallToolResponse{
Content: result.Content,
IsError: result.IsError,
Content: finalResult.Content,
IsError: false,
})
s.logger.Info("工具执行完成",
zap.String("toolName", req.Name),
zap.Bool("isError", result.IsError),
zap.Bool("isError", finalResult.IsError),
)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
@@ -359,6 +430,25 @@ func (s *Server) handleCallTool(msg *Message) *Message {
// updateStats 更新统计信息
func (s *Server) updateStats(toolName string, failed bool) {
now := time.Now()
if s.storage != nil {
totalCalls := 1
successCalls := 0
failedCalls := 0
if failed {
failedCalls = 1
} else {
successCalls = 1
}
if err := s.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil {
s.logger.Warn("保存统计信息到数据库失败", zap.Error(err))
}
return
}
s.mu.Lock()
defer s.mu.Unlock()
if s.stats[toolName] == nil {
s.stats[toolName] = &ToolStats{
ToolName: toolName,
@@ -367,7 +457,6 @@ func (s *Server) updateStats(toolName string, failed bool) {
stats := s.stats[toolName]
stats.TotalCalls++
now := time.Now()
stats.LastCallTime = &now
if failed {
@@ -377,41 +466,129 @@ func (s *Server) updateStats(toolName string, failed bool) {
}
}
// GetExecution 获取执行记录
// GetExecution 获取执行记录（先从内存查找，再从数据库查找）
func (s *Server) GetExecution(id string) (*ToolExecution, bool) {
s.mu.RLock()
defer s.mu.RUnlock()
exec, exists := s.executions[id]
return exec, exists
s.mu.RUnlock()
if exists {
return exec, true
}
if s.storage != nil {
exec, err := s.storage.GetToolExecution(id)
if err == nil {
return exec, true
}
}
return nil, false
}
// GetAllExecutions 获取所有执行记录
// loadHistoricalData 从数据库加载历史数据
func (s *Server) loadHistoricalData() {
if s.storage == nil {
return
}
// 加载历史执行记录（最近1000条）
executions, err := s.storage.LoadToolExecutions()
if err != nil {
s.logger.Warn("加载历史执行记录失败", zap.Error(err))
} else {
s.mu.Lock()
for _, exec := range executions {
// 只加载最近1000条，避免内存占用过大
if len(s.executions) < 1000 {
s.executions[exec.ID] = exec
}
}
s.mu.Unlock()
s.logger.Info("加载历史执行记录", zap.Int("count", len(executions)))
}
// 加载历史统计信息
stats, err := s.storage.LoadToolStats()
if err != nil {
s.logger.Warn("加载历史统计信息失败", zap.Error(err))
} else {
s.mu.Lock()
for k, v := range stats {
s.stats[k] = v
}
s.mu.Unlock()
s.logger.Info("加载历史统计信息", zap.Int("count", len(stats)))
}
}
// GetAllExecutions 获取所有执行记录（合并内存和数据库）
func (s *Server) GetAllExecutions() []*ToolExecution {
if s.storage != nil {
dbExecutions, err := s.storage.LoadToolExecutions()
if err == nil {
execMap := make(map[string]*ToolExecution)
for _, exec := range dbExecutions {
if _, exists := execMap[exec.ID]; !exists {
execMap[exec.ID] = exec
}
}
s.mu.RLock()
for id, exec := range s.executions {
if _, exists := execMap[id]; !exists {
execMap[id] = exec
}
}
s.mu.RUnlock()
result := make([]*ToolExecution, 0, len(execMap))
for _, exec := range execMap {
result = append(result, exec)
}
return result
} else {
s.logger.Warn("从数据库加载执行记录失败", zap.Error(err))
}
}
s.mu.RLock()
defer s.mu.RUnlock()
executions := make([]*ToolExecution, 0, len(s.executions))
memExecutions := make([]*ToolExecution, 0, len(s.executions))
for _, exec := range s.executions {
executions = append(executions, exec)
memExecutions = append(memExecutions, exec)
}
return executions
return memExecutions
}
// GetStats 获取统计信息
// GetStats 获取统计信息（合并内存和数据库）
func (s *Server) GetStats() map[string]*ToolStats {
if s.storage != nil {
dbStats, err := s.storage.LoadToolStats()
if err == nil {
return dbStats
}
s.logger.Warn("从数据库加载统计信息失败", zap.Error(err))
}
s.mu.RLock()
defer s.mu.RUnlock()
stats := make(map[string]*ToolStats)
memStats := make(map[string]*ToolStats)
for k, v := range s.stats {
stats[k] = v
statCopy := *v
memStats[k] = &statCopy
}
return stats
return memStats
}
// GetAllTools 获取所有已注册的工具（用于Agent动态获取工具列表）
func (s *Server) GetAllTools() []Tool {
s.mu.RLock()
defer s.mu.RUnlock()
tools := make([]Tool, 0, len(s.toolDefs))
for _, tool := range s.toolDefs {
tools = append(tools, tool)
@@ -443,37 +620,80 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
s.executions[executionID] = execution
s.mu.Unlock()
// 更新统计
s.updateStats(toolName, false)
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
// 执行工具
result, err := handler(ctx, args)
s.mu.Lock()
now := time.Now()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
var failed bool
var finalResult *ToolResult
if err != nil {
execution.Status = "failed"
execution.Error = err.Error()
s.updateStats(toolName, true)
s.mu.Unlock()
return nil, executionID, err
failed = true
} else if result != nil && result.IsError {
execution.Status = "failed"
if len(result.Content) > 0 {
execution.Error = result.Content[0].Text
} else {
execution.Error = "工具执行返回错误结果"
}
execution.Result = result
failed = true
finalResult = result
} else {
execution.Status = "completed"
if result == nil {
result = &ToolResult{
Content: []Content{
{Type: "text", Text: "工具执行完成，但未返回结果"},
},
}
}
execution.Result = result
s.updateStats(toolName, false)
s.mu.Unlock()
return result, executionID, nil
finalResult = result
failed = false
}
if finalResult == nil {
finalResult = execution.Result
}
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
s.updateStats(toolName, failed)
if s.storage != nil {
s.mu.Lock()
delete(s.executions, executionID)
s.mu.Unlock()
}
if err != nil {
return nil, executionID, err
}
return finalResult, executionID, nil
}
// initDefaultPrompts 初始化默认提示词模板
func (s *Server) initDefaultPrompts() {
s.mu.Lock()
defer s.mu.Unlock()
// 网络安全测试提示词
s.prompts["security_scan"] = &Prompt{
Name: "security_scan",
@@ -483,7 +703,7 @@ func (s *Server) initDefaultPrompts() {
{Name: "scan_type", Description: "扫描类型（port, vuln, web等）", Required: false},
},
}
// 渗透测试提示词
s.prompts["penetration_test"] = &Prompt{
Name: "penetration_test",
@@ -509,7 +729,7 @@ func (s *Server) handleListPrompts(msg *Message) *Message {
prompts = append(prompts, *prompt)
}
s.mu.RUnlock()
response := ListPromptsResponse{
Prompts: prompts,
}
@@ -533,11 +753,11 @@ func (s *Server) handleGetPrompt(msg *Message) *Message {
Error: &Error{Code: -32602, Message: "Invalid params"},
}
}
s.mu.RLock()
prompt, exists := s.prompts[req.Name]
s.mu.RUnlock()
if !exists {
return &Message{
ID: msg.ID,
@@ -546,10 +766,10 @@ func (s *Server) handleGetPrompt(msg *Message) *Message {
Error: &Error{Code: -32601, Message: "Prompt not found"},
}
}
// 根据提示词名称生成消息
messages := s.generatePromptMessages(prompt, req.Arguments)
response := GetPromptResponse{
Messages: messages,
}
@@ -565,7 +785,7 @@ func (s *Server) handleGetPrompt(msg *Message) *Message {
// generatePromptMessages 生成提示词消息
func (s *Server) generatePromptMessages(prompt *Prompt, args map[string]interface{}) []PromptMessage {
messages := []PromptMessage{}
switch prompt.Name {
case "security_scan":
target, _ := args["target"].(string)
@@ -573,40 +793,40 @@ func (s *Server) generatePromptMessages(prompt *Prompt, args map[string]interfac
if scanType == "" {
scanType = "comprehensive"
}
content := fmt.Sprintf(`请对目标 %s 执行%s安全扫描。包括：
1. 端口扫描和服务识别
2. 漏洞检测
3. Web应用安全测试
4. 生成详细的安全报告`, target, scanType)
messages = append(messages, PromptMessage{
Role: "user",
Content: content,
})
case "penetration_test":
target, _ := args["target"].(string)
scope, _ := args["scope"].(string)
content := fmt.Sprintf(`请对目标 %s 执行渗透测试。`, target)
if scope != "" {
content += fmt.Sprintf("测试范围：%s", scope)
}
content += "\n请按照OWASP Top 10进行全面的安全测试。"
messages = append(messages, PromptMessage{
Role: "user",
Content: content,
})
default:
messages = append(messages, PromptMessage{
Role: "user",
Content: "请执行安全测试任务",
})
}
return messages
}
@@ -618,7 +838,7 @@ func (s *Server) handleListResources(msg *Message) *Message {
resources = append(resources, *resource)
}
s.mu.RUnlock()
response := ListResourcesResponse{
Resources: resources,
}
@@ -642,11 +862,11 @@ func (s *Server) handleReadResource(msg *Message) *Message {
Error: &Error{Code: -32602, Message: "Invalid params"},
}
}
s.mu.RLock()
resource, exists := s.resources[req.URI]
s.mu.RUnlock()
if !exists {
return &Message{
ID: msg.ID,
@@ -655,10 +875,10 @@ func (s *Server) handleReadResource(msg *Message) *Message {
Error: &Error{Code: -32601, Message: "Resource not found"},
}
}
// 生成资源内容
content := s.generateResourceContent(resource)
response := ReadResourceResponse{
Contents: []ResourceContent{content},
}
@@ -677,7 +897,7 @@ func (s *Server) generateResourceContent(resource *Resource) ResourceContent {
URI: resource.URI,
MimeType: resource.MimeType,
}
// 如果是工具资源，生成详细文档
if strings.HasPrefix(resource.URI, "tool://") {
toolName := strings.TrimPrefix(resource.URI, "tool://")
@@ -686,116 +906,36 @@ func (s *Server) generateResourceContent(resource *Resource) ResourceContent {
// 其他资源使用描述或默认内容
content.Text = resource.Description
}
return content
}
// generateToolDocumentation 生成工具文档
// 注意：硬编码的工具文档已移除，现在只使用工具定义中的信息
func (s *Server) generateToolDocumentation(toolName string, resource *Resource) string {
// 获取工具定义以获取更详细的信息
s.mu.RLock()
tool, hasTool := s.toolDefs[toolName]
s.mu.RUnlock()
// 为常见工具生成详细文档
switch toolName {
case "nmap":
return `Nmap (Network Mapper) 是一个强大的网络扫描工具。
主要功能：
- 端口扫描：发现目标主机开放的端口
- 服务识别：识别运行在端口上的服务
- 版本检测：检测服务版本信息
- 操作系统检测：识别目标操作系统
常用命令：
- nmap -sT target # TCP连接扫描
- nmap -sV target # 版本检测
- nmap -sC target # 默认脚本扫描
- nmap -p 1-1000 target # 扫描指定端口范围
参数说明：
- target: 目标IP地址或域名（必需）
- ports: 端口范围，例如: 1-1000（可选）`
case "sqlmap":
return `SQLMap 是一个自动化的SQL注入检测和利用工具。
主要功能：
- 自动检测SQL注入漏洞
- 数据库指纹识别
- 数据提取
- 文件系统访问
常用命令：
- sqlmap -u "http://target.com/page?id=1" # 检测URL参数
- sqlmap -u "http://target.com" --forms # 检测表单
- sqlmap -u "http://target.com" --dbs # 列出数据库
参数说明：
- url: 目标URL（必需）`
case "nikto":
return `Nikto 是一个Web服务器扫描工具。
主要功能：
- Web服务器漏洞扫描
- 检测过时的服务器软件
- 检测危险文件和程序
- 检测服务器配置问题
常用命令：
- nikto -h target # 扫描目标主机
- nikto -h target -p 80,443 # 扫描指定端口
参数说明：
- target: 目标URL（必需）`
case "dirb":
return `Dirb 是一个Web内容扫描器。
主要功能：
- 扫描Web目录和文件
- 发现隐藏的目录和文件
- 支持自定义字典
常用命令：
- dirb url # 扫描目标URL
- dirb url -w wordlist.txt # 使用自定义字典
参数说明：
- target: 目标URL（必需）`
case "exec":
return `Exec 工具用于执行系统命令。
⚠️ 警告：此工具可以执行任意系统命令，请谨慎使用！
参数说明：
- command: 要执行的系统命令（必需）
- shell: 使用的shell，默认为sh（可选）
- workdir: 工作目录（可选）`
default:
// 对于其他工具，使用工具定义中的描述信息
if hasTool {
doc := fmt.Sprintf("%s\n\n", resource.Description)
if tool.InputSchema != nil {
if props, ok := tool.InputSchema["properties"].(map[string]interface{}); ok {
doc += "参数说明：\n"
for paramName, paramInfo := range props {
if paramMap, ok := paramInfo.(map[string]interface{}); ok {
if desc, ok := paramMap["description"].(string); ok {
doc += fmt.Sprintf("- %s: %s\n", paramName, desc)
}
// 使用工具定义中的描述信息
if hasTool {
doc := fmt.Sprintf("%s\n\n", resource.Description)
if tool.InputSchema != nil {
if props, ok := tool.InputSchema["properties"].(map[string]interface{}); ok {
doc += "参数说明：\n"
for paramName, paramInfo := range props {
if paramMap, ok := paramInfo.(map[string]interface{}); ok {
if desc, ok := paramMap["description"].(string); ok {
doc += fmt.Sprintf("- %s: %s\n", paramName, desc)
}
}
}
}
return doc
}
return resource.Description
return doc
}
return resource.Description
}
// handleSamplingRequest 处理采样请求
@@ -809,13 +949,13 @@ func (s *Server) handleSamplingRequest(msg *Message) *Message {
Error: &Error{Code: -32602, Message: "Invalid params"},
}
}
// 注意：采样功能通常需要连接到实际的LLM服务
// 这里返回一个占位符响应，实际实现需要集成LLM API
s.logger.Warn("Sampling request received but not fully implemented",
zap.Any("request", req),
)
response := SamplingResponse{
Content: []SamplingContent{
{
@@ -908,4 +1048,3 @@ func (s *Server) sendError(w http.ResponseWriter, id interface{}, code int, mess
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(response)
}
internal/security/executor.go
+5 -69
View File
@@ -722,76 +722,12 @@ type Vulnerability struct {
}
// AnalyzeResults 分析工具执行结果，提取漏洞信息
// 注意：硬编码的漏洞解析逻辑已移除，此函数现在返回空数组
// 漏洞检测应该由工具本身或专门的漏洞扫描工具来完成
func (e *Executor) AnalyzeResults(toolName string, result *mcp.ToolResult) []Vulnerability {
vulnerabilities := []Vulnerability{}
if result.IsError {
return vulnerabilities
}
// 分析输出内容
for _, content := range result.Content {
if content.Type == "text" {
vulns := e.parseToolOutput(toolName, content.Text)
vulnerabilities = append(vulnerabilities, vulns...)
}
}
return vulnerabilities
}
// parseToolOutput 解析工具输出
func (e *Executor) parseToolOutput(toolName, output string) []Vulnerability {
vulnerabilities := []Vulnerability{}
// 简单的漏洞检测逻辑
outputLower := strings.ToLower(output)
// SQL注入检测
if strings.Contains(outputLower, "sql injection") || strings.Contains(outputLower, "sqli") {
vulnerabilities = append(vulnerabilities, Vulnerability{
ID: fmt.Sprintf("sql-%d", time.Now().Unix()),
Type: "SQL Injection",
Severity: "high",
Title: "SQL注入漏洞",
Description: "检测到潜在的SQL注入漏洞",
FoundAt: time.Now(),
Details: output,
})
}
// XSS检测
if strings.Contains(outputLower, "xss") || strings.Contains(outputLower, "cross-site scripting") {
vulnerabilities = append(vulnerabilities, Vulnerability{
ID: fmt.Sprintf("xss-%d", time.Now().Unix()),
Type: "XSS",
Severity: "medium",
Title: "跨站脚本攻击漏洞",
Description: "检测到潜在的XSS漏洞",
FoundAt: time.Now(),
Details: output,
})
}
// 开放端口检测
if toolName == "nmap" {
lines := strings.Split(output, "\n")
for _, line := range lines {
if strings.Contains(line, "open") && strings.Contains(line, "port") {
vulnerabilities = append(vulnerabilities, Vulnerability{
ID: fmt.Sprintf("port-%d", time.Now().Unix()),
Type: "Open Port",
Severity: "low",
Title: "开放端口",
Description: fmt.Sprintf("发现开放端口: %s", line),
FoundAt: time.Now(),
Details: line,
})
}
}
}
return vulnerabilities
// 不再进行硬编码的漏洞解析
// 漏洞检测应该由工具本身（如sqlmap、nmap等）的输出结果来体现
return []Vulnerability{}
}
// GetVulnerabilityReport 生成漏洞报告