diff --git a/internal/app/app.go b/internal/app/app.go index af19d041..d6dbb65f 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -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 diff --git a/internal/app/c2_tools.go b/internal/app/c2_tools.go index 20a85f19..14853da2 100644 --- a/internal/app/c2_tools.go +++ b/internal/app/c2_tools.go @@ -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")), diff --git a/internal/app/mcp_authorization.go b/internal/app/mcp_authorization.go index 1bfeb0a0..d9947c23 100644 --- a/internal/app/mcp_authorization.go +++ b/internal/app/mcp_authorization.go @@ -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) { diff --git a/internal/app/mcp_authorization_test.go b/internal/app/mcp_authorization_test.go index e8b29d6e..e4bfdc44 100644 --- a/internal/app/mcp_authorization_test.go +++ b/internal/app/mcp_authorization_test.go @@ -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 { diff --git a/internal/app/mcp_project_scope.go b/internal/app/mcp_project_scope.go new file mode 100644 index 00000000..171acbab --- /dev/null +++ b/internal/app/mcp_project_scope.go @@ -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 "" +} diff --git a/internal/database/c2.go b/internal/database/c2.go index b8eda549..fee5184f 100644 --- a/internal/database/c2.go +++ b/internal/database/c2.go @@ -46,6 +46,7 @@ func validC2TextIDForDelete(id string) bool { // C2Listener 监听器实体 type C2Listener struct { ID string `json:"id"` + ProjectID string `json:"project_id,omitempty"` Name string `json:"name"` Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns BindHost string `json:"bindHost"` // 默认 127.0.0.1 @@ -165,12 +166,12 @@ func (db *DB) CreateC2Listener(l *C2Listener) error { l.ConfigJSON = "{}" } query := ` - INSERT INTO c2_listeners (id, name, type, bind_host, bind_port, profile_id, encryption_key, + INSERT INTO c2_listeners (id, project_id, name, type, bind_host, bind_port, profile_id, encryption_key, implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ` _, err := db.Exec(query, - l.ID, l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey, + l.ID, strings.TrimSpace(l.ProjectID), l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey, l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError, ) if err != nil { @@ -190,12 +191,12 @@ func (db *DB) UpdateC2Listener(l *C2Listener) error { } query := ` UPDATE c2_listeners SET - name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?, + project_id = ?, name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?, implant_token = ?, status = ?, config_json = ?, remark = ?, owner_user_id = ?, started_at = ?, last_error = ? WHERE id = ? ` res, err := db.Exec(query, - l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey, + strings.TrimSpace(l.ProjectID), l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey, l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID, ) if err != nil { @@ -229,7 +230,7 @@ func (db *DB) SetC2ListenerStatus(id, status, lastError string, startedAt *time. // GetC2Listener 单条查询 func (db *DB) GetC2Listener(id string) (*C2Listener, error) { query := ` - SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''), + SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''), COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status, COALESCE(config_json, '{}'), COALESCE(remark, ''), COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '') @@ -238,7 +239,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) { var l C2Listener var startedAt sql.NullTime err := db.QueryRow(query, id).Scan( - &l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID, + &l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID, &l.EncryptionKey, &l.ImplantToken, &l.Status, &l.ConfigJSON, &l.Remark, &l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError, @@ -259,7 +260,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) { // ListC2Listeners 全量列表，按创建时间倒序 func (db *DB) ListC2Listeners() ([]*C2Listener, error) { query := ` - SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''), + SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''), COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status, COALESCE(config_json, '{}'), COALESCE(remark, ''), COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '') @@ -275,7 +276,7 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) { var l C2Listener var startedAt sql.NullTime if err := rows.Scan( - &l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID, + &l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID, &l.EncryptionKey, &l.ImplantToken, &l.Status, &l.ConfigJSON, &l.Remark, &l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError, @@ -293,12 +294,18 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) { } // ListC2ListenersForAccess lists listeners visible to the resolved RBAC scope. -func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, error) { +func (db *DB) ListC2ListenersForAccess(access RBACListAccess, projectID string) ([]*C2Listener, error) { conditions := []string{"1=1"} args := []interface{}{} + if projectID = strings.TrimSpace(projectID); projectID == ProjectFilterUnbound { + conditions = append(conditions, "COALESCE(project_id, '') = ''") + } else if projectID != "" { + conditions = append(conditions, "COALESCE(project_id, '') = ?") + args = append(args, projectID) + } appendC2ListenerAccessFilter(&conditions, &args, access) query := ` - SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''), + SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''), COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status, COALESCE(config_json, '{}'), COALESCE(remark, ''), COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '') @@ -316,7 +323,7 @@ func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, er var l C2Listener var startedAt sql.NullTime if err := rows.Scan( - &l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID, + &l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID, &l.EncryptionKey, &l.ImplantToken, &l.Status, &l.ConfigJSON, &l.Remark, &l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError, @@ -535,6 +542,7 @@ func (db *DB) queryC2SessionWhere(whereClause string, args ...interface{}) (*C2S // ListC2SessionsFilter 列表过滤参数 type ListC2SessionsFilter struct { ListenerID string + ProjectID string Status string // active|sleeping|dead|killed；空表示全部 OS string Search string // 模糊匹配 hostname/username/internal_ip @@ -550,6 +558,18 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error) conditions = append(conditions, "listener_id = ?") args = append(args, filter.ListenerID) } + if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound { + conditions = append(conditions, `EXISTS ( + SELECT 1 FROM c2_listeners l + WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = '' + )`) + } else if strings.TrimSpace(filter.ProjectID) != "" { + conditions = append(conditions, `EXISTS ( + SELECT 1 FROM c2_listeners l + WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ? + )`) + args = append(args, strings.TrimSpace(filter.ProjectID)) + } if filter.Status != "" { conditions = append(conditions, "status = ?") args = append(args, filter.Status) @@ -645,6 +665,18 @@ func buildC2SessionsWhere(filter ListC2SessionsFilter) ([]string, []interface{}) conditions = append(conditions, "listener_id = ?") args = append(args, filter.ListenerID) } + if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound { + conditions = append(conditions, `EXISTS ( + SELECT 1 FROM c2_listeners l + WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = '' + )`) + } else if strings.TrimSpace(filter.ProjectID) != "" { + conditions = append(conditions, `EXISTS ( + SELECT 1 FROM c2_listeners l + WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ? + )`) + args = append(args, strings.TrimSpace(filter.ProjectID)) + } if filter.Status != "" { conditions = append(conditions, "status = ?") args = append(args, filter.Status) @@ -947,7 +979,10 @@ func (db *DB) GetC2Task(id string) (*C2Task, error) { // ListC2TasksFilter 任务过滤 type ListC2TasksFilter struct { SessionID string + ProjectID string Status string + TaskType string + Since *time.Time Limit int Offset int } @@ -959,10 +994,32 @@ func buildC2TasksWhere(filter ListC2TasksFilter) (where string, args []interface conditions = append(conditions, "session_id = ?") args = append(args, filter.SessionID) } + if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound { + conditions = append(conditions, `EXISTS ( + SELECT 1 FROM c2_sessions s + JOIN c2_listeners l ON l.id = s.listener_id + WHERE s.id = c2_tasks.session_id AND COALESCE(l.project_id, '') = '' + )`) + } else if strings.TrimSpace(filter.ProjectID) != "" { + conditions = append(conditions, `EXISTS ( + SELECT 1 FROM c2_sessions s + JOIN c2_listeners l ON l.id = s.listener_id + WHERE s.id = c2_tasks.session_id AND COALESCE(l.project_id, '') = ? + )`) + args = append(args, strings.TrimSpace(filter.ProjectID)) + } if filter.Status != "" { conditions = append(conditions, "status = ?") args = append(args, filter.Status) } + if strings.TrimSpace(filter.TaskType) != "" { + conditions = append(conditions, "task_type = ?") + args = append(args, strings.TrimSpace(filter.TaskType)) + } + if filter.Since != nil { + conditions = append(conditions, sqliteEpochGE("created_at", ">=")) + args = append(args, formatSQLiteUTC(*filter.Since)) + } return strings.Join(conditions, " AND "), args } @@ -1016,6 +1073,43 @@ func (db *DB) CountC2TasksForAccess(filter ListC2TasksFilter, access RBACListAcc return n, err } +// CountC2TasksByStatusForAccess 与 ListC2Tasks 相同过滤条件下按状态统计 +func (db *DB) CountC2TasksByStatusForAccess(filter ListC2TasksFilter, access RBACListAccess) (map[string]int64, error) { + where, args := buildC2TasksWhereForAccess(filter, access) + query := `SELECT status, COUNT(*) FROM c2_tasks WHERE ` + where + ` GROUP BY status` + rows, err := db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + counts := map[string]int64{ + "queued": 0, + "sent": 0, + "running": 0, + "success": 0, + "failed": 0, + "cancelled": 0, + "pending": 0, + } + var legacyPending int64 + for rows.Next() { + var status string + var n int64 + if err := rows.Scan(&status, &n); err != nil { + continue + } + if status == "pending" { + legacyPending = n + continue + } + if _, ok := counts[status]; ok { + counts[status] = n + } + } + counts["pending"] = counts["queued"] + counts["sent"] + counts["running"] + legacyPending + return counts, rows.Err() +} + // CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数（仪表盘「待审任务」） func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) { conditions := []string{"status IN ('queued', 'pending')"} @@ -1030,8 +1124,8 @@ func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) { return n, err } -func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID string, access RBACListAccess) (int64, error) { - filter := ListC2TasksFilter{SessionID: sessionID} +func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID, projectID string, access RBACListAccess) (int64, error) { + filter := ListC2TasksFilter{SessionID: sessionID, ProjectID: projectID} where, args := buildC2TasksWhereForAccess(filter, access) query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where var n int64 @@ -1412,6 +1506,7 @@ func (db *DB) AppendC2Event(e *C2Event) error { type ListC2EventsFilter struct { Level string Category string + ProjectID string SessionID string TaskID string Since *time.Time @@ -1430,6 +1525,49 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa conditions = append(conditions, "category = ?") args = append(args, filter.Category) } + if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound { + conditions = append(conditions, `( + EXISTS ( + SELECT 1 FROM c2_sessions s + JOIN c2_listeners l ON l.id = s.listener_id + WHERE s.id = c2_events.session_id AND COALESCE(l.project_id, '') = '' + ) + OR EXISTS ( + SELECT 1 FROM c2_tasks t + JOIN c2_sessions s ON s.id = t.session_id + JOIN c2_listeners l ON l.id = s.listener_id + WHERE t.id = c2_events.task_id AND COALESCE(l.project_id, '') = '' + ) + OR EXISTS ( + SELECT 1 FROM c2_listeners l + WHERE json_valid(c2_events.data_json) + AND l.id = json_extract(c2_events.data_json, '$.listener_id') + AND COALESCE(l.project_id, '') = '' + ) + )`) + } else if strings.TrimSpace(filter.ProjectID) != "" { + conditions = append(conditions, `( + EXISTS ( + SELECT 1 FROM c2_sessions s + JOIN c2_listeners l ON l.id = s.listener_id + WHERE s.id = c2_events.session_id AND COALESCE(l.project_id, '') = ? + ) + OR EXISTS ( + SELECT 1 FROM c2_tasks t + JOIN c2_sessions s ON s.id = t.session_id + JOIN c2_listeners l ON l.id = s.listener_id + WHERE t.id = c2_events.task_id AND COALESCE(l.project_id, '') = ? + ) + OR EXISTS ( + SELECT 1 FROM c2_listeners l + WHERE json_valid(c2_events.data_json) + AND l.id = json_extract(c2_events.data_json, '$.listener_id') + AND COALESCE(l.project_id, '') = ? + ) + )`) + pid := strings.TrimSpace(filter.ProjectID) + args = append(args, pid, pid, pid) + } if filter.SessionID != "" { conditions = append(conditions, "session_id = ?") args = append(args, filter.SessionID) diff --git a/internal/database/database.go b/internal/database/database.go index 1f9ecc05..056d4791 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -494,6 +494,7 @@ func (db *DB) initTables() error { createWebshellConnectionsTable := ` CREATE TABLE IF NOT EXISTS webshell_connections ( id TEXT PRIMARY KEY, + project_id TEXT, url TEXT NOT NULL, password TEXT NOT NULL DEFAULT '', type TEXT NOT NULL DEFAULT 'php', @@ -520,6 +521,7 @@ func (db *DB) initTables() error { createC2ListenersTable := ` CREATE TABLE IF NOT EXISTS c2_listeners ( id TEXT PRIMARY KEY, + project_id TEXT, name TEXT NOT NULL, type TEXT NOT NULL, bind_host TEXT NOT NULL DEFAULT '127.0.0.1', @@ -756,8 +758,10 @@ func (db *DB) initTables() error { CREATE INDEX IF NOT EXISTS idx_batch_task_queues_created_at ON batch_task_queues(created_at); CREATE INDEX IF NOT EXISTS idx_batch_task_queues_title ON batch_task_queues(title); CREATE INDEX IF NOT EXISTS idx_webshell_connections_created_at ON webshell_connections(created_at); + CREATE INDEX IF NOT EXISTS idx_webshell_connections_project_id ON webshell_connections(project_id); CREATE INDEX IF NOT EXISTS idx_webshell_connection_states_updated_at ON webshell_connection_states(updated_at); CREATE INDEX IF NOT EXISTS idx_c2_listeners_created_at ON c2_listeners(created_at); + CREATE INDEX IF NOT EXISTS idx_c2_listeners_project_id ON c2_listeners(project_id); CREATE INDEX IF NOT EXISTS idx_c2_listeners_status ON c2_listeners(status); CREATE INDEX IF NOT EXISTS idx_c2_sessions_listener ON c2_sessions(listener_id); CREATE INDEX IF NOT EXISTS idx_c2_sessions_status ON c2_sessions(status); @@ -956,6 +960,9 @@ func (db *DB) initTables() error { db.logger.Warn("迁移webshell_connections表失败", zap.Error(err)) // 不返回错误，允许继续运行 } + if err := db.migrateC2ListenersTable(); err != nil { + db.logger.Warn("迁移c2_listeners表失败", zap.Error(err)) + } if err := db.migrateWorkflowRunsTable(); err != nil { db.logger.Warn("迁移workflow_runs表失败", zap.Error(err)) } @@ -1610,6 +1617,7 @@ func (db *DB) migrateWebshellConnectionsTable() error { name string stmt string }{ + {name: "project_id", stmt: "ALTER TABLE webshell_connections ADD COLUMN project_id TEXT"}, {name: "encoding", stmt: "ALTER TABLE webshell_connections ADD COLUMN encoding TEXT NOT NULL DEFAULT ''"}, {name: "os", stmt: "ALTER TABLE webshell_connections ADD COLUMN os TEXT NOT NULL DEFAULT ''"}, } @@ -1635,6 +1643,10 @@ func (db *DB) migrateWebshellConnectionsTable() error { return nil } +func (db *DB) migrateC2ListenersTable() error { + return db.addColumnIfMissing("c2_listeners", "project_id", "ALTER TABLE c2_listeners ADD COLUMN project_id TEXT") +} + // NewKnowledgeDB 创建知识库数据库连接（只包含知识库相关的表） func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) { sqlDB, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL") diff --git a/internal/database/project.go b/internal/database/project.go index 75be4d69..791c3bee 100644 --- a/internal/database/project.go +++ b/internal/database/project.go @@ -263,7 +263,7 @@ func (db *DB) UpdateProject(p *Project) error { return nil } -// DeleteProject 删除项目（级联删除事实；对话 project_id 置空由 FK 处理；漏洞 project_id 置空）。 +// DeleteProject 删除项目（级联删除事实；对话 project_id 置空由 FK 处理；其他资源 project_id 置空）。 func (db *DB) DeleteProject(id string) error { if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil { return fmt.Errorf("解除漏洞项目关联失败: %w", err) @@ -271,6 +271,12 @@ func (db *DB) DeleteProject(id string) error { if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil { return fmt.Errorf("解除资产项目关联失败: %w", err) } + if _, err := db.Exec(`UPDATE webshell_connections SET project_id = NULL WHERE project_id = ?`, id); err != nil { + return fmt.Errorf("解除 WebShell 项目关联失败: %w", err) + } + if _, err := db.Exec(`UPDATE c2_listeners SET project_id = NULL WHERE project_id = ?`, id); err != nil { + return fmt.Errorf("解除 C2 监听器项目关联失败: %w", err) + } _, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id) if err != nil { return fmt.Errorf("删除项目失败: %w", err) diff --git a/internal/database/rbac_access_test.go b/internal/database/rbac_access_test.go index dcdbfdaf..4ecef6b6 100644 --- a/internal/database/rbac_access_test.go +++ b/internal/database/rbac_access_test.go @@ -373,22 +373,46 @@ func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) { func TestRBACWebshellAndBatchListAccess(t *testing.T) { db := newRBACTestDB(t) - ws1 := WebShellConnection{ID: "ws_visible", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()} - ws2 := WebShellConnection{ID: "ws_hidden", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()} + ws1 := WebShellConnection{ID: "ws_visible", ProjectID: "p1", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()} + ws2 := WebShellConnection{ID: "ws_hidden", ProjectID: "p2", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()} + ws3 := WebShellConnection{ID: "ws_other_project", ProjectID: "p2", URL: "http://c", Type: "php", Method: "post", CreatedAt: time.Now()} + ws4 := WebShellConnection{ID: "ws_unbound", URL: "http://d", Type: "php", Method: "post", CreatedAt: time.Now()} if err := db.CreateWebshellConnection(&ws1); err != nil { t.Fatal(err) } if err := db.CreateWebshellConnection(&ws2); err != nil { t.Fatal(err) } + if err := db.CreateWebshellConnection(&ws3); err != nil { + t.Fatal(err) + } + if err := db.CreateWebshellConnection(&ws4); err != nil { + t.Fatal(err) + } _ = db.SetResourceOwner("webshell", ws1.ID, "u1") _ = db.SetResourceOwner("webshell", ws2.ID, "u2") - webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn) + _ = db.SetResourceOwner("webshell", ws3.ID, "u1") + _ = db.SetResourceOwner("webshell", ws4.ID, "u1") + webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "") + if err != nil { + t.Fatal(err) + } + if len(webshells) != 3 { + t.Fatalf("webshells = %#v, want 3 owned webshells including unbound", webshells) + } + webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "p1") if err != nil { t.Fatal(err) } if len(webshells) != 1 || webshells[0].ID != ws1.ID { - t.Fatalf("webshells = %#v, want only %s", webshells, ws1.ID) + t.Fatalf("webshells scoped to p1 = %#v, want only %s", webshells, ws1.ID) + } + webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, ProjectFilterUnbound) + if err != nil { + t.Fatal(err) + } + if len(webshells) != 1 || webshells[0].ID != ws4.ID { + t.Fatalf("unbound webshells = %#v, want only %s", webshells, ws4.ID) } if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil { @@ -411,61 +435,143 @@ func TestRBACWebshellAndBatchListAccess(t *testing.T) { func TestRBACC2AccessInheritsListener(t *testing.T) { db := newRBACTestDB(t) now := time.Now() - l1 := &C2Listener{ID: "l_visible", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now} - l2 := &C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now} + l1 := &C2Listener{ID: "l_visible", ProjectID: "p1", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now} + l2 := &C2Listener{ID: "l_hidden", ProjectID: "p2", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now} + l3 := &C2Listener{ID: "l_other_project", ProjectID: "p2", Name: "other project", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9003, OwnerUserID: "u1", CreatedAt: now} + l4 := &C2Listener{ID: "l_unbound", Name: "unbound", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9004, OwnerUserID: "u1", CreatedAt: now} if err := db.CreateC2Listener(l1); err != nil { t.Fatal(err) } if err := db.CreateC2Listener(l2); err != nil { t.Fatal(err) } + if err := db.CreateC2Listener(l3); err != nil { + t.Fatal(err) + } + if err := db.CreateC2Listener(l4); err != nil { + t.Fatal(err) + } if err := db.UpsertC2Session(&C2Session{ID: "s_visible", ListenerID: l1.ID, ImplantUUID: "implant-visible", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil { t.Fatal(err) } if err := db.UpsertC2Session(&C2Session{ID: "s_hidden", ListenerID: l2.ID, ImplantUUID: "implant-hidden", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil { t.Fatal(err) } + if err := db.UpsertC2Session(&C2Session{ID: "s_other_project", ListenerID: l3.ID, ImplantUUID: "implant-other-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil { + t.Fatal(err) + } + if err := db.UpsertC2Session(&C2Session{ID: "s_unbound", ListenerID: l4.ID, ImplantUUID: "implant-unbound", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil { + t.Fatal(err) + } if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil { t.Fatal(err) } if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil { t.Fatal(err) } + if err := db.CreateC2Task(&C2Task{ID: "t_other_project", SessionID: "s_other_project", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil { + t.Fatal(err) + } + if err := db.CreateC2Task(&C2Task{ID: "t_unbound", SessionID: "s_unbound", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil { + t.Fatal(err) + } if err := db.AppendC2Event(&C2Event{ID: "e_visible", Level: "info", Category: "task", SessionID: "s_visible", TaskID: "t_visible", Message: "visible", CreatedAt: now}); err != nil { t.Fatal(err) } if err := db.AppendC2Event(&C2Event{ID: "e_hidden", Level: "info", Category: "task", SessionID: "s_hidden", TaskID: "t_hidden", Message: "hidden", CreatedAt: now}); err != nil { t.Fatal(err) } + if err := db.AppendC2Event(&C2Event{ID: "e_other_project", Level: "info", Category: "task", SessionID: "s_other_project", TaskID: "t_other_project", Message: "other project", CreatedAt: now}); err != nil { + t.Fatal(err) + } + if err := db.AppendC2Event(&C2Event{ID: "e_unbound", Level: "info", Category: "task", SessionID: "s_unbound", TaskID: "t_unbound", Message: "unbound", CreatedAt: now}); err != nil { + t.Fatal(err) + } access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn} - listeners, err := db.ListC2ListenersForAccess(access) + listeners, err := db.ListC2ListenersForAccess(access, "") + if err != nil { + t.Fatal(err) + } + if len(listeners) != 3 { + t.Fatalf("listeners = %#v, want 3 owned listeners including unbound", listeners) + } + listeners, err = db.ListC2ListenersForAccess(access, "p1") if err != nil { t.Fatal(err) } if len(listeners) != 1 || listeners[0].ID != l1.ID { - t.Fatalf("listeners = %#v, want only %s", listeners, l1.ID) + t.Fatalf("listeners scoped to p1 = %#v, want only %s", listeners, l1.ID) + } + listeners, err = db.ListC2ListenersForAccess(access, ProjectFilterUnbound) + if err != nil { + t.Fatal(err) + } + if len(listeners) != 1 || listeners[0].ID != l4.ID { + t.Fatalf("unbound listeners = %#v, want only %s", listeners, l4.ID) } sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access) if err != nil { t.Fatal(err) } + if len(sessions) != 3 { + t.Fatalf("sessions = %#v, want 3 owned sessions including unbound", sessions) + } + sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: "p1"}, access) + if err != nil { + t.Fatal(err) + } if len(sessions) != 1 || sessions[0].ID != "s_visible" { - t.Fatalf("sessions = %#v, want only s_visible", sessions) + t.Fatalf("sessions scoped to p1 = %#v, want only s_visible", sessions) + } + sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: ProjectFilterUnbound}, access) + if err != nil { + t.Fatal(err) + } + if len(sessions) != 1 || sessions[0].ID != "s_unbound" { + t.Fatalf("unbound sessions = %#v, want only s_unbound", sessions) } tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access) if err != nil { t.Fatal(err) } + if len(tasks) != 3 { + t.Fatalf("tasks = %#v, want 3 owned tasks including unbound", tasks) + } + tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: "p1"}, access) + if err != nil { + t.Fatal(err) + } if len(tasks) != 1 || tasks[0].ID != "t_visible" { - t.Fatalf("tasks = %#v, want only t_visible", tasks) + t.Fatalf("tasks scoped to p1 = %#v, want only t_visible", tasks) + } + tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: ProjectFilterUnbound}, access) + if err != nil { + t.Fatal(err) + } + if len(tasks) != 1 || tasks[0].ID != "t_unbound" { + t.Fatalf("unbound tasks = %#v, want only t_unbound", tasks) } events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access) if err != nil { t.Fatal(err) } + if len(events) != 3 { + t.Fatalf("events = %#v, want 3 owned events including unbound", events) + } + events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: "p1"}, access) + if err != nil { + t.Fatal(err) + } if len(events) != 1 || events[0].ID != "e_visible" { - t.Fatalf("events = %#v, want only e_visible", events) + t.Fatalf("events scoped to p1 = %#v, want only e_visible", events) + } + events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: ProjectFilterUnbound}, access) + if err != nil { + t.Fatal(err) + } + if len(events) != 1 || events[0].ID != "e_unbound" { + t.Fatalf("unbound events = %#v, want only e_unbound", events) } if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") { t.Fatalf("expected listener ownership to allow task detail") diff --git a/internal/database/webshell.go b/internal/database/webshell.go index 424d61bc..e5335626 100644 --- a/internal/database/webshell.go +++ b/internal/database/webshell.go @@ -11,6 +11,7 @@ import ( // WebShellConnection WebShell 连接配置 type WebShellConnection struct { ID string `json:"id"` + ProjectID string `json:"project_id,omitempty"` URL string `json:"url"` Password string `json:"password"` Type string `json:"type"` @@ -60,17 +61,24 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro // ListWebshellConnections 列出所有 WebShell 连接，按创建时间倒序 func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) { - return db.ListWebshellConnectionsForAccess("", "") + return db.ListWebshellConnectionsForAccess("", "", "") } -func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShellConnection, error) { +func (db *DB) ListWebshellConnectionsForAccess(userID, scope, projectID string) ([]WebShellConnection, error) { query := ` - SELECT id, url, password, type, method, cmd_param, remark, + SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark, COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at FROM webshell_connections WHERE 1=1 ` args := []interface{}{} + projectID = strings.TrimSpace(projectID) + if projectID == ProjectFilterUnbound { + query += ` AND COALESCE(project_id, '') = ''` + } else if projectID != "" { + query += ` AND COALESCE(project_id, '') = ?` + args = append(args, projectID) + } userID = strings.TrimSpace(userID) if userID != "" && scope != RBACScopeAll { query += ` AND ( @@ -93,7 +101,7 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell var list []WebShellConnection for rows.Next() { var c WebShellConnection - err := rows.Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt) + err := rows.Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt) if err != nil { db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err)) continue @@ -106,12 +114,12 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell // GetWebshellConnection 根据 ID 获取一条连接 func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) { query := ` - SELECT id, url, password, type, method, cmd_param, remark, + SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark, COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at FROM webshell_connections WHERE id = ? ` var c WebShellConnection - err := db.QueryRow(query, id).Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt) + err := db.QueryRow(query, id).Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt) if err == sql.ErrNoRows { return nil, nil } @@ -125,10 +133,10 @@ func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) { // CreateWebshellConnection 创建 WebShell 连接 func (db *DB) CreateWebshellConnection(c *WebShellConnection) error { query := ` - INSERT INTO webshell_connections (id, url, password, type, method, cmd_param, remark, encoding, os, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + INSERT INTO webshell_connections (id, project_id, url, password, type, method, cmd_param, remark, encoding, os, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ` - _, err := db.Exec(query, c.ID, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt) + _, err := db.Exec(query, c.ID, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt) if err != nil { db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID)) return err @@ -140,10 +148,10 @@ func (db *DB) CreateWebshellConnection(c *WebShellConnection) error { func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error { query := ` UPDATE webshell_connections - SET url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ? + SET project_id = ?, url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ? WHERE id = ? ` - result, err := db.Exec(query, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID) + result, err := db.Exec(query, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID) if err != nil { db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID)) return err