diff --git a/internal/database/audit.go b/internal/database/audit.go index a4bfe6cb..57d0347f 100644 --- a/internal/database/audit.go +++ b/internal/database/audit.go @@ -78,8 +78,8 @@ func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) { } if q := strings.TrimSpace(filter.Query); q != "" { like := "%" + q + "%" - conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ?)") - args = append(args, like, like, like, like) + conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ? OR detail_json LIKE ?)") + args = append(args, like, like, like, like, like) } return strings.Join(conditions, " AND "), args } diff --git a/internal/database/batch_task.go b/internal/database/batch_task.go index b59a9acb..0be6cac2 100644 --- a/internal/database/batch_task.go +++ b/internal/database/batch_task.go @@ -137,7 +137,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) { // GetAllBatchQueues 获取所有批量任务队列 func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) { rows, err := db.Query( - "SELECT "+batchQueueSelectColumns+" FROM batch_task_queues ORDER BY created_at DESC", + "SELECT " + batchQueueSelectColumns + " FROM batch_task_queues ORDER BY created_at DESC", ) if err != nil { return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err) @@ -168,6 +168,10 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) { // ListBatchQueues 列出批量任务队列（支持筛选和分页） func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) { + return db.ListBatchQueuesForAccess(limit, offset, status, keyword, "", "") +} + +func (db *DB) ListBatchQueuesForAccess(limit, offset int, status, keyword, userID, scope string) ([]*BatchTaskQueueRow, error) { query := "SELECT " + batchQueueSelectColumns + " FROM batch_task_queues WHERE 1=1" args := []interface{}{} @@ -182,6 +186,26 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat query += " AND (id LIKE ? OR title LIKE ?)" args = append(args, "%"+keyword+"%", "%"+keyword+"%") } + userID = strings.TrimSpace(userID) + if userID != "" && scope != RBACScopeAll { + query += ` AND ( + owner_user_id = ? + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id + ) + OR ( + project_id IS NOT NULL AND project_id <> '' AND ( + EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments pra + WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id + ) + ) + ) + )` + args = append(args, userID, userID, userID, userID) + } query += " ORDER BY created_at DESC LIMIT ? OFFSET ?" args = append(args, limit, offset) @@ -216,6 +240,10 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat // CountBatchQueues 统计批量任务队列总数（支持筛选条件） func (db *DB) CountBatchQueues(status, keyword string) (int, error) { + return db.CountBatchQueuesForAccess(status, keyword, "", "") +} + +func (db *DB) CountBatchQueuesForAccess(status, keyword, userID, scope string) (int, error) { query := "SELECT COUNT(*) FROM batch_task_queues WHERE 1=1" args := []interface{}{} @@ -230,6 +258,26 @@ func (db *DB) CountBatchQueues(status, keyword string) (int, error) { query += " AND (id LIKE ? OR title LIKE ?)" args = append(args, "%"+keyword+"%", "%"+keyword+"%") } + userID = strings.TrimSpace(userID) + if userID != "" && scope != RBACScopeAll { + query += ` AND ( + owner_user_id = ? + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id + ) + OR ( + project_id IS NOT NULL AND project_id <> '' AND ( + EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments pra + WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id + ) + ) + ) + )` + args = append(args, userID, userID, userID, userID) + } var count int err := db.QueryRow(query, args...).Scan(&count) diff --git a/internal/database/c2.go b/internal/database/c2.go index 4290ea7b..38e361f6 100644 --- a/internal/database/c2.go +++ b/internal/database/c2.go @@ -45,20 +45,21 @@ func validC2TextIDForDelete(id string) bool { // C2Listener 监听器实体 type C2Listener struct { - ID string `json:"id"` - Name string `json:"name"` - Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns - BindHost string `json:"bindHost"` // 默认 127.0.0.1 - BindPort int `json:"bindPort"` // 1-65535 - ProfileID string `json:"profileId"` // 可空：关联 c2_profiles.id - EncryptionKey string `json:"-"` // base64(AES-256)，前端不返回 - ImplantToken string `json:"-"` // beacon 携带的鉴权 token，前端不返回 - Status string `json:"status"` // stopped|running|error - ConfigJSON string `json:"configJson"` // TLS 证书路径 / URI 模式 / 上限并发 等 - Remark string `json:"remark"` - CreatedAt time.Time `json:"createdAt"` + ID string `json:"id"` + Name string `json:"name"` + Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns + BindHost string `json:"bindHost"` // 默认 127.0.0.1 + BindPort int `json:"bindPort"` // 1-65535 + ProfileID string `json:"profileId"` // 可空：关联 c2_profiles.id + EncryptionKey string `json:"-"` // base64(AES-256)，前端不返回 + ImplantToken string `json:"-"` // beacon 携带的鉴权 token，前端不返回 + Status string `json:"status"` // stopped|running|error + ConfigJSON string `json:"configJson"` // TLS 证书路径 / URI 模式 / 上限并发 等 + Remark string `json:"remark"` + OwnerUserID string `json:"ownerUserId,omitempty"` + CreatedAt time.Time `json:"createdAt"` StartedAt *time.Time `json:"startedAt,omitempty"` - LastError string `json:"lastError,omitempty"` + LastError string `json:"lastError,omitempty"` } // C2Session 已上线会话 @@ -132,17 +133,17 @@ type C2Event struct { // C2Profile Malleable Profile type C2Profile struct { - ID string `json:"id"` - Name string `json:"name"` - UserAgent string `json:"userAgent"` - URIs []string `json:"uris"` - RequestHeaders map[string]string `json:"requestHeaders,omitempty"` - ResponseHeaders map[string]string `json:"responseHeaders,omitempty"` - BodyTemplate string `json:"bodyTemplate"` - JitterMinMS int `json:"jitterMinMs"` - JitterMaxMS int `json:"jitterMaxMs"` - Extra map[string]interface{} `json:"extra,omitempty"` - CreatedAt time.Time `json:"createdAt"` + ID string `json:"id"` + Name string `json:"name"` + UserAgent string `json:"userAgent"` + URIs []string `json:"uris"` + RequestHeaders map[string]string `json:"requestHeaders,omitempty"` + ResponseHeaders map[string]string `json:"responseHeaders,omitempty"` + BodyTemplate string `json:"bodyTemplate"` + JitterMinMS int `json:"jitterMinMs"` + JitterMaxMS int `json:"jitterMaxMs"` + Extra map[string]interface{} `json:"extra,omitempty"` + CreatedAt time.Time `json:"createdAt"` } // ---------------------------------------------------------------------------- @@ -165,12 +166,12 @@ func (db *DB) CreateC2Listener(l *C2Listener) error { } query := ` INSERT INTO c2_listeners (id, name, type, bind_host, bind_port, profile_id, encryption_key, - implant_token, status, config_json, remark, created_at, started_at, last_error) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ` _, err := db.Exec(query, l.ID, l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey, - l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.CreatedAt, l.StartedAt, l.LastError, + l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError, ) if err != nil { db.logger.Error("创建 C2 监听器失败", zap.Error(err), zap.String("id", l.ID)) @@ -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 = ?, - implant_token = ?, status = ?, config_json = ?, remark = ?, started_at = ?, last_error = ? + 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, - l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.StartedAt, l.LastError, l.ID, + l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID, ) if err != nil { db.logger.Error("更新 C2 监听器失败", zap.Error(err), zap.String("id", l.ID)) @@ -231,7 +232,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) { SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''), COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status, COALESCE(config_json, '{}'), COALESCE(remark, ''), - created_at, started_at, COALESCE(last_error, '') + COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '') FROM c2_listeners WHERE id = ? ` var l C2Listener @@ -240,7 +241,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) { &l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID, &l.EncryptionKey, &l.ImplantToken, &l.Status, &l.ConfigJSON, &l.Remark, - &l.CreatedAt, &startedAt, &l.LastError, + &l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError, ) if err == sql.ErrNoRows { return nil, nil @@ -261,7 +262,7 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) { SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''), COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status, COALESCE(config_json, '{}'), COALESCE(remark, ''), - created_at, started_at, COALESCE(last_error, '') + COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '') FROM c2_listeners ORDER BY created_at DESC ` rows, err := db.Query(query) @@ -277,6 +278,47 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) { &l.ID, &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, + ); err != nil { + db.logger.Warn("扫描 c2_listeners 行失败", zap.Error(err)) + continue + } + if startedAt.Valid { + t := startedAt.Time + l.StartedAt = &t + } + list = append(list, &l) + } + return list, rows.Err() +} + +// ListC2ListenersForAccess lists listeners visible to the resolved RBAC scope. +func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, error) { + conditions := []string{"1=1"} + args := []interface{}{} + appendC2ListenerAccessFilter(&conditions, &args, access) + query := ` + SELECT 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, '') + FROM c2_listeners + WHERE ` + strings.Join(conditions, " AND ") + ` + ORDER BY created_at DESC + ` + rows, err := db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var list []*C2Listener + for rows.Next() { + var l C2Listener + var startedAt sql.NullTime + if err := rows.Scan( + &l.ID, &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, ); err != nil { db.logger.Warn("扫描 c2_listeners 行失败", zap.Error(err)) @@ -291,6 +333,26 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) { return list, rows.Err() } +func appendC2ListenerAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) { + if access.Scope == RBACScopeAll { + return + } + if access.UserID == "" { + *conditions = append(*conditions, "1=0") + return + } + clauses := []string{"owner_user_id = ?"} + *args = append(*args, access.UserID) + if access.Scope == RBACScopeAssigned { + clauses = append(clauses, `EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'c2_listener' AND ra.resource_id = c2_listeners.id + )`) + *args = append(*args, access.UserID) + } + *conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")") +} + // DeleteC2Listener 级联删除（会话/任务/文件/事件随之消失） func (db *DB) DeleteC2Listener(id string) error { res, err := db.Exec(`DELETE FROM c2_listeners WHERE id = ?`, id) @@ -550,6 +612,109 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error) return list, rows.Err() } +// ListC2SessionsForAccess lists sessions whose parent listener is visible. +func (db *DB) ListC2SessionsForAccess(filter ListC2SessionsFilter, access RBACListAccess) ([]*C2Session, error) { + conditions, args := buildC2SessionsWhere(filter) + appendC2SessionAccessFilter(&conditions, &args, access) + query := ` + SELECT id, listener_id, implant_uuid, COALESCE(hostname,''), COALESCE(username,''), + COALESCE(os,''), COALESCE(arch,''), COALESCE(pid, 0), COALESCE(process_name,''), + COALESCE(is_admin, 0), COALESCE(internal_ip,''), COALESCE(external_ip,''), + COALESCE(user_agent,''), COALESCE(sleep_seconds, 5), COALESCE(jitter_percent, 0), + status, first_seen_at, last_check_in, COALESCE(metadata_json, '{}'), + COALESCE(note, '') + FROM c2_sessions + WHERE ` + strings.Join(conditions, " AND ") + ` + ORDER BY last_check_in DESC + ` + if filter.Limit > 0 { + query += fmt.Sprintf(" LIMIT %d", filter.Limit) + } + rows, err := db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + return db.scanC2SessionRows(rows) +} + +func buildC2SessionsWhere(filter ListC2SessionsFilter) ([]string, []interface{}) { + conditions := []string{"1=1"} + args := []interface{}{} + if filter.ListenerID != "" { + conditions = append(conditions, "listener_id = ?") + args = append(args, filter.ListenerID) + } + if filter.Status != "" { + conditions = append(conditions, "status = ?") + args = append(args, filter.Status) + } + if filter.OS != "" { + conditions = append(conditions, "os = ?") + args = append(args, filter.OS) + } + if filter.Search != "" { + conditions = append(conditions, "(hostname LIKE ? OR username LIKE ? OR internal_ip LIKE ?)") + kw := "%" + filter.Search + "%" + args = append(args, kw, kw, kw) + } + if filter.Suspicious { + conditions = append(conditions, `status = 'dead' AND ( + hostname LIKE 'tcp_%' OR LOWER(COALESCE(username,'')) = 'unknown' OR COALESCE(pid, 0) = 0 + )`) + } + return conditions, args +} + +func (db *DB) scanC2SessionRows(rows *sql.Rows) ([]*C2Session, error) { + var list []*C2Session + for rows.Next() { + var s C2Session + var isAdminInt int + var metadataJSON string + if err := rows.Scan( + &s.ID, &s.ListenerID, &s.ImplantUUID, &s.Hostname, &s.Username, + &s.OS, &s.Arch, &s.PID, &s.ProcessName, + &isAdminInt, &s.InternalIP, &s.ExternalIP, + &s.UserAgent, &s.SleepSeconds, &s.JitterPercent, + &s.Status, &s.FirstSeenAt, &s.LastCheckIn, &metadataJSON, + &s.Note, + ); err != nil { + db.logger.Warn("扫描 c2_sessions 行失败", zap.Error(err)) + continue + } + s.IsAdmin = isAdminInt != 0 + if metadataJSON != "" && metadataJSON != "{}" { + _ = json.Unmarshal([]byte(metadataJSON), &s.Metadata) + } + list = append(list, &s) + } + return list, rows.Err() +} + +func appendC2SessionAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) { + if access.Scope == RBACScopeAll { + return + } + if access.UserID == "" { + *conditions = append(*conditions, "1=0") + return + } + clauses := []string{`EXISTS ( + SELECT 1 FROM c2_listeners + WHERE c2_listeners.id = c2_sessions.listener_id AND c2_listeners.owner_user_id = ? + )`} + *args = append(*args, access.UserID) + if access.Scope == RBACScopeAssigned { + clauses = append(clauses, `EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'c2_listener' AND ra.resource_id = c2_sessions.listener_id + )`) + *args = append(*args, access.UserID) + } + *conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")") +} + // DeleteC2Session 级联删除其 tasks/files func (db *DB) DeleteC2Session(id string) error { res, err := db.Exec(`DELETE FROM c2_sessions WHERE id = ?`, id) @@ -601,6 +766,29 @@ func (db *DB) DeleteC2SessionsByIDs(ids []string) (int64, error) { return res.RowsAffected() } +func (db *DB) DeleteC2SessionsByIDsForAccess(ids []string, access RBACListAccess) (int64, error) { + if access.Scope == RBACScopeAll { + return db.DeleteC2SessionsByIDs(ids) + } + clean := cleanC2IDs(ids) + if len(clean) == 0 { + return 0, ErrNoValidC2SessionIDs + } + placeholders := strings.Repeat("?,", len(clean)-1) + "?" + args := make([]interface{}, 0, len(clean)+2) + for _, id := range clean { + args = append(args, id) + } + conditions := []string{"id IN (" + placeholders + ")"} + appendC2SessionAccessFilter(&conditions, &args, access) + query := `DELETE FROM c2_sessions WHERE ` + strings.Join(conditions, " AND ") + res, err := db.Exec(query, args...) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + // ---------------------------------------------------------------------------- // CRUD：C2 任务 // ---------------------------------------------------------------------------- @@ -778,6 +966,39 @@ func buildC2TasksWhere(filter ListC2TasksFilter) (where string, args []interface return strings.Join(conditions, " AND "), args } +func appendC2TaskAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) { + if access.Scope == RBACScopeAll { + return + } + if access.UserID == "" { + *conditions = append(*conditions, "1=0") + return + } + clauses := []string{`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 l.owner_user_id = ? + )`} + *args = append(*args, access.UserID) + if access.Scope == RBACScopeAssigned { + clauses = append(clauses, `EXISTS ( + SELECT 1 FROM c2_sessions s + JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id + WHERE s.id = c2_tasks.session_id + AND ra.user_id = ? AND ra.resource_type = 'c2_listener' + )`) + *args = append(*args, access.UserID) + } + *conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")") +} + +func buildC2TasksWhereForAccess(filter ListC2TasksFilter, access RBACListAccess) (string, []interface{}) { + where, args := buildC2TasksWhere(filter) + conditions := []string{where} + appendC2TaskAccessFilter(&conditions, &args, access) + return strings.Join(conditions, " AND "), args +} + // CountC2Tasks 与 ListC2Tasks 相同过滤条件下的记录总数 func (db *DB) CountC2Tasks(filter ListC2TasksFilter) (int64, error) { where, args := buildC2TasksWhere(filter) @@ -787,6 +1008,14 @@ func (db *DB) CountC2Tasks(filter ListC2TasksFilter) (int64, error) { return n, err } +func (db *DB) CountC2TasksForAccess(filter ListC2TasksFilter, access RBACListAccess) (int64, error) { + where, args := buildC2TasksWhereForAccess(filter, access) + query := `SELECT COUNT(*) FROM c2_tasks WHERE ` + where + var n int64 + err := db.QueryRow(query, args...).Scan(&n) + return n, err +} + // CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数（仪表盘「待审任务」） func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) { conditions := []string{"status IN ('queued', 'pending')"} @@ -801,6 +1030,15 @@ 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} + where, args := buildC2TasksWhereForAccess(filter, access) + query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where + var n int64 + err := db.QueryRow(query, args...).Scan(&n) + return n, err +} + // ListC2Tasks 任务列表，按创建时间倒序 func (db *DB) ListC2Tasks(filter ListC2TasksFilter) ([]*C2Task, error) { where, args := buildC2TasksWhere(filter) @@ -866,6 +1104,74 @@ func (db *DB) ListC2Tasks(filter ListC2TasksFilter) ([]*C2Task, error) { return list, rows.Err() } +func (db *DB) ListC2TasksForAccess(filter ListC2TasksFilter, access RBACListAccess) ([]*C2Task, error) { + where, args := buildC2TasksWhereForAccess(filter, access) + query := ` + SELECT id, session_id, task_type, COALESCE(payload_json, '{}'), + status, COALESCE(result_text, ''), COALESCE(result_blob_path, ''), + COALESCE(error, ''), COALESCE(source, 'manual'), + COALESCE(conversation_id, ''), COALESCE(approval_status, ''), + created_at, sent_at, started_at, completed_at, COALESCE(duration_ms, 0) + FROM c2_tasks + WHERE ` + where + ` + ORDER BY created_at DESC + ` + limit := filter.Limit + offset := filter.Offset + if offset < 0 { + offset = 0 + } + if limit > 0 { + if limit > 1000 { + limit = 1000 + } + query += ` LIMIT ? OFFSET ?` + args = append(args, limit, offset) + } + rows, err := db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + return db.scanC2TaskRows(rows) +} + +func (db *DB) scanC2TaskRows(rows *sql.Rows) ([]*C2Task, error) { + var list []*C2Task + for rows.Next() { + var t C2Task + var payloadJSON string + var sentAt, startedAt, completedAt sql.NullTime + if err := rows.Scan( + &t.ID, &t.SessionID, &t.TaskType, &payloadJSON, + &t.Status, &t.ResultText, &t.ResultBlobPath, + &t.Error, &t.Source, + &t.ConversationID, &t.ApprovalStatus, + &t.CreatedAt, &sentAt, &startedAt, &completedAt, &t.DurationMS, + ); err != nil { + db.logger.Warn("扫描 c2_tasks 行失败", zap.Error(err)) + continue + } + if payloadJSON != "" && payloadJSON != "{}" { + _ = json.Unmarshal([]byte(payloadJSON), &t.Payload) + } + if sentAt.Valid { + x := sentAt.Time + t.SentAt = &x + } + if startedAt.Valid { + x := startedAt.Time + t.StartedAt = &x + } + if completedAt.Valid { + x := completedAt.Time + t.CompletedAt = &x + } + list = append(list, &t) + } + return list, rows.Err() +} + // PopQueuedC2Tasks 取出某会话所有 queued/approved 任务（用于 beacon 拉取），原子置为 sent func (db *DB) PopQueuedC2Tasks(sessionID string, limit int) ([]*C2Task, error) { if limit <= 0 { @@ -978,6 +1284,29 @@ func (db *DB) DeleteC2TasksByIDs(ids []string) (int64, error) { return res.RowsAffected() } +func (db *DB) DeleteC2TasksByIDsForAccess(ids []string, access RBACListAccess) (int64, error) { + if access.Scope == RBACScopeAll { + return db.DeleteC2TasksByIDs(ids) + } + clean := cleanC2IDs(ids) + if len(clean) == 0 { + return 0, ErrNoValidC2TaskIDs + } + placeholders := strings.Repeat("?,", len(clean)-1) + "?" + args := make([]interface{}, 0, len(clean)+2) + for _, id := range clean { + args = append(args, id) + } + conditions := []string{"id IN (" + placeholders + ")"} + appendC2TaskAccessFilter(&conditions, &args, access) + query := `DELETE FROM c2_tasks WHERE ` + strings.Join(conditions, " AND ") + res, err := db.Exec(query, args...) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + // ---------------------------------------------------------------------------- // CRUD：C2 文件 // ---------------------------------------------------------------------------- @@ -1024,6 +1353,27 @@ func (db *DB) ListC2FilesBySession(sessionID string) ([]*C2File, error) { return list, rows.Err() } +func cleanC2IDs(ids []string) []string { + const maxBatch = 500 + if len(ids) > maxBatch { + ids = ids[:maxBatch] + } + clean := make([]string, 0, len(ids)) + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + id = strings.TrimSpace(id) + if !validC2TextIDForDelete(id) { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + clean = append(clean, id) + } + return clean +} + // ---------------------------------------------------------------------------- // CRUD：C2 事件审计 // ---------------------------------------------------------------------------- @@ -1093,6 +1443,56 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa return strings.Join(conditions, " AND "), args } +func appendC2EventAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) { + if access.Scope == RBACScopeAll { + return + } + if access.UserID == "" { + *conditions = append(*conditions, "1=0") + return + } + clauses := []string{`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 l.owner_user_id = ? + )`} + *args = append(*args, access.UserID) + if access.Scope == RBACScopeAssigned { + clauses = append(clauses, `EXISTS ( + SELECT 1 FROM c2_sessions s + JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id + WHERE s.id = c2_events.session_id + AND ra.user_id = ? AND ra.resource_type = 'c2_listener' + )`) + *args = append(*args, access.UserID) + } + clauses = append(clauses, `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 l.owner_user_id = ? + )`) + *args = append(*args, access.UserID) + if access.Scope == RBACScopeAssigned { + clauses = append(clauses, `EXISTS ( + SELECT 1 FROM c2_tasks t + JOIN c2_sessions s ON s.id = t.session_id + JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id + WHERE t.id = c2_events.task_id + AND ra.user_id = ? AND ra.resource_type = 'c2_listener' + )`) + *args = append(*args, access.UserID) + } + *conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")") +} + +func buildC2EventsWhereForAccess(filter ListC2EventsFilter, access RBACListAccess) (string, []interface{}) { + where, args := buildC2EventsWhere(filter) + conditions := []string{where} + appendC2EventAccessFilter(&conditions, &args, access) + return strings.Join(conditions, " AND "), args +} + // CountC2Events 与 ListC2Events 相同过滤条件下的记录总数 func (db *DB) CountC2Events(filter ListC2EventsFilter) (int64, error) { where, args := buildC2EventsWhere(filter) @@ -1102,6 +1502,14 @@ func (db *DB) CountC2Events(filter ListC2EventsFilter) (int64, error) { return n, err } +func (db *DB) CountC2EventsForAccess(filter ListC2EventsFilter, access RBACListAccess) (int64, error) { + where, args := buildC2EventsWhereForAccess(filter, access) + query := `SELECT COUNT(*) FROM c2_events WHERE ` + where + var n int64 + err := db.QueryRow(query, args...).Scan(&n) + return n, err +} + // ListC2Events 事件查询，按创建时间倒序 func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) { where, args := buildC2EventsWhere(filter) @@ -1143,6 +1551,50 @@ func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) { return list, rows.Err() } +func (db *DB) ListC2EventsForAccess(filter ListC2EventsFilter, access RBACListAccess) ([]*C2Event, error) { + where, args := buildC2EventsWhereForAccess(filter, access) + limit := filter.Limit + if limit <= 0 || limit > 1000 { + limit = 200 + } + offset := filter.Offset + if offset < 0 { + offset = 0 + } + query := ` + SELECT id, level, category, COALESCE(session_id, ''), COALESCE(task_id, ''), + message, COALESCE(data_json, ''), created_at + FROM c2_events + WHERE ` + where + ` + ORDER BY created_at DESC + LIMIT ? OFFSET ? + ` + args = append(args, limit, offset) + rows, err := db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + return scanC2EventRows(rows) +} + +func scanC2EventRows(rows *sql.Rows) ([]*C2Event, error) { + var list []*C2Event + for rows.Next() { + var e C2Event + var dataJSON string + if err := rows.Scan(&e.ID, &e.Level, &e.Category, &e.SessionID, &e.TaskID, + &e.Message, &dataJSON, &e.CreatedAt); err != nil { + continue + } + if dataJSON != "" { + _ = json.Unmarshal([]byte(dataJSON), &e.Data) + } + list = append(list, &e) + } + return list, rows.Err() +} + // DeleteC2EventsByIDs 按主键批量删除事件，返回实际删除行数 func (db *DB) DeleteC2EventsByIDs(ids []string) (int64, error) { if len(ids) == 0 { @@ -1181,6 +1633,29 @@ func (db *DB) DeleteC2EventsByIDs(ids []string) (int64, error) { return res.RowsAffected() } +func (db *DB) DeleteC2EventsByIDsForAccess(ids []string, access RBACListAccess) (int64, error) { + if access.Scope == RBACScopeAll { + return db.DeleteC2EventsByIDs(ids) + } + clean := cleanC2IDs(ids) + if len(clean) == 0 { + return 0, ErrNoValidC2EventIDs + } + placeholders := strings.Repeat("?,", len(clean)-1) + "?" + args := make([]interface{}, 0, len(clean)+4) + for _, id := range clean { + args = append(args, id) + } + conditions := []string{"id IN (" + placeholders + ")"} + appendC2EventAccessFilter(&conditions, &args, access) + query := `DELETE FROM c2_events WHERE ` + strings.Join(conditions, " AND ") + res, err := db.Exec(query, args...) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + // ---------------------------------------------------------------------------- // CRUD：C2 Malleable Profile // ---------------------------------------------------------------------------- diff --git a/internal/database/conversation.go b/internal/database/conversation.go index 631431f7..5be515f3 100644 --- a/internal/database/conversation.go +++ b/internal/database/conversation.go @@ -384,6 +384,31 @@ func appendConversationProjectFilter(where string, args []interface{}, projectID return where + fmt.Sprintf(" AND %s = ?", col), append(args, pid) } +func appendConversationAccessFilter(where string, args []interface{}, userID, scope, alias string) (string, []interface{}) { + userID = strings.TrimSpace(userID) + if userID == "" || scope == RBACScopeAll { + return where, args + } + prefix := "" + if alias != "" { + prefix = alias + "." + } + where += fmt.Sprintf(` AND (%sowner_user_id = ? OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'conversation' AND ra.resource_id = %sid + ) OR EXISTS ( + SELECT 1 FROM projects p + WHERE p.id = %sproject_id AND ( + p.owner_user_id = ? OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments pra + WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = p.id + ) + ) + ))`, prefix, prefix, prefix) + args = append(args, userID, userID, userID, userID) + return where, args +} + // CountConversations 统计对话数量。 func (db *DB) CountConversations(search, projectID string) (int, error) { var count int @@ -410,6 +435,33 @@ func (db *DB) CountConversations(search, projectID string) (int, error) { return count, nil } +func (db *DB) CountConversationsForAccess(search, projectID, userID, scope string) (int, error) { + var count int + var err error + if search != "" { + searchPattern := "%" + search + "%" + where := ` WHERE (c.title LIKE ? + OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))` + args := []interface{}{searchPattern, searchPattern} + where, args = appendConversationProjectFilter(where, args, projectID, "c") + where, args = appendConversationAccessFilter(where, args, userID, scope, "c") + err = db.QueryRow(`SELECT COUNT(*) FROM conversations c`+where, args...).Scan(&count) + } else { + where := "" + args := []interface{}{} + where, args = appendConversationProjectFilter(where, args, projectID, "") + where, args = appendConversationAccessFilter(where, args, userID, scope, "") + if where != "" { + where = " WHERE" + strings.TrimPrefix(where, " AND") + } + err = db.QueryRow(`SELECT COUNT(*) FROM conversations`+where, args...).Scan(&count) + } + if err != nil { + return 0, fmt.Errorf("统计对话失败: %w", err) + } + return count, nil +} + func conversationOrderClause(sortBy, tableAlias string) string { col := "updated_at" if strings.TrimSpace(strings.ToLower(sortBy)) == "created_at" { @@ -503,6 +555,81 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str return conversations, nil } +func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) { + if scope == RBACScopeAll || strings.TrimSpace(userID) == "" { + return db.ListConversations(limit, offset, search, sortBy, projectID) + } + var rows *sql.Rows + var err error + if search != "" { + searchPattern := "%" + search + "%" + orderClause := conversationOrderClause(sortBy, "c") + where := ` WHERE (c.title LIKE ? + OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))` + args := []interface{}{searchPattern, searchPattern} + where, args = appendConversationProjectFilter(where, args, projectID, "c") + where, args = appendConversationAccessFilter(where, args, userID, scope, "c") + args = append(args, limit, offset) + rows, err = db.Query( + `SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id + FROM conversations c`+where+` + `+orderClause+` + LIMIT ? OFFSET ?`, args...) + } else { + orderClause := conversationOrderClause(sortBy, "") + where := "" + args := []interface{}{} + where, args = appendConversationProjectFilter(where, args, projectID, "") + where, args = appendConversationAccessFilter(where, args, userID, scope, "") + if where != "" { + where = " WHERE" + strings.TrimPrefix(where, " AND") + } + args = append(args, limit, offset) + rows, err = db.Query( + "SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?", + args...) + } + if err != nil { + return nil, fmt.Errorf("查询对话列表失败: %w", err) + } + defer rows.Close() + return scanConversationRows(rows) +} + +func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) { + var conversations []*Conversation + for rows.Next() { + var conv Conversation + var createdAt, updatedAt string + var pinned int + var projectID sql.NullString + if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil { + return nil, fmt.Errorf("扫描对话失败: %w", err) + } + if projectID.Valid { + conv.ProjectID = strings.TrimSpace(projectID.String) + } + var err1, err2 error + conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt) + if err1 != nil { + conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt) + } + if err1 != nil { + conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt) + } + conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt) + if err2 != nil { + conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt) + } + if err2 != nil { + conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt) + } + conv.Pinned = pinned != 0 + conversations = append(conversations, &conv) + } + return conversations, rows.Err() +} + const ungroupedConversationsSQL = ` FROM conversations c WHERE NOT EXISTS ( @@ -521,6 +648,18 @@ func (db *DB) CountUngroupedConversations(projectID string) (int, error) { return count, nil } +func (db *DB) CountUngroupedConversationsForAccess(projectID, userID, scope string) (int, error) { + where := ungroupedConversationsSQL + args := []interface{}{} + where, args = appendConversationProjectFilter(where, args, projectID, "c") + where, args = appendConversationAccessFilter(where, args, userID, scope, "c") + var count int + if err := db.QueryRow(`SELECT COUNT(*) `+where, args...).Scan(&count); err != nil { + return 0, fmt.Errorf("统计未分组对话失败: %w", err) + } + return count, nil +} + // ListUngroupedConversations 列出不在任何分组中的对话（最近对话侧栏）。 func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID string) ([]*Conversation, error) { orderClause := conversationOrderClause(sortBy, "c") @@ -578,6 +717,30 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st return conversations, rows.Err() } +func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) { + if scope == RBACScopeAll || strings.TrimSpace(userID) == "" { + return db.ListUngroupedConversations(limit, offset, sortBy, projectID) + } + orderClause := conversationOrderClause(sortBy, "c") + where := ungroupedConversationsSQL + args := []interface{}{} + where, args = appendConversationProjectFilter(where, args, projectID, "c") + where, args = appendConversationAccessFilter(where, args, userID, scope, "c") + args = append(args, limit, offset) + rows, err := db.Query( + `SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+ + where+` + `+orderClause+` + LIMIT ? OFFSET ?`, + args..., + ) + if err != nil { + return nil, fmt.Errorf("查询未分组对话失败: %w", err) + } + defer rows.Close() + return scanConversationRows(rows) +} + // GetConversationTitle 获取对话标题（轻量查询，不加载消息） func (db *DB) GetConversationTitle(id string) (string, error) { var title string @@ -1280,8 +1443,13 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary return nil, fmt.Errorf("查询工具执行摘要失败: %w", err) } seenExecIDs := make(map[string]bool) - toolIndexByCallID := make(map[string]int) - unmatchedToolIdx := 0 + // A provider may reuse a fallback toolCallId across streaming rounds. Keep a + // FIFO per ID instead of a single index so every persisted call gets at most + // one result. Results without an ID fall back to the oldest unmatched call. + toolIndexesByCallID := make(map[string][]int) + lastMatchedToolIndexByCallID := make(map[string]int) + matchedToolIndexes := make([]bool, 0) + nextUnmatchedToolIdx := 0 for execRows.Next() { var detailID string var eventType string @@ -1320,24 +1488,52 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary ProcessDetailID: strings.TrimSpace(detailID), ToolName: toolName, ToolCallID: toolCallID, - Status: "running", + // This summary is reconstructed from persisted history, not live + // execution state. Until a matching result is found the honest state + // is "result_missing", never "running". + Status: "result_missing", }) + matchedToolIndexes = append(matchedToolIndexes, false) if toolCallID != "" { - toolIndexByCallID[toolCallID] = len(summary.ToolExecutions) - 1 + toolIndexesByCallID[toolCallID] = append(toolIndexesByCallID[toolCallID], len(summary.ToolExecutions)-1) } } if eventType == "tool_result" { idx := -1 if toolCallID != "" { - if found, ok := toolIndexByCallID[toolCallID]; ok { - idx = found + queue := toolIndexesByCallID[toolCallID] + for len(queue) > 0 { + candidate := queue[0] + queue = queue[1:] + if candidate >= 0 && candidate < len(matchedToolIndexes) && !matchedToolIndexes[candidate] { + idx = candidate + break + } + } + toolIndexesByCallID[toolCallID] = queue + if idx < 0 { + // Multiple persisted result events for one call (for example an + // agent-facing reduced result replacing an earlier preview) update + // that call instead of consuming an unrelated FIFO entry. + if previous, ok := lastMatchedToolIndexByCallID[toolCallID]; ok { + idx = previous + } } } - if idx < 0 && unmatchedToolIdx < len(summary.ToolExecutions) { - idx = unmatchedToolIdx - unmatchedToolIdx++ + if idx < 0 { + for nextUnmatchedToolIdx < len(matchedToolIndexes) && matchedToolIndexes[nextUnmatchedToolIdx] { + nextUnmatchedToolIdx++ + } + if nextUnmatchedToolIdx < len(matchedToolIndexes) { + idx = nextUnmatchedToolIdx + nextUnmatchedToolIdx++ + } } if idx >= 0 && idx < len(summary.ToolExecutions) { + matchedToolIndexes[idx] = true + if toolCallID != "" { + lastMatchedToolIndexByCallID[toolCallID] = idx + } if summary.ToolExecutions[idx].ToolName == "" { summary.ToolExecutions[idx].ToolName = toolName } @@ -1356,6 +1552,7 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary ExecutionID: execID, Status: status, }) + matchedToolIndexes = append(matchedToolIndexes, true) } } if execID != "" && !seenExecIDs[execID] { diff --git a/internal/database/database.go b/internal/database/database.go index abe1f786..e41f23bb 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -478,6 +478,7 @@ func (db *DB) initTables() error { status TEXT NOT NULL DEFAULT 'stopped', config_json TEXT NOT NULL DEFAULT '{}', remark TEXT NOT NULL DEFAULT '', + owner_user_id TEXT, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, started_at DATETIME, last_error TEXT @@ -783,6 +784,10 @@ func (db *DB) initTables() error { return fmt.Errorf("创建audit_logs表失败: %w", err) } + if err := db.initRBACTables(); err != nil { + return fmt.Errorf("创建RBAC表失败: %w", err) + } + for tableName, ddl := range map[string]string{ "workflow_definitions": createWorkflowDefinitionsTable, "workflow_runs": createWorkflowRunsTable, @@ -853,6 +858,9 @@ func (db *DB) initTables() error { if err := db.migrateWorkflowRunsTable(); err != nil { db.logger.Warn("迁移workflow_runs表失败", zap.Error(err)) } + if err := db.migrateRBACOwnershipColumns(); err != nil { + db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err)) + } if _, err := db.Exec(createIndexes); err != nil { return fmt.Errorf("创建索引失败: %w", err) diff --git a/internal/database/process_details_summary_test.go b/internal/database/process_details_summary_test.go new file mode 100644 index 00000000..0ff8d9ab --- /dev/null +++ b/internal/database/process_details_summary_test.go @@ -0,0 +1,108 @@ +package database + +import ( + "path/filepath" + "testing" + + "go.uber.org/zap" +) + +func TestProcessDetailsSummaryPairsMixedIdentifiedAndIDLessResults(t *testing.T) { + db, conversationID, messageID := setupProcessDetailsSummaryTest(t) + for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} { + if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{ + "toolName": "http-framework-test", "toolCallId": id, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_call): %v", err) + } + } + results := []map[string]interface{}{ + {"toolName": "http-framework-test", "toolCallId": "call-1", "success": true}, + {"toolName": "http-framework-test", "toolCallId": "call-2", "success": true}, + {"toolName": "http-framework-test", "success": true}, + {"toolName": "http-framework-test", "success": true}, + } + for _, result := range results { + if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", result); err != nil { + t.Fatalf("AddProcessDetail(tool_result): %v", err) + } + } + + summary, err := db.GetProcessDetailsSummary(messageID) + if err != nil { + t.Fatalf("GetProcessDetailsSummary: %v", err) + } + if len(summary.ToolExecutions) != 4 { + t.Fatalf("tool executions = %d, want 4", len(summary.ToolExecutions)) + } + for i, execution := range summary.ToolExecutions { + if execution.Status != "completed" { + t.Fatalf("execution %d status = %q, want completed", i, execution.Status) + } + } +} + +func TestProcessDetailsSummaryPairsRepeatedToolCallIDsFIFO(t *testing.T) { + db, conversationID, messageID := setupProcessDetailsSummaryTest(t) + for i := 0; i < 2; i++ { + if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{ + "toolName": "execute", "toolCallId": "legacy-reused-id", + }); err != nil { + t.Fatalf("AddProcessDetail(tool_call): %v", err) + } + } + for i := 0; i < 2; i++ { + if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", map[string]interface{}{ + "toolName": "execute", "toolCallId": "legacy-reused-id", "success": true, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_result): %v", err) + } + } + + summary, err := db.GetProcessDetailsSummary(messageID) + if err != nil { + t.Fatalf("GetProcessDetailsSummary: %v", err) + } + if len(summary.ToolExecutions) != 2 { + t.Fatalf("tool executions = %d, want 2", len(summary.ToolExecutions)) + } + for i, execution := range summary.ToolExecutions { + if execution.Status != "completed" { + t.Fatalf("execution %d status = %q, want completed", i, execution.Status) + } + } +} + +func TestProcessDetailsSummaryDoesNotReportPersistedOrphanAsRunning(t *testing.T) { + db, conversationID, messageID := setupProcessDetailsSummaryTest(t) + if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{ + "toolName": "execute", "toolCallId": "orphan", + }); err != nil { + t.Fatalf("AddProcessDetail(tool_call): %v", err) + } + summary, err := db.GetProcessDetailsSummary(messageID) + if err != nil { + t.Fatalf("GetProcessDetailsSummary: %v", err) + } + if len(summary.ToolExecutions) != 1 || summary.ToolExecutions[0].Status != "result_missing" { + t.Fatalf("tool executions = %#v, want result_missing", summary.ToolExecutions) + } +} + +func setupProcessDetailsSummaryTest(t *testing.T) (*DB, string, string) { + t.Helper() + db, err := NewDB(filepath.Join(t.TempDir(), "process-details.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + conversation, err := db.CreateConversation("process details", ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + message, err := db.AddMessage(conversation.ID, "assistant", "done", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + return db, conversation.ID, message.ID +} diff --git a/internal/database/project.go b/internal/database/project.go index b5ad9cfe..d163aee8 100644 --- a/internal/database/project.go +++ b/internal/database/project.go @@ -58,11 +58,11 @@ type ProjectFact struct { // ProjectFactListFilter 事实列表筛选。 type ProjectFactListFilter struct { - Category string - Confidence string - Search string - RelatedVulnerabilityID string - ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated + Category string + Confidence string + Search string + RelatedVulnerabilityID string + ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated } // CreateProject 创建项目。 @@ -143,6 +143,19 @@ func appendProjectListFilters(query string, args []interface{}, status, search s return query, args } +func appendProjectAccessFilter(query string, args []interface{}, userID, scope string) (string, []interface{}) { + userID = strings.TrimSpace(userID) + if userID == "" || scope == RBACScopeAll { + return query, args + } + query += ` AND (owner_user_id = ? OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = projects.id + ))` + args = append(args, userID, userID) + return query, args +} + // CountProjects 统计项目数量。 func (db *DB) CountProjects(status, search string) (int, error) { query := `SELECT COUNT(*) FROM projects WHERE 1=1` @@ -155,6 +168,18 @@ func (db *DB) CountProjects(status, search string) (int, error) { return count, nil } +func (db *DB) CountProjectsForAccess(status, search, userID, scope string) (int, error) { + query := `SELECT COUNT(*) FROM projects WHERE 1=1` + args := []interface{}{} + query, args = appendProjectListFilters(query, args, status, search) + query, args = appendProjectAccessFilter(query, args, userID, scope) + var count int + if err := db.QueryRow(query, args...).Scan(&count); err != nil { + return 0, fmt.Errorf("统计项目失败: %w", err) + } + return count, nil +} + // ListProjects 列出项目。 func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project, error) { if limit <= 0 { @@ -189,6 +214,42 @@ func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project return out, rows.Err() } +func (db *DB) ListProjectsForAccess(status, search string, limit, offset int, userID, scope string) ([]*Project, error) { + if scope == RBACScopeAll || strings.TrimSpace(userID) == "" { + return db.ListProjects(status, search, limit, offset) + } + if limit <= 0 { + limit = 50 + } + query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at + FROM projects WHERE 1=1` + args := []interface{}{} + query, args = appendProjectListFilters(query, args, status, search) + query, args = appendProjectAccessFilter(query, args, userID, scope) + query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?" + args = append(args, limit, offset) + + rows, err := db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("列出项目失败: %w", err) + } + defer rows.Close() + var out []*Project + for rows.Next() { + var p Project + var pinned int + var createdAt, updatedAt string + if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt); err != nil { + return nil, err + } + p.Pinned = pinned != 0 + p.CreatedAt = parseDBTime(createdAt) + p.UpdatedAt = parseDBTime(updatedAt) + out = append(out, &p) + } + return out, rows.Err() +} + // UpdateProject 更新项目。 func (db *DB) UpdateProject(p *Project) error { p.UpdatedAt = time.Now() diff --git a/internal/database/project_dashboard.go b/internal/database/project_dashboard.go index e4408fdf..0a3bdbda 100644 --- a/internal/database/project_dashboard.go +++ b/internal/database/project_dashboard.go @@ -33,6 +33,10 @@ type ProjectDashboardSummary struct { // GetProjectDashboardSummary 聚合跨项目近期事实（仅活跃项目、排除 deprecated）。 func (db *DB) GetProjectDashboardSummary(factLimit int) (*ProjectDashboardSummary, error) { + return db.GetProjectDashboardSummaryForAccess(factLimit, "", "") +} + +func (db *DB) GetProjectDashboardSummaryForAccess(factLimit int, userID, scope string) (*ProjectDashboardSummary, error) { if factLimit <= 0 { factLimit = 5 } @@ -44,25 +48,42 @@ func (db *DB) GetProjectDashboardSummary(factLimit int) (*ProjectDashboardSummar RecentFacts: []ProjectDashboardFact{}, } - if err := db.QueryRow(`SELECT COUNT(*) FROM projects WHERE status = 'active'`).Scan(&out.Totals.ActiveProjects); err != nil { + projectAccess := "" + args := []interface{}{} + userID = strings.TrimSpace(userID) + if userID != "" && scope != RBACScopeAll { + projectAccess = ` AND ( + p.owner_user_id = ? + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = p.id + ) + )` + args = append(args, userID, userID) + } + + if err := db.QueryRow(`SELECT COUNT(*) FROM projects p WHERE p.status = 'active'`+projectAccess, args...).Scan(&out.Totals.ActiveProjects); err != nil { return nil, fmt.Errorf("统计活跃项目失败: %w", err) } if err := db.QueryRow( `SELECT COUNT(*) FROM project_facts f INNER JOIN projects p ON p.id = f.project_id - WHERE f.confidence != 'deprecated' AND p.status = 'active'`, + WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess, + args..., ).Scan(&out.Totals.TotalFacts); err != nil { return nil, fmt.Errorf("统计事实失败: %w", err) } + queryArgs := append([]interface{}{}, args...) + queryArgs = append(queryArgs, factLimit) rows, err := db.Query( `SELECT f.id, f.project_id, p.name, f.fact_key, f.category, f.summary, f.confidence, f.pinned, f.updated_at FROM project_facts f INNER JOIN projects p ON p.id = f.project_id - WHERE f.confidence != 'deprecated' AND p.status = 'active' + WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess+` ORDER BY f.pinned DESC, f.updated_at DESC LIMIT ?`, - factLimit, + queryArgs..., ) if err != nil { return nil, fmt.Errorf("查询近期事实失败: %w", err) diff --git a/internal/database/rbac.go b/internal/database/rbac.go new file mode 100644 index 00000000..a5c27486 --- /dev/null +++ b/internal/database/rbac.go @@ -0,0 +1,1111 @@ +package database + +import ( + "database/sql" + "errors" + "fmt" + "strings" + "time" + + "github.com/google/uuid" +) + +const ( + RBACSystemRoleAdmin = "admin" + RBACSystemRoleOperator = "operator" + RBACSystemRoleAuditor = "auditor" + RBACSystemRoleViewer = "viewer" + + RBACScopeAll = "all" + RBACScopeAssigned = "assigned" + RBACScopeOwn = "own" + + RBACMaxBatchResourceAssignments = 100 +) + +var rbacAssignableResourceTables = map[string]string{ + "project": "projects", + "conversation": "conversations", + "vulnerability": "vulnerabilities", + "webshell": "webshell_connections", + "batch_task": "batch_task_queues", + "c2_listener": "c2_listeners", +} + +// RBACUser is a local platform account. +type RBACUser struct { + ID string `json:"id"` + Username string `json:"username"` + DisplayName string `json:"displayName,omitempty"` + PasswordHash string `json:"-"` + Enabled bool `json:"enabled"` + IsBuiltin bool `json:"isBuiltin"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// RBACRole groups permissions and a resource visibility scope. +type RBACRole struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + Scope string `json:"scope"` + IsSystem bool `json:"isSystem"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// RBACResourceAssignment grants a user access to one resource. +type RBACResourceAssignment struct { + ID string `json:"id"` + UserID string `json:"userId"` + ResourceType string `json:"resourceType"` + ResourceID string `json:"resourceId"` + ResourceLabel string `json:"resourceLabel,omitempty"` + ResourceDetail string `json:"resourceDetail,omitempty"` + CreatedAt time.Time `json:"createdAt"` +} + +// RBACResourceOption is a safe, minimal projection used by the assignment picker. +// It intentionally excludes resource contents and credentials. +type RBACResourceOption struct { + ID string `json:"id"` + Label string `json:"label"` + Detail string `json:"detail,omitempty"` +} + +// RBACAccess is the resolved authorization profile for one user. +type RBACAccess struct { + User RBACUser `json:"user"` + Roles []RBACRole `json:"roles"` + Permissions map[string]bool `json:"permissions"` + Scope string `json:"scope"` +} + +func (db *DB) initRBACTables() error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS rbac_users ( + id TEXT PRIMARY KEY, + username TEXT NOT NULL UNIQUE, + display_name TEXT NOT NULL DEFAULT '', + password_hash TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + is_builtin INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL + );`, + `CREATE TABLE IF NOT EXISTS rbac_roles ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL UNIQUE, + description TEXT NOT NULL DEFAULT '', + scope TEXT NOT NULL DEFAULT 'assigned', + is_system INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL + );`, + `CREATE TABLE IF NOT EXISTS rbac_permissions ( + key TEXT PRIMARY KEY, + description TEXT NOT NULL DEFAULT '', + created_at DATETIME NOT NULL + );`, + `CREATE TABLE IF NOT EXISTS rbac_role_permissions ( + role_id TEXT NOT NULL, + permission_key TEXT NOT NULL, + created_at DATETIME NOT NULL, + PRIMARY KEY (role_id, permission_key), + FOREIGN KEY (role_id) REFERENCES rbac_roles(id) ON DELETE CASCADE, + FOREIGN KEY (permission_key) REFERENCES rbac_permissions(key) ON DELETE CASCADE + );`, + `CREATE TABLE IF NOT EXISTS rbac_user_roles ( + user_id TEXT NOT NULL, + role_id TEXT NOT NULL, + created_at DATETIME NOT NULL, + PRIMARY KEY (user_id, role_id), + FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, + FOREIGN KEY (role_id) REFERENCES rbac_roles(id) ON DELETE CASCADE + );`, + `CREATE TABLE IF NOT EXISTS rbac_resource_assignments ( + id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + resource_type TEXT NOT NULL, + resource_id TEXT NOT NULL, + created_at DATETIME NOT NULL, + FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, + UNIQUE(user_id, resource_type, resource_id) + );`, + `CREATE INDEX IF NOT EXISTS idx_rbac_user_roles_user ON rbac_user_roles(user_id);`, + `CREATE INDEX IF NOT EXISTS idx_rbac_role_permissions_role ON rbac_role_permissions(role_id);`, + `CREATE INDEX IF NOT EXISTS idx_rbac_assignments_user_resource ON rbac_resource_assignments(user_id, resource_type, resource_id);`, + `CREATE INDEX IF NOT EXISTS idx_rbac_assignments_resource ON rbac_resource_assignments(resource_type, resource_id);`, + } + for _, stmt := range stmts { + if _, err := db.Exec(stmt); err != nil { + return err + } + } + return nil +} + +func (db *DB) migrateRBACOwnershipColumns() error { + for _, col := range []struct { + table string + name string + stmt string + }{ + {"projects", "owner_user_id", "ALTER TABLE projects ADD COLUMN owner_user_id TEXT"}, + {"conversations", "owner_user_id", "ALTER TABLE conversations ADD COLUMN owner_user_id TEXT"}, + {"vulnerabilities", "owner_user_id", "ALTER TABLE vulnerabilities ADD COLUMN owner_user_id TEXT"}, + {"webshell_connections", "owner_user_id", "ALTER TABLE webshell_connections ADD COLUMN owner_user_id TEXT"}, + {"batch_task_queues", "owner_user_id", "ALTER TABLE batch_task_queues ADD COLUMN owner_user_id TEXT"}, + {"c2_listeners", "owner_user_id", "ALTER TABLE c2_listeners ADD COLUMN owner_user_id TEXT"}, + } { + if err := db.addColumnIfMissing(col.table, col.name, col.stmt); err != nil { + return err + } + } + return nil +} + +func (db *DB) addColumnIfMissing(table, name, stmt string) error { + var count int + err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?", table, name).Scan(&count) + if err != nil || count == 0 { + if _, addErr := db.Exec(stmt); addErr != nil { + msg := strings.ToLower(addErr.Error()) + if !strings.Contains(msg, "duplicate column") && !strings.Contains(msg, "already exists") { + return fmt.Errorf("添加%s.%s字段失败: %w", table, name, addErr) + } + } + } + return nil +} + +// BootstrapRBAC seeds the local admin account and system roles. +func (db *DB) BootstrapRBAC(adminPasswordHash string, permissions map[string]string) error { + if strings.TrimSpace(adminPasswordHash) == "" { + return errors.New("admin password hash is required") + } + now := time.Now() + tx, err := db.Begin() + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + + for key, desc := range permissions { + key = strings.TrimSpace(key) + if key == "" { + continue + } + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_permissions (key, description, created_at) VALUES (?, ?, ?)`, key, desc, now); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE rbac_permissions SET description = ? WHERE key = ?`, desc, key); err != nil { + return err + } + } + + systemRoles := []RBACRole{ + {ID: RBACSystemRoleAdmin, Name: "管理员", Description: "全局管理权限", Scope: RBACScopeAll, IsSystem: true}, + {ID: RBACSystemRoleOperator, Name: "操作员", Description: "可执行日常安全工作流，不能管理账号与核心配置", Scope: RBACScopeAssigned, IsSystem: true}, + {ID: RBACSystemRoleAuditor, Name: "审计员", Description: "只读查看审计、监控与资产", Scope: RBACScopeAll, IsSystem: true}, + {ID: RBACSystemRoleViewer, Name: "只读用户", Description: "只读查看被授权资源", Scope: RBACScopeAssigned, IsSystem: true}, + } + for _, role := range systemRoles { + if _, err := tx.Exec(` + INSERT INTO rbac_roles (id, name, description, scope, is_system, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET name=excluded.name, description=excluded.description, scope=excluded.scope, is_system=excluded.is_system, updated_at=excluded.updated_at + `, role.ID, role.Name, role.Description, role.Scope, boolToInt(role.IsSystem), now, now); err != nil { + return err + } + } + + var userCount int + if err := tx.QueryRow(`SELECT COUNT(*) FROM rbac_users`).Scan(&userCount); err != nil { + return err + } + if userCount == 0 { + if _, err := tx.Exec(` + INSERT INTO rbac_users (id, username, display_name, password_hash, enabled, is_builtin, created_at, updated_at) + VALUES (?, 'admin', '管理员', ?, 1, 1, ?, ?) + `, "admin", adminPasswordHash, now, now); err != nil { + return err + } + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_user_roles (user_id, role_id, created_at) VALUES ('admin', ?, ?)`, RBACSystemRoleAdmin, now); err != nil { + return err + } + } else { + if _, err := tx.Exec(`UPDATE rbac_users SET password_hash = ?, updated_at = ? WHERE username = 'admin' AND is_builtin = 1 AND (password_hash = '' OR password_hash IS NULL)`, adminPasswordHash, now); err != nil { + return err + } + } + + if err := grantSystemRolePermissions(tx, permissions); err != nil { + return err + } + + return tx.Commit() +} + +func grantSystemRolePermissions(tx *sql.Tx, permissions map[string]string) error { + now := time.Now() + for key := range permissions { + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_role_permissions (role_id, permission_key, created_at) VALUES (?, ?, ?)`, RBACSystemRoleAdmin, key, now); err != nil { + return err + } + switch { + case strings.HasSuffix(key, ":read"), key == "auth:self": + for _, roleID := range []string{RBACSystemRoleOperator, RBACSystemRoleAuditor, RBACSystemRoleViewer} { + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_role_permissions (role_id, permission_key, created_at) VALUES (?, ?, ?)`, roleID, key, now); err != nil { + return err + } + } + case strings.HasPrefix(key, "rbac:"), strings.HasPrefix(key, "config:"), strings.HasPrefix(key, "terminal:"), strings.HasPrefix(key, "audit:"): + continue + default: + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_role_permissions (role_id, permission_key, created_at) VALUES (?, ?, ?)`, RBACSystemRoleOperator, key, now); err != nil { + return err + } + } + } + return nil +} + +func (db *DB) GetRBACUserByUsername(username string) (*RBACUser, error) { + username = strings.TrimSpace(strings.ToLower(username)) + if username == "" { + return nil, sql.ErrNoRows + } + return db.scanRBACUser(db.QueryRow(` + SELECT id, username, display_name, password_hash, enabled, is_builtin, created_at, updated_at + FROM rbac_users WHERE username = ? + `, username)) +} + +func (db *DB) GetRBACUserByID(id string) (*RBACUser, error) { + id = strings.TrimSpace(id) + if id == "" { + return nil, sql.ErrNoRows + } + return db.scanRBACUser(db.QueryRow(` + SELECT id, username, display_name, password_hash, enabled, is_builtin, created_at, updated_at + FROM rbac_users WHERE id = ? + `, id)) +} + +func (db *DB) scanRBACUser(row *sql.Row) (*RBACUser, error) { + var u RBACUser + var enabled, builtin int + var createdAt, updatedAt string + if err := row.Scan(&u.ID, &u.Username, &u.DisplayName, &u.PasswordHash, &enabled, &builtin, &createdAt, &updatedAt); err != nil { + return nil, err + } + u.Enabled = enabled != 0 + u.IsBuiltin = builtin != 0 + u.CreatedAt = parseDBTime(createdAt) + u.UpdatedAt = parseDBTime(updatedAt) + return &u, nil +} + +func (db *DB) ResolveRBACAccess(userID string) (*RBACAccess, error) { + u, err := db.GetRBACUserByID(userID) + if err != nil { + return nil, err + } + rows, err := db.Query(` + SELECT r.id, r.name, r.description, r.scope, r.is_system, r.created_at, r.updated_at + FROM rbac_roles r + JOIN rbac_user_roles ur ON ur.role_id = r.id + WHERE ur.user_id = ? + ORDER BY r.is_system DESC, r.name ASC + `, userID) + if err != nil { + return nil, err + } + defer rows.Close() + + access := &RBACAccess{User: *u, Permissions: map[string]bool{}, Scope: RBACScopeOwn} + for rows.Next() { + var role RBACRole + var isSystem int + var createdAt, updatedAt string + if err := rows.Scan(&role.ID, &role.Name, &role.Description, &role.Scope, &isSystem, &createdAt, &updatedAt); err != nil { + return nil, err + } + role.IsSystem = isSystem != 0 + role.CreatedAt = parseDBTime(createdAt) + role.UpdatedAt = parseDBTime(updatedAt) + access.Roles = append(access.Roles, role) + access.Scope = mergeRBACScope(access.Scope, role.Scope) + } + if err := rows.Err(); err != nil { + return nil, err + } + + prows, err := db.Query(` + SELECT DISTINCT rp.permission_key + FROM rbac_role_permissions rp + JOIN rbac_user_roles ur ON ur.role_id = rp.role_id + WHERE ur.user_id = ? + `, userID) + if err != nil { + return nil, err + } + defer prows.Close() + for prows.Next() { + var key string + if err := prows.Scan(&key); err != nil { + return nil, err + } + access.Permissions[key] = true + } + return access, prows.Err() +} + +func mergeRBACScope(a, b string) string { + if a == RBACScopeAll || b == RBACScopeAll { + return RBACScopeAll + } + if a == RBACScopeAssigned || b == RBACScopeAssigned { + return RBACScopeAssigned + } + return RBACScopeOwn +} + +func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID string) bool { + userID = strings.TrimSpace(userID) + resourceType = strings.TrimSpace(resourceType) + resourceID = strings.TrimSpace(resourceID) + if userID == "" || resourceType == "" || resourceID == "" { + return false + } + if scope == RBACScopeAll { + return true + } + if scope == RBACScopeOwn { + if db.userOwnsResource(userID, resourceType, resourceID) { + return true + } + } + var n int + err := db.QueryRow(`SELECT COUNT(*) FROM rbac_resource_assignments WHERE user_id = ? AND resource_type = ? AND resource_id = ?`, userID, resourceType, resourceID).Scan(&n) + if err == nil && n > 0 { + return true + } + if resourceType == "vulnerability" { + return db.userCanAccessVulnerabilityViaParent(userID, scope, resourceID) + } + if resourceType == "conversation" { + return db.userCanAccessConversationViaParent(userID, scope, resourceID) + } + if strings.HasPrefix(resourceType, "c2_") { + return db.userCanAccessC2ViaParent(userID, scope, resourceType, resourceID) + } + return false +} + +func (db *DB) userCanAccessConversationViaParent(userID, scope, conversationID string) bool { + var projectID sql.NullString + if err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&projectID); err != nil { + return false + } + return projectID.Valid && strings.TrimSpace(projectID.String) != "" && + db.UserCanAccessResource(userID, scope, "project", strings.TrimSpace(projectID.String)) +} + +func (db *DB) userCanAccessVulnerabilityViaParent(userID, scope, vulnerabilityID string) bool { + var projectID, conversationID sql.NullString + err := db.QueryRow(`SELECT project_id, conversation_id FROM vulnerabilities WHERE id = ?`, vulnerabilityID).Scan(&projectID, &conversationID) + if err != nil { + return false + } + if projectID.Valid && strings.TrimSpace(projectID.String) != "" && db.UserCanAccessResource(userID, scope, "project", strings.TrimSpace(projectID.String)) { + return true + } + if conversationID.Valid && strings.TrimSpace(conversationID.String) != "" && db.UserCanAccessResource(userID, scope, "conversation", strings.TrimSpace(conversationID.String)) { + return true + } + return false +} + +func (db *DB) UserCanAccessMessage(userID, scope, messageID string) bool { + var conversationID string + err := db.QueryRow(`SELECT conversation_id FROM messages WHERE id = ?`, strings.TrimSpace(messageID)).Scan(&conversationID) + if err != nil { + return false + } + return db.UserCanAccessResource(userID, scope, "conversation", conversationID) +} + +func (db *DB) UserCanAccessProcessDetail(userID, scope, processDetailID string) bool { + var conversationID string + err := db.QueryRow(`SELECT conversation_id FROM process_details WHERE id = ?`, strings.TrimSpace(processDetailID)).Scan(&conversationID) + if err != nil { + return false + } + return db.UserCanAccessResource(userID, scope, "conversation", conversationID) +} + +func (db *DB) userCanAccessC2ViaParent(userID, scope, resourceType, resourceID string) bool { + switch resourceType { + case "c2_session": + var listenerID string + if err := db.QueryRow(`SELECT listener_id FROM c2_sessions WHERE id = ?`, resourceID).Scan(&listenerID); err != nil { + return false + } + return db.UserCanAccessResource(userID, scope, "c2_listener", listenerID) + case "c2_task": + var sessionID string + if err := db.QueryRow(`SELECT session_id FROM c2_tasks WHERE id = ?`, resourceID).Scan(&sessionID); err != nil { + return false + } + return db.UserCanAccessResource(userID, scope, "c2_session", sessionID) + case "c2_file": + var sessionID string + if err := db.QueryRow(`SELECT session_id FROM c2_files WHERE id = ?`, resourceID).Scan(&sessionID); err != nil { + return false + } + return db.UserCanAccessResource(userID, scope, "c2_session", sessionID) + case "c2_event": + var sessionID, taskID sql.NullString + if err := db.QueryRow(`SELECT session_id, task_id FROM c2_events WHERE id = ?`, resourceID).Scan(&sessionID, &taskID); err != nil { + return false + } + if sessionID.Valid && strings.TrimSpace(sessionID.String) != "" { + return db.UserCanAccessResource(userID, scope, "c2_session", strings.TrimSpace(sessionID.String)) + } + if taskID.Valid && strings.TrimSpace(taskID.String) != "" { + return db.UserCanAccessResource(userID, scope, "c2_task", strings.TrimSpace(taskID.String)) + } + } + return false +} + +func (db *DB) userOwnsResource(userID, resourceType, resourceID string) bool { + table := "" + switch resourceType { + case "project": + table = "projects" + case "conversation": + table = "conversations" + case "vulnerability": + table = "vulnerabilities" + case "webshell": + table = "webshell_connections" + case "batch_task": + table = "batch_task_queues" + case "c2_listener": + table = "c2_listeners" + default: + return false + } + var n int + err := db.QueryRow(`SELECT COUNT(*) FROM `+table+` WHERE id = ? AND owner_user_id = ?`, resourceID, userID).Scan(&n) + return err == nil && n > 0 +} + +func (db *DB) SetResourceOwner(resourceType, resourceID, userID string) error { + userID = strings.TrimSpace(userID) + if userID == "" { + return nil + } + table := "" + switch resourceType { + case "project": + table = "projects" + case "conversation": + table = "conversations" + case "vulnerability": + table = "vulnerabilities" + case "webshell": + table = "webshell_connections" + case "batch_task": + table = "batch_task_queues" + case "c2_listener": + table = "c2_listeners" + default: + return nil + } + _, err := db.Exec(`UPDATE `+table+` SET owner_user_id = COALESCE(NULLIF(owner_user_id, ''), ?) WHERE id = ?`, userID, resourceID) + return err +} + +func (db *DB) AssignResourceToUser(userID, resourceType, resourceID string) error { + _, err := db.AssignResourcesToUser(userID, resourceType, []string{resourceID}) + return err +} + +// ListAssignableRBACResources returns real resources for the admin assignment +// picker without exposing full records or secret-bearing fields. +func (db *DB) ListAssignableRBACResources(resourceType, search string, limit int) ([]RBACResourceOption, error) { + return db.ListAssignableRBACResourcesPage(resourceType, search, limit, 0) +} + +// ListAssignableRBACResourcesPage returns one stable page for the assignment +// picker. Callers can request limit+1 rows to determine whether another page +// exists without running a separate COUNT query. +func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit, offset int) ([]RBACResourceOption, error) { + resourceType = strings.TrimSpace(resourceType) + if _, ok := rbacAssignableResourceTables[resourceType]; !ok { + return nil, fmt.Errorf("不支持的资源类型: %s", resourceType) + } + if limit <= 0 || limit > 100 { + limit = 50 + } + if offset < 0 { + offset = 0 + } + pattern := "%" + strings.ToLower(strings.NewReplacer( + `\`, `\\`, + `%`, `\%`, + `_`, `\_`, + ).Replace(strings.TrimSpace(search))) + "%" + + var query string + switch resourceType { + case "project": + query = `SELECT id, name, status FROM projects + WHERE LOWER(name) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' + ORDER BY updated_at DESC LIMIT ? OFFSET ?` + case "conversation": + query = `SELECT id, COALESCE(NULLIF(TRIM(title), ''), '未命名对话'), COALESCE(project_id, '') FROM conversations + WHERE LOWER(COALESCE(NULLIF(TRIM(title), ''), id)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' + ORDER BY updated_at DESC LIMIT ? OFFSET ?` + case "vulnerability": + query = `SELECT id, title, severity FROM vulnerabilities + WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' + ORDER BY updated_at DESC LIMIT ? OFFSET ?` + case "webshell": + query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections + WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' + ORDER BY created_at DESC LIMIT ? OFFSET ?` + case "batch_task": + query = `SELECT id, COALESCE(NULLIF(title, ''), id), status FROM batch_task_queues + WHERE LOWER(COALESCE(NULLIF(title, ''), id)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' + ORDER BY created_at DESC LIMIT ? OFFSET ?` + case "c2_listener": + query = `SELECT id, name, type || ' · ' || status FROM c2_listeners + WHERE LOWER(name) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' + ORDER BY created_at DESC LIMIT ? OFFSET ?` + } + + rows, err := db.Query(query, pattern, pattern, limit, offset) + if err != nil { + return nil, err + } + defer rows.Close() + options := make([]RBACResourceOption, 0) + for rows.Next() { + var option RBACResourceOption + if err := rows.Scan(&option.ID, &option.Label, &option.Detail); err != nil { + return nil, err + } + option.Label = normalizeRBACResourceLabel(option.Label, option.ID) + options = append(options, option) + } + return options, rows.Err() +} + +func normalizeRBACResourceLabel(label, id string) string { + label = strings.TrimSpace(label) + if label == "" { + return "资源 " + shortRBACResourceID(id) + } + if isWeakRBACResourceLabel(label) { + return label + " · " + shortRBACResourceID(id) + } + return label +} + +func isWeakRBACResourceLabel(label string) bool { + runes := []rune(strings.TrimSpace(label)) + if len(runes) <= 1 { + return true + } + if len(runes) <= 3 { + numeric := true + for _, r := range runes { + if r < '0' || r > '9' { + numeric = false + break + } + } + return numeric + } + return false +} + +func shortRBACResourceID(id string) string { + id = strings.TrimSpace(id) + if len(id) <= 12 { + return id + } + return id[:8] + "…" +} + +func (db *DB) lookupRBACResourceOptionsByIDs(resourceType string, ids []string) (map[string]RBACResourceOption, error) { + resourceType = strings.TrimSpace(resourceType) + if _, ok := rbacAssignableResourceTables[resourceType]; !ok { + return nil, fmt.Errorf("不支持的资源类型: %s", resourceType) + } + unique := make([]string, 0, len(ids)) + seen := make(map[string]struct{}, len(ids)) + for _, rawID := range ids { + id := strings.TrimSpace(rawID) + if id == "" { + continue + } + if _, exists := seen[id]; exists { + continue + } + seen[id] = struct{}{} + unique = append(unique, id) + } + out := make(map[string]RBACResourceOption, len(unique)) + if len(unique) == 0 { + return out, nil + } + + placeholders := strings.TrimRight(strings.Repeat("?,", len(unique)), ",") + args := make([]interface{}, 0, len(unique)) + for _, id := range unique { + args = append(args, id) + } + + var query string + switch resourceType { + case "project": + query = `SELECT id, name, status FROM projects WHERE id IN (` + placeholders + `)` + case "conversation": + query = `SELECT id, COALESCE(NULLIF(TRIM(title), ''), '未命名对话'), COALESCE(project_id, '') FROM conversations WHERE id IN (` + placeholders + `)` + case "vulnerability": + query = `SELECT id, title, severity FROM vulnerabilities WHERE id IN (` + placeholders + `)` + case "webshell": + query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections WHERE id IN (` + placeholders + `)` + case "batch_task": + query = `SELECT id, COALESCE(NULLIF(title, ''), id), status FROM batch_task_queues WHERE id IN (` + placeholders + `)` + case "c2_listener": + query = `SELECT id, name, type || ' · ' || status FROM c2_listeners WHERE id IN (` + placeholders + `)` + default: + return out, nil + } + + rows, err := db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var option RBACResourceOption + if err := rows.Scan(&option.ID, &option.Label, &option.Detail); err != nil { + return nil, err + } + option.Label = normalizeRBACResourceLabel(option.Label, option.ID) + out[option.ID] = option + } + return out, rows.Err() +} + +func enrichRBACAssignmentLabels(rows []RBACResourceAssignment, lookup func(resourceType string, ids []string) (map[string]RBACResourceOption, error)) error { + if lookup == nil || len(rows) == 0 { + return nil + } + idsByType := make(map[string][]string) + for _, row := range rows { + idsByType[row.ResourceType] = append(idsByType[row.ResourceType], row.ResourceID) + } + labelsByType := make(map[string]map[string]RBACResourceOption, len(idsByType)) + for resourceType, ids := range idsByType { + options, err := lookup(resourceType, ids) + if err != nil { + return err + } + labelsByType[resourceType] = options + } + for i := range rows { + options := labelsByType[rows[i].ResourceType] + if options == nil { + continue + } + if option, ok := options[rows[i].ResourceID]; ok { + rows[i].ResourceLabel = option.Label + rows[i].ResourceDetail = option.Detail + } + } + return nil +} + +// AssignResourcesToUser validates the complete request before writing anything, +// then inserts all grants in one transaction. Existing grants are idempotent. +func (db *DB) AssignResourcesToUser(userID, resourceType string, resourceIDs []string) (int64, error) { + userID = strings.TrimSpace(userID) + resourceType = strings.TrimSpace(resourceType) + if userID == "" || resourceType == "" || len(resourceIDs) == 0 { + return 0, errors.New("user_id, resource_type and resource_ids are required") + } + if len(resourceIDs) > RBACMaxBatchResourceAssignments { + return 0, fmt.Errorf("一次最多授权 %d 个资源", RBACMaxBatchResourceAssignments) + } + table, ok := rbacAssignableResourceTables[resourceType] + if !ok { + return 0, fmt.Errorf("不支持的资源类型: %s", resourceType) + } + + uniqueIDs := make([]string, 0, len(resourceIDs)) + seen := make(map[string]struct{}, len(resourceIDs)) + for _, rawID := range resourceIDs { + id := strings.TrimSpace(rawID) + if id == "" { + return 0, errors.New("资源 ID 不能为空") + } + if _, exists := seen[id]; exists { + continue + } + seen[id] = struct{}{} + uniqueIDs = append(uniqueIDs, id) + } + if len(uniqueIDs) == 0 { + return 0, errors.New("资源 ID 不能为空") + } + + tx, err := db.Begin() + if err != nil { + return 0, err + } + defer func() { _ = tx.Rollback() }() + + var userExists int + if err := tx.QueryRow(`SELECT COUNT(*) FROM rbac_users WHERE id = ?`, userID).Scan(&userExists); err != nil { + return 0, err + } + if userExists == 0 { + return 0, errors.New("用户不存在") + } + for _, resourceID := range uniqueIDs { + var exists int + if err := tx.QueryRow(`SELECT COUNT(*) FROM `+table+` WHERE id = ?`, resourceID).Scan(&exists); err != nil { + return 0, err + } + if exists == 0 { + return 0, fmt.Errorf("资源不存在: %s/%s", resourceType, resourceID) + } + } + + var created int64 + for _, resourceID := range uniqueIDs { + result, err := tx.Exec(` + INSERT OR IGNORE INTO rbac_resource_assignments (id, user_id, resource_type, resource_id, created_at) + VALUES (?, ?, ?, ?, ?) + `, uuid.NewString(), userID, resourceType, resourceID, time.Now()) + if err != nil { + return 0, err + } + if n, err := result.RowsAffected(); err == nil { + created += n + } + } + if err := tx.Commit(); err != nil { + return 0, err + } + return created, nil +} + +func (db *DB) ListRBACUsers() ([]RBACUser, error) { + rows, err := db.Query(`SELECT id, username, display_name, password_hash, enabled, is_builtin, created_at, updated_at FROM rbac_users ORDER BY username ASC`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []RBACUser + for rows.Next() { + var u RBACUser + var enabled, builtin int + var createdAt, updatedAt string + if err := rows.Scan(&u.ID, &u.Username, &u.DisplayName, &u.PasswordHash, &enabled, &builtin, &createdAt, &updatedAt); err != nil { + return nil, err + } + u.Enabled = enabled != 0 + u.IsBuiltin = builtin != 0 + u.CreatedAt = parseDBTime(createdAt) + u.UpdatedAt = parseDBTime(updatedAt) + out = append(out, u) + } + return out, rows.Err() +} + +func (db *DB) ListRBACRoles() ([]RBACRole, error) { + rows, err := db.Query(`SELECT id, name, description, scope, is_system, created_at, updated_at FROM rbac_roles ORDER BY is_system DESC, name ASC`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []RBACRole + for rows.Next() { + var r RBACRole + var system int + var createdAt, updatedAt string + if err := rows.Scan(&r.ID, &r.Name, &r.Description, &r.Scope, &system, &createdAt, &updatedAt); err != nil { + return nil, err + } + r.IsSystem = system != 0 + r.CreatedAt = parseDBTime(createdAt) + r.UpdatedAt = parseDBTime(updatedAt) + out = append(out, r) + } + return out, rows.Err() +} + +func (db *DB) GetRBACRoleByID(id string) (*RBACRole, error) { + id = strings.TrimSpace(id) + if id == "" { + return nil, sql.ErrNoRows + } + var r RBACRole + var system int + var createdAt, updatedAt string + err := db.QueryRow(`SELECT id, name, description, scope, is_system, created_at, updated_at FROM rbac_roles WHERE id = ?`, id). + Scan(&r.ID, &r.Name, &r.Description, &r.Scope, &system, &createdAt, &updatedAt) + if err != nil { + return nil, err + } + r.IsSystem = system != 0 + r.CreatedAt = parseDBTime(createdAt) + r.UpdatedAt = parseDBTime(updatedAt) + return &r, nil +} + +func (db *DB) UpsertRBACRole(id, name, description, scope string, permissionKeys []string) (*RBACRole, error) { + id = strings.TrimSpace(id) + name = strings.TrimSpace(name) + scope = strings.TrimSpace(scope) + if name == "" { + return nil, errors.New("role name is required") + } + if scope != RBACScopeAll && scope != RBACScopeAssigned && scope != RBACScopeOwn { + scope = RBACScopeAssigned + } + if id == "" { + id = uuid.NewString() + } + now := time.Now() + tx, err := db.Begin() + if err != nil { + return nil, err + } + defer func() { _ = tx.Rollback() }() + var isSystem int + _ = tx.QueryRow(`SELECT is_system FROM rbac_roles WHERE id = ?`, id).Scan(&isSystem) + if _, err := tx.Exec(` + INSERT INTO rbac_roles (id, name, description, scope, is_system, created_at, updated_at) + VALUES (?, ?, ?, ?, 0, ?, ?) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + description = excluded.description, + scope = excluded.scope, + updated_at = excluded.updated_at + `, id, name, strings.TrimSpace(description), scope, now, now); err != nil { + return nil, err + } + if _, err := tx.Exec(`DELETE FROM rbac_role_permissions WHERE role_id = ?`, id); err != nil { + return nil, err + } + for _, key := range permissionKeys { + key = strings.TrimSpace(key) + if key == "" { + continue + } + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_permissions (key, description, created_at) VALUES (?, '', ?)`, key, now); err != nil { + return nil, err + } + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_role_permissions (role_id, permission_key, created_at) VALUES (?, ?, ?)`, id, key, now); err != nil { + return nil, err + } + } + if err := tx.Commit(); err != nil { + return nil, err + } + return db.GetRBACRoleByID(id) +} + +func (db *DB) DeleteRBACRole(id string) error { + id = strings.TrimSpace(id) + if id == "" { + return errors.New("role id is required") + } + if id == RBACSystemRoleAdmin || id == RBACSystemRoleOperator || id == RBACSystemRoleAuditor || id == RBACSystemRoleViewer { + return errors.New("system role cannot be deleted") + } + _, err := db.Exec(`DELETE FROM rbac_roles WHERE id = ? AND is_system = 0`, id) + return err +} + +func (db *DB) UpdateRBACUserPassword(userID, passwordHash string) error { + userID = strings.TrimSpace(userID) + passwordHash = strings.TrimSpace(passwordHash) + if userID == "" || passwordHash == "" { + return errors.New("user_id and password_hash are required") + } + _, err := db.Exec(`UPDATE rbac_users SET password_hash = ?, updated_at = ? WHERE id = ?`, passwordHash, time.Now(), userID) + return err +} + +func (db *DB) UpdateRBACAdminPassword(passwordHash string) error { + return db.UpdateRBACUserPassword("admin", passwordHash) +} + +func (db *DB) CreateRBACUser(username, displayName, passwordHash string, enabled bool, roleIDs []string) (*RBACUser, error) { + username = strings.TrimSpace(strings.ToLower(username)) + if username == "" || passwordHash == "" { + return nil, errors.New("username and password are required") + } + id := uuid.NewString() + now := time.Now() + tx, err := db.Begin() + if err != nil { + return nil, err + } + defer func() { _ = tx.Rollback() }() + if _, err := tx.Exec(` + INSERT INTO rbac_users (id, username, display_name, password_hash, enabled, is_builtin, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, 0, ?, ?) + `, id, username, strings.TrimSpace(displayName), passwordHash, boolToInt(enabled), now, now); err != nil { + return nil, err + } + for _, roleID := range roleIDs { + roleID = strings.TrimSpace(roleID) + if roleID == "" { + continue + } + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_user_roles (user_id, role_id, created_at) VALUES (?, ?, ?)`, id, roleID, now); err != nil { + return nil, err + } + } + if err := tx.Commit(); err != nil { + return nil, err + } + return db.GetRBACUserByID(id) +} + +func (db *DB) UpdateRBACUser(userID, displayName string, enabled *bool, roleIDs *[]string) error { + userID = strings.TrimSpace(userID) + if userID == "" { + return errors.New("user_id is required") + } + tx, err := db.Begin() + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + if enabled != nil { + if _, err := tx.Exec(`UPDATE rbac_users SET display_name = ?, enabled = ?, updated_at = ? WHERE id = ?`, strings.TrimSpace(displayName), boolToInt(*enabled), time.Now(), userID); err != nil { + return err + } + } else { + if _, err := tx.Exec(`UPDATE rbac_users SET display_name = ?, updated_at = ? WHERE id = ?`, strings.TrimSpace(displayName), time.Now(), userID); err != nil { + return err + } + } + if roleIDs != nil { + if _, err := tx.Exec(`DELETE FROM rbac_user_roles WHERE user_id = ?`, userID); err != nil { + return err + } + for _, roleID := range *roleIDs { + roleID = strings.TrimSpace(roleID) + if roleID == "" { + continue + } + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_user_roles (user_id, role_id, created_at) VALUES (?, ?, ?)`, userID, roleID, time.Now()); err != nil { + return err + } + } + } + return tx.Commit() +} + +func (db *DB) DeleteRBACUser(userID string) error { + userID = strings.TrimSpace(userID) + if userID == "" || userID == "admin" { + return errors.New("cannot delete this user") + } + _, err := db.Exec(`DELETE FROM rbac_users WHERE id = ? AND is_builtin = 0`, userID) + return err +} + +func (db *DB) ListRBACUserRoleIDs(userID string) ([]string, error) { + rows, err := db.Query(`SELECT role_id FROM rbac_user_roles WHERE user_id = ? ORDER BY role_id ASC`, userID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []string + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + return nil, err + } + out = append(out, id) + } + return out, rows.Err() +} + +func (db *DB) ListRBACRolePermissionKeys(roleID string) ([]string, error) { + rows, err := db.Query(`SELECT permission_key FROM rbac_role_permissions WHERE role_id = ? ORDER BY permission_key ASC`, roleID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []string + for rows.Next() { + var key string + if err := rows.Scan(&key); err != nil { + return nil, err + } + out = append(out, key) + } + return out, rows.Err() +} + +func (db *DB) ListRBACResourceAssignments(userID string) ([]RBACResourceAssignment, error) { + query := `SELECT id, user_id, resource_type, resource_id, created_at FROM rbac_resource_assignments WHERE 1=1` + args := []interface{}{} + if strings.TrimSpace(userID) != "" { + query += ` AND user_id = ?` + args = append(args, strings.TrimSpace(userID)) + } + query += ` ORDER BY created_at DESC` + rows, err := db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []RBACResourceAssignment + for rows.Next() { + var row RBACResourceAssignment + var createdAt string + if err := rows.Scan(&row.ID, &row.UserID, &row.ResourceType, &row.ResourceID, &createdAt); err != nil { + return nil, err + } + row.CreatedAt = parseDBTime(createdAt) + out = append(out, row) + } + if err := enrichRBACAssignmentLabels(out, db.lookupRBACResourceOptionsByIDs); err != nil { + return nil, err + } + return out, rows.Err() +} + +func (db *DB) DeleteRBACResourceAssignment(id string) error { + id = strings.TrimSpace(id) + if id == "" { + return errors.New("assignment id is required") + } + result, err := db.Exec(`DELETE FROM rbac_resource_assignments WHERE id = ?`, id) + if err != nil { + return err + } + if affected, err := result.RowsAffected(); err == nil && affected == 0 { + return errors.New("资源授权不存在或已撤销") + } + return nil +} diff --git a/internal/database/rbac_access_test.go b/internal/database/rbac_access_test.go new file mode 100644 index 00000000..8e93067e --- /dev/null +++ b/internal/database/rbac_access_test.go @@ -0,0 +1,401 @@ +package database + +import ( + "path/filepath" + "strings" + "testing" + "time" + + "go.uber.org/zap" +) + +func newRBACTestDB(t *testing.T) *DB { + t.Helper() + db, err := NewDB(filepath.Join(t.TempDir(), "rbac.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + return db +} + +func TestRBACProjectAndConversationListAccess(t *testing.T) { + db := newRBACTestDB(t) + p1, _ := db.CreateProject(&Project{Name: "visible"}) + p2, _ := db.CreateProject(&Project{Name: "hidden"}) + if err := db.SetResourceOwner("project", p1.ID, "u1"); err != nil { + t.Fatal(err) + } + c1, _ := db.CreateConversation("visible conv", ConversationCreateMeta{ProjectID: p1.ID}) + c2, _ := db.CreateConversation("hidden conv", ConversationCreateMeta{ProjectID: p2.ID}) + _ = db.SetResourceOwner("conversation", c1.ID, "u1") + _ = db.SetResourceOwner("conversation", c2.ID, "u2") + + projects, err := db.ListProjectsForAccess("", "", 50, 0, "u1", RBACScopeOwn) + if err != nil { + t.Fatal(err) + } + if len(projects) != 1 || projects[0].ID != p1.ID { + t.Fatalf("projects = %#v, want only %s", projects, p1.ID) + } + + convs, err := db.ListConversationsForAccess(50, 0, "", "", "", "u1", RBACScopeOwn) + if err != nil { + t.Fatal(err) + } + if len(convs) != 1 || convs[0].ID != c1.ID { + t.Fatalf("conversations = %#v, want only %s", convs, c1.ID) + } +} + +func TestRBACVulnerabilityAccessInheritsProject(t *testing.T) { + db := newRBACTestDB(t) + user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + p1, _ := db.CreateProject(&Project{Name: "visible"}) + p2, _ := db.CreateProject(&Project{Name: "hidden"}) + if err := db.AssignResourceToUser(user.ID, "project", p1.ID); err != nil { + t.Fatal(err) + } + v1, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p1.ID, Title: "v1", Severity: "high"}) + v2, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p2.ID, Title: "v2", Severity: "high"}) + + items, err := db.ListVulnerabilitiesForAccess(50, 0, VulnerabilityListFilter{}, RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned}) + if err != nil { + t.Fatal(err) + } + if len(items) != 1 || items[0].ID != v1.ID { + t.Fatalf("vulnerabilities = %#v, want only %s; hidden %s", items, v1.ID, v2.ID) + } + if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v1.ID) { + t.Fatalf("expected project assignment to allow vulnerability detail") + } + if db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v2.ID) { + t.Fatalf("unexpected access to hidden vulnerability") + } +} + +func TestRBACConversationAccessInheritsProject(t *testing.T) { + db := newRBACTestDB(t) + user, err := db.CreateRBACUser("project-member", "Project Member", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + project, err := db.CreateProject(&Project{Name: "assigned project"}) + if err != nil { + t.Fatal(err) + } + conversation, err := db.CreateConversation("project conversation", ConversationCreateMeta{ProjectID: project.ID}) + if err != nil { + t.Fatal(err) + } + if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil { + t.Fatal(err) + } + + rows, err := db.ListConversationsForAccess(50, 0, "", "", "", user.ID, RBACScopeAssigned) + if err != nil { + t.Fatal(err) + } + if len(rows) != 1 || rows[0].ID != conversation.ID { + t.Fatalf("conversations = %#v, want project conversation %s", rows, conversation.ID) + } + if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "conversation", conversation.ID) { + t.Fatal("expected project assignment to allow conversation detail") + } +} + +func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) { + db := newRBACTestDB(t) + user, err := db.CreateRBACUser("batch-member", "Batch Member", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + p1, err := db.CreateProject(&Project{Name: "p1"}) + if err != nil { + t.Fatal(err) + } + p2, err := db.CreateProject(&Project{Name: "p2"}) + if err != nil { + t.Fatal(err) + } + p3, err := db.CreateProject(&Project{Name: "p3"}) + if err != nil { + t.Fatal(err) + } + options, err := db.ListAssignableRBACResources("project", "p1", 50) + if err != nil { + t.Fatal(err) + } + if len(options) != 1 || options[0].ID != p1.ID || options[0].Label != "p1" { + t.Fatalf("resource options = %#v, want p1", options) + } + firstPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 0) + if err != nil { + t.Fatal(err) + } + secondPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 2) + if err != nil { + t.Fatal(err) + } + if len(firstPage) != 2 || len(secondPage) != 1 { + t.Fatalf("paged resource options = %d + %d, want 2 + 1", len(firstPage), len(secondPage)) + } + seen := map[string]bool{} + for _, option := range append(firstPage, secondPage...) { + seen[option.ID] = true + } + if !seen[p1.ID] || !seen[p2.ID] || !seen[p3.ID] { + t.Fatalf("paged resource options missed resources: %#v", seen) + } + if _, err := db.ListAssignableRBACResources("secret_table", "", 50); err == nil { + t.Fatal("expected unsupported picker resource type to fail") + } + + if _, err := db.AssignResourcesToUser(user.ID, "unknown_type", []string{p1.ID}); err == nil { + t.Fatal("expected unsupported resource type to fail") + } + if _, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, "missing-project"}); err == nil { + t.Fatal("expected missing resource to fail the entire batch") + } + rows, err := db.ListRBACResourceAssignments(user.ID) + if err != nil { + t.Fatal(err) + } + if len(rows) != 0 { + t.Fatalf("partial grants persisted after failed batch: %#v", rows) + } + + created, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p1.ID, p2.ID}) + if err != nil { + t.Fatal(err) + } + if created != 2 { + t.Fatalf("created = %d, want 2 unique grants", created) + } + created, err = db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p2.ID}) + if err != nil { + t.Fatal(err) + } + if created != 0 { + t.Fatalf("idempotent retry created = %d, want 0", created) + } + rows, err = db.ListRBACResourceAssignments(user.ID) + if err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("assignment count = %d, want 2", len(rows)) + } +} + +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()} + if err := db.CreateWebshellConnection(&ws1); err != nil { + t.Fatal(err) + } + if err := db.CreateWebshellConnection(&ws2); err != nil { + t.Fatal(err) + } + _ = db.SetResourceOwner("webshell", ws1.ID, "u1") + _ = db.SetResourceOwner("webshell", ws2.ID, "u2") + webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn) + 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) + } + + if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil { + t.Fatal(err) + } + if err := db.CreateBatchQueue("q_hidden", "hidden", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t2", "message": "b"}}); err != nil { + t.Fatal(err) + } + _ = db.SetResourceOwner("batch_task", "q_visible", "u1") + _ = db.SetResourceOwner("batch_task", "q_hidden", "u2") + queues, err := db.ListBatchQueuesForAccess(50, 0, "all", "", "u1", RBACScopeOwn) + if err != nil { + t.Fatal(err) + } + if len(queues) != 1 || queues[0].ID != "q_visible" { + t.Fatalf("queues = %#v, want only q_visible", queues) + } +} + +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} + if err := db.CreateC2Listener(l1); err != nil { + t.Fatal(err) + } + if err := db.CreateC2Listener(l2); 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.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.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) + } + + access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn} + listeners, err := db.ListC2ListenersForAccess(access) + 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) + } + sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, 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) + } + tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, 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) + } + events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, 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) + } + if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") { + t.Fatalf("expected listener ownership to allow task detail") + } + if db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_hidden") { + t.Fatalf("unexpected access to hidden task") + } +} + +func TestRBACC2AssignedDeleteIsScoped(t *testing.T) { + db := newRBACTestDB(t) + user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + now := time.Now() + if err := db.CreateC2Listener(&C2Listener{ID: "l_assigned", Name: "assigned", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, CreatedAt: now}); err != nil { + t.Fatal(err) + } + if err := db.CreateC2Listener(&C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, CreatedAt: now}); err != nil { + t.Fatal(err) + } + if err := db.AssignResourceToUser(user.ID, "c2_listener", "l_assigned"); err != nil { + t.Fatal(err) + } + for _, row := range []struct { + sessionID string + listener string + taskID string + eventID string + }{ + {"s_assigned", "l_assigned", "t_assigned", "e_assigned"}, + {"s_hidden", "l_hidden", "t_hidden", "e_hidden"}, + } { + if err := db.UpsertC2Session(&C2Session{ID: row.sessionID, ListenerID: row.listener, ImplantUUID: row.sessionID + "_uuid", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil { + t.Fatal(err) + } + if err := db.CreateC2Task(&C2Task{ID: row.taskID, SessionID: row.sessionID, TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil { + t.Fatal(err) + } + if err := db.AppendC2Event(&C2Event{ID: row.eventID, Level: "info", Category: "task", SessionID: row.sessionID, TaskID: row.taskID, Message: row.eventID, CreatedAt: now}); err != nil { + t.Fatal(err) + } + } + access := RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned} + n, err := db.DeleteC2TasksByIDsForAccess([]string{"t_assigned", "t_hidden"}, access) + if err != nil { + t.Fatal(err) + } + if n != 1 { + t.Fatalf("deleted tasks = %d, want 1", n) + } + if task, _ := db.GetC2Task("t_hidden"); task == nil { + t.Fatalf("hidden task was deleted") + } + n, err = db.DeleteC2EventsByIDsForAccess([]string{"e_assigned", "e_hidden"}, access) + if err != nil { + t.Fatal(err) + } + if n != 1 { + t.Fatalf("deleted events = %d, want 1", n) + } + hiddenEvents, err := db.ListC2Events(ListC2EventsFilter{TaskID: "t_hidden"}) + if err != nil { + t.Fatal(err) + } + if len(hiddenEvents) != 1 { + t.Fatalf("hidden event count = %d, want 1", len(hiddenEvents)) + } +} + +func TestRBACAssignmentLabelsAndWeakTitles(t *testing.T) { + db := newRBACTestDB(t) + user, err := db.CreateRBACUser("label-member", "Label Member", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + project, err := db.CreateProject(&Project{Name: "Alpha Project"}) + if err != nil { + t.Fatal(err) + } + conversation, err := db.CreateConversation("1", ConversationCreateMeta{}) + if err != nil { + t.Fatal(err) + } + if _, err := db.AssignResourcesToUser(user.ID, "project", []string{project.ID}); err != nil { + t.Fatal(err) + } + + options, err := db.ListAssignableRBACResources("conversation", "", 10) + if err != nil { + t.Fatal(err) + } + if len(options) == 0 { + t.Fatal("expected conversation options") + } + for _, option := range options { + if option.ID == conversation.ID && !strings.Contains(option.Label, "1 ·") { + t.Fatalf("weak conversation label = %q, want suffix with short id", option.Label) + } + } + + rows, err := db.ListRBACResourceAssignments(user.ID) + if err != nil { + t.Fatal(err) + } + if len(rows) != 1 { + t.Fatalf("assignments = %#v, want 1", rows) + } + if rows[0].ResourceLabel != "Alpha Project" { + t.Fatalf("assignment label = %q, want Alpha Project", rows[0].ResourceLabel) + } +} diff --git a/internal/database/vulnerability.go b/internal/database/vulnerability.go index c7bef02a..51d2a3db 100644 --- a/internal/database/vulnerability.go +++ b/internal/database/vulnerability.go @@ -23,6 +23,11 @@ type VulnerabilityListFilter struct { TaskTag string } +type RBACListAccess struct { + UserID string + Scope string +} + func escapeVulnerabilityLikePattern(s string) string { s = strings.ReplaceAll(s, `\`, `\\`) s = strings.ReplaceAll(s, `%`, `\%`) @@ -89,6 +94,40 @@ func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) ( return query, args } +func appendVulnerabilityAccessFilter(query string, args []interface{}, access RBACListAccess) (string, []interface{}) { + userID := strings.TrimSpace(access.UserID) + if userID == "" || access.Scope == RBACScopeAll { + return query, args + } + query += ` AND ( + owner_user_id = ? + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'vulnerability' AND ra.resource_id = vulnerabilities.id + ) + OR ( + project_id IS NOT NULL AND project_id <> '' AND ( + EXISTS (SELECT 1 FROM projects p WHERE p.id = vulnerabilities.project_id AND p.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments pra + WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = vulnerabilities.project_id + ) + ) + ) + OR ( + conversation_id IS NOT NULL AND conversation_id <> '' AND ( + EXISTS (SELECT 1 FROM conversations c WHERE c.id = vulnerabilities.conversation_id AND c.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments cra + WHERE cra.user_id = ? AND cra.resource_type = 'conversation' AND cra.resource_id = vulnerabilities.conversation_id + ) + ) + ) + )` + args = append(args, userID, userID, userID, userID, userID, userID) + return query, args +} + // Vulnerability 漏洞 type Vulnerability struct { ID string `json:"id"` @@ -190,6 +229,10 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) { // ListVulnerabilities 列出漏洞 func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) { + return db.ListVulnerabilitiesForAccess(limit, offset, filter, RBACListAccess{}) +} + +func (db *DB) ListVulnerabilitiesForAccess(limit, offset int, filter VulnerabilityListFilter, access RBACListAccess) ([]*Vulnerability, error) { query := ` SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag, vulnerability_type, target, @@ -203,6 +246,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFil ` args := []interface{}{} query, args = filter.appendWhere(query, args) + query, args = appendVulnerabilityAccessFilter(query, args, access) query += " ORDER BY created_at DESC LIMIT ? OFFSET ?" args = append(args, limit, offset) @@ -235,9 +279,14 @@ func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFil // CountVulnerabilities 统计漏洞总数（支持筛选条件） func (db *DB) CountVulnerabilities(filter VulnerabilityListFilter) (int, error) { + return db.CountVulnerabilitiesForAccess(filter, RBACListAccess{}) +} + +func (db *DB) CountVulnerabilitiesForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int, error) { query := "SELECT COUNT(*) FROM vulnerabilities WHERE 1=1" args := []interface{}{} query, args = filter.appendWhere(query, args) + query, args = appendVulnerabilityAccessFilter(query, args, access) var count int err := db.QueryRow(query, args...).Scan(&count) @@ -275,6 +324,10 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error { // DeleteVulnerabilitiesByFilter 按筛选条件批量删除漏洞，返回实际删除条数 func (db *DB) DeleteVulnerabilitiesByFilter(filter VulnerabilityListFilter) (int64, error) { + return db.DeleteVulnerabilitiesByFilterForAccess(filter, RBACListAccess{}) +} + +func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int64, error) { tx, err := db.Begin() if err != nil { return 0, fmt.Errorf("开启事务失败: %w", err) @@ -284,6 +337,7 @@ func (db *DB) DeleteVulnerabilitiesByFilter(filter VulnerabilityListFilter) (int where := "WHERE 1=1" args := []interface{}{} where, args = filter.appendWhere(where, args) + where, args = appendVulnerabilityAccessFilter(where, args, access) clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)` @@ -329,11 +383,16 @@ func (db *DB) DeleteVulnerability(id string) error { // GetVulnerabilityStats 获取漏洞统计（筛选条件与 ListVulnerabilities / CountVulnerabilities 一致） func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) { + return db.GetVulnerabilityStatsForAccess(filter, RBACListAccess{}) +} + +func (db *DB) GetVulnerabilityStatsForAccess(filter VulnerabilityListFilter, access RBACListAccess) (map[string]interface{}, error) { stats := make(map[string]interface{}) where := "WHERE 1=1" args := []interface{}{} where, args = filter.appendWhere(where, args) + where, args = appendVulnerabilityAccessFilter(where, args, access) // 总漏洞数 var totalCount int @@ -389,6 +448,10 @@ func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string] // GetVulnerabilityFilterOptions 获取漏洞筛选建议项 func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) { + return db.GetVulnerabilityFilterOptionsForAccess(RBACListAccess{}) +} + +func (db *DB) GetVulnerabilityFilterOptionsForAccess(access RBACListAccess) (map[string][]string, error) { collect := func(query string, args ...interface{}) ([]string, error) { rows, err := db.Query(query, args...) if err != nil { @@ -409,31 +472,35 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) { return items, nil } - vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities ORDER BY created_at DESC LIMIT 500`) + where := "WHERE 1=1" + accessArgs := []interface{}{} + where, accessArgs = appendVulnerabilityAccessFilter(where, accessArgs, access) + + vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities `+where+` ORDER BY created_at DESC LIMIT 500`, accessArgs...) if err != nil { return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err) } - conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`) + conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities `+where+` AND conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`, accessArgs...) if err != nil { return nil, fmt.Errorf("查询会话ID建议失败: %w", err) } - taskIDs, err := collect(`SELECT DISTINCT id FROM batch_tasks WHERE id <> '' ORDER BY rowid DESC LIMIT 500`) + taskIDs, err := collect(`SELECT DISTINCT bt.id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...) if err != nil { return nil, fmt.Errorf("查询任务ID建议失败: %w", err) } - queueIDs, err := collect(`SELECT DISTINCT queue_id FROM batch_tasks WHERE queue_id <> '' ORDER BY rowid DESC LIMIT 500`) + queueIDs, err := collect(`SELECT DISTINCT bt.queue_id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.queue_id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...) if err != nil { return nil, fmt.Errorf("查询队列ID建议失败: %w", err) } - conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities WHERE conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`) + conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities `+where+` AND conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`, accessArgs...) if err != nil { return nil, fmt.Errorf("查询对话标签建议失败: %w", err) } - taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities WHERE task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`) + taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities `+where+` AND task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`, accessArgs...) if err != nil { return nil, fmt.Errorf("查询任务标签建议失败: %w", err) } - projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities WHERE project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`) + projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities `+where+` AND project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`, accessArgs...) if err != nil { return nil, fmt.Errorf("查询项目ID建议失败: %w", err) } diff --git a/internal/database/webshell.go b/internal/database/webshell.go index db4e912f..424d61bc 100644 --- a/internal/database/webshell.go +++ b/internal/database/webshell.go @@ -2,6 +2,7 @@ package database import ( "database/sql" + "strings" "time" "go.uber.org/zap" @@ -59,13 +60,30 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro // ListWebshellConnections 列出所有 WebShell 连接，按创建时间倒序 func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) { + return db.ListWebshellConnectionsForAccess("", "") +} + +func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShellConnection, error) { query := ` SELECT id, url, password, type, method, cmd_param, remark, COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at FROM webshell_connections - ORDER BY created_at DESC + WHERE 1=1 ` - rows, err := db.Query(query) + args := []interface{}{} + userID = strings.TrimSpace(userID) + if userID != "" && scope != RBACScopeAll { + query += ` AND ( + owner_user_id = ? + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'webshell' AND ra.resource_id = webshell_connections.id + ) + )` + args = append(args, userID, userID) + } + query += ` ORDER BY created_at DESC` + rows, err := db.Query(query, args...) if err != nil { db.logger.Error("查询 WebShell 连接列表失败", zap.Error(err)) return nil, err diff --git a/internal/workflow/nodes.go b/internal/workflow/nodes.go index 5c342296..eb0342e2 100644 --- a/internal/workflow/nodes.go +++ b/internal/workflow/nodes.go @@ -4,8 +4,10 @@ import ( "context" "fmt" "strings" + "unicode/utf8" "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/multiagent" ) @@ -85,6 +87,11 @@ func runToolNode(ctx context.Context, args RunArgs, node graphNode, state *Workf executionID = result.ExecutionID isError = result.IsError } + maxToolOutputBytes := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective() + if args.AppCfg != nil { + maxToolOutputBytes = args.AppCfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective() + } + output = truncateWorkflowToolOutput(output, maxToolOutputBytes, executionID) out := toolOutputMap(node, output, toolName, toolArgs, executionID, isError) if key := cfgString(node.Config, "output_key"); key != "" { state.Outputs[key] = output @@ -99,6 +106,30 @@ func runToolNode(ctx context.Context, args RunArgs, node graphNode, state *Workf return out, true, "completed", "" } +func truncateWorkflowToolOutput(output string, maxBytes int, executionID string) string { + if maxBytes <= 0 || len(output) <= maxBytes { + return output + } + marker := fmt.Sprintf("\n\n...[workflow tool output truncated; full result is stored in execution %s]...\n\n", strings.TrimSpace(executionID)) + if strings.TrimSpace(executionID) == "" { + marker = "\n\n...[workflow tool output truncated; full result remains in the tool execution record]...\n\n" + } + budget := maxBytes - len(marker) + if budget <= 0 { + return marker + } + head := budget / 2 + tail := budget - head + for head > 0 && !utf8.RuneStart(output[head]) { + head-- + } + tailStart := len(output) - tail + for tailStart < len(output) && !utf8.RuneStart(output[tailStart]) { + tailStart++ + } + return output[:head] + marker + output[tailStart:] +} + func runAgentNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) { if args.AppCfg == nil || args.Agent == nil { errText := "Agent 节点执行失败：应用配置或 Agent 为空" diff --git a/internal/workflow/tool_output_budget_test.go b/internal/workflow/tool_output_budget_test.go new file mode 100644 index 00000000..d66d27eb --- /dev/null +++ b/internal/workflow/tool_output_budget_test.go @@ -0,0 +1,23 @@ +package workflow + +import ( + "strings" + "testing" +) + +func TestTruncateWorkflowToolOutputBoundsBytesAndKeepsExecutionReference(t *testing.T) { + out := truncateWorkflowToolOutput(strings.Repeat("响应正文", 1000), 256, "exec-123") + if len(out) > 256 { + t.Fatalf("workflow output bytes=%d, want <=256", len(out)) + } + if !strings.Contains(out, "exec-123") || !strings.Contains(out, "truncated") { + t.Fatalf("missing truncation reference: %q", out) + } +} + +func TestTruncateWorkflowToolOutputLeavesBoundedContentUntouched(t *testing.T) { + const want = "small-result" + if got := truncateWorkflowToolOutput(want, 256, "exec-123"); got != want { + t.Fatalf("got %q want %q", got, want) + } +}