mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-13 22:50:24 +02:00
Add files via upload
This commit is contained in:
公明
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user