Add files via upload

This commit is contained in:
公明
2026-07-23 19:02:37 +08:00
committed by GitHub
parent 0cd4df3518
commit d304ff5c99
10 changed files with 607 additions and 52 deletions
+83 -9
View File
@@ -1313,6 +1313,7 @@ func setupRoutes(
c2Routes.GET("/sessions/:id", c2Handler.GetSession)
c2Routes.DELETE("/sessions/:id", c2Handler.DeleteSession)
c2Routes.PUT("/sessions/:id/sleep", c2Handler.SetSessionSleep)
c2Routes.PUT("/sessions/:id/note", c2Handler.SetSessionNote)
c2Routes.GET("/tasks", c2Handler.ListTasks)
c2Routes.DELETE("/tasks", c2Handler.DeleteTasks)
c2Routes.GET("/tasks/:id", c2Handler.GetTask)
@@ -1574,22 +1575,62 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
logger.Warn("跳过 WebShell 管理工具注册:db 为空")
return
}
projectIDFromToolArgs := func(ctx context.Context, args map[string]interface{}) string {
projectID, _ := args["project_id"].(string)
projectID = strings.TrimSpace(projectID)
if projectID == "" {
projectID = strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx))
}
return projectID
}
explicitProjectIDFromToolArgs := func(args map[string]interface{}) string {
projectID, _ := args["project_id"].(string)
return strings.TrimSpace(projectID)
}
authorizeWebshellToolProject := func(principal authctx.Principal, permission, projectID string) *mcp.ToolResult {
projectID = strings.TrimSpace(projectID)
if projectID == "" {
return nil
}
if projectID == database.ProjectFilterUnbound {
return nil
}
if !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
return &mcp.ToolResult{
Content: []mcp.Content{{Type: "text", Text: "无权访问项目: " + projectID}},
IsError: true,
}
}
return nil
}
// manage_webshell_list - 列出所有 webshell 连接
listTool := mcp.Tool{
Name: builtin.ToolManageWebshellList,
Description: "列出所有已保存的 WebShell 连接,返回连接ID、URL、类型、备注等信息。",
Description: "列出已保存的 WebShell 连接,返回连接ID、URL、类型、所属项目、备注等信息。默认按当前对话项目边界过滤:项目对话看本项目,未绑定项目的对话看未绑定连接;显式传 project_id 时按指定项目过滤。",
ShortDescription: "列出所有 WebShell 连接",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{},
"type": "object",
"properties": map[string]interface{}{
"project_id": map[string]interface{}{
"type": "string",
"description": "项目 ID;不填时在项目会话中默认使用当前项目。",
},
},
},
}
listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
connections := []database.WebShellConnection{}
var err error
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"))
projectID := explicitProjectIDFromToolArgs(args)
if projectID == "" {
projectID = mcpEffectiveProjectFilter(ctx, db)
}
if result := authorizeWebshellToolProject(principal, "webshell:read", projectID); result != nil {
return result, nil
}
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"), projectID)
} else {
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
}
@@ -1613,6 +1654,11 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
sb.WriteString(fmt.Sprintf(" 类型: %s\n", conn.Type))
sb.WriteString(fmt.Sprintf(" 请求方式: %s\n", conn.Method))
sb.WriteString(fmt.Sprintf(" 命令参数: %s\n", conn.CmdParam))
if conn.ProjectID != "" {
sb.WriteString(fmt.Sprintf(" 项目ID: %s\n", conn.ProjectID))
} else {
sb.WriteString(" 项目: 未绑定\n")
}
if conn.Remark != "" {
sb.WriteString(fmt.Sprintf(" 备注: %s\n", conn.Remark))
}
@@ -1687,6 +1733,14 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
cmdParam = "cmd"
}
remark, _ := args["remark"].(string)
principal, ok := authctx.PrincipalFromContext(ctx)
if !ok {
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
}
projectID := projectIDFromToolArgs(ctx, args)
if result := authorizeWebshellToolProject(principal, "webshell:write", projectID); result != nil {
return result, nil
}
// 生成连接ID
connID := "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12]
@@ -1698,6 +1752,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
Method: strings.ToLower(method),
CmdParam: cmdParam,
Remark: remark,
ProjectID: projectID,
CreatedAt: time.Now(),
}
@@ -1707,15 +1762,17 @@ 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)
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
projectLine := "项目: 未绑定"
if conn.ProjectID != "" {
projectLine = "项目ID: " + conn.ProjectID
}
return &mcp.ToolResult{
Content: []mcp.Content{{
Type: "text",
Text: fmt.Sprintf("WebShell 连接添加成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s", conn.ID, conn.URL, conn.Type, conn.Method, conn.CmdParam),
Text: fmt.Sprintf("WebShell 连接添加成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n%s", conn.ID, conn.URL, conn.Type, conn.Method, conn.CmdParam, projectLine),
}},
IsError: false,
}, nil
@@ -1760,6 +1817,10 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
"type": "string",
"description": "新的备注",
},
"project_id": map[string]interface{}{
"type": "string",
"description": "新的所属项目 ID;传空字符串可取消绑定。",
},
},
"required": []string{"connection_id"},
},
@@ -1801,6 +1862,19 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
if remark, ok := args["remark"].(string); ok {
existing.Remark = remark
}
if projectID, ok := args["project_id"].(string); ok {
projectID = strings.TrimSpace(projectID)
if projectID != "" {
principal, ok := authctx.PrincipalFromContext(ctx)
if !ok {
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
}
if result := authorizeWebshellToolProject(principal, "webshell:write", projectID); result != nil {
return result, nil
}
}
existing.ProjectID = projectID
}
if err := db.UpdateWebshellConnection(existing); err != nil {
return &mcp.ToolResult{
@@ -1812,7 +1886,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
return &mcp.ToolResult{
Content: []mcp.Content{{
Type: "text",
Text: fmt.Sprintf("WebShell 连接更新成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n备注: %s", existing.ID, existing.URL, existing.Type, existing.Method, existing.CmdParam, existing.Remark),
Text: fmt.Sprintf("WebShell 连接更新成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n项目ID: %s\n备注: %s", existing.ID, existing.URL, existing.Type, existing.Method, existing.CmdParam, existing.ProjectID, existing.Remark),
}},
IsError: false,
}, nil
+5 -1
View File
@@ -87,7 +87,7 @@ tcp_reverse 默认仅接受 CSB1 加密 BeaconAES-GCM + ImplantToken)才登
switch action {
case "list":
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx))
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx), mcpEffectiveProjectFilter(ctx, m.DB()))
if err != nil {
return makeC2Result(nil, err)
}
@@ -123,6 +123,7 @@ tcp_reverse 默认仅接受 CSB1 加密 BeaconAES-GCM + ImplantToken)才登
BindPort: int(getFloat64(params, "bind_port")),
ProfileID: getString(params, "profile_id"),
Remark: getString(params, "remark"),
ProjectID: strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx)),
Config: cfg,
CallbackHost: getString(params, "callback_host"),
}
@@ -260,6 +261,7 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
case "list":
filter := database.ListC2SessionsFilter{
ListenerID: getString(params, "listener_id"),
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
Status: getString(params, "status"),
OS: getString(params, "os"),
Search: getString(params, "search"),
@@ -495,6 +497,7 @@ func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
case "list":
filter := database.ListC2TasksFilter{
SessionID: getString(params, "session_id"),
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
Status: getString(params, "status"),
}
if limit := int(getFloat64(params, "limit")); limit > 0 {
@@ -645,6 +648,7 @@ func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
filter := database.ListC2EventsFilter{
Level: getString(params, "level"),
Category: getString(params, "category"),
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
SessionID: getString(params, "session_id"),
TaskID: getString(params, "task_id"),
Limit: int(getFloat64(params, "limit")),
+101 -5
View File
@@ -32,6 +32,9 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
return fmt.Errorf("no access to %s %s", resourceType, id)
}
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
return err
}
return nil
}
toolExecutionResource := func(permission string) error {
@@ -141,20 +144,26 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
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")
return authorizeC2Action(ctx, 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(ctx, principal, db, args, "c2_task", "task_id")
}
return authorizeC2Action(principal, db, args, "c2_session", "session_id")
return authorizeC2Action(ctx, principal, db, args, "c2_session", "session_id")
case builtin.ToolC2TaskManage:
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
return authorizeC2Action(ctx, 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 id := mcpAuthorizationString(args, "task_id"); id != "" {
return resource("c2:read", "c2_task", "task_id")
}
if filter := mcpEffectiveProjectFilter(ctx, db); filter != "" {
return require("c2:read")
}
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
return fmt.Errorf("unfiltered C2 event list requires global scope")
}
@@ -207,7 +216,7 @@ func externalMCPToolAuthorizer() func(context.Context, string, map[string]interf
}
}
func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
func authorizeC2Action(ctx context.Context, 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" {
@@ -228,6 +237,9 @@ func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[st
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
return fmt.Errorf("no access to %s %s", resourceType, candidate)
}
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, candidate); err != nil {
return err
}
}
return nil
}
@@ -240,9 +252,93 @@ func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[st
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
return fmt.Errorf("no access to %s %s", resourceType, id)
}
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
return err
}
return nil
}
func authorizeMCPProjectResourceBoundary(ctx context.Context, db *database.DB, resourceType, resourceID string) error {
filter := mcpEffectiveProjectFilter(ctx, db)
if filter == "" || db == nil {
return nil
}
projectID, ok, err := mcpResourceProjectID(db, resourceType, resourceID)
if err != nil {
return err
}
if !ok {
return nil
}
if filter == database.ProjectFilterUnbound {
if projectID != "" {
return fmt.Errorf("resource %s %s belongs to project %s, current conversation is unbound", resourceType, resourceID, projectID)
}
return nil
}
if projectID != filter {
if projectID == "" {
return fmt.Errorf("resource %s %s is unbound, current conversation project is %s", resourceType, resourceID, filter)
}
return fmt.Errorf("resource %s %s belongs to project %s, current conversation project is %s", resourceType, resourceID, projectID, filter)
}
return nil
}
func mcpResourceProjectID(db *database.DB, resourceType, resourceID string) (string, bool, error) {
switch resourceType {
case "webshell":
conn, err := db.GetWebshellConnection(resourceID)
if err != nil {
return "", true, err
}
if conn == nil {
return "", true, fmt.Errorf("webshell not found")
}
return strings.TrimSpace(conn.ProjectID), true, nil
case "c2_listener":
listener, err := db.GetC2Listener(resourceID)
if err != nil {
return "", true, err
}
if listener == nil {
return "", true, fmt.Errorf("listener not found")
}
return strings.TrimSpace(listener.ProjectID), true, nil
case "c2_session":
session, err := db.GetC2Session(resourceID)
if err != nil {
return "", true, err
}
if session == nil {
return "", true, fmt.Errorf("session not found")
}
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
case "c2_task":
task, err := db.GetC2Task(resourceID)
if err != nil {
return "", true, err
}
if task == nil {
return "", true, fmt.Errorf("task not found")
}
return mcpResourceProjectIDFromC2Session(db, task.SessionID)
default:
return "", false, nil
}
}
func mcpResourceProjectIDFromC2Session(db *database.DB, sessionID string) (string, bool, error) {
session, err := db.GetC2Session(sessionID)
if err != nil {
return "", true, err
}
if session == nil {
return "", true, fmt.Errorf("session not found")
}
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
}
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
values := []string{}
switch raw := args[key].(type) {
+85
View File
@@ -49,6 +49,91 @@ func TestMCPToolAuthorizerEnforcesPermissionAndResource(t *testing.T) {
}
}
func TestMCPToolAuthorizerEnforcesConversationProjectBoundary(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-project-boundary.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
user, err := db.CreateRBACUser("boundary-user", "Boundary User", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
project, err := db.CreateProject(&database.Project{Name: "Project 123"})
if err != nil {
t.Fatal(err)
}
projectConv, err := db.CreateConversation("project conversation", database.ConversationCreateMeta{ProjectID: project.ID})
if err != nil {
t.Fatal(err)
}
unboundConv, err := db.CreateConversation("unbound conversation", database.ConversationCreateMeta{})
if err != nil {
t.Fatal(err)
}
wsProject := database.WebShellConnection{ID: "ws_project", ProjectID: project.ID, URL: "http://127.0.0.1/project.php", Type: "php", Method: "post", CreatedAt: time.Now()}
wsUnbound := database.WebShellConnection{ID: "ws_unbound", URL: "http://127.0.0.1/unbound.php", Type: "php", Method: "post", CreatedAt: time.Now()}
if err := db.CreateWebshellConnection(&wsProject); err != nil {
t.Fatal(err)
}
if err := db.CreateWebshellConnection(&wsUnbound); err != nil {
t.Fatal(err)
}
for _, id := range []string{wsProject.ID, wsUnbound.ID} {
if err := db.AssignResourceToUser(user.ID, "webshell", id); err != nil {
t.Fatal(err)
}
}
now := time.Now()
listener := &database.C2Listener{ID: "l_project", ProjectID: project.ID, Name: "project listener", Type: "tcp_reverse", BindHost: "127.0.0.1", BindPort: 5555, OwnerUserID: user.ID, CreatedAt: now}
if err := db.CreateC2Listener(listener); err != nil {
t.Fatal(err)
}
if err := db.AssignResourceToUser(user.ID, "c2_listener", listener.ID); err != nil {
t.Fatal(err)
}
session := &database.C2Session{ID: "s_project", ListenerID: listener.ID, ImplantUUID: "implant-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}
if err := db.UpsertC2Session(session); err != nil {
t.Fatal(err)
}
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{
"webshell:read": true, "webshell:write": true,
"c2:read": true, "c2:write": true,
})
authorize := mcpToolAuthorizer(db)
unboundCtx := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), unboundConv.ID), principal)
projectCtx := authctx.WithPrincipal(mcp.WithMCPProjectID(mcp.WithMCPConversationID(context.Background(), projectConv.ID), project.ID), principal)
projectCtxFromConversationOnly := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), projectConv.ID), principal)
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err == nil {
t.Fatal("unbound conversation was allowed to use project-bound webshell")
}
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err != nil {
t.Fatalf("unbound webshell denied in unbound conversation: %v", err)
}
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
t.Fatalf("project webshell denied in project conversation: %v", err)
}
if err := authorize(projectCtxFromConversationOnly, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
t.Fatalf("project webshell denied when only conversation id is present: %v", err)
}
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err == nil {
t.Fatal("project conversation was allowed to use unbound webshell by id")
}
if err := authorize(unboundCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err == nil {
t.Fatal("unbound conversation was allowed to use project-bound c2 session")
}
if err := authorize(projectCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
t.Fatalf("project c2 session denied in project conversation: %v", err)
}
if err := authorize(projectCtxFromConversationOnly, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
t.Fatalf("project c2 session denied when only conversation id is present: %v", err)
}
}
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
if err != nil {
+26
View File
@@ -0,0 +1,26 @@
package app
import (
"context"
"strings"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
)
func mcpEffectiveProjectFilter(ctx context.Context, db *database.DB) string {
if projectID := strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx)); projectID != "" {
return projectID
}
if conversationID := mcpAuthorizationConversationID(ctx); conversationID != "" {
if db != nil {
if projectID, err := db.GetConversationProjectID(conversationID); err == nil {
if projectID = strings.TrimSpace(projectID); projectID != "" {
return projectID
}
}
}
return database.ProjectFilterUnbound
}
return ""
}