Files
CyberStrikeAI/internal/database/rbac.go
T
2026-07-10 16:42:35 +08:00

1112 lines
36 KiB
Go

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
}