CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-07-19 10:28:09 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-16 19:37:01 +08:00
committed by GitHub
parent 15cc08a95b
commit d03cc47ac3
10 changed files with 1979 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/database/rbac.go
+39 -2
View File
@@ -27,6 +27,7 @@ var rbacAssignableResourceTables = map[string]string{
"project": "projects",
"conversation": "conversations",
"vulnerability": "vulnerabilities",
"asset": "assets",
"webshell": "webshell_connections",
"batch_task": "batch_task_queues",
"c2_listener": "c2_listeners",
@@ -528,6 +529,9 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri
if resourceType == "vulnerability" {
return db.userCanAccessVulnerabilityViaParent(userID, scope, resourceID)
}
if resourceType == "asset" {
return db.userCanAccessAssetViaParent(userID, scope, resourceID)
}
if resourceType == "conversation" {
return db.userCanAccessConversationViaParent(userID, scope, resourceID)
}
@@ -537,6 +541,15 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri
return false
}
func (db *DB) userCanAccessAssetViaParent(userID, scope, assetID string) bool {
var projectID sql.NullString
if err := db.QueryRow(`SELECT project_id FROM assets WHERE id = ?`, assetID).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) 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 {
@@ -623,6 +636,8 @@ func (db *DB) userOwnsResource(userID, resourceType, resourceID string) bool {
table = "conversations"
case "vulnerability":
table = "vulnerabilities"
case "asset":
table = "assets"
case "webshell":
table = "webshell_connections"
case "batch_task":
@@ -650,6 +665,8 @@ func (db *DB) SetResourceOwner(resourceType, resourceID, userID string) error {
table = "conversations"
case "vulnerability":
table = "vulnerabilities"
case "asset":
table = "assets"
case "webshell":
table = "webshell_connections"
case "batch_task":
@@ -672,6 +689,8 @@ func (db *DB) GetResourceOwner(resourceType, resourceID string) string {
table = "conversations"
case "vulnerability":
table = "vulnerabilities"
case "asset":
table = "assets"
case "webshell":
table = "webshell_connections"
case "batch_task":
@@ -733,6 +752,10 @@ func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit
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 "asset":
query = `SELECT id, COALESCE(NULLIF(host,''),NULLIF(domain,''),NULLIF(ip,''),id), protocol || CASE WHEN port>0 THEN ':' || port ELSE '' END FROM assets
WHERE LOWER(host) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\' OR LOWER(ip) 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 '\'
@@ -747,7 +770,12 @@ func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit
ORDER BY created_at DESC LIMIT ? OFFSET ?`
}
rows, err := db.Query(query, pattern, pattern, limit, offset)
queryArgs := []interface{}{pattern, pattern}
if resourceType == "asset" {
queryArgs = append(queryArgs, pattern)
}
queryArgs = append(queryArgs, limit, offset)
rows, err := db.Query(query, queryArgs...)
if err != nil {
return nil, err
}
@@ -781,6 +809,8 @@ func (db *DB) CountAssignableRBACResources(resourceType, search string) (int, er
query = `SELECT COUNT(*) FROM conversations WHERE LOWER(COALESCE(NULLIF(TRIM(title), ''), id)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
case "vulnerability":
query = `SELECT COUNT(*) FROM vulnerabilities WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
case "asset":
query = `SELECT COUNT(*) FROM assets WHERE LOWER(host) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\'`
case "webshell":
query = `SELECT COUNT(*) FROM webshell_connections WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
case "batch_task":
@@ -789,7 +819,11 @@ func (db *DB) CountAssignableRBACResources(resourceType, search string) (int, er
query = `SELECT COUNT(*) FROM c2_listeners WHERE LOWER(name) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
}
var total int
if err := db.QueryRow(query, pattern, pattern).Scan(&total); err != nil {
queryArgs := []interface{}{pattern, pattern}
if resourceType == "asset" {
queryArgs = append(queryArgs, pattern)
}
if err := db.QueryRow(query, queryArgs...).Scan(&total); err != nil {
return 0, err
}
return total, nil
@@ -869,6 +903,8 @@ func (db *DB) lookupRBACResourceOptionsByIDs(resourceType string, ids []string)
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 "asset":
query = `SELECT id, COALESCE(NULLIF(host,''),NULLIF(domain,''),NULLIF(ip,''),id), protocol || CASE WHEN port>0 THEN ':' || port ELSE '' END FROM assets WHERE id IN (` + placeholders + `)`
case "webshell":
query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections WHERE id IN (` + placeholders + `)`
case "batch_task":
@@ -1039,6 +1075,7 @@ func (db *DB) AssignResourcesToUserAuto(userID string, resourceIDs []string) (in
typeTablePairs := []struct{ resourceType, table string }{
{"project", "projects"}, {"conversation", "conversations"},
{"vulnerability", "vulnerabilities"}, {"webshell", "webshell_connections"},
{"asset", "assets"},
{"batch_task", "batch_task_queues"}, {"c2_listener", "c2_listeners"},
}
detected := make(map[string]string, len(uniqueIDs))