mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-05 02:18:46 +02:00
Add files via upload
This commit is contained in:
+83
-9
@@ -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
|
||||
|
||||
@@ -87,7 +87,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-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 加密 Beacon(AES-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")),
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 ""
|
||||
}
|
||||
Reference in New Issue
Block a user