mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-13 23:47:24 +02:00
Add files via upload
This commit is contained in:
+61
-14
@@ -8,6 +8,7 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -16,6 +17,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
@@ -123,6 +125,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
|
||||
// 创建MCP服务器(带数据库持久化)
|
||||
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
|
||||
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
|
||||
|
||||
// 创建安全工具执行器
|
||||
@@ -146,6 +149,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
|
||||
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
|
||||
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
|
||||
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
|
||||
if cfg.ExternalMCP.Servers != nil {
|
||||
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
|
||||
// 启动所有启用的外部MCP客户端
|
||||
@@ -372,7 +376,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
vulnerabilityHandler.SetAudit(auditSvc)
|
||||
webshellHandler := handler.NewWebShellHandler(log.Logger, db)
|
||||
webshellHandler.SetAudit(auditSvc)
|
||||
chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger)
|
||||
chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger, db)
|
||||
chatUploadsHandler.SetAudit(auditSvc)
|
||||
registerWebshellTools(mcpServer, db, webshellHandler, log.Logger)
|
||||
registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger)
|
||||
@@ -541,6 +545,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
terminalHandler,
|
||||
app.c2Handler,
|
||||
auditHandler,
|
||||
auditSvc,
|
||||
rbacHandler,
|
||||
mcpServer,
|
||||
authManager,
|
||||
@@ -554,17 +559,30 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
// mcpHandlerWithAuth 在鉴权通过后转发到 MCP 处理;若配置了 auth_header 则校验请求头,否则直接放行
|
||||
func (a *App) mcpHandlerWithAuth(w http.ResponseWriter, r *http.Request) {
|
||||
cfg := a.config.MCP
|
||||
if cfg.AuthHeader != "" {
|
||||
actual := []byte(r.Header.Get(cfg.AuthHeader))
|
||||
expected := []byte(cfg.AuthHeaderValue)
|
||||
if subtle.ConstantTimeCompare(actual, expected) != 1 {
|
||||
a.logger.Logger.Debug("MCP 鉴权失败:header 缺失或值不匹配", zap.String("header", cfg.AuthHeader))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
w.Write([]byte(`{"error":"unauthorized"}`))
|
||||
if authHeader := strings.TrimSpace(r.Header.Get("Authorization")); len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "Bearer ") {
|
||||
if session, ok := a.auth.ValidateToken(strings.TrimSpace(authHeader[7:])); ok && session.Permissions["mcp:execute"] {
|
||||
principal := authctx.NewPrincipalWithScopes(session.UserID, session.Username, session.Scope, session.Permissions, session.PermissionScopes)
|
||||
a.mcpServer.HandleHTTP(w, r.WithContext(authctx.WithPrincipal(r.Context(), principal)))
|
||||
return
|
||||
}
|
||||
}
|
||||
if !cfg.AllowGlobalAccess || strings.TrimSpace(cfg.AuthHeader) == "" || strings.TrimSpace(cfg.AuthHeaderValue) == "" {
|
||||
http.Error(w, "use an authorized user bearer token; global MCP service access is disabled", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if subtle.ConstantTimeCompare([]byte(r.Header.Get(cfg.AuthHeader)), []byte(cfg.AuthHeaderValue)) != 1 {
|
||||
a.logger.Logger.Debug("MCP 鉴权失败:header 缺失或值不匹配", zap.String("header", cfg.AuthHeader))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
w.Write([]byte(`{"error":"unauthorized"}`))
|
||||
return
|
||||
}
|
||||
permissions := make(map[string]bool, len(security.PermissionCatalog))
|
||||
for permission := range security.PermissionCatalog {
|
||||
permissions[permission] = true
|
||||
}
|
||||
principal := authctx.NewPrincipal("service:mcp", "mcp-service", database.RBACScopeAll, permissions)
|
||||
r = r.WithContext(authctx.WithPrincipal(r.Context(), principal))
|
||||
a.mcpServer.HandleHTTP(w, r)
|
||||
}
|
||||
|
||||
@@ -825,6 +843,7 @@ func setupRoutes(
|
||||
terminalHandler *handler.TerminalHandler,
|
||||
c2Handler *handler.C2Handler,
|
||||
auditHandler *handler.AuditHandler,
|
||||
auditSvc *audit.Service,
|
||||
rbacHandler *handler.RBACHandler,
|
||||
mcpServer *mcp.Server,
|
||||
authManager *security.AuthManager,
|
||||
@@ -835,8 +854,9 @@ func setupRoutes(
|
||||
|
||||
// 认证相关路由
|
||||
authRoutes := api.Group("/auth")
|
||||
loginRL := security.NewRateLimiter(10, 1*time.Minute)
|
||||
{
|
||||
authRoutes.POST("/login", authHandler.Login)
|
||||
authRoutes.POST("/login", security.RateLimitMiddleware(loginRL), authHandler.Login)
|
||||
authRoutes.POST("/logout", security.AuthMiddleware(authManager), authHandler.Logout)
|
||||
authRoutes.POST("/change-password", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), authHandler.ChangePassword)
|
||||
authRoutes.GET("/validate", security.AuthMiddleware(authManager), authHandler.Validate)
|
||||
@@ -856,7 +876,15 @@ func setupRoutes(
|
||||
|
||||
protected := api.Group("")
|
||||
protected.Use(security.AuthMiddleware(authManager))
|
||||
protected.Use(security.RBACMiddleware(app.db))
|
||||
protected.Use(security.RBACMiddlewareWithDenyHook(app.db, func(c *gin.Context, reason, permission string) {
|
||||
if auditSvc != nil {
|
||||
auditSvc.Record(c, audit.Entry{
|
||||
Level: "warn", Category: "rbac", Action: "access_denied", Result: "failure",
|
||||
Message: "RBAC 拒绝访问", ResourceType: "route", ResourceID: c.FullPath(),
|
||||
Detail: map[string]interface{}{"reason": reason, "permission": permission, "method": c.Request.Method},
|
||||
})
|
||||
}
|
||||
}))
|
||||
{
|
||||
protected.GET("/rbac/me", rbacHandler.Me)
|
||||
protected.GET("/rbac/metadata", rbacHandler.Metadata)
|
||||
@@ -1490,7 +1518,13 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
},
|
||||
}
|
||||
listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
connections, err := db.ListWebshellConnections()
|
||||
connections := []database.WebShellConnection{}
|
||||
var err error
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"))
|
||||
} else {
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: "获取连接列表失败: " + err.Error()}},
|
||||
@@ -1605,6 +1639,10 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
||||
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
||||
}
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{
|
||||
@@ -2000,8 +2038,17 @@ func initializeKnowledge(
|
||||
// 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")
|
||||
origin := strings.TrimSpace(c.GetHeader("Origin"))
|
||||
if origin != "" {
|
||||
parsed, err := url.Parse(origin)
|
||||
if err != nil || parsed.Host == "" || !strings.EqualFold(parsed.Host, c.Request.Host) {
|
||||
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "cross-origin request denied"})
|
||||
return
|
||||
}
|
||||
c.Writer.Header().Set("Access-Control-Allow-Origin", origin)
|
||||
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
|
||||
c.Writer.Header().Add("Vary", "Origin")
|
||||
}
|
||||
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")
|
||||
|
||||
|
||||
+29
-12
@@ -4,11 +4,13 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
@@ -66,16 +68,16 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(get/update/start/stop/delete 需要)"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "监听器名称(create/update)"},
|
||||
"type": map[string]interface{}{"type": "string", "description": "监听器类型(create)", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}},
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(get/update/start/stop/delete 需要)"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "监听器名称(create/update)"},
|
||||
"type": map[string]interface{}{"type": "string", "description": "监听器类型(create)", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}},
|
||||
"bind_host": map[string]interface{}{"type": "string", "description": "绑定地址,默认 127.0.0.1;外网监听常用 0.0.0.0"},
|
||||
"callback_host": map[string]interface{}{"type": "string", "description": "可选:植入端/Payload 回连主机名(公网 IP 或域名)。写入 config_json;生成 oneliner/beacon 时优先于 bind_host。update 时传入空字符串可清除"},
|
||||
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
||||
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
||||
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"},
|
||||
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
||||
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
||||
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
@@ -85,7 +87,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
listeners, err := m.DB().ListC2Listeners()
|
||||
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx))
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
@@ -128,6 +130,10 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = m.DB().SetResourceOwner("c2_listener", listener.ID, principal.UserID)
|
||||
_ = m.DB().AssignResourceToUser(principal.UserID, "c2_listener", listener.ID)
|
||||
}
|
||||
implantToken := listener.ImplantToken
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
@@ -264,7 +270,7 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
if v, ok := params["suspicious"].(bool); ok && v {
|
||||
filter.Suspicious = true
|
||||
}
|
||||
sessions, err := m.DB().ListC2Sessions(filter)
|
||||
sessions, err := m.DB().ListC2SessionsForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"sessions": sessions, "count": len(sessions)}, err)
|
||||
|
||||
case "get":
|
||||
@@ -494,7 +500,7 @@ func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||
filter.Limit = limit
|
||||
}
|
||||
tasks, err := m.DB().ListC2Tasks(filter)
|
||||
tasks, err := m.DB().ListC2TasksForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"tasks": tasks, "count": len(tasks)}, err)
|
||||
|
||||
case "cancel":
|
||||
@@ -602,6 +608,9 @@ func registerC2PayloadTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListe
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = m.DB().RecordC2PayloadArtifact(filepath.Base(result.OutputPath), result.PayloadID, result.ListenerID, principal.UserID)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"payload_id": result.PayloadID, "download_path": result.DownloadPath,
|
||||
"os": result.OS, "arch": result.Arch, "size_bytes": result.SizeBytes,
|
||||
@@ -648,11 +657,19 @@ func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
events, err := m.DB().ListC2Events(filter)
|
||||
events, err := m.DB().ListC2EventsForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"events": events, "count": len(events)}, err)
|
||||
})
|
||||
}
|
||||
|
||||
func c2ToolAccess(ctx context.Context) database.RBACListAccess {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return database.RBACListAccess{Scope: database.RBACScopeAssigned}
|
||||
}
|
||||
return database.RBACListAccess{UserID: principal.UserID, Scope: principal.ScopeFor("c2:read")}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_profile — Malleable Profile 管理工具(新增)
|
||||
// ============================================================================
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestCORSMiddlewareAllowsSameOriginAndRejectsForeignOrigin(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(corsMiddleware())
|
||||
router.GET("/test", func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||||
|
||||
same := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil)
|
||||
same.Host = "app.example"
|
||||
same.Header.Set("Origin", "http://app.example")
|
||||
sameW := httptest.NewRecorder()
|
||||
router.ServeHTTP(sameW, same)
|
||||
if sameW.Code != http.StatusNoContent || sameW.Header().Get("Access-Control-Allow-Origin") != "http://app.example" {
|
||||
t.Fatalf("same-origin response = %d, allow-origin=%q", sameW.Code, sameW.Header().Get("Access-Control-Allow-Origin"))
|
||||
}
|
||||
|
||||
foreign := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil)
|
||||
foreign.Host = "app.example"
|
||||
foreign.Header.Set("Origin", "https://evil.example")
|
||||
foreignW := httptest.NewRecorder()
|
||||
router.ServeHTTP(foreignW, foreign)
|
||||
if foreignW.Code != http.StatusForbidden {
|
||||
t.Fatalf("foreign-origin response = %d, want %d", foreignW.Code, http.StatusForbidden)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
)
|
||||
|
||||
func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string]interface{}) error {
|
||||
return func(ctx context.Context, toolName string, args map[string]interface{}) error {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return fmt.Errorf("missing authenticated principal")
|
||||
}
|
||||
require := func(permission string) error {
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
resource := func(permission, resourceType, argument string) error {
|
||||
if err := require(permission); err != nil {
|
||||
return err
|
||||
}
|
||||
id := mcpAuthorizationString(args, argument)
|
||||
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
switch toolName {
|
||||
case builtin.ToolWebshellExec, builtin.ToolWebshellFileWrite:
|
||||
return resource("webshell:write", "webshell", "connection_id")
|
||||
case builtin.ToolWebshellFileList, builtin.ToolWebshellFileRead:
|
||||
return resource("webshell:read", "webshell", "connection_id")
|
||||
case builtin.ToolManageWebshellList:
|
||||
return require("webshell:read")
|
||||
case builtin.ToolManageWebshellAdd:
|
||||
return require("webshell:write")
|
||||
case builtin.ToolManageWebshellUpdate, builtin.ToolManageWebshellTest:
|
||||
return resource("webshell:write", "webshell", "connection_id")
|
||||
case builtin.ToolManageWebshellDelete:
|
||||
return resource("webshell:delete", "webshell", "connection_id")
|
||||
case builtin.ToolRecordVulnerability:
|
||||
if err := require("vulnerability:write"); err != nil {
|
||||
return err
|
||||
}
|
||||
conversationID := mcpAuthorizationString(args, "conversation_id")
|
||||
if conversationID == "" {
|
||||
conversationID = mcpAuthorizationConversationID(ctx)
|
||||
}
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:write"), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolListVulnerabilities:
|
||||
if err := require("vulnerability:read"); err != nil {
|
||||
return err
|
||||
}
|
||||
conversationID := mcpAuthorizationConversationID(ctx)
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:read"), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolGetVulnerability:
|
||||
return resource("vulnerability:read", "vulnerability", "id")
|
||||
case builtin.ToolUpsertProjectFact, builtin.ToolDeprecateProjectFact, builtin.ToolRestoreProjectFact:
|
||||
return authorizeProjectTool(ctx, principal, db, "project:write")
|
||||
case builtin.ToolGetProjectFact, builtin.ToolListProjectFacts, builtin.ToolSearchProjectFacts:
|
||||
return authorizeProjectTool(ctx, principal, db, "project:read")
|
||||
case builtin.ToolListKnowledgeRiskTypes, builtin.ToolSearchKnowledgeBase:
|
||||
return require("knowledge:read")
|
||||
case builtin.ToolAnalyzeImage:
|
||||
return require("agent:execute")
|
||||
case builtin.ToolBatchTaskList:
|
||||
return require("tasks:read")
|
||||
case builtin.ToolBatchTaskGet:
|
||||
return resource("tasks:read", "batch_task", "queue_id")
|
||||
case builtin.ToolBatchTaskCreate:
|
||||
if err := require("tasks:write"); err != nil {
|
||||
return err
|
||||
}
|
||||
if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("tasks:write"), "project", projectID)) {
|
||||
return fmt.Errorf("no access to project %s", projectID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolBatchTaskDelete, builtin.ToolBatchTaskRemove:
|
||||
return resource("tasks:delete", "batch_task", "queue_id")
|
||||
case builtin.ToolBatchTaskStart, builtin.ToolBatchTaskRerun, builtin.ToolBatchTaskPause,
|
||||
builtin.ToolBatchTaskUpdateMetadata, builtin.ToolBatchTaskUpdateSchedule,
|
||||
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
||||
return resource("tasks:write", "batch_task", "queue_id")
|
||||
case builtin.ToolC2Listener:
|
||||
return authorizeC2Action(principal, db, args, "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File:
|
||||
if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" {
|
||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
||||
}
|
||||
return authorizeC2Action(principal, db, args, "c2_session", "session_id")
|
||||
case builtin.ToolC2TaskManage:
|
||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
||||
case builtin.ToolC2Payload:
|
||||
return resource("c2:write", "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Event:
|
||||
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
||||
return resource("c2:read", "c2_session", "session_id")
|
||||
}
|
||||
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
|
||||
return fmt.Errorf("unfiltered C2 event list requires global scope")
|
||||
}
|
||||
return require("c2:read")
|
||||
case builtin.ToolC2Profile:
|
||||
// Profiles are process-global and do not yet have an owner. Writes are
|
||||
// therefore reserved for global scope; reads require c2:read.
|
||||
if mcpAuthorizationString(args, "action") == "list" || mcpAuthorizationString(args, "action") == "get" {
|
||||
return require("c2:read")
|
||||
}
|
||||
permission := "c2:write"
|
||||
if mcpAuthorizationString(args, "action") == "delete" {
|
||||
permission = "c2:delete"
|
||||
}
|
||||
if principal.ScopeFor(permission) != database.RBACScopeAll {
|
||||
return fmt.Errorf("C2 profile mutation requires global scope")
|
||||
}
|
||||
if mcpAuthorizationString(args, "action") == "delete" {
|
||||
return require("c2:delete")
|
||||
}
|
||||
return require("c2:write")
|
||||
default:
|
||||
if builtin.IsBuiltinTool(toolName) {
|
||||
return fmt.Errorf("no authorization policy registered for builtin tool %s", toolName)
|
||||
}
|
||||
if principal.HasPermission("agent:local-execute") {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("missing agent:local-execute")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func externalMCPToolAuthorizer() func(context.Context, string, map[string]interface{}) error {
|
||||
return func(ctx context.Context, toolName string, _ map[string]interface{}) error {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return fmt.Errorf("missing authenticated principal")
|
||||
}
|
||||
if !principal.HasPermission("mcp:external:execute") {
|
||||
return fmt.Errorf("missing permission mcp:external:execute")
|
||||
}
|
||||
if principal.ScopeFor("mcp:external:execute") != database.RBACScopeAll {
|
||||
return fmt.Errorf("external MCP invocation requires global scope")
|
||||
}
|
||||
if strings.TrimSpace(toolName) == "" {
|
||||
return fmt.Errorf("missing external tool name")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
|
||||
action := mcpAuthorizationString(args, "action")
|
||||
permission := "c2:write"
|
||||
if action == "list" || action == "get" || action == "get_result" || action == "wait" {
|
||||
permission = "c2:read"
|
||||
} else if action == "delete" || action == "delete_batch" {
|
||||
permission = "c2:delete"
|
||||
}
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
id := mcpAuthorizationString(args, argument)
|
||||
if action == "delete_batch" {
|
||||
ids := mcpAuthorizationStrings(args, argument+"s")
|
||||
if len(ids) == 0 {
|
||||
return fmt.Errorf("missing resource identifiers %ss", argument)
|
||||
}
|
||||
for _, candidate := range ids {
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, candidate)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if id == "" {
|
||||
if action == "create" || action == "list" {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("missing resource identifier %s", argument)
|
||||
}
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
|
||||
values := []string{}
|
||||
switch raw := args[key].(type) {
|
||||
case []string:
|
||||
for _, value := range raw {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
values = append(values, value)
|
||||
}
|
||||
}
|
||||
case []interface{}:
|
||||
for _, item := range raw {
|
||||
if value, ok := item.(string); ok {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
values = append(values, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
func authorizeProjectTool(ctx context.Context, principal authctx.Principal, db *database.DB, permission string) error {
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
conversationID := mcpAuthorizationConversationID(ctx)
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
projectID, err := db.GetConversationProjectID(conversationID)
|
||||
if err != nil || strings.TrimSpace(projectID) == "" || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
|
||||
return fmt.Errorf("no access to project %s", projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mcpAuthorizationConversationID(ctx context.Context) string {
|
||||
if id := strings.TrimSpace(agent.ConversationIDFromContext(ctx)); id != "" {
|
||||
return id
|
||||
}
|
||||
return strings.TrimSpace(mcp.MCPConversationIDFromContext(ctx))
|
||||
}
|
||||
|
||||
func mcpAuthorizationString(args map[string]interface{}, key string) string {
|
||||
value, _ := args[key].(string)
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestMCPToolAuthorizerEnforcesPermissionAndResource(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-authz.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
user, err := db.CreateRBACUser("mcp-user", "MCP User", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, id := range []string{"ws_allowed", "ws_hidden"} {
|
||||
if err := db.CreateWebshellConnection(&database.WebShellConnection{ID: id, URL: "http://127.0.0.1/" + id, Type: "php", Method: "post", CmdParam: "cmd", CreatedAt: time.Now()}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "webshell", "ws_allowed"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"mcp:write": true, "webshell:write": true})
|
||||
ctx := authctx.WithPrincipal(context.Background(), principal)
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_allowed"}); err != nil {
|
||||
t.Fatalf("allowed resource denied: %v", err)
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_hidden"}); err == nil {
|
||||
t.Fatal("foreign webshell resource was allowed")
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolManageWebshellDelete, map[string]interface{}{"connection_id": "ws_allowed"}); err == nil {
|
||||
t.Fatal("delete without webshell:delete was allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
permissions := map[string]bool{}
|
||||
for permission := range security.PermissionCatalog {
|
||||
permissions[permission] = true
|
||||
}
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("admin", "admin", database.RBACScopeAll, permissions))
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
args := map[string]interface{}{
|
||||
"action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x",
|
||||
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x",
|
||||
}
|
||||
for _, toolName := range builtin.GetAllBuiltinTools() {
|
||||
err := authorize(ctx, toolName, args)
|
||||
if err != nil && strings.Contains(err.Error(), "no authorization policy registered") {
|
||||
t.Errorf("builtin tool %s has no explicit policy", toolName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPRequiresDedicatedPermission(t *testing.T) {
|
||||
authorize := externalMCPToolAuthorizer()
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
||||
if err := authorize(ctx, "server::tool", nil); err == nil {
|
||||
t.Fatal("agent:execute alone authorized an external MCP tool")
|
||||
}
|
||||
ctx = authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAll, map[string]bool{"mcp:external:execute": true}))
|
||||
if err := authorize(ctx, "server::tool", nil); err != nil {
|
||||
t.Fatalf("dedicated external MCP permission rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfiguredCommandToolRequiresLocalExecutePermission(t *testing.T) {
|
||||
authorize := mcpToolAuthorizer(nil)
|
||||
agentOnly := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
||||
if err := authorize(agentOnly, "nmap_scan", nil); err == nil {
|
||||
t.Fatal("agent:execute alone authorized a configured command tool")
|
||||
}
|
||||
local := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:local-execute": true}))
|
||||
if err := authorize(local, "nmap_scan", nil); err != nil {
|
||||
t.Fatalf("agent:local-execute rejected: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestStandaloneMCPPrefersUserRBACAndDisablesGlobalTokenByDefault(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-http-auth.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
auth, err := security.NewAuthManager("admin-secret", 12)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := auth.AttachRBACStore(db); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
token, _, err := auth.Authenticate("admin", "admin-secret")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
server := mcp.NewServer(zap.NewNop())
|
||||
server.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
a := &App{config: &config.Config{MCP: config.MCPConfig{AuthHeader: "X-MCP-Token", AuthHeaderValue: "static-secret"}}, auth: auth, mcpServer: server}
|
||||
body := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)
|
||||
|
||||
userReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body))
|
||||
userReq.Header.Set("Authorization", "Bearer "+token)
|
||||
userW := httptest.NewRecorder()
|
||||
a.mcpHandlerWithAuth(userW, userReq)
|
||||
if userW.Code != http.StatusOK {
|
||||
t.Fatalf("user bearer status = %d: %s", userW.Code, userW.Body.String())
|
||||
}
|
||||
|
||||
staticReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body))
|
||||
staticReq.Header.Set("X-MCP-Token", "static-secret")
|
||||
staticW := httptest.NewRecorder()
|
||||
a.mcpHandlerWithAuth(staticW, staticReq)
|
||||
if staticW.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("global static token status = %d, want 401", staticW.Code)
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
@@ -313,6 +314,10 @@ func registerRecordVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, log
|
||||
}
|
||||
return textResult(fmt.Sprintf("记录漏洞失败: %v", err), true), nil
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = db.SetResourceOwner("vulnerability", created.ID, principal.UserID)
|
||||
_ = db.AssignResourceToUser(principal.UserID, "vulnerability", created.ID)
|
||||
}
|
||||
|
||||
if logger != nil {
|
||||
logger.Info("漏洞记录成功",
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
package authctx
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Principal is the immutable authorization identity propagated beyond the
|
||||
// transport layer into Agent, MCP and background task contexts.
|
||||
type Principal struct {
|
||||
UserID string
|
||||
Username string
|
||||
Permissions map[string]bool
|
||||
PermissionScopes map[string]string
|
||||
Scope string
|
||||
}
|
||||
|
||||
type principalContextKey struct{}
|
||||
|
||||
func NewPrincipal(userID, username, scope string, permissions map[string]bool) Principal {
|
||||
return NewPrincipalWithScopes(userID, username, scope, permissions, nil)
|
||||
}
|
||||
|
||||
func NewPrincipalWithScopes(userID, username, scope string, permissions map[string]bool, permissionScopes map[string]string) Principal {
|
||||
permissionCopy := make(map[string]bool, len(permissions))
|
||||
scopeCopy := make(map[string]string, len(permissionScopes))
|
||||
for permission, allowed := range permissions {
|
||||
if allowed {
|
||||
permissionCopy[permission] = true
|
||||
if permissionScope := strings.TrimSpace(permissionScopes[permission]); permissionScope != "" {
|
||||
scopeCopy[permission] = permissionScope
|
||||
}
|
||||
}
|
||||
}
|
||||
return Principal{
|
||||
UserID: strings.TrimSpace(userID), Username: strings.TrimSpace(username),
|
||||
Scope: strings.TrimSpace(scope), Permissions: permissionCopy, PermissionScopes: scopeCopy,
|
||||
}
|
||||
}
|
||||
|
||||
func WithPrincipal(ctx context.Context, principal Principal) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if strings.TrimSpace(principal.UserID) == "" {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, principalContextKey{}, principal)
|
||||
}
|
||||
|
||||
func PrincipalFromContext(ctx context.Context) (Principal, bool) {
|
||||
if ctx == nil {
|
||||
return Principal{}, false
|
||||
}
|
||||
principal, ok := ctx.Value(principalContextKey{}).(Principal)
|
||||
return principal, ok && strings.TrimSpace(principal.UserID) != ""
|
||||
}
|
||||
|
||||
func (p Principal) HasPermission(permission string) bool {
|
||||
return p.Permissions[strings.TrimSpace(permission)]
|
||||
}
|
||||
|
||||
// ScopeFor returns the scope attached to the permission that authorizes the
|
||||
// current action. Falling back to Scope keeps explicit service principals and
|
||||
// legacy callers compatible without reintroducing cross-role scope widening.
|
||||
func (p Principal) ScopeFor(permission string) string {
|
||||
permission = strings.TrimSpace(permission)
|
||||
if scope := strings.TrimSpace(p.PermissionScopes[permission]); scope != "" {
|
||||
return scope
|
||||
}
|
||||
return strings.TrimSpace(p.Scope)
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -37,29 +38,30 @@ type listToolsInflight struct {
|
||||
|
||||
// ExternalMCPManager 外部MCP管理器
|
||||
type ExternalMCPManager struct {
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
listToolsInflight map[string]*listToolsInflight
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
runningCancels map[string]context.CancelFunc
|
||||
abortUserNotes map[string]string
|
||||
reconnectMu sync.Mutex
|
||||
reconnecting map[string]bool
|
||||
reconnectLastTry map[string]time.Time
|
||||
reconnectAttempts map[string]int
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
}
|
||||
|
||||
// NewExternalMCPManager 创建外部MCP管理器
|
||||
@@ -67,16 +69,24 @@ func NewExternalMCPManager(logger *zap.Logger) *ExternalMCPManager {
|
||||
return NewExternalMCPManagerWithStorage(logger, nil)
|
||||
}
|
||||
|
||||
// SetToolAuthorizer installs the policy decision point for all external MCP
|
||||
// invocations. App wiring configures this before any Agent can call a tool.
|
||||
func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
|
||||
m.mu.Lock()
|
||||
m.toolAuthorizer = authorizer
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
|
||||
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
|
||||
manager := &ExternalMCPManager{
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
toolCounts: make(map[string]int),
|
||||
toolCache: make(map[string]toolListCacheEntry),
|
||||
listToolsInflight: make(map[string]*listToolsInflight),
|
||||
@@ -554,6 +564,17 @@ func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
|
||||
|
||||
// CallTool 调用外部MCP工具(返回执行ID)
|
||||
func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
|
||||
_, authenticated := authctx.PrincipalFromContext(ctx)
|
||||
m.mu.RLock()
|
||||
authorizer := m.toolAuthorizer
|
||||
m.mu.RUnlock()
|
||||
if authorizer != nil {
|
||||
if err := authorizer(ctx, toolName, args); err != nil {
|
||||
return nil, "", fmt.Errorf("external tool authorization denied: %w", err)
|
||||
}
|
||||
} else if authenticated {
|
||||
return nil, "", fmt.Errorf("external tool authorization policy is not configured")
|
||||
}
|
||||
// 解析工具名称:name::toolName
|
||||
var mcpName, actualToolName string
|
||||
if idx := findSubstring(toolName, "::"); idx > 0 {
|
||||
@@ -591,6 +612,10 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
m.mu.Lock()
|
||||
m.executions[executionID] = execution
|
||||
|
||||
@@ -2,14 +2,30 @@ package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
t.Cleanup(manager.StopAll)
|
||||
manager.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error {
|
||||
return errors.New("denied by policy")
|
||||
})
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true}))
|
||||
_, _, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{})
|
||||
if err == nil || !strings.Contains(err.Error(), "authorization denied") {
|
||||
t.Fatalf("external call bypassed authorizer: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
manager := NewExternalMCPManager(logger)
|
||||
|
||||
+74
-17
@@ -14,6 +14,8 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -49,6 +51,18 @@ type Server struct {
|
||||
// nil 表示未配置,沿用默认 30 分钟;指向 0 表示不限制;>0 为分钟数。
|
||||
httpToolTimeoutMinutes *int
|
||||
httpToolTimeoutMu sync.RWMutex
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
}
|
||||
|
||||
// SetToolAuthorizer installs the common policy decision point for every
|
||||
// user-attributed tool call, whether it originates from HTTP or an Agent.
|
||||
func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.toolAuthorizer = authorizer
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
type sseClient struct {
|
||||
@@ -111,21 +125,24 @@ func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) {
|
||||
s.httpToolTimeoutMinutes = &v
|
||||
}
|
||||
|
||||
func (s *Server) effectiveHTTPToolCallDeadline() (context.Context, context.CancelFunc) {
|
||||
func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) {
|
||||
const defaultDur = 30 * time.Minute
|
||||
if parent == nil {
|
||||
parent = context.Background()
|
||||
}
|
||||
if s == nil {
|
||||
return context.WithTimeout(context.Background(), defaultDur)
|
||||
return context.WithTimeout(parent, defaultDur)
|
||||
}
|
||||
s.httpToolTimeoutMu.RLock()
|
||||
mPtr := s.httpToolTimeoutMinutes
|
||||
s.httpToolTimeoutMu.RUnlock()
|
||||
if mPtr == nil {
|
||||
return context.WithTimeout(context.Background(), defaultDur)
|
||||
return context.WithTimeout(parent, defaultDur)
|
||||
}
|
||||
if *mPtr <= 0 {
|
||||
return context.WithCancel(context.Background())
|
||||
return context.WithCancel(parent)
|
||||
}
|
||||
return context.WithTimeout(context.Background(), time.Duration(*mPtr)*time.Minute)
|
||||
return context.WithTimeout(parent, time.Duration(*mPtr)*time.Minute)
|
||||
}
|
||||
|
||||
// RegisterTool 注册工具
|
||||
@@ -196,7 +213,7 @@ func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
response := s.handleMessage(&msg)
|
||||
response := s.handleMessage(r.Context(), &msg)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(response)
|
||||
}
|
||||
@@ -223,7 +240,7 @@ func (s *Server) serveSSESessionMessage(w http.ResponseWriter, r *http.Request,
|
||||
return
|
||||
}
|
||||
|
||||
response := s.handleMessage(&msg)
|
||||
response := s.handleMessage(r.Context(), &msg)
|
||||
if response == nil {
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
return
|
||||
@@ -317,7 +334,7 @@ func (s *Server) removeSSEClient(id string) {
|
||||
}
|
||||
|
||||
// handleMessage 处理MCP消息
|
||||
func (s *Server) handleMessage(msg *Message) *Message {
|
||||
func (s *Server) handleMessage(ctx context.Context, msg *Message) *Message {
|
||||
// 检查是否是通知(notification)- 通知没有id字段,不需要响应
|
||||
isNotification := msg.ID.Value() == nil || msg.ID.String() == ""
|
||||
|
||||
@@ -332,7 +349,7 @@ func (s *Server) handleMessage(msg *Message) *Message {
|
||||
case "tools/list":
|
||||
return s.handleListTools(msg)
|
||||
case "tools/call":
|
||||
return s.handleCallTool(msg)
|
||||
return s.handleCallTool(ctx, msg)
|
||||
case "prompts/list":
|
||||
return s.handleListPrompts(msg)
|
||||
case "prompts/get":
|
||||
@@ -433,7 +450,7 @@ func (s *Server) handleListTools(msg *Message) *Message {
|
||||
}
|
||||
|
||||
// handleCallTool 处理工具调用请求
|
||||
func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Message {
|
||||
var req CallToolRequest
|
||||
if err := json.Unmarshal(msg.Params, &req); err != nil {
|
||||
return &Message{
|
||||
@@ -443,6 +460,17 @@ func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
Error: &Error{Code: -32602, Message: "Invalid params"},
|
||||
}
|
||||
}
|
||||
_, authenticated := authctx.PrincipalFromContext(requestCtx)
|
||||
s.mu.RLock()
|
||||
authorizer := s.toolAuthorizer
|
||||
s.mu.RUnlock()
|
||||
if authorizer != nil {
|
||||
if err := authorizer(requestCtx, req.Name, req.Arguments); err != nil {
|
||||
return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Forbidden", Data: err.Error()}}
|
||||
}
|
||||
} else if authenticated {
|
||||
return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Tool authorization policy is not configured"}}
|
||||
}
|
||||
|
||||
executionID := uuid.New().String()
|
||||
execution := &ToolExecution{
|
||||
@@ -452,6 +480,10 @@ func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(requestCtx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(requestCtx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
@@ -495,7 +527,7 @@ func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
}
|
||||
}
|
||||
|
||||
baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline()
|
||||
baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline(requestCtx)
|
||||
defer timeoutCancel()
|
||||
execCtx, runCancel := context.WithCancel(baseCtx)
|
||||
s.registerRunningCancel(executionID, runCancel)
|
||||
@@ -809,6 +841,17 @@ func (s *Server) GetAllTools() []Tool {
|
||||
|
||||
// CallTool 直接调用工具(用于内部调用)
|
||||
func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
|
||||
_, authenticated := authctx.PrincipalFromContext(ctx)
|
||||
s.mu.RLock()
|
||||
authorizer := s.toolAuthorizer
|
||||
s.mu.RUnlock()
|
||||
if authorizer != nil {
|
||||
if err := authorizer(ctx, toolName, args); err != nil {
|
||||
return nil, "", fmt.Errorf("tool authorization denied: %w", err)
|
||||
}
|
||||
} else if authenticated {
|
||||
return nil, "", errors.New("tool authorization policy is not configured")
|
||||
}
|
||||
s.mu.RLock()
|
||||
handler, exists := s.tools[toolName]
|
||||
s.mu.RUnlock()
|
||||
@@ -826,6 +869,10 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
@@ -922,7 +969,7 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
||||
}
|
||||
|
||||
// BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。
|
||||
func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}) string {
|
||||
func (s *Server) BeginToolExecution(ctx context.Context, toolName string, args map[string]interface{}) string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
@@ -937,6 +984,10 @@ func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
@@ -952,7 +1003,7 @@ func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}
|
||||
}
|
||||
|
||||
// FinishToolExecution 完成先前 BeginToolExecution 创建的记录;executionID 为空时等同 RecordCompletedToolInvocation。
|
||||
func (s *Server) FinishToolExecution(executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
@@ -961,7 +1012,7 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri
|
||||
}
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" {
|
||||
return s.RecordCompletedToolInvocation(toolName, args, resultText, invokeErr)
|
||||
id = uuid.New().String()
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
@@ -984,6 +1035,12 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri
|
||||
if len(args) > 0 {
|
||||
exec.Arguments = args
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
exec.OwnerUserID = principal.UserID
|
||||
}
|
||||
if conversationID := MCPConversationIDFromContext(ctx); conversationID != "" {
|
||||
exec.ConversationID = conversationID
|
||||
}
|
||||
exec.EndTime = &now
|
||||
if exec.StartTime.IsZero() {
|
||||
exec.StartTime = now
|
||||
@@ -1027,8 +1084,8 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri
|
||||
|
||||
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
|
||||
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
|
||||
func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
return s.FinishToolExecution("", toolName, args, resultText, invokeErr)
|
||||
func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
return s.FinishToolExecution(ctx, "", toolName, args, resultText, invokeErr)
|
||||
}
|
||||
|
||||
// UpdateToolExecutionResult 将监控库中的工具结果更新为送入模型的展示正文(如 reduction 后的 persisted-output)。
|
||||
@@ -1519,7 +1576,7 @@ func (s *Server) HandleStdio() error {
|
||||
}
|
||||
|
||||
// 处理消息
|
||||
response := s.handleMessage(&msg)
|
||||
response := s.handleMessage(context.Background(), &msg)
|
||||
|
||||
// 如果是通知(response 为 nil),不需要发送响应
|
||||
if response == nil {
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.RegisterTool(Tool{Name: "echo", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil
|
||||
})
|
||||
server.SetToolAuthorizer(func(ctx context.Context, toolName string, args map[string]interface{}) error {
|
||||
if _, ok := authctx.PrincipalFromContext(ctx); !ok {
|
||||
return errors.New("principal required")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if _, _, err := server.CallTool(context.Background(), "echo", nil); err == nil {
|
||||
t.Fatal("tool call without principal was allowed")
|
||||
}
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true}))
|
||||
_, executionID, err := server.CallTool(ctx, "echo", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
execution, ok := server.GetExecution(executionID)
|
||||
if !ok || execution.OwnerUserID != "u1" {
|
||||
t.Fatalf("execution owner = %#v, want u1", execution)
|
||||
}
|
||||
}
|
||||
@@ -201,6 +201,7 @@ type ToolExecution struct {
|
||||
Duration time.Duration `json:"duration,omitempty"`
|
||||
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
OwnerUserID string `json:"-"`
|
||||
}
|
||||
|
||||
// ToolStats 工具统计信息
|
||||
|
||||
Reference in New Issue
Block a user