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

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-10 16:42:35 +08:00
committed by GitHub
parent 225d9ba0d9
commit a2744d6936
14 changed files with 2631 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/database/vulnerability.go
+74 -7
View File
@@ -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)
}