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

Add files via upload

Browse Source
This commit is contained in:
公明
2025-11-08 18:56:23 +08:00
committed by GitHub
commit add33e1cf7
24 changed files with 5228 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/agent/agent.go
+576
View File
@@ -0,0 +1,576 @@
package agent
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
// Agent AI代理
type Agent struct {
openAIClient *http.Client
config *config.OpenAIConfig
mcpServer *mcp.Server
logger *zap.Logger
}
// NewAgent 创建新的Agent
func NewAgent(cfg *config.OpenAIConfig, mcpServer *mcp.Server, logger *zap.Logger) *Agent {
return &Agent{
openAIClient: &http.Client{Timeout: 5 * time.Minute},
config: cfg,
mcpServer: mcpServer,
logger: logger,
}
}
// ChatMessage 聊天消息
type ChatMessage struct {
Role string `json:"role"`
Content string `json:"content,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
}
// MarshalJSON 自定义JSON序列化，将tool_calls中的arguments转换为JSON字符串
func (cm ChatMessage) MarshalJSON() ([]byte, error) {
// 构建序列化结构
aux := map[string]interface{}{
"role": cm.Role,
}
// 添加content（如果存在）
if cm.Content != "" {
aux["content"] = cm.Content
}
// 添加tool_call_id（如果存在）
if cm.ToolCallID != "" {
aux["tool_call_id"] = cm.ToolCallID
}
// 转换tool_calls，将arguments转换为JSON字符串
if len(cm.ToolCalls) > 0 {
toolCallsJSON := make([]map[string]interface{}, len(cm.ToolCalls))
for i, tc := range cm.ToolCalls {
// 将arguments转换为JSON字符串
argsJSON := ""
if tc.Function.Arguments != nil {
argsBytes, err := json.Marshal(tc.Function.Arguments)
if err != nil {
return nil, err
}
argsJSON = string(argsBytes)
}
toolCallsJSON[i] = map[string]interface{}{
"id": tc.ID,
"type": tc.Type,
"function": map[string]interface{}{
"name": tc.Function.Name,
"arguments": argsJSON,
},
}
}
aux["tool_calls"] = toolCallsJSON
}
return json.Marshal(aux)
}
// OpenAIRequest OpenAI API请求
type OpenAIRequest struct {
Model string `json:"model"`
Messages []ChatMessage `json:"messages"`
Tools []Tool `json:"tools,omitempty"`
}
// OpenAIResponse OpenAI API响应
type OpenAIResponse struct {
ID string `json:"id"`
Choices []Choice `json:"choices"`
Error *Error `json:"error,omitempty"`
}
// Choice 选择
type Choice struct {
Message MessageWithTools `json:"message"`
FinishReason string `json:"finish_reason"`
}
// MessageWithTools 带工具调用的消息
type MessageWithTools struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
// Tool OpenAI工具定义
type Tool struct {
Type string `json:"type"`
Function FunctionDefinition `json:"function"`
}
// FunctionDefinition 函数定义
type FunctionDefinition struct {
Name string `json:"name"`
Description string `json:"description"`
Parameters map[string]interface{} `json:"parameters"`
}
// Error OpenAI错误
type Error struct {
Message string `json:"message"`
Type string `json:"type"`
}
// ToolCall 工具调用
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function FunctionCall `json:"function"`
}
// FunctionCall 函数调用
type FunctionCall struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments"`
}
// UnmarshalJSON 自定义JSON解析，处理arguments可能是字符串或对象的情况
func (fc *FunctionCall) UnmarshalJSON(data []byte) error {
type Alias FunctionCall
aux := &struct {
Name string `json:"name"`
Arguments interface{} `json:"arguments"`
*Alias
}{
Alias: (*Alias)(fc),
}
if err := json.Unmarshal(data, &aux); err != nil {
return err
}
fc.Name = aux.Name
// 处理arguments可能是字符串或对象的情况
switch v := aux.Arguments.(type) {
case map[string]interface{}:
fc.Arguments = v
case string:
// 如果是字符串，尝试解析为JSON
if err := json.Unmarshal([]byte(v), &fc.Arguments); err != nil {
// 如果解析失败，创建一个包含原始字符串的map
fc.Arguments = map[string]interface{}{
"raw": v,
}
}
case nil:
fc.Arguments = make(map[string]interface{})
default:
// 其他类型，尝试转换为map
fc.Arguments = map[string]interface{}{
"value": v,
}
}
return nil
}
// AgentLoopResult Agent Loop执行结果
type AgentLoopResult struct {
Response string
MCPExecutionIDs []string
}
// AgentLoop 执行Agent循环
func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages []ChatMessage) (*AgentLoopResult, error) {
messages := []ChatMessage{
{
Role: "system",
Content: "你是一个专业的网络安全渗透测试专家。你可以使用各种安全工具进行自主渗透测试。分析目标并选择最佳测试策略。当需要执行工具时，使用提供的工具函数。",
},
}
// 添加历史消息（数据库只保存user和assistant消息）
a.logger.Info("处理历史消息",
zap.Int("count", len(historyMessages)),
)
addedCount := 0
for i, msg := range historyMessages {
// 只添加有内容的消息
if msg.Content != "" {
messages = append(messages, ChatMessage{
Role: msg.Role,
Content: msg.Content,
})
addedCount++
contentPreview := msg.Content
if len(contentPreview) > 50 {
contentPreview = contentPreview[:50] + "..."
}
a.logger.Info("添加历史消息到上下文",
zap.Int("index", i),
zap.String("role", msg.Role),
zap.String("content", contentPreview),
)
}
}
a.logger.Info("构建消息数组",
zap.Int("historyMessages", len(historyMessages)),
zap.Int("addedMessages", addedCount),
zap.Int("totalMessages", len(messages)),
)
// 添加当前用户消息
messages = append(messages, ChatMessage{
Role: "user",
Content: userInput,
})
result := &AgentLoopResult{
MCPExecutionIDs: make([]string, 0),
}
maxIterations := 10
for i := 0; i < maxIterations; i++ {
// 获取可用工具
tools := a.getAvailableTools()
// 记录每次调用OpenAI
if i == 0 {
a.logger.Info("调用OpenAI",
zap.Int("iteration", i+1),
zap.Int("messagesCount", len(messages)),
)
// 记录前几条消息的内容（用于调试）
for j, msg := range messages {
if j >= 5 { // 只记录前5条
break
}
contentPreview := msg.Content
if len(contentPreview) > 100 {
contentPreview = contentPreview[:100] + "..."
}
a.logger.Debug("消息内容",
zap.Int("index", j),
zap.String("role", msg.Role),
zap.String("content", contentPreview),
)
}
} else {
a.logger.Info("调用OpenAI",
zap.Int("iteration", i+1),
zap.Int("messagesCount", len(messages)),
)
}
// 调用OpenAI
response, err := a.callOpenAI(ctx, messages, tools)
if err != nil {
result.Response = ""
return result, fmt.Errorf("调用OpenAI失败: %w", err)
}
if response.Error != nil {
result.Response = ""
return result, fmt.Errorf("OpenAI错误: %s", response.Error.Message)
}
if len(response.Choices) == 0 {
result.Response = ""
return result, fmt.Errorf("没有收到响应")
}
choice := response.Choices[0]
// 检查是否有工具调用
if len(choice.Message.ToolCalls) > 0 {
// 添加assistant消息（包含工具调用）
messages = append(messages, ChatMessage{
Role: "assistant",
Content: choice.Message.Content,
ToolCalls: choice.Message.ToolCalls,
})
// 执行所有工具调用
for _, toolCall := range choice.Message.ToolCalls {
// 执行工具
execResult, err := a.executeToolViaMCP(ctx, toolCall.Function.Name, toolCall.Function.Arguments)
if err != nil {
messages = append(messages, ChatMessage{
Role: "tool",
ToolCallID: toolCall.ID,
Content: fmt.Sprintf("工具执行失败: %v", err),
})
} else {
messages = append(messages, ChatMessage{
Role: "tool",
ToolCallID: toolCall.ID,
Content: execResult.Result,
})
// 收集执行ID
if execResult.ExecutionID != "" {
result.MCPExecutionIDs = append(result.MCPExecutionIDs, execResult.ExecutionID)
}
}
}
continue
}
// 添加assistant响应
messages = append(messages, ChatMessage{
Role: "assistant",
Content: choice.Message.Content,
})
// 如果完成，返回结果
if choice.FinishReason == "stop" {
result.Response = choice.Message.Content
return result, nil
}
}
result.Response = "达到最大迭代次数"
return result, nil
}
// getAvailableTools 获取可用工具
func (a *Agent) getAvailableTools() []Tool {
// 从MCP服务器获取工具列表
executions := a.mcpServer.GetAllExecutions()
toolNames := make(map[string]bool)
for _, exec := range executions {
toolNames[exec.ToolName] = true
}
tools := []Tool{
{
Type: "function",
Function: FunctionDefinition{
Name: "nmap",
Description: "使用nmap进行网络扫描，发现开放端口和服务。支持IP地址、域名或URL（会自动提取域名）。使用TCP连接扫描，不需要root权限。",
Parameters: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"target": map[string]interface{}{
"type": "string",
"description": "目标IP地址、域名或URL（如 https://example.com）。如果是URL，会自动提取域名部分。",
},
"ports": map[string]interface{}{
"type": "string",
"description": "要扫描的端口范围，例如: 1-1000 或 80,443,8080。如果不指定，将扫描常用端口。",
},
},
"required": []string{"target"},
},
},
},
{
Type: "function",
Function: FunctionDefinition{
Name: "sqlmap",
Description: "使用sqlmap检测SQL注入漏洞",
Parameters: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"url": map[string]interface{}{
"type": "string",
"description": "目标URL",
},
},
"required": []string{"url"},
},
},
},
{
Type: "function",
Function: FunctionDefinition{
Name: "nikto",
Description: "使用nikto扫描Web服务器漏洞",
Parameters: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"target": map[string]interface{}{
"type": "string",
"description": "目标URL",
},
},
"required": []string{"target"},
},
},
},
{
Type: "function",
Function: FunctionDefinition{
Name: "dirb",
Description: "使用dirb进行目录扫描",
Parameters: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"url": map[string]interface{}{
"type": "string",
"description": "目标URL",
},
},
"required": []string{"url"},
},
},
},
{
Type: "function",
Function: FunctionDefinition{
Name: "exec",
Description: "执行系统命令（谨慎使用，仅用于必要的系统操作）",
Parameters: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"command": map[string]interface{}{
"type": "string",
"description": "要执行的系统命令",
},
"shell": map[string]interface{}{
"type": "string",
"description": "使用的shell（可选，默认为sh）",
},
"workdir": map[string]interface{}{
"type": "string",
"description": "工作目录（可选）",
},
},
"required": []string{"command"},
},
},
},
}
return tools
}
// callOpenAI 调用OpenAI API
func (a *Agent) callOpenAI(ctx context.Context, messages []ChatMessage, tools []Tool) (*OpenAIResponse, error) {
reqBody := OpenAIRequest{
Model: a.config.Model,
Messages: messages,
}
if len(tools) > 0 {
reqBody.Tools = tools
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, "POST", a.config.BaseURL+"/chat/completions", bytes.NewBuffer(jsonData))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+a.config.APIKey)
resp, err := a.openAIClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
// 记录响应内容（用于调试）
if resp.StatusCode != http.StatusOK {
a.logger.Warn("OpenAI API返回非200状态码",
zap.Int("status", resp.StatusCode),
zap.String("body", string(body)),
)
}
var response OpenAIResponse
if err := json.Unmarshal(body, &response); err != nil {
a.logger.Error("解析OpenAI响应失败",
zap.Error(err),
zap.String("body", string(body)),
)
return nil, fmt.Errorf("解析响应失败: %w, 响应内容: %s", err, string(body))
}
return &response, nil
}
// parseToolCall 解析工具调用
func (a *Agent) parseToolCall(content string) (map[string]interface{}, error) {
// 简单解析，实际应该更复杂
// 格式: [TOOL_CALL]tool_name:arg1=value1,arg2=value2
if !strings.HasPrefix(content, "[TOOL_CALL]") {
return nil, fmt.Errorf("不是有效的工具调用格式")
}
parts := strings.Split(content[len("[TOOL_CALL]"):], ":")
if len(parts) < 2 {
return nil, fmt.Errorf("工具调用格式错误")
}
toolName := strings.TrimSpace(parts[0])
argsStr := strings.TrimSpace(parts[1])
args := make(map[string]interface{})
argPairs := strings.Split(argsStr, ",")
for _, pair := range argPairs {
kv := strings.Split(pair, "=")
if len(kv) == 2 {
args[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1])
}
}
args["_tool_name"] = toolName
return args, nil
}
// ToolExecutionResult 工具执行结果
type ToolExecutionResult struct {
Result string
ExecutionID string
}
// executeToolViaMCP 通过MCP执行工具
func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
a.logger.Info("通过MCP执行工具",
zap.String("tool", toolName),
zap.Any("args", args),
)
// 通过MCP服务器调用工具
result, executionID, err := a.mcpServer.CallTool(ctx, toolName, args)
if err != nil {
return nil, fmt.Errorf("工具执行失败: %w", err)
}
// 格式化结果
var resultText strings.Builder
for _, content := range result.Content {
resultText.WriteString(content.Text)
resultText.WriteString("\n")
}
return &ToolExecutionResult{
Result: resultText.String(),
ExecutionID: executionID,
}, nil
}
internal/app/app.go
+163
View File
@@ -0,0 +1,163 @@
package app
import (
"fmt"
"net/http"
"os"
"path/filepath"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/handler"
"cyberstrike-ai/internal/logger"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// App 应用
type App struct {
config *config.Config
logger *logger.Logger
router *gin.Engine
mcpServer *mcp.Server
agent *agent.Agent
executor *security.Executor
db *database.DB
}
// New 创建新应用
func New(cfg *config.Config, log *logger.Logger) (*App, error) {
gin.SetMode(gin.ReleaseMode)
router := gin.Default()
// CORS中间件
router.Use(corsMiddleware())
// 初始化数据库
dbPath := cfg.Database.Path
if dbPath == "" {
dbPath = "data/conversations.db"
}
// 确保目录存在
if err := os.MkdirAll(filepath.Dir(dbPath), 0755); err != nil {
return nil, fmt.Errorf("创建数据库目录失败: %w", err)
}
db, err := database.NewDB(dbPath, log.Logger)
if err != nil {
return nil, fmt.Errorf("初始化数据库失败: %w", err)
}
// 创建MCP服务器
mcpServer := mcp.NewServer(log.Logger)
// 创建安全工具执行器
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
// 注册工具
executor.RegisterTools(mcpServer)
// 创建Agent
agent := agent.NewAgent(&cfg.OpenAI, mcpServer, log.Logger)
// 创建处理器
agentHandler := handler.NewAgentHandler(agent, db, log.Logger)
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, log.Logger)
conversationHandler := handler.NewConversationHandler(db, log.Logger)
// 设置路由
setupRoutes(router, agentHandler, monitorHandler, conversationHandler, mcpServer)
return &App{
config: cfg,
logger: log,
router: router,
mcpServer: mcpServer,
agent: agent,
executor: executor,
db: db,
}, nil
}
// Run 启动应用
func (a *App) Run() error {
// 启动MCP服务器（如果启用）
if a.config.MCP.Enabled {
go func() {
mcpAddr := fmt.Sprintf("%s:%d", a.config.MCP.Host, a.config.MCP.Port)
a.logger.Info("启动MCP服务器", zap.String("address", mcpAddr))
mux := http.NewServeMux()
mux.HandleFunc("/mcp", a.mcpServer.HandleHTTP)
if err := http.ListenAndServe(mcpAddr, mux); err != nil {
a.logger.Error("MCP服务器启动失败", zap.Error(err))
}
}()
}
// 启动主服务器
addr := fmt.Sprintf("%s:%d", a.config.Server.Host, a.config.Server.Port)
a.logger.Info("启动HTTP服务器", zap.String("address", addr))
return a.router.Run(addr)
}
// setupRoutes 设置路由
func setupRoutes(router *gin.Engine, agentHandler *handler.AgentHandler, monitorHandler *handler.MonitorHandler, conversationHandler *handler.ConversationHandler, mcpServer *mcp.Server) {
// API路由
api := router.Group("/api")
{
// Agent Loop
api.POST("/agent-loop", agentHandler.AgentLoop)
// 对话历史
api.POST("/conversations", conversationHandler.CreateConversation)
api.GET("/conversations", conversationHandler.ListConversations)
api.GET("/conversations/:id", conversationHandler.GetConversation)
api.DELETE("/conversations/:id", conversationHandler.DeleteConversation)
// 监控
api.GET("/monitor", monitorHandler.Monitor)
api.GET("/monitor/execution/:id", monitorHandler.GetExecution)
api.GET("/monitor/stats", monitorHandler.GetStats)
api.GET("/monitor/vulnerabilities", monitorHandler.GetVulnerabilities)
// MCP端点
api.POST("/mcp", func(c *gin.Context) {
mcpServer.HandleHTTP(c.Writer, c.Request)
})
}
// 静态文件
router.Static("/static", "./web/static")
router.LoadHTMLGlob("web/templates/*")
// 前端页面
router.GET("/", func(c *gin.Context) {
c.HTML(http.StatusOK, "index.html", nil)
})
}
// corsMiddleware CORS中间件
func corsMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With")
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE")
if c.Request.Method == "OPTIONS" {
c.AbortWithStatus(204)
return
}
c.Next()
}
}
internal/config/config.go
+114
View File
@@ -0,0 +1,114 @@
package config
import (
"fmt"
"os"
"gopkg.in/yaml.v3"
)
type Config struct {
Server ServerConfig `yaml:"server"`
Log LogConfig `yaml:"log"`
MCP MCPConfig `yaml:"mcp"`
OpenAI OpenAIConfig `yaml:"openai"`
Security SecurityConfig `yaml:"security"`
Database DatabaseConfig `yaml:"database"`
}
type ServerConfig struct {
Host string `yaml:"host"`
Port int `yaml:"port"`
}
type LogConfig struct {
Level string `yaml:"level"`
Output string `yaml:"output"`
}
type MCPConfig struct {
Enabled bool `yaml:"enabled"`
Host string `yaml:"host"`
Port int `yaml:"port"`
}
type OpenAIConfig struct {
APIKey string `yaml:"api_key"`
BaseURL string `yaml:"base_url"`
Model string `yaml:"model"`
}
type SecurityConfig struct {
Tools []ToolConfig `yaml:"tools"`
}
type DatabaseConfig struct {
Path string `yaml:"path"`
}
type ToolConfig struct {
Name string `yaml:"name"`
Command string `yaml:"command"`
Args []string `yaml:"args,omitempty"` // 固定参数（可选）
Description string `yaml:"description"`
Enabled bool `yaml:"enabled"`
Parameters []ParameterConfig `yaml:"parameters,omitempty"` // 参数定义（可选）
ArgMapping string `yaml:"arg_mapping,omitempty"` // 参数映射方式: "auto", "manual", "template"（可选）
}
// ParameterConfig 参数配置
type ParameterConfig struct {
Name string `yaml:"name"` // 参数名称
Type string `yaml:"type"` // 参数类型: string, int, bool, array
Description string `yaml:"description"` // 参数描述
Required bool `yaml:"required,omitempty"` // 是否必需
Default interface{} `yaml:"default,omitempty"` // 默认值
Flag string `yaml:"flag,omitempty"` // 命令行标志，如 "-u", "--url", "-p"
Position *int `yaml:"position,omitempty"` // 位置参数的位置（从0开始）
Format string `yaml:"format,omitempty"` // 参数格式: "flag", "positional", "combined" (flag=value), "template"
Template string `yaml:"template,omitempty"` // 模板字符串，如 "{flag} {value}" 或 "{value}"
Options []string `yaml:"options,omitempty"` // 可选值列表（用于枚举）
}
func Load(path string) (*Config, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("读取配置文件失败: %w", err)
}
var cfg Config
if err := yaml.Unmarshal(data, &cfg); err != nil {
return nil, fmt.Errorf("解析配置文件失败: %w", err)
}
return &cfg, nil
}
func Default() *Config {
return &Config{
Server: ServerConfig{
Host: "0.0.0.0",
Port: 8080,
},
Log: LogConfig{
Level: "info",
Output: "stdout",
},
MCP: MCPConfig{
Enabled: true,
Host: "0.0.0.0",
Port: 8081,
},
OpenAI: OpenAIConfig{
BaseURL: "https://api.openai.com/v1",
Model: "gpt-4",
},
Security: SecurityConfig{
Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 加载，不在此硬编码
},
Database: DatabaseConfig{
Path: "data/conversations.db",
},
}
}
internal/database/conversation.go
+256
View File
@@ -0,0 +1,256 @@
package database
import (
"database/sql"
"encoding/json"
"fmt"
"time"
"github.com/google/uuid"
"go.uber.org/zap"
)
// Conversation 对话
type Conversation struct {
ID string `json:"id"`
Title string `json:"title"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
Messages []Message `json:"messages,omitempty"`
}
// Message 消息
type Message struct {
ID string `json:"id"`
ConversationID string `json:"conversationId"`
Role string `json:"role"`
Content string `json:"content"`
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"`
CreatedAt time.Time `json:"createdAt"`
}
// CreateConversation 创建新对话
func (db *DB) CreateConversation(title string) (*Conversation, error) {
id := uuid.New().String()
now := time.Now()
_, err := db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
id, title, now, now,
)
if err != nil {
return nil, fmt.Errorf("创建对话失败: %w", err)
}
return &Conversation{
ID: id,
Title: title,
CreatedAt: now,
UpdatedAt: now,
}, nil
}
// GetConversation 获取对话
func (db *DB) GetConversation(id string) (*Conversation, error) {
var conv Conversation
var createdAt, updatedAt string
err := db.QueryRow(
"SELECT id, title, created_at, updated_at FROM conversations WHERE id = ?",
id,
).Scan(&conv.ID, &conv.Title, &createdAt, &updatedAt)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("对话不存在")
}
return nil, fmt.Errorf("查询对话失败: %w", err)
}
// 尝试多种时间格式解析
var err1, err2 error
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err1 != nil {
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err1 != nil {
conv.CreatedAt, err1 = time.Parse(time.RFC3339, createdAt)
}
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
if err2 != nil {
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
}
if err2 != nil {
conv.UpdatedAt, err2 = time.Parse(time.RFC3339, updatedAt)
}
// 加载消息
messages, err := db.GetMessages(id)
if err != nil {
return nil, fmt.Errorf("加载消息失败: %w", err)
}
conv.Messages = messages
return &conv, nil
}
// ListConversations 列出所有对话
func (db *DB) ListConversations(limit, offset int) ([]*Conversation, error) {
rows, err := db.Query(
"SELECT id, title, created_at, updated_at FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?",
limit, offset,
)
if err != nil {
return nil, fmt.Errorf("查询对话列表失败: %w", err)
}
defer rows.Close()
var conversations []*Conversation
for rows.Next() {
var conv Conversation
var createdAt, updatedAt string
if err := rows.Scan(&conv.ID, &conv.Title, &createdAt, &updatedAt); err != nil {
return nil, fmt.Errorf("扫描对话失败: %w", err)
}
// 尝试多种时间格式解析
var err1, err2 error
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err1 != nil {
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err1 != nil {
conv.CreatedAt, err1 = time.Parse(time.RFC3339, createdAt)
}
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
if err2 != nil {
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
}
if err2 != nil {
conv.UpdatedAt, err2 = time.Parse(time.RFC3339, updatedAt)
}
conversations = append(conversations, &conv)
}
return conversations, nil
}
// UpdateConversationTitle 更新对话标题
func (db *DB) UpdateConversationTitle(id, title string) error {
_, err := db.Exec(
"UPDATE conversations SET title = ?, updated_at = ? WHERE id = ?",
title, time.Now(), id,
)
if err != nil {
return fmt.Errorf("更新对话标题失败: %w", err)
}
return nil
}
// UpdateConversationTime 更新对话时间
func (db *DB) UpdateConversationTime(id string) error {
_, err := db.Exec(
"UPDATE conversations SET updated_at = ? WHERE id = ?",
time.Now(), id,
)
if err != nil {
return fmt.Errorf("更新对话时间失败: %w", err)
}
return nil
}
// DeleteConversation 删除对话
func (db *DB) DeleteConversation(id string) error {
_, err := db.Exec("DELETE FROM conversations WHERE id = ?", id)
if err != nil {
return fmt.Errorf("删除对话失败: %w", err)
}
return nil
}
// AddMessage 添加消息
func (db *DB) AddMessage(conversationID, role, content string, mcpExecutionIDs []string) (*Message, error) {
id := uuid.New().String()
var mcpIDsJSON string
if len(mcpExecutionIDs) > 0 {
jsonData, err := json.Marshal(mcpExecutionIDs)
if err != nil {
db.logger.Warn("序列化MCP执行ID失败", zap.Error(err))
} else {
mcpIDsJSON = string(jsonData)
}
}
_, err := db.Exec(
"INSERT INTO messages (id, conversation_id, role, content, mcp_execution_ids, created_at) VALUES (?, ?, ?, ?, ?, ?)",
id, conversationID, role, content, mcpIDsJSON, time.Now(),
)
if err != nil {
return nil, fmt.Errorf("添加消息失败: %w", err)
}
// 更新对话时间
if err := db.UpdateConversationTime(conversationID); err != nil {
db.logger.Warn("更新对话时间失败", zap.Error(err))
}
message := &Message{
ID: id,
ConversationID: conversationID,
Role: role,
Content: content,
MCPExecutionIDs: mcpExecutionIDs,
CreatedAt: time.Now(),
}
return message, nil
}
// GetMessages 获取对话的所有消息
func (db *DB) GetMessages(conversationID string) ([]Message, error) {
rows, err := db.Query(
"SELECT id, conversation_id, role, content, mcp_execution_ids, created_at FROM messages WHERE conversation_id = ? ORDER BY created_at ASC",
conversationID,
)
if err != nil {
return nil, fmt.Errorf("查询消息失败: %w", err)
}
defer rows.Close()
var messages []Message
for rows.Next() {
var msg Message
var mcpIDsJSON sql.NullString
var createdAt string
if err := rows.Scan(&msg.ID, &msg.ConversationID, &msg.Role, &msg.Content, &mcpIDsJSON, &createdAt); err != nil {
return nil, fmt.Errorf("扫描消息失败: %w", err)
}
// 尝试多种时间格式解析
var err error
msg.CreatedAt, err = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err != nil {
msg.CreatedAt, err = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err != nil {
msg.CreatedAt, err = time.Parse(time.RFC3339, createdAt)
}
// 解析MCP执行ID
if mcpIDsJSON.Valid && mcpIDsJSON.String != "" {
if err := json.Unmarshal([]byte(mcpIDsJSON.String), &msg.MCPExecutionIDs); err != nil {
db.logger.Warn("解析MCP执行ID失败", zap.Error(err))
}
}
messages = append(messages, msg)
}
return messages, nil
}
internal/database/database.go
+90
View File
@@ -0,0 +1,90 @@
package database
import (
"database/sql"
"fmt"
_ "github.com/mattn/go-sqlite3"
"go.uber.org/zap"
)
// DB 数据库连接
type DB struct {
*sql.DB
logger *zap.Logger
}
// NewDB 创建数据库连接
func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1")
if err != nil {
return nil, fmt.Errorf("打开数据库失败: %w", err)
}
if err := db.Ping(); err != nil {
return nil, fmt.Errorf("连接数据库失败: %w", err)
}
database := &DB{
DB: db,
logger: logger,
}
// 初始化表
if err := database.initTables(); err != nil {
return nil, fmt.Errorf("初始化表失败: %w", err)
}
return database, nil
}
// initTables 初始化数据库表
func (db *DB) initTables() error {
// 创建对话表
createConversationsTable := `
CREATE TABLE IF NOT EXISTS conversations (
id TEXT PRIMARY KEY,
title TEXT NOT NULL,
created_at DATETIME NOT NULL,
updated_at DATETIME NOT NULL
);`
// 创建消息表
createMessagesTable := `
CREATE TABLE IF NOT EXISTS messages (
id TEXT PRIMARY KEY,
conversation_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
mcp_execution_ids TEXT,
created_at DATETIME NOT NULL,
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
);`
// 创建索引
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);
`
if _, err := db.Exec(createConversationsTable); err != nil {
return fmt.Errorf("创建conversations表失败: %w", err)
}
if _, err := db.Exec(createMessagesTable); err != nil {
return fmt.Errorf("创建messages表失败: %w", err)
}
if _, err := db.Exec(createIndexes); err != nil {
return fmt.Errorf("创建索引失败: %w", err)
}
db.logger.Info("数据库表初始化完成")
return nil
}
// Close 关闭数据库连接
func (db *DB) Close() error {
return db.DB.Close()
}
internal/handler/agent.go
+134
View File
@@ -0,0 +1,134 @@
package handler
import (
"net/http"
"time"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/database"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// AgentHandler Agent处理器
type AgentHandler struct {
agent *agent.Agent
db *database.DB
logger *zap.Logger
}
// NewAgentHandler 创建新的Agent处理器
func NewAgentHandler(agent *agent.Agent, db *database.DB, logger *zap.Logger) *AgentHandler {
return &AgentHandler{
agent: agent,
db: db,
logger: logger,
}
}
// ChatRequest 聊天请求
type ChatRequest struct {
Message string `json:"message" binding:"required"`
ConversationID string `json:"conversationId,omitempty"`
}
// ChatResponse 聊天响应
type ChatResponse struct {
Response string `json:"response"`
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表
ConversationID string `json:"conversationId"` // 对话ID
Time time.Time `json:"time"`
}
// AgentLoop 处理Agent Loop请求
func (h *AgentHandler) AgentLoop(c *gin.Context) {
var req ChatRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
h.logger.Info("收到Agent Loop请求",
zap.String("message", req.Message),
zap.String("conversationId", req.ConversationID),
)
// 如果没有对话ID，创建新对话
conversationID := req.ConversationID
if conversationID == "" {
title := req.Message
if len(title) > 50 {
title = title[:50] + "..."
}
conv, err := h.db.CreateConversation(title)
if err != nil {
h.logger.Error("创建对话失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
conversationID = conv.ID
}
// 获取历史消息（排除当前消息，因为还没保存）
historyMessages, err := h.db.GetMessages(conversationID)
if err != nil {
h.logger.Warn("获取历史消息失败", zap.Error(err))
historyMessages = []database.Message{}
}
h.logger.Info("获取历史消息",
zap.String("conversationId", conversationID),
zap.Int("count", len(historyMessages)),
)
// 将数据库消息转换为Agent消息格式
agentHistoryMessages := make([]agent.ChatMessage, 0, len(historyMessages))
for i, msg := range historyMessages {
agentHistoryMessages = append(agentHistoryMessages, agent.ChatMessage{
Role: msg.Role,
Content: msg.Content,
})
contentPreview := msg.Content
if len(contentPreview) > 50 {
contentPreview = contentPreview[:50] + "..."
}
h.logger.Info("添加历史消息",
zap.Int("index", i),
zap.String("role", msg.Role),
zap.String("content", contentPreview),
)
}
h.logger.Info("历史消息转换完成",
zap.Int("originalCount", len(historyMessages)),
zap.Int("convertedCount", len(agentHistoryMessages)),
)
// 保存用户消息
_, err = h.db.AddMessage(conversationID, "user", req.Message, nil)
if err != nil {
h.logger.Error("保存用户消息失败", zap.Error(err))
}
// 执行Agent Loop，传入历史消息
result, err := h.agent.AgentLoop(c.Request.Context(), req.Message, agentHistoryMessages)
if err != nil {
h.logger.Error("Agent Loop执行失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
// 保存助手回复
_, err = h.db.AddMessage(conversationID, "assistant", result.Response, result.MCPExecutionIDs)
if err != nil {
h.logger.Error("保存助手消息失败", zap.Error(err))
}
c.JSON(http.StatusOK, ChatResponse{
Response: result.Response,
MCPExecutionIDs: result.MCPExecutionIDs,
ConversationID: conversationID,
Time: time.Now(),
})
}
internal/handler/conversation.go
+102
View File
@@ -0,0 +1,102 @@
package handler
import (
"net/http"
"strconv"
"cyberstrike-ai/internal/database"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// ConversationHandler 对话处理器
type ConversationHandler struct {
db *database.DB
logger *zap.Logger
}
// NewConversationHandler 创建新的对话处理器
func NewConversationHandler(db *database.DB, logger *zap.Logger) *ConversationHandler {
return &ConversationHandler{
db: db,
logger: logger,
}
}
// CreateConversationRequest 创建对话请求
type CreateConversationRequest struct {
Title string `json:"title"`
}
// CreateConversation 创建新对话
func (h *ConversationHandler) CreateConversation(c *gin.Context) {
var req CreateConversationRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
title := req.Title
if title == "" {
title = "新对话"
}
conv, err := h.db.CreateConversation(title)
if err != nil {
h.logger.Error("创建对话失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, conv)
}
// ListConversations 列出对话
func (h *ConversationHandler) ListConversations(c *gin.Context) {
limitStr := c.DefaultQuery("limit", "50")
offsetStr := c.DefaultQuery("offset", "0")
limit, _ := strconv.Atoi(limitStr)
offset, _ := strconv.Atoi(offsetStr)
if limit <= 0 || limit > 100 {
limit = 50
}
conversations, err := h.db.ListConversations(limit, offset)
if err != nil {
h.logger.Error("获取对话列表失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, conversations)
}
// GetConversation 获取对话
func (h *ConversationHandler) GetConversation(c *gin.Context) {
id := c.Param("id")
conv, err := h.db.GetConversation(id)
if err != nil {
h.logger.Error("获取对话失败", zap.Error(err))
c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"})
return
}
c.JSON(http.StatusOK, conv)
}
// DeleteConversation 删除对话
func (h *ConversationHandler) DeleteConversation(c *gin.Context) {
id := c.Param("id")
if err := h.db.DeleteConversation(id); err != nil {
h.logger.Error("删除对话失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
}
internal/handler/monitor.go
+92
View File
@@ -0,0 +1,92 @@
package handler
import (
"net/http"
"time"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// MonitorHandler 监控处理器
type MonitorHandler struct {
mcpServer *mcp.Server
executor *security.Executor
logger *zap.Logger
vulns []security.Vulnerability
}
// NewMonitorHandler 创建新的监控处理器
func NewMonitorHandler(mcpServer *mcp.Server, executor *security.Executor, logger *zap.Logger) *MonitorHandler {
return &MonitorHandler{
mcpServer: mcpServer,
executor: executor,
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"`
}
// 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)
c.JSON(http.StatusOK, MonitorResponse{
Executions: executions,
Stats: stats,
Vulnerabilities: h.vulns,
Report: report,
Timestamp: time.Now(),
})
}
// 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": "执行记录未找到"})
return
}
c.JSON(http.StatusOK, exec)
}
// GetStats 获取统计信息
func (h *MonitorHandler) GetStats(c *gin.Context) {
stats := h.mcpServer.GetStats()
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/logger/logger.go
+60
View File
@@ -0,0 +1,60 @@
package logger
import (
"os"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
)
type Logger struct {
*zap.Logger
}
func New(level, output string) *Logger {
var zapLevel zapcore.Level
switch level {
case "debug":
zapLevel = zapcore.DebugLevel
case "info":
zapLevel = zapcore.InfoLevel
case "warn":
zapLevel = zapcore.WarnLevel
case "error":
zapLevel = zapcore.ErrorLevel
default:
zapLevel = zapcore.InfoLevel
}
config := zap.NewProductionConfig()
config.Level = zap.NewAtomicLevelAt(zapLevel)
config.EncoderConfig.TimeKey = "timestamp"
config.EncoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
var writeSyncer zapcore.WriteSyncer
if output == "stdout" {
writeSyncer = zapcore.AddSync(os.Stdout)
} else {
file, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
if err != nil {
writeSyncer = zapcore.AddSync(os.Stdout)
} else {
writeSyncer = zapcore.AddSync(file)
}
}
core := zapcore.NewCore(
zapcore.NewJSONEncoder(config.EncoderConfig),
writeSyncer,
zapLevel,
)
logger := zap.New(core, zap.AddCaller(), zap.AddStacktrace(zapcore.ErrorLevel))
return &Logger{Logger: logger}
}
func (l *Logger) Fatal(msg string, fields ...interface{}) {
l.Logger.Fatal(msg, zap.Any("fields", fields))
}
internal/mcp/server.go
+798
View File
@@ -0,0 +1,798 @@
package mcp
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"sync"
"time"
"github.com/google/uuid"
"go.uber.org/zap"
)
// Server MCP服务器
type Server struct {
tools map[string]ToolHandler
toolDefs map[string]Tool // 工具定义
executions map[string]*ToolExecution
stats map[string]*ToolStats
prompts map[string]*Prompt // 提示词模板
resources map[string]*Resource // 资源
mu sync.RWMutex
logger *zap.Logger
}
// ToolHandler 工具处理函数
type ToolHandler func(ctx context.Context, args map[string]interface{}) (*ToolResult, error)
// NewServer 创建新的MCP服务器
func NewServer(logger *zap.Logger) *Server {
s := &Server{
tools: make(map[string]ToolHandler),
toolDefs: make(map[string]Tool),
executions: make(map[string]*ToolExecution),
stats: make(map[string]*ToolStats),
prompts: make(map[string]*Prompt),
resources: make(map[string]*Resource),
logger: logger,
}
// 初始化默认提示词和资源
s.initDefaultPrompts()
s.initDefaultResources()
return s
}
// RegisterTool 注册工具
func (s *Server) RegisterTool(tool Tool, handler ToolHandler) {
s.mu.Lock()
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{
URI: resourceURI,
Name: fmt.Sprintf("%s工具文档", tool.Name),
Description: tool.Description,
MimeType: "text/plain",
}
}
// HandleHTTP 处理HTTP请求
func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
body, err := io.ReadAll(r.Body)
if err != nil {
s.sendError(w, nil, -32700, "Parse error", err.Error())
return
}
var msg Message
if err := json.Unmarshal(body, &msg); err != nil {
s.sendError(w, nil, -32700, "Parse error", err.Error())
return
}
// 处理消息
response := s.handleMessage(&msg)
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(response)
}
// handleMessage 处理MCP消息
func (s *Server) handleMessage(msg *Message) *Message {
if msg.ID == "" {
msg.ID = uuid.New().String()
}
switch msg.Method {
case "initialize":
return s.handleInitialize(msg)
case "tools/list":
return s.handleListTools(msg)
case "tools/call":
return s.handleCallTool(msg)
case "prompts/list":
return s.handleListPrompts(msg)
case "prompts/get":
return s.handleGetPrompt(msg)
case "resources/list":
return s.handleListResources(msg)
case "resources/read":
return s.handleReadResource(msg)
case "sampling/request":
return s.handleSamplingRequest(msg)
default:
return &Message{
ID: msg.ID,
Type: MessageTypeError,
Error: &Error{Code: -32601, Message: "Method not found"},
}
}
}
// handleInitialize 处理初始化请求
func (s *Server) handleInitialize(msg *Message) *Message {
var req InitializeRequest
if err := json.Unmarshal(msg.Params, &req); err != nil {
return &Message{
ID: msg.ID,
Type: MessageTypeError,
Error: &Error{Code: -32602, Message: "Invalid params"},
}
}
response := InitializeResponse{
ProtocolVersion: ProtocolVersion,
Capabilities: ServerCapabilities{
Tools: map[string]interface{}{
"listChanged": true,
},
Prompts: map[string]interface{}{
"listChanged": true,
},
Resources: map[string]interface{}{
"subscribe": true,
"listChanged": true,
},
Sampling: map[string]interface{}{},
},
ServerInfo: ServerInfo{
Name: "CyberStrikeAI",
Version: "1.0.0",
},
}
result, _ := json.Marshal(response)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: result,
}
}
// handleListTools 处理列出工具请求
func (s *Server) handleListTools(msg *Message) *Message {
s.mu.RLock()
tools := make([]Tool, 0, len(s.toolDefs))
for _, tool := range s.toolDefs {
tools = append(tools, tool)
}
s.mu.RUnlock()
response := ListToolsResponse{Tools: tools}
result, _ := json.Marshal(response)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: result,
}
}
// handleCallTool 处理工具调用请求
func (s *Server) handleCallTool(msg *Message) *Message {
var req CallToolRequest
if err := json.Unmarshal(msg.Params, &req); err != nil {
return &Message{
ID: msg.ID,
Type: MessageTypeError,
Error: &Error{Code: -32602, Message: "Invalid params"},
}
}
// 创建执行记录
executionID := uuid.New().String()
execution := &ToolExecution{
ID: executionID,
ToolName: req.Name,
Arguments: req.Arguments,
Status: "running",
StartTime: time.Now(),
}
s.mu.Lock()
s.executions[executionID] = execution
s.mu.Unlock()
// 更新统计
s.updateStats(req.Name, false)
// 执行工具
s.mu.RLock()
handler, exists := s.tools[req.Name]
s.mu.RUnlock()
if !exists {
execution.Status = "failed"
execution.Error = "Tool not found"
now := time.Now()
execution.EndTime = &now
return &Message{
ID: msg.ID,
Type: MessageTypeError,
Error: &Error{Code: -32601, Message: "Tool not found"},
}
}
// 同步执行所有工具，确保错误能正确返回
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
defer cancel()
s.logger.Info("开始执行工具",
zap.String("toolName", req.Name),
zap.Any("arguments", req.Arguments),
)
result, err := handler(ctx, req.Arguments)
s.mu.Lock()
now := time.Now()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
if err != nil {
execution.Status = "failed"
execution.Error = err.Error()
s.updateStats(req.Name, true)
s.mu.Unlock()
s.logger.Error("工具执行失败",
zap.String("toolName", req.Name),
zap.Error(err),
)
// 返回错误结果
errorResult, _ := json.Marshal(CallToolResponse{
Content: []Content{
{Type: "text", Text: fmt.Sprintf("工具执行失败: %v", err)},
},
IsError: true,
})
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: errorResult,
}
}
// 检查result是否为错误
if result != nil && result.IsError {
execution.Status = "failed"
if len(result.Content) > 0 {
execution.Error = result.Content[0].Text
}
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{
Content: []Content{
{Type: "text", Text: "工具执行完成，但未返回结果"},
},
}
}
resultJSON, _ := json.Marshal(CallToolResponse{
Content: result.Content,
IsError: result.IsError,
})
s.logger.Info("工具执行完成",
zap.String("toolName", req.Name),
zap.Bool("isError", result.IsError),
)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: resultJSON,
}
}
// updateStats 更新统计信息
func (s *Server) updateStats(toolName string, failed bool) {
if s.stats[toolName] == nil {
s.stats[toolName] = &ToolStats{
ToolName: toolName,
}
}
stats := s.stats[toolName]
stats.TotalCalls++
now := time.Now()
stats.LastCallTime = &now
if failed {
stats.FailedCalls++
} else {
stats.SuccessCalls++
}
}
// GetExecution 获取执行记录
func (s *Server) GetExecution(id string) (*ToolExecution, bool) {
s.mu.RLock()
defer s.mu.RUnlock()
exec, exists := s.executions[id]
return exec, exists
}
// GetAllExecutions 获取所有执行记录
func (s *Server) GetAllExecutions() []*ToolExecution {
s.mu.RLock()
defer s.mu.RUnlock()
executions := make([]*ToolExecution, 0, len(s.executions))
for _, exec := range s.executions {
executions = append(executions, exec)
}
return executions
}
// GetStats 获取统计信息
func (s *Server) GetStats() map[string]*ToolStats {
s.mu.RLock()
defer s.mu.RUnlock()
stats := make(map[string]*ToolStats)
for k, v := range s.stats {
stats[k] = v
}
return stats
}
// CallTool 直接调用工具（用于内部调用）
func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
s.mu.RLock()
handler, exists := s.tools[toolName]
s.mu.RUnlock()
if !exists {
return nil, "", fmt.Errorf("工具 %s 未找到", toolName)
}
// 创建执行记录
executionID := uuid.New().String()
execution := &ToolExecution{
ID: executionID,
ToolName: toolName,
Arguments: args,
Status: "running",
StartTime: time.Now(),
}
s.mu.Lock()
s.executions[executionID] = execution
s.mu.Unlock()
// 更新统计
s.updateStats(toolName, false)
// 执行工具
result, err := handler(ctx, args)
s.mu.Lock()
now := time.Now()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
if err != nil {
execution.Status = "failed"
execution.Error = err.Error()
s.updateStats(toolName, true)
s.mu.Unlock()
return nil, executionID, err
} else {
execution.Status = "completed"
execution.Result = result
s.updateStats(toolName, false)
s.mu.Unlock()
return result, executionID, nil
}
}
// initDefaultPrompts 初始化默认提示词模板
func (s *Server) initDefaultPrompts() {
s.mu.Lock()
defer s.mu.Unlock()
// 网络安全测试提示词
s.prompts["security_scan"] = &Prompt{
Name: "security_scan",
Description: "生成网络安全扫描任务的提示词",
Arguments: []PromptArgument{
{Name: "target", Description: "扫描目标（IP地址或域名）", Required: true},
{Name: "scan_type", Description: "扫描类型（port, vuln, web等）", Required: false},
},
}
// 渗透测试提示词
s.prompts["penetration_test"] = &Prompt{
Name: "penetration_test",
Description: "生成渗透测试任务的提示词",
Arguments: []PromptArgument{
{Name: "target", Description: "测试目标", Required: true},
{Name: "scope", Description: "测试范围", Required: false},
},
}
}
// initDefaultResources 初始化默认资源
// 注意：工具资源现在在 RegisterTool 时自动创建，此函数保留用于其他非工具资源
func (s *Server) initDefaultResources() {
// 工具资源已改为在 RegisterTool 时自动创建，无需在此硬编码
}
// handleListPrompts 处理列出提示词请求
func (s *Server) handleListPrompts(msg *Message) *Message {
s.mu.RLock()
prompts := make([]Prompt, 0, len(s.prompts))
for _, prompt := range s.prompts {
prompts = append(prompts, *prompt)
}
s.mu.RUnlock()
response := ListPromptsResponse{
Prompts: prompts,
}
result, _ := json.Marshal(response)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: result,
}
}
// handleGetPrompt 处理获取提示词请求
func (s *Server) handleGetPrompt(msg *Message) *Message {
var req GetPromptRequest
if err := json.Unmarshal(msg.Params, &req); err != nil {
return &Message{
ID: msg.ID,
Type: MessageTypeError,
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,
Type: MessageTypeError,
Error: &Error{Code: -32601, Message: "Prompt not found"},
}
}
// 根据提示词名称生成消息
messages := s.generatePromptMessages(prompt, req.Arguments)
response := GetPromptResponse{
Messages: messages,
}
result, _ := json.Marshal(response)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: result,
}
}
// generatePromptMessages 生成提示词消息
func (s *Server) generatePromptMessages(prompt *Prompt, args map[string]interface{}) []PromptMessage {
messages := []PromptMessage{}
switch prompt.Name {
case "security_scan":
target, _ := args["target"].(string)
scanType, _ := args["scan_type"].(string)
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
}
// handleListResources 处理列出资源请求
func (s *Server) handleListResources(msg *Message) *Message {
s.mu.RLock()
resources := make([]Resource, 0, len(s.resources))
for _, resource := range s.resources {
resources = append(resources, *resource)
}
s.mu.RUnlock()
response := ListResourcesResponse{
Resources: resources,
}
result, _ := json.Marshal(response)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: result,
}
}
// handleReadResource 处理读取资源请求
func (s *Server) handleReadResource(msg *Message) *Message {
var req ReadResourceRequest
if err := json.Unmarshal(msg.Params, &req); err != nil {
return &Message{
ID: msg.ID,
Type: MessageTypeError,
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,
Type: MessageTypeError,
Error: &Error{Code: -32601, Message: "Resource not found"},
}
}
// 生成资源内容
content := s.generateResourceContent(resource)
response := ReadResourceResponse{
Contents: []ResourceContent{content},
}
result, _ := json.Marshal(response)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: result,
}
}
// generateResourceContent 生成资源内容
func (s *Server) generateResourceContent(resource *Resource) ResourceContent {
content := ResourceContent{
URI: resource.URI,
MimeType: resource.MimeType,
}
// 如果是工具资源，生成详细文档
if strings.HasPrefix(resource.URI, "tool://") {
toolName := strings.TrimPrefix(resource.URI, "tool://")
content.Text = s.generateToolDocumentation(toolName, resource)
} else {
// 其他资源使用描述或默认内容
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)
}
}
}
}
}
return doc
}
return resource.Description
}
}
// handleSamplingRequest 处理采样请求
func (s *Server) handleSamplingRequest(msg *Message) *Message {
var req SamplingRequest
if err := json.Unmarshal(msg.Params, &req); err != nil {
return &Message{
ID: msg.ID,
Type: MessageTypeError,
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{
{
Type: "text",
Text: "采样功能需要配置LLM服务。请使用Agent Loop API进行AI对话。",
},
},
StopReason: "length",
}
result, _ := json.Marshal(response)
return &Message{
ID: msg.ID,
Type: MessageTypeResponse,
Version: "2.0",
Result: result,
}
}
// RegisterPrompt 注册提示词模板
func (s *Server) RegisterPrompt(prompt *Prompt) {
s.mu.Lock()
defer s.mu.Unlock()
s.prompts[prompt.Name] = prompt
}
// RegisterResource 注册资源
func (s *Server) RegisterResource(resource *Resource) {
s.mu.Lock()
defer s.mu.Unlock()
s.resources[resource.URI] = resource
}
// sendError 发送错误响应
func (s *Server) sendError(w http.ResponseWriter, id interface{}, code int, message, data string) {
response := Message{
ID: fmt.Sprintf("%v", id),
Type: MessageTypeError,
Error: &Error{Code: code, Message: message, Data: data},
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(response)
}
internal/mcp/types.go
+232
View File
@@ -0,0 +1,232 @@
package mcp
import (
"encoding/json"
"time"
)
// MCP消息类型
const (
MessageTypeRequest = "request"
MessageTypeResponse = "response"
MessageTypeError = "error"
MessageTypeNotify = "notify"
)
// MCP协议版本
const ProtocolVersion = "2024-11-05"
// Message 表示MCP消息
type Message struct {
ID string `json:"id,omitempty"`
Type string `json:"type"`
Method string `json:"method,omitempty"`
Params json.RawMessage `json:"params,omitempty"`
Result json.RawMessage `json:"result,omitempty"`
Error *Error `json:"error,omitempty"`
Version string `json:"jsonrpc,omitempty"`
}
// Error 表示MCP错误
type Error struct {
Code int `json:"code"`
Message string `json:"message"`
Data interface{} `json:"data,omitempty"`
}
// Tool 表示MCP工具定义
type Tool struct {
Name string `json:"name"`
Description string `json:"description"`
InputSchema map[string]interface{} `json:"inputSchema"`
}
// ToolCall 表示工具调用
type ToolCall struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments"`
}
// ToolResult 表示工具执行结果
type ToolResult struct {
Content []Content `json:"content"`
IsError bool `json:"isError,omitempty"`
}
// Content 表示内容
type Content struct {
Type string `json:"type"`
Text string `json:"text"`
}
// InitializeRequest 初始化请求
type InitializeRequest struct {
ProtocolVersion string `json:"protocolVersion"`
Capabilities map[string]interface{} `json:"capabilities"`
ClientInfo ClientInfo `json:"clientInfo"`
}
// ClientInfo 客户端信息
type ClientInfo struct {
Name string `json:"name"`
Version string `json:"version"`
}
// InitializeResponse 初始化响应
type InitializeResponse struct {
ProtocolVersion string `json:"protocolVersion"`
Capabilities ServerCapabilities `json:"capabilities"`
ServerInfo ServerInfo `json:"serverInfo"`
}
// ServerCapabilities 服务器能力
type ServerCapabilities struct {
Tools map[string]interface{} `json:"tools,omitempty"`
Prompts map[string]interface{} `json:"prompts,omitempty"`
Resources map[string]interface{} `json:"resources,omitempty"`
Sampling map[string]interface{} `json:"sampling,omitempty"`
}
// ServerInfo 服务器信息
type ServerInfo struct {
Name string `json:"name"`
Version string `json:"version"`
}
// ListToolsRequest 列出工具请求
type ListToolsRequest struct{}
// ListToolsResponse 列出工具响应
type ListToolsResponse struct {
Tools []Tool `json:"tools"`
}
// ListPromptsResponse 列出提示词响应
type ListPromptsResponse struct {
Prompts []Prompt `json:"prompts"`
}
// ListResourcesResponse 列出资源响应
type ListResourcesResponse struct {
Resources []Resource `json:"resources"`
}
// CallToolRequest 调用工具请求
type CallToolRequest struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments"`
}
// CallToolResponse 调用工具响应
type CallToolResponse struct {
Content []Content `json:"content"`
IsError bool `json:"isError,omitempty"`
}
// ToolExecution 工具执行记录
type ToolExecution struct {
ID string `json:"id"`
ToolName string `json:"toolName"`
Arguments map[string]interface{} `json:"arguments"`
Status string `json:"status"` // pending, running, completed, failed
Result *ToolResult `json:"result,omitempty"`
Error string `json:"error,omitempty"`
StartTime time.Time `json:"startTime"`
EndTime *time.Time `json:"endTime,omitempty"`
Duration time.Duration `json:"duration,omitempty"`
}
// ToolStats 工具统计信息
type ToolStats struct {
ToolName string `json:"toolName"`
TotalCalls int `json:"totalCalls"`
SuccessCalls int `json:"successCalls"`
FailedCalls int `json:"failedCalls"`
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
}
// Prompt 提示词模板
type Prompt struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Arguments []PromptArgument `json:"arguments,omitempty"`
}
// PromptArgument 提示词参数
type PromptArgument struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Required bool `json:"required,omitempty"`
}
// GetPromptRequest 获取提示词请求
type GetPromptRequest struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments,omitempty"`
}
// GetPromptResponse 获取提示词响应
type GetPromptResponse struct {
Messages []PromptMessage `json:"messages"`
}
// PromptMessage 提示词消息
type PromptMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
// Resource 资源
type Resource struct {
URI string `json:"uri"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
MimeType string `json:"mimeType,omitempty"`
}
// ReadResourceRequest 读取资源请求
type ReadResourceRequest struct {
URI string `json:"uri"`
}
// ReadResourceResponse 读取资源响应
type ReadResourceResponse struct {
Contents []ResourceContent `json:"contents"`
}
// ResourceContent 资源内容
type ResourceContent struct {
URI string `json:"uri"`
MimeType string `json:"mimeType,omitempty"`
Text string `json:"text,omitempty"`
Blob string `json:"blob,omitempty"`
}
// SamplingRequest 采样请求
type SamplingRequest struct {
Messages []SamplingMessage `json:"messages"`
Model string `json:"model,omitempty"`
MaxTokens int `json:"maxTokens,omitempty"`
Temperature float64 `json:"temperature,omitempty"`
TopP float64 `json:"topP,omitempty"`
}
// SamplingMessage 采样消息
type SamplingMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
// SamplingResponse 采样响应
type SamplingResponse struct {
Content []SamplingContent `json:"content"`
Model string `json:"model,omitempty"`
StopReason string `json:"stopReason,omitempty"`
}
// SamplingContent 采样内容
type SamplingContent struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
}
internal/security/executor.go
+730
View File
@@ -0,0 +1,730 @@
package security
import (
"context"
"fmt"
"os/exec"
"strings"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
// Executor 安全工具执行器
type Executor struct {
config *config.SecurityConfig
mcpServer *mcp.Server
logger *zap.Logger
}
// NewExecutor 创建新的执行器
func NewExecutor(cfg *config.SecurityConfig, mcpServer *mcp.Server, logger *zap.Logger) *Executor {
return &Executor{
config: cfg,
mcpServer: mcpServer,
logger: logger,
}
}
// ExecuteTool 执行安全工具
func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[string]interface{}) (*mcp.ToolResult, error) {
e.logger.Info("ExecuteTool被调用",
zap.String("toolName", toolName),
zap.Any("args", args),
)
// 特殊处理：exec工具直接执行系统命令
if toolName == "exec" {
e.logger.Info("执行exec工具")
return e.executeSystemCommand(ctx, args)
}
// 查找工具配置
var toolConfig *config.ToolConfig
for i := range e.config.Tools {
if e.config.Tools[i].Name == toolName && e.config.Tools[i].Enabled {
toolConfig = &e.config.Tools[i]
break
}
}
if toolConfig == nil {
e.logger.Error("工具未找到或未启用",
zap.String("toolName", toolName),
zap.Int("totalTools", len(e.config.Tools)),
)
return nil, fmt.Errorf("工具 %s 未找到或未启用", toolName)
}
e.logger.Info("找到工具配置",
zap.String("toolName", toolName),
zap.String("command", toolConfig.Command),
zap.Strings("args", toolConfig.Args),
)
// 构建命令 - 根据工具类型使用不同的参数格式
cmdArgs := e.buildCommandArgs(toolName, toolConfig, args)
e.logger.Info("构建命令参数完成",
zap.String("toolName", toolName),
zap.Strings("cmdArgs", cmdArgs),
zap.Int("argsCount", len(cmdArgs)),
)
// 验证命令参数
if len(cmdArgs) == 0 {
e.logger.Warn("命令参数为空",
zap.String("toolName", toolName),
zap.Any("inputArgs", args),
)
return &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: fmt.Sprintf("错误: 工具 %s 缺少必需的参数。接收到的参数: %v", toolName, args),
},
},
IsError: true,
}, nil
}
// 执行命令
cmd := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
e.logger.Info("执行安全工具",
zap.String("tool", toolName),
zap.Strings("args", cmdArgs),
)
output, err := cmd.CombinedOutput()
if err != nil {
e.logger.Error("工具执行失败",
zap.String("tool", toolName),
zap.Error(err),
zap.String("output", string(output)),
)
return &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: fmt.Sprintf("工具执行失败: %v\n输出: %s", err, string(output)),
},
},
IsError: true,
}, nil
}
e.logger.Info("工具执行成功",
zap.String("tool", toolName),
zap.String("output", string(output)),
)
return &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: string(output),
},
},
IsError: false,
}, nil
}
// RegisterTools 注册工具到MCP服务器
func (e *Executor) RegisterTools(mcpServer *mcp.Server) {
e.logger.Info("开始注册工具",
zap.Int("totalTools", len(e.config.Tools)),
)
for i, toolConfig := range e.config.Tools {
if !toolConfig.Enabled {
e.logger.Debug("跳过未启用的工具",
zap.String("tool", toolConfig.Name),
)
continue
}
// 创建工具配置的副本，避免闭包问题
toolName := toolConfig.Name
toolConfigCopy := toolConfig
tool := mcp.Tool{
Name: toolConfigCopy.Name,
Description: toolConfigCopy.Description,
InputSchema: e.buildInputSchema(&toolConfigCopy),
}
handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
e.logger.Info("工具handler被调用",
zap.String("toolName", toolName),
zap.Any("args", args),
)
return e.ExecuteTool(ctx, toolName, args)
}
mcpServer.RegisterTool(tool, handler)
e.logger.Info("注册安全工具成功",
zap.String("tool", toolConfigCopy.Name),
zap.String("command", toolConfigCopy.Command),
zap.Int("index", i),
)
}
e.logger.Info("工具注册完成",
zap.Int("registeredCount", len(e.config.Tools)),
)
}
// buildCommandArgs 构建命令参数
func (e *Executor) buildCommandArgs(toolName string, toolConfig *config.ToolConfig, args map[string]interface{}) []string {
cmdArgs := make([]string, 0)
// 如果配置中定义了参数映射，使用配置中的映射规则
if len(toolConfig.Parameters) > 0 {
// 先添加固定参数
cmdArgs = append(cmdArgs, toolConfig.Args...)
// 按位置参数排序
positionalParams := make([]config.ParameterConfig, 0)
flagParams := make([]config.ParameterConfig, 0)
for _, param := range toolConfig.Parameters {
if param.Position != nil {
positionalParams = append(positionalParams, param)
} else {
flagParams = append(flagParams, param)
}
}
// 对位置参数按位置排序
for i := 0; i < len(positionalParams); i++ {
for _, param := range positionalParams {
if param.Position != nil && *param.Position == i {
value := e.getParamValue(args, param)
if value == nil {
if param.Required {
// 必需参数缺失，返回空数组让上层处理错误
e.logger.Warn("缺少必需的位置参数",
zap.String("tool", toolName),
zap.String("param", param.Name),
zap.Int("position", *param.Position),
)
return []string{}
}
break
}
cmdArgs = append(cmdArgs, e.formatParamValue(param, value))
break
}
}
}
// 处理标志参数
for _, param := range flagParams {
value := e.getParamValue(args, param)
if value == nil {
if param.Required {
// 必需参数缺失，返回空数组让上层处理错误
e.logger.Warn("缺少必需的标志参数",
zap.String("tool", toolName),
zap.String("param", param.Name),
)
return []string{}
}
continue
}
// 布尔值特殊处理：如果为 false，跳过；如果为 true，只添加标志
if param.Type == "bool" {
if boolVal, ok := value.(bool); ok {
if !boolVal {
continue // false 时不添加任何参数
}
// true 时只添加标志，不添加值
if param.Flag != "" {
cmdArgs = append(cmdArgs, param.Flag)
}
continue
}
}
format := param.Format
if format == "" {
format = "flag" // 默认格式
}
switch format {
case "flag":
// --flag value 或 -f value
if param.Flag != "" {
cmdArgs = append(cmdArgs, param.Flag)
}
formattedValue := e.formatParamValue(param, value)
if formattedValue != "" {
cmdArgs = append(cmdArgs, formattedValue)
}
case "combined":
// --flag=value 或 -f=value
if param.Flag != "" {
cmdArgs = append(cmdArgs, fmt.Sprintf("%s=%s", param.Flag, e.formatParamValue(param, value)))
} else {
cmdArgs = append(cmdArgs, e.formatParamValue(param, value))
}
case "template":
// 使用模板字符串
if param.Template != "" {
template := param.Template
template = strings.ReplaceAll(template, "{flag}", param.Flag)
template = strings.ReplaceAll(template, "{value}", e.formatParamValue(param, value))
template = strings.ReplaceAll(template, "{name}", param.Name)
cmdArgs = append(cmdArgs, strings.Fields(template)...)
} else {
// 如果没有模板，使用默认格式
if param.Flag != "" {
cmdArgs = append(cmdArgs, param.Flag)
}
cmdArgs = append(cmdArgs, e.formatParamValue(param, value))
}
case "positional":
// 位置参数（已在上面处理）
cmdArgs = append(cmdArgs, e.formatParamValue(param, value))
default:
// 默认：直接添加值
cmdArgs = append(cmdArgs, e.formatParamValue(param, value))
}
}
return cmdArgs
}
// 向后兼容：如果没有定义参数，使用旧的硬编码逻辑
switch toolName {
case "nmap":
// nmap -sT -sV -sC target [ports]
// 使用 -sT (TCP连接扫描) 而不是 -sS (SYN扫描)，因为 -sS 需要root权限
e.logger.Debug("处理nmap参数",
zap.Any("args", args),
)
// 尝试多种方式获取target参数
var target string
var ok bool
// 方式1: 直接获取target
if target, ok = args["target"].(string); !ok || target == "" {
// 方式2: 尝试从tool字段获取（兼容某些格式）
if toolVal, exists := args["tool"]; exists {
if toolMap, ok := toolVal.(map[string]interface{}); ok {
if t, ok := toolMap["target"].(string); ok {
target = t
}
}
}
}
if target == "" {
e.logger.Warn("nmap缺少target参数",
zap.Any("args", args),
)
return cmdArgs // 返回空数组，让上层处理错误
}
e.logger.Debug("提取到target",
zap.String("target", target),
)
// 处理URL格式的目标（提取域名）
if strings.HasPrefix(target, "http://") || strings.HasPrefix(target, "https://") {
// 提取域名部分
target = strings.TrimPrefix(target, "http://")
target = strings.TrimPrefix(target, "https://")
// 移除路径部分
if idx := strings.Index(target, "/"); idx != -1 {
target = target[:idx]
}
}
// 添加扫描选项：-sT (TCP连接扫描，不需要root权限), -sV (版本检测), -sC (默认脚本)
cmdArgs = append(cmdArgs, "-sT", "-sV", "-sC")
// 添加端口范围（如果指定）
if ports, ok := args["ports"].(string); ok && ports != "" {
cmdArgs = append(cmdArgs, "-p", ports)
}
// 添加目标
cmdArgs = append(cmdArgs, target)
e.logger.Debug("nmap命令参数构建完成",
zap.Strings("cmdArgs", cmdArgs),
)
case "sqlmap":
// sqlmap -u url
if url, ok := args["url"].(string); ok {
cmdArgs = append(cmdArgs, "-u", url, "--batch", "--level=3", "--risk=2")
}
case "nikto":
// nikto -h target
if target, ok := args["target"].(string); ok {
cmdArgs = append(cmdArgs, "-h", target)
}
case "dirb":
// dirb url
if url, ok := args["url"].(string); ok {
cmdArgs = append(cmdArgs, url)
}
default:
// 通用处理
cmdArgs = append(cmdArgs, toolConfig.Args...)
for key, value := range args {
if key == "_tool_name" {
continue
}
cmdArgs = append(cmdArgs, fmt.Sprintf("--%s", key))
if strValue, ok := value.(string); ok {
cmdArgs = append(cmdArgs, strValue)
} else {
cmdArgs = append(cmdArgs, fmt.Sprintf("%v", value))
}
}
}
return cmdArgs
}
// getParamValue 获取参数值，支持默认值
func (e *Executor) getParamValue(args map[string]interface{}, param config.ParameterConfig) interface{} {
// 从参数中获取值
if value, ok := args[param.Name]; ok && value != nil {
return value
}
// 如果参数是必需的但没有提供，返回 nil（让上层处理错误）
if param.Required {
return nil
}
// 返回默认值
return param.Default
}
// formatParamValue 格式化参数值
func (e *Executor) formatParamValue(param config.ParameterConfig, value interface{}) string {
switch param.Type {
case "bool":
// 布尔值应该在上层处理，这里不应该被调用
if boolVal, ok := value.(bool); ok {
return fmt.Sprintf("%v", boolVal)
}
return "false"
case "array":
// 数组：转换为逗号分隔的字符串
if arr, ok := value.([]interface{}); ok {
strs := make([]string, 0, len(arr))
for _, item := range arr {
strs = append(strs, fmt.Sprintf("%v", item))
}
return strings.Join(strs, ",")
}
return fmt.Sprintf("%v", value)
default:
return fmt.Sprintf("%v", value)
}
}
// executeSystemCommand 执行系统命令
func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
// 获取命令
command, ok := args["command"].(string)
if !ok {
return &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: "错误: 缺少command参数",
},
},
IsError: true,
}, nil
}
if command == "" {
return &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: "错误: command参数不能为空",
},
},
IsError: true,
}, nil
}
// 安全检查：记录执行的命令
e.logger.Warn("执行系统命令",
zap.String("command", command),
)
// 获取shell类型（可选，默认为sh）
shell := "sh"
if s, ok := args["shell"].(string); ok && s != "" {
shell = s
}
// 获取工作目录（可选）
workDir := ""
if wd, ok := args["workdir"].(string); ok && wd != "" {
workDir = wd
}
// 构建命令
var cmd *exec.Cmd
if workDir != "" {
cmd = exec.CommandContext(ctx, shell, "-c", command)
cmd.Dir = workDir
} else {
cmd = exec.CommandContext(ctx, shell, "-c", command)
}
// 执行命令
e.logger.Info("执行系统命令",
zap.String("command", command),
zap.String("shell", shell),
zap.String("workdir", workDir),
)
output, err := cmd.CombinedOutput()
if err != nil {
e.logger.Error("系统命令执行失败",
zap.String("command", command),
zap.Error(err),
zap.String("output", string(output)),
)
return &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: fmt.Sprintf("命令执行失败: %v\n输出: %s", err, string(output)),
},
},
IsError: true,
}, nil
}
e.logger.Info("系统命令执行成功",
zap.String("command", command),
zap.String("output_length", fmt.Sprintf("%d", len(output))),
)
return &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: string(output),
},
},
IsError: false,
}, nil
}
// buildInputSchema 构建输入模式
func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]interface{} {
schema := map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{},
"required": []string{},
}
// 如果配置中定义了参数，优先使用配置中的参数定义
if len(toolConfig.Parameters) > 0 {
properties := make(map[string]interface{})
required := []string{}
for _, param := range toolConfig.Parameters {
prop := map[string]interface{}{
"type": param.Type,
"description": param.Description,
}
// 添加默认值
if param.Default != nil {
prop["default"] = param.Default
}
// 添加枚举选项
if len(param.Options) > 0 {
prop["enum"] = param.Options
}
properties[param.Name] = prop
// 添加到必需参数列表
if param.Required {
required = append(required, param.Name)
}
}
schema["properties"] = properties
schema["required"] = required
return schema
}
// 向后兼容：如果没有定义参数，使用旧的硬编码逻辑
switch toolConfig.Name {
case "nmap":
schema["properties"] = map[string]interface{}{
"target": map[string]interface{}{
"type": "string",
"description": "目标IP地址或域名",
},
"ports": map[string]interface{}{
"type": "string",
"description": "端口范围，例如: 1-1000",
},
}
schema["required"] = []string{"target"}
case "sqlmap":
schema["properties"] = map[string]interface{}{
"url": map[string]interface{}{
"type": "string",
"description": "目标URL",
},
}
schema["required"] = []string{"url"}
case "nikto", "dirb":
schema["properties"] = map[string]interface{}{
"target": map[string]interface{}{
"type": "string",
"description": "目标URL",
},
}
schema["required"] = []string{"target"}
case "exec":
schema["properties"] = map[string]interface{}{
"command": map[string]interface{}{
"type": "string",
"description": "要执行的系统命令",
},
"shell": map[string]interface{}{
"type": "string",
"description": "使用的shell（可选，默认为sh）",
},
"workdir": map[string]interface{}{
"type": "string",
"description": "工作目录（可选）",
},
}
schema["required"] = []string{"command"}
}
return schema
}
// Vulnerability 漏洞信息
type Vulnerability struct {
ID string `json:"id"`
Type string `json:"type"`
Severity string `json:"severity"` // low, medium, high, critical
Title string `json:"title"`
Description string `json:"description"`
Target string `json:"target"`
FoundAt time.Time `json:"foundAt"`
Details string `json:"details"`
}
// 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
}
// GetVulnerabilityReport 生成漏洞报告
func (e *Executor) GetVulnerabilityReport(vulnerabilities []Vulnerability) map[string]interface{} {
severityCount := map[string]int{
"critical": 0,
"high": 0,
"medium": 0,
"low": 0,
}
for _, vuln := range vulnerabilities {
severityCount[vuln.Severity]++
}
return map[string]interface{}{
"total": len(vulnerabilities),
"severityCount": severityCount,
"vulnerabilities": vulnerabilities,
"generatedAt": time.Now(),
}
}