mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-30 13:09:40 +02:00
Add files via upload
This commit is contained in:
44 files changed
+10016
-72
No files matched your search
@@ -0,0 +1,227 @@
|
||||
package skills
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Manager Skills管理器
|
||||
type Manager struct {
|
||||
skillsDir string
|
||||
logger *zap.Logger
|
||||
skills map[string]*Skill // 缓存已加载的skills
|
||||
}
|
||||
|
||||
// Skill Skill定义
|
||||
type Skill struct {
|
||||
Name string // Skill名称
|
||||
Description string // Skill描述
|
||||
Content string // Skill内容（从SKILL.md中提取）
|
||||
Path string // Skill路径
|
||||
}
|
||||
|
||||
// NewManager 创建新的Skills管理器
|
||||
func NewManager(skillsDir string, logger *zap.Logger) *Manager {
|
||||
return &Manager{
|
||||
skillsDir: skillsDir,
|
||||
logger: logger,
|
||||
skills: make(map[string]*Skill),
|
||||
}
|
||||
}
|
||||
|
||||
// LoadSkill 加载单个skill
|
||||
func (m *Manager) LoadSkill(skillName string) (*Skill, error) {
|
||||
// 检查缓存
|
||||
if skill, exists := m.skills[skillName]; exists {
|
||||
return skill, nil
|
||||
}
|
||||
|
||||
// 构建skill路径
|
||||
skillPath := filepath.Join(m.skillsDir, skillName)
|
||||
|
||||
// 检查目录是否存在
|
||||
if _, err := os.Stat(skillPath); os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("skill %s not found", skillName)
|
||||
}
|
||||
|
||||
// 查找SKILL.md文件
|
||||
skillFile := filepath.Join(skillPath, "SKILL.md")
|
||||
if _, err := os.Stat(skillFile); os.IsNotExist(err) {
|
||||
// 尝试其他可能的文件名
|
||||
alternatives := []string{
|
||||
filepath.Join(skillPath, "skill.md"),
|
||||
filepath.Join(skillPath, "README.md"),
|
||||
filepath.Join(skillPath, "readme.md"),
|
||||
}
|
||||
found := false
|
||||
for _, alt := range alternatives {
|
||||
if _, err := os.Stat(alt); err == nil {
|
||||
skillFile = alt
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
return nil, fmt.Errorf("skill file not found for %s", skillName)
|
||||
}
|
||||
}
|
||||
|
||||
// 读取skill文件
|
||||
content, err := os.ReadFile(skillFile)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read skill file: %w", err)
|
||||
}
|
||||
|
||||
// 解析skill内容
|
||||
skill := m.parseSkillContent(string(content), skillName, skillPath)
|
||||
|
||||
// 缓存skill
|
||||
m.skills[skillName] = skill
|
||||
|
||||
return skill, nil
|
||||
}
|
||||
|
||||
// LoadSkills 批量加载skills
|
||||
func (m *Manager) LoadSkills(skillNames []string) ([]*Skill, error) {
|
||||
var skills []*Skill
|
||||
var errors []string
|
||||
|
||||
for _, name := range skillNames {
|
||||
skill, err := m.LoadSkill(name)
|
||||
if err != nil {
|
||||
errors = append(errors, fmt.Sprintf("failed to load skill %s: %v", name, err))
|
||||
m.logger.Warn("加载skill失败", zap.String("skill", name), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
skills = append(skills, skill)
|
||||
}
|
||||
|
||||
if len(errors) > 0 && len(skills) == 0 {
|
||||
return nil, fmt.Errorf("failed to load any skills: %s", strings.Join(errors, "; "))
|
||||
}
|
||||
|
||||
return skills, nil
|
||||
}
|
||||
|
||||
// ListSkills 列出所有可用的skills
|
||||
func (m *Manager) ListSkills() ([]string, error) {
|
||||
if _, err := os.Stat(m.skillsDir); os.IsNotExist(err) {
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(m.skillsDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read skills directory: %w", err)
|
||||
}
|
||||
|
||||
var skills []string
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
|
||||
skillName := entry.Name()
|
||||
// 检查是否有SKILL.md文件
|
||||
skillFile := filepath.Join(m.skillsDir, skillName, "SKILL.md")
|
||||
if _, err := os.Stat(skillFile); err == nil {
|
||||
skills = append(skills, skillName)
|
||||
continue
|
||||
}
|
||||
|
||||
// 尝试其他可能的文件名
|
||||
alternatives := []string{
|
||||
filepath.Join(m.skillsDir, skillName, "skill.md"),
|
||||
filepath.Join(m.skillsDir, skillName, "README.md"),
|
||||
filepath.Join(m.skillsDir, skillName, "readme.md"),
|
||||
}
|
||||
for _, alt := range alternatives {
|
||||
if _, err := os.Stat(alt); err == nil {
|
||||
skills = append(skills, skillName)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return skills, nil
|
||||
}
|
||||
|
||||
// parseSkillContent 解析skill内容
|
||||
// 支持YAML front matter格式，类似goskills
|
||||
func (m *Manager) parseSkillContent(content, skillName, skillPath string) *Skill {
|
||||
skill := &Skill{
|
||||
Name: skillName,
|
||||
Path: skillPath,
|
||||
}
|
||||
|
||||
// 检查是否有YAML front matter
|
||||
if strings.HasPrefix(content, "---") {
|
||||
parts := strings.SplitN(content, "---", 3)
|
||||
if len(parts) >= 3 {
|
||||
// 解析front matter（简单实现，只提取name和description）
|
||||
frontMatter := parts[1]
|
||||
lines := strings.Split(frontMatter, "\n")
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
if strings.HasPrefix(line, "name:") {
|
||||
name := strings.TrimSpace(strings.TrimPrefix(line, "name:"))
|
||||
name = strings.Trim(name, `"'"`)
|
||||
if name != "" {
|
||||
skill.Name = name
|
||||
}
|
||||
} else if strings.HasPrefix(line, "description:") {
|
||||
desc := strings.TrimSpace(strings.TrimPrefix(line, "description:"))
|
||||
desc = strings.Trim(desc, `"'"`)
|
||||
skill.Description = desc
|
||||
}
|
||||
}
|
||||
// 剩余部分是内容
|
||||
if len(parts) == 3 {
|
||||
skill.Content = strings.TrimSpace(parts[2])
|
||||
}
|
||||
} else {
|
||||
// 没有front matter，整个内容就是skill内容
|
||||
skill.Content = content
|
||||
}
|
||||
} else {
|
||||
// 没有front matter，整个内容就是skill内容
|
||||
skill.Content = content
|
||||
}
|
||||
|
||||
// 如果内容为空，使用描述作为内容
|
||||
if skill.Content == "" {
|
||||
skill.Content = skill.Description
|
||||
}
|
||||
|
||||
return skill
|
||||
}
|
||||
|
||||
// GetSkillContent 获取skill的完整内容（用于注入到系统提示词）
|
||||
func (m *Manager) GetSkillContent(skillNames []string) (string, error) {
|
||||
skills, err := m.LoadSkills(skillNames)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if len(skills) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
var builder strings.Builder
|
||||
builder.WriteString("## 可用Skills\n\n")
|
||||
builder.WriteString("在执行任务前，请仔细阅读以下skills内容，这些内容包含了相关的专业知识和方法：\n\n")
|
||||
|
||||
for _, skill := range skills {
|
||||
builder.WriteString(fmt.Sprintf("### Skill: %s\n", skill.Name))
|
||||
if skill.Description != "" {
|
||||
builder.WriteString(fmt.Sprintf("**描述**: %s\n\n", skill.Description))
|
||||
}
|
||||
builder.WriteString(skill.Content)
|
||||
builder.WriteString("\n\n---\n\n")
|
||||
}
|
||||
|
||||
return builder.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
package skills
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RegisterSkillsTool 注册Skills工具到MCP服务器
|
||||
func RegisterSkillsTool(
|
||||
mcpServer *mcp.Server,
|
||||
manager *Manager,
|
||||
logger *zap.Logger,
|
||||
) {
|
||||
RegisterSkillsToolWithStorage(mcpServer, manager, nil, logger)
|
||||
}
|
||||
|
||||
// RegisterSkillsToolWithStorage 注册Skills工具到MCP服务器（带存储支持）
|
||||
func RegisterSkillsToolWithStorage(
|
||||
mcpServer *mcp.Server,
|
||||
manager *Manager,
|
||||
storage SkillStatsStorage,
|
||||
logger *zap.Logger,
|
||||
) {
|
||||
// 注册第一个工具：获取所有可用的skills列表
|
||||
listSkillsTool := mcp.Tool{
|
||||
Name: builtin.ToolListSkills,
|
||||
Description: "获取所有可用的skills列表。Skills是专业知识文档，可以在执行任务前阅读以获取相关专业知识。使用此工具可以查看系统中所有可用的skills，然后使用read_skill工具读取特定skill的内容。",
|
||||
ShortDescription: "获取所有可用的skills列表",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{},
|
||||
"required": []string{},
|
||||
},
|
||||
}
|
||||
|
||||
listSkillsHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
skills, err := manager.ListSkills()
|
||||
if err != nil {
|
||||
logger.Error("获取skills列表失败", zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("获取skills列表失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if len(skills) == 0 {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "当前没有可用的skills。\n\nSkills是专业知识文档，可以在执行任务前阅读以获取相关专业知识。你可以在skills目录下创建新的skill。",
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
var result strings.Builder
|
||||
result.WriteString(fmt.Sprintf("共有 %d 个可用的skills：\n\n", len(skills)))
|
||||
for i, skill := range skills {
|
||||
result.WriteString(fmt.Sprintf("%d. %s\n", i+1, skill))
|
||||
}
|
||||
result.WriteString("\n使用 read_skill 工具可以读取特定skill的详细内容。\n")
|
||||
result.WriteString("例如：read_skill(skill_name=\"sql-injection-testing\")")
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: result.String(),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(listSkillsTool, listSkillsHandler)
|
||||
logger.Info("注册skills列表工具成功")
|
||||
|
||||
// 注册第二个工具：读取特定skill的内容
|
||||
readSkillTool := mcp.Tool{
|
||||
Name: builtin.ToolReadSkill,
|
||||
Description: "读取指定skill的详细内容。Skills是专业知识文档，包含测试方法、工具使用、最佳实践等。在执行相关任务前，可以调用此工具读取相关skill的内容，以获取专业知识和指导。",
|
||||
ShortDescription: "读取指定skill的详细内容",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"skill_name": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "要读取的skill名称（必需）。可以使用list_skills工具获取所有可用的skill名称。",
|
||||
},
|
||||
},
|
||||
"required": []string{"skill_name"},
|
||||
},
|
||||
}
|
||||
|
||||
readSkillHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
skillName, ok := args["skill_name"].(string)
|
||||
if !ok || skillName == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: skill_name 参数必需且不能为空。请使用list_skills工具获取所有可用的skill名称。",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
skill, err := manager.LoadSkill(skillName)
|
||||
failed := err != nil
|
||||
now := time.Now()
|
||||
|
||||
// 记录调用统计
|
||||
if storage != nil {
|
||||
totalCalls := 1
|
||||
successCalls := 0
|
||||
failedCalls := 0
|
||||
if failed {
|
||||
failedCalls = 1
|
||||
} else {
|
||||
successCalls = 1
|
||||
}
|
||||
if err := storage.UpdateSkillStats(skillName, totalCalls, successCalls, failedCalls, &now); err != nil {
|
||||
logger.Warn("保存Skills统计信息失败", zap.String("skill", skillName), zap.Error(err))
|
||||
} else {
|
||||
logger.Info("Skills统计信息已更新",
|
||||
zap.String("skill", skillName),
|
||||
zap.Int("totalCalls", totalCalls),
|
||||
zap.Int("successCalls", successCalls),
|
||||
zap.Int("failedCalls", failedCalls))
|
||||
}
|
||||
} else {
|
||||
logger.Warn("Skills统计存储未配置，无法记录调用统计", zap.String("skill", skillName))
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
logger.Warn("读取skill失败", zap.String("skill", skillName), zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("读取skill失败: %v\n\n请使用list_skills工具确认skill名称是否正确。", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
var result strings.Builder
|
||||
result.WriteString(fmt.Sprintf("## Skill: %s\n\n", skill.Name))
|
||||
if skill.Description != "" {
|
||||
result.WriteString(fmt.Sprintf("**描述**: %s\n\n", skill.Description))
|
||||
}
|
||||
result.WriteString("---\n\n")
|
||||
result.WriteString(skill.Content)
|
||||
result.WriteString("\n\n---\n\n")
|
||||
result.WriteString(fmt.Sprintf("*Skill路径: %s*", skill.Path))
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: result.String(),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(readSkillTool, readSkillHandler)
|
||||
logger.Info("注册skill读取工具成功")
|
||||
}
|
||||
|
||||
// SkillStatsStorage Skills统计存储接口
|
||||
type SkillStatsStorage interface {
|
||||
UpdateSkillStats(skillName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error
|
||||
LoadSkillStats() (map[string]*SkillStats, error)
|
||||
}
|
||||
|
||||
// SkillStats Skills统计信息
|
||||
type SkillStats struct {
|
||||
SkillName string
|
||||
TotalCalls int
|
||||
SuccessCalls int
|
||||
FailedCalls int
|
||||
LastCallTime *time.Time
|
||||
}
|
||||
Reference in new issue
Block a user