mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-11-14 00:07:10 +01:00
754 lines
29 KiB
Go
754 lines
29 KiB
Go
package database
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"net/netip"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
"unicode"
|
|
|
|
"github.com/google/uuid"
|
|
"go.uber.org/zap"
|
|
"golang.org/x/net/idna"
|
|
"golang.org/x/net/publicsuffix"
|
|
)
|
|
|
|
// VulnerabilityListFilter 列表/统计/导出共用的筛选条件
|
|
type VulnerabilityListFilter struct {
|
|
ID string
|
|
Search string // 关键词模糊匹配（标题、描述、类型、目标等）
|
|
HostQuery string
|
|
TaskQuery string
|
|
CaseID string
|
|
SummaryOnly bool
|
|
ConversationID string
|
|
ProjectID string
|
|
Severity string
|
|
Status string
|
|
TaskID string
|
|
ConversationTag string
|
|
TaskTag string
|
|
}
|
|
|
|
type RBACListAccess struct {
|
|
UserID string
|
|
Scope string
|
|
}
|
|
|
|
func escapeVulnerabilityLikePattern(s string) string {
|
|
s = strings.ReplaceAll(s, `\`, `\\`)
|
|
s = strings.ReplaceAll(s, `%`, `\%`)
|
|
s = strings.ReplaceAll(s, `_`, `\_`)
|
|
return "%" + s + "%"
|
|
}
|
|
|
|
func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (string, []interface{}) {
|
|
if host := strings.TrimSpace(f.HostQuery); host != "" {
|
|
query += ` AND EXISTS (SELECT 1 FROM json_each(targets_json) jt WHERE json_extract(jt.value,'$.host') LIKE ? ESCAPE '\')`
|
|
args = append(args, escapeVulnerabilityLikePattern(strings.ToLower(host)))
|
|
}
|
|
if f.CaseID != "" {
|
|
query += " AND case_id=?"
|
|
args = append(args, f.CaseID)
|
|
}
|
|
if q := strings.TrimSpace(f.TaskQuery); q != "" {
|
|
query += ` AND (COALESCE(task_tag,'') LIKE ? ESCAPE '\' OR EXISTS (
|
|
SELECT 1 FROM batch_tasks bt LEFT JOIN batch_task_queues bq ON bq.id=bt.queue_id
|
|
WHERE bt.conversation_id=vulnerabilities.conversation_id
|
|
AND (bt.id LIKE ? ESCAPE '\' OR bt.queue_id LIKE ? ESCAPE '\' OR COALESCE(bq.title,'') LIKE ? ESCAPE '\')))`
|
|
pattern := escapeVulnerabilityLikePattern(q)
|
|
args = append(args, pattern, pattern, pattern, pattern)
|
|
}
|
|
if f.ID != "" {
|
|
query += " AND vulnerabilities.id = ?"
|
|
args = append(args, f.ID)
|
|
}
|
|
if f.ConversationID != "" {
|
|
query += " AND conversation_id = ?"
|
|
args = append(args, f.ConversationID)
|
|
}
|
|
if f.ProjectID != "" {
|
|
query += " AND project_id = ?"
|
|
args = append(args, f.ProjectID)
|
|
}
|
|
if f.TaskID != "" {
|
|
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
|
args = append(args, f.TaskID, f.TaskID)
|
|
}
|
|
if f.ConversationTag != "" {
|
|
query += " AND conversation_tag = ?"
|
|
args = append(args, f.ConversationTag)
|
|
}
|
|
if f.TaskTag != "" {
|
|
query += " AND task_tag = ?"
|
|
args = append(args, f.TaskTag)
|
|
}
|
|
if f.Severity != "" {
|
|
query += " AND severity = ?"
|
|
args = append(args, f.Severity)
|
|
}
|
|
if f.Status != "" {
|
|
query += " AND status = ?"
|
|
args = append(args, f.Status)
|
|
}
|
|
search := strings.TrimSpace(f.Search)
|
|
if search != "" {
|
|
pattern := escapeVulnerabilityLikePattern(search)
|
|
query += ` AND (
|
|
LOWER(vulnerabilities.id) LIKE LOWER(?) OR
|
|
LOWER(title) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(description, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(vulnerability_type, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(target, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(preconditions, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(reproduction_steps, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(evidence, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(impact, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(recommendation, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(retest_notes, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(conversation_id, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(conversation_tag, '')) LIKE LOWER(?) OR
|
|
LOWER(COALESCE(task_tag, '')) LIKE LOWER(?)
|
|
)`
|
|
for i := 0; i < 14; i++ {
|
|
args = append(args, pattern)
|
|
}
|
|
}
|
|
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"`
|
|
ConversationID string `json:"conversation_id"`
|
|
ProjectID string `json:"project_id,omitempty"`
|
|
ConversationTag string `json:"conversation_tag,omitempty"`
|
|
TaskTag string `json:"task_tag,omitempty"`
|
|
TaskID string `json:"task_id,omitempty"`
|
|
TaskQueueID string `json:"task_queue_id,omitempty"`
|
|
TaskQueueTitle string `json:"task_queue_title,omitempty"`
|
|
Title string `json:"title"`
|
|
Description string `json:"description"`
|
|
Severity string `json:"severity"` // critical, high, medium, low, info
|
|
Status string `json:"status"` // open, confirmed, fixed, false_positive, ignored
|
|
Type string `json:"type"`
|
|
Target string `json:"target"`
|
|
TargetHost *string `json:"target_host,omitempty"`
|
|
TargetDomain string `json:"target_domain,omitempty"`
|
|
Targets []VulnerabilityTarget `json:"targets"`
|
|
CaseID string `json:"case_id,omitempty"`
|
|
CaseTitle string `json:"case_title,omitempty"`
|
|
SummaryOnly bool `json:"summary_only,omitempty"`
|
|
Preconditions string `json:"preconditions"`
|
|
ReproSteps string `json:"reproduction_steps"`
|
|
Evidence string `json:"evidence"`
|
|
Impact string `json:"impact"`
|
|
Recommendation string `json:"recommendation"`
|
|
RetestNotes string `json:"retest_notes"`
|
|
RetestLog string `json:"retest_log"` // append-only retest evidence; kept out of retest_notes because the edit form rewrites that field wholesale
|
|
CreatedAt time.Time `json:"created_at"`
|
|
UpdatedAt time.Time `json:"updated_at"`
|
|
}
|
|
|
|
// CreateVulnerability 创建漏洞
|
|
func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
|
if vuln.ID == "" {
|
|
vuln.ID = uuid.New().String()
|
|
}
|
|
if vuln.Status == "" {
|
|
vuln.Status = "open"
|
|
}
|
|
now := time.Now()
|
|
if vuln.CreatedAt.IsZero() {
|
|
vuln.CreatedAt = now
|
|
}
|
|
vuln.UpdatedAt = now
|
|
|
|
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID != "" {
|
|
if pid, err := db.GetConversationProjectID(vuln.ConversationID); err == nil {
|
|
vuln.ProjectID = pid
|
|
}
|
|
}
|
|
|
|
query := `
|
|
INSERT INTO vulnerabilities (
|
|
id, conversation_id, project_id, conversation_tag, task_tag, title, description, severity, status,
|
|
vulnerability_type, target, preconditions, reproduction_steps, evidence, impact, recommendation, retest_notes,
|
|
created_at, updated_at, targets_json, targets_version
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`
|
|
|
|
_, err := db.Exec(
|
|
query,
|
|
vuln.ID, nullIfEmpty(vuln.ConversationID), nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
|
|
vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
|
|
vuln.Preconditions, vuln.ReproSteps, vuln.Evidence, vuln.Impact, vuln.Recommendation, vuln.RetestNotes,
|
|
vuln.CreatedAt, vuln.UpdatedAt, vulnerabilityTargetsJSON(vuln.Target), vulnerabilityTargetsVersion,
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("创建漏洞失败: %w", err)
|
|
}
|
|
db.refreshAssetRiskCacheForConversationsBestEffort(vuln.ConversationID)
|
|
return db.GetVulnerability(vuln.ID)
|
|
}
|
|
|
|
// GetVulnerability 获取漏洞
|
|
func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
|
var vuln Vulnerability
|
|
var targets string
|
|
query := `
|
|
SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, COALESCE(description,''), severity, status,
|
|
COALESCE(conversation_tag,''), COALESCE(task_tag,''), COALESCE(vulnerability_type,''), COALESCE(target,''),
|
|
COALESCE(preconditions,''), COALESCE(reproduction_steps,''), COALESCE(evidence,''),
|
|
COALESCE(impact,''), COALESCE(recommendation,''), COALESCE(retest_notes,''), COALESCE(retest_log,''),
|
|
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
|
|
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
|
|
created_at, updated_at, targets_json, COALESCE(case_id,''),
|
|
COALESCE((SELECT title FROM vulnerability_cases vc WHERE vc.id=vulnerabilities.case_id),'')
|
|
FROM vulnerabilities
|
|
WHERE id = ?
|
|
`
|
|
|
|
err := db.QueryRow(query, id).Scan(
|
|
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
|
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
|
&vuln.Preconditions, &vuln.ReproSteps, &vuln.Evidence, &vuln.Impact, &vuln.Recommendation, &vuln.RetestNotes,
|
|
&vuln.RetestLog, &vuln.TaskID, &vuln.TaskQueueID,
|
|
&vuln.CreatedAt, &vuln.UpdatedAt, &targets, &vuln.CaseID, &vuln.CaseTitle,
|
|
)
|
|
if err != nil {
|
|
if err == sql.ErrNoRows {
|
|
return nil, fmt.Errorf("漏洞不存在")
|
|
}
|
|
return nil, fmt.Errorf("获取漏洞失败: %w", err)
|
|
}
|
|
|
|
return &vuln, scanVulnerabilityTargets(targets, &vuln)
|
|
}
|
|
|
|
// 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) {
|
|
return db.ListVulnerabilitiesByGroupForAccess(limit, offset, filter, access, "", "")
|
|
}
|
|
|
|
func vulnerabilityTargetHost(target string) string {
|
|
target = strings.TrimSpace(target)
|
|
// Historical findings can append endpoint notes or secondary targets to
|
|
// the primary URL or host. Use the leading target for a single group.
|
|
if end := strings.IndexFunc(target, func(r rune) bool {
|
|
return unicode.IsSpace(r) || strings.ContainsRune("（），；、(,;", r)
|
|
}); end >= 0 {
|
|
target = target[:end]
|
|
}
|
|
if ip, err := netip.ParseAddr(target); err == nil {
|
|
return ip.String()
|
|
}
|
|
scheme, _, hasScheme := strings.Cut(target, "://")
|
|
explicitURL := strings.HasPrefix(target, "//") || (hasScheme && !strings.ContainsAny(scheme, "/?#"))
|
|
if explicitURL {
|
|
if strings.HasPrefix(target, "//") {
|
|
target = strings.TrimPrefix(target, "//")
|
|
} else {
|
|
target = target[len(scheme)+3:]
|
|
}
|
|
}
|
|
// Paths and queries can contain raw request fragments or malformed escape
|
|
// sequences. Host extraction only needs the authority.
|
|
if end := strings.IndexAny(target, "/?#"); end >= 0 {
|
|
target = target[:end]
|
|
}
|
|
parsed, err := url.Parse("//" + target)
|
|
if err != nil || parsed.Host == "" {
|
|
return ""
|
|
}
|
|
host := parsed.Hostname()
|
|
if ip, err := netip.ParseAddr(host); err == nil {
|
|
return ip.String()
|
|
}
|
|
host, err = idna.Lookup.ToASCII(host)
|
|
host = strings.ToLower(strings.TrimSuffix(host, "."))
|
|
if err != nil || !validAssetDomain(host) {
|
|
return ""
|
|
}
|
|
// A relative endpoint such as "api/login" does not identify a host.
|
|
if !explicitURL && parsed.Port() == "" && host != "localhost" && !strings.Contains(host, ".") {
|
|
return ""
|
|
}
|
|
return host
|
|
}
|
|
|
|
const vulnerabilityTaskIDExpr = `COALESCE((SELECT bt.id FROM batch_tasks bt
|
|
WHERE bt.conversation_id = vulnerabilities.conversation_id ORDER BY bt.id LIMIT 1), '')`
|
|
|
|
const vulnerabilitySeverityRankExpr = `CASE severity
|
|
WHEN 'critical' THEN 0 WHEN 'high' THEN 1 WHEN 'medium' THEN 2
|
|
WHEN 'low' THEN 3 WHEN 'info' THEN 4 ELSE 5 END`
|
|
|
|
func vulnerabilityGroupQuery(filter VulnerabilityListFilter, access RBACListAccess, groupBy string) (string, []interface{}) {
|
|
var key string
|
|
join := ""
|
|
isIP := "0"
|
|
switch groupBy {
|
|
case "host":
|
|
key = "COALESCE(json_extract(jt.value,'$.host'),'')"
|
|
join = " LEFT JOIN json_each(targets_json) jt ON json_extract(jt.value,'$.group_host')=1"
|
|
isIP = "COALESCE(json_extract(jt.value,'$.kind')='ip',0)"
|
|
case "case":
|
|
key = "CASE WHEN case_id IS NOT NULL THEN 'case:' || case_id ELSE 'row:' || vulnerabilities.id END"
|
|
case "task":
|
|
key = `CASE
|
|
WHEN ` + vulnerabilityTaskIDExpr + ` <> '' THEN 'id:' || ` + vulnerabilityTaskIDExpr + `
|
|
WHEN TRIM(task_tag) <> '' THEN 'tag:' || TRIM(task_tag)
|
|
ELSE '' END`
|
|
default:
|
|
return "", nil
|
|
}
|
|
query := "SELECT DISTINCT vulnerabilities.id, created_at, " + isIP + " AS group_is_ip, " + vulnerabilitySeverityRankExpr + " AS severity_rank, " +
|
|
key + " AS group_key FROM vulnerabilities" + join + " WHERE 1=1"
|
|
query, args := filter.appendWhere(query, nil)
|
|
if groupBy == "host" && strings.TrimSpace(filter.HostQuery) != "" {
|
|
query += ` AND json_extract(jt.value,'$.host') LIKE ? ESCAPE '\'`
|
|
args = append(args, escapeVulnerabilityLikePattern(strings.ToLower(strings.TrimSpace(filter.HostQuery))))
|
|
}
|
|
return appendVulnerabilityAccessFilter(query, args, access)
|
|
}
|
|
|
|
func (db *DB) CountVulnerabilityGroupsForAccess(filter VulnerabilityListFilter, access RBACListAccess, groupBy string) (int, error) {
|
|
query, args := vulnerabilityGroupQuery(filter, access, groupBy)
|
|
if query == "" {
|
|
return db.CountVulnerabilitiesForAccess(filter, access)
|
|
}
|
|
var count int
|
|
err := db.QueryRow("SELECT COUNT(DISTINCT group_key) FROM ("+query+")", args...).Scan(&count)
|
|
return count, err
|
|
}
|
|
|
|
func (db *DB) ListVulnerabilitiesByGroupForAccess(limit, offset int, filter VulnerabilityListFilter, access RBACListAccess, groupBy, sortBy string) ([]*Vulnerability, error) {
|
|
hostExpr, order := "NULL", "created_at DESC"
|
|
itemOrder := "created_at DESC, id"
|
|
groupOrder := "group_key = '', group_key"
|
|
if groupBy == "host" {
|
|
groupOrder = "group_key = '', group_is_ip, group_key"
|
|
}
|
|
switch sortBy {
|
|
case "severity":
|
|
itemOrder = vulnerabilitySeverityRankExpr + ", created_at DESC, id"
|
|
groupOrder = "group_severity, group_latest DESC, " + groupOrder
|
|
case "latest":
|
|
groupOrder = "group_latest DESC, " + groupOrder
|
|
}
|
|
groupQuery, args := vulnerabilityGroupQuery(filter, access, groupBy)
|
|
var prefix, join string
|
|
if groupQuery != "" {
|
|
prefix = `WITH scoped AS (` + groupQuery + `),
|
|
page_groups AS (
|
|
SELECT group_key, MIN(group_is_ip) AS group_is_ip, MIN(severity_rank) AS group_severity, MAX(created_at) AS group_latest
|
|
FROM scoped GROUP BY group_key
|
|
ORDER BY ` + groupOrder + ` LIMIT ? OFFSET ?
|
|
),
|
|
page_items AS (
|
|
SELECT scoped.id, scoped.group_key, page_groups.group_is_ip, page_groups.group_severity, page_groups.group_latest
|
|
FROM scoped JOIN page_groups USING (group_key)
|
|
)`
|
|
args = append(args, limit, offset)
|
|
join = " JOIN page_items USING (id)"
|
|
order = groupOrder + ", " + itemOrder
|
|
if groupBy == "host" {
|
|
hostExpr = "group_key"
|
|
}
|
|
} else if groupBy == "none" || sortBy != "" {
|
|
order = itemOrder
|
|
}
|
|
description := "COALESCE(description,'')"
|
|
details := `COALESCE(preconditions,''), COALESCE(reproduction_steps,''), COALESCE(evidence,''),
|
|
COALESCE(impact,''), COALESCE(recommendation,''), COALESCE(retest_notes,''), COALESCE(retest_log,'')`
|
|
if filter.SummaryOnly {
|
|
description = "SUBSTR(COALESCE(description,''),1,240)"
|
|
details = "'','','','','','',''"
|
|
}
|
|
query := prefix + `
|
|
SELECT vulnerabilities.id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, ` + description + `, severity, status, COALESCE(conversation_tag,''), COALESCE(task_tag,''),
|
|
COALESCE(vulnerability_type,''), COALESCE(target,''),
|
|
` + details + `,
|
|
` + vulnerabilityTaskIDExpr + ` AS task_id,
|
|
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id ORDER BY bt.id LIMIT 1), '') AS task_queue_id,
|
|
COALESCE((SELECT bq.title FROM batch_tasks bt JOIN batch_task_queues bq ON bq.id=bt.queue_id WHERE bt.conversation_id=vulnerabilities.conversation_id ORDER BY bt.id LIMIT 1),''),
|
|
created_at, updated_at, ` + hostExpr + ` AS target_host, targets_json, COALESCE(case_id,''),
|
|
COALESCE((SELECT title FROM vulnerability_cases vc WHERE vc.id=vulnerabilities.case_id),'')
|
|
FROM vulnerabilities` + join + `
|
|
WHERE 1=1
|
|
`
|
|
if groupQuery == "" {
|
|
query, args = filter.appendWhere(query, args)
|
|
query, args = appendVulnerabilityAccessFilter(query, args, access)
|
|
}
|
|
query += " ORDER BY " + order
|
|
if groupQuery == "" {
|
|
query += " 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()
|
|
|
|
vulnerabilities := make([]*Vulnerability, 0)
|
|
for rows.Next() {
|
|
var vuln Vulnerability
|
|
var targets string
|
|
err := rows.Scan(
|
|
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
|
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
|
&vuln.Preconditions, &vuln.ReproSteps, &vuln.Evidence, &vuln.Impact, &vuln.Recommendation, &vuln.RetestNotes,
|
|
&vuln.RetestLog, &vuln.TaskID, &vuln.TaskQueueID, &vuln.TaskQueueTitle,
|
|
&vuln.CreatedAt, &vuln.UpdatedAt, &vuln.TargetHost, &targets, &vuln.CaseID, &vuln.CaseTitle,
|
|
)
|
|
if err != nil {
|
|
db.logger.Warn("扫描漏洞记录失败", zap.Error(err))
|
|
continue
|
|
}
|
|
if vuln.TargetHost != nil {
|
|
vuln.TargetDomain = *vuln.TargetHost
|
|
if _, err := netip.ParseAddr(*vuln.TargetHost); err != nil {
|
|
if domain, err := publicsuffix.EffectiveTLDPlusOne(*vuln.TargetHost); err == nil {
|
|
vuln.TargetDomain = domain
|
|
}
|
|
}
|
|
}
|
|
if err := scanVulnerabilityTargets(targets, &vuln); err != nil {
|
|
return nil, err
|
|
}
|
|
vuln.SummaryOnly = filter.SummaryOnly
|
|
vulnerabilities = append(vulnerabilities, &vuln)
|
|
}
|
|
|
|
return vulnerabilities, rows.Err()
|
|
}
|
|
|
|
// 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)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("统计漏洞总数失败: %w", err)
|
|
}
|
|
|
|
return count, nil
|
|
}
|
|
|
|
// UpdateVulnerability 更新漏洞
|
|
func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
|
vuln.UpdatedAt = time.Now()
|
|
var oldConversationID string
|
|
_ = db.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&oldConversationID)
|
|
|
|
query := `
|
|
UPDATE vulnerabilities
|
|
SET project_id = ?, conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
|
|
vulnerability_type = ?, target = ?, preconditions = ?, reproduction_steps = ?, evidence = ?, impact = ?,
|
|
recommendation = ?, retest_notes = ?, updated_at = ?, targets_json = ?, targets_version=?
|
|
WHERE id = ?
|
|
`
|
|
|
|
tx, err := db.Begin()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback()
|
|
if err := detachVulnerabilityCaseOnProjectChange(tx, id, vuln.ProjectID); err != nil {
|
|
return err
|
|
}
|
|
_, err = tx.Exec(
|
|
query,
|
|
nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
|
|
vuln.Type, vuln.Target, vuln.Preconditions, vuln.ReproSteps, vuln.Evidence, vuln.Impact,
|
|
vuln.Recommendation, vuln.RetestNotes, vuln.UpdatedAt, vulnerabilityTargetsJSON(vuln.Target), vulnerabilityTargetsVersion, id,
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("更新漏洞失败: %w", err)
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return err
|
|
}
|
|
|
|
db.refreshAssetRiskCacheForConversationsBestEffort(oldConversationID, vuln.ConversationID)
|
|
return nil
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
|
|
where := "WHERE 1=1"
|
|
args := []interface{}{}
|
|
where, args = filter.appendWhere(where, args)
|
|
where, args = appendVulnerabilityAccessFilter(where, args, access)
|
|
affectedConversations, err := collectVulnerabilityConversationIDs(tx, where, args)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL
|
|
WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)`
|
|
if _, err := tx.Exec(clearQuery, args...); err != nil {
|
|
return 0, fmt.Errorf("清理事实漏洞关联失败: %w", err)
|
|
}
|
|
|
|
deleteQuery := `DELETE FROM vulnerabilities ` + where
|
|
result, err := tx.Exec(deleteQuery, args...)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("批量删除漏洞失败: %w", err)
|
|
}
|
|
deleted, err := result.RowsAffected()
|
|
if err != nil {
|
|
return 0, fmt.Errorf("获取删除条数失败: %w", err)
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return 0, fmt.Errorf("提交事务失败: %w", err)
|
|
}
|
|
db.refreshAssetRiskCacheForConversationsBestEffort(affectedConversations...)
|
|
return deleted, nil
|
|
}
|
|
|
|
// DeleteVulnerability 删除漏洞
|
|
func (db *DB) DeleteVulnerability(id string) error {
|
|
tx, err := db.Begin()
|
|
if err != nil {
|
|
return fmt.Errorf("开启事务失败: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
var conversationID string
|
|
_ = tx.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&conversationID)
|
|
|
|
// 删除漏洞前先解除项目事实中的关联，避免前端继续显示已删除漏洞的短 ID。
|
|
if _, err := tx.Exec("UPDATE project_facts SET related_vulnerability_id = NULL WHERE related_vulnerability_id = ?", id); err != nil {
|
|
return fmt.Errorf("清理事实漏洞关联失败: %w", err)
|
|
}
|
|
if _, err := tx.Exec("DELETE FROM vulnerabilities WHERE id = ?", id); err != nil {
|
|
return fmt.Errorf("删除漏洞失败: %w", err)
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("提交事务失败: %w", err)
|
|
}
|
|
db.refreshAssetRiskCacheForConversationsBestEffort(conversationID)
|
|
return nil
|
|
}
|
|
|
|
func collectVulnerabilityConversationIDs(tx *sql.Tx, where string, args []interface{}) ([]string, error) {
|
|
rows, err := tx.Query(`SELECT DISTINCT COALESCE(conversation_id,'') FROM vulnerabilities `+where, args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("查询受影响漏洞会话失败: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
ids := []string{}
|
|
for rows.Next() {
|
|
var id string
|
|
if err := rows.Scan(&id); err != nil {
|
|
return nil, err
|
|
}
|
|
if strings.TrimSpace(id) != "" {
|
|
ids = append(ids, id)
|
|
}
|
|
}
|
|
return ids, rows.Err()
|
|
}
|
|
|
|
// 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
|
|
query := "SELECT COUNT(*) FROM vulnerabilities " + where
|
|
err := db.QueryRow(query, args...).Scan(&totalCount)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("获取总漏洞数失败: %w", err)
|
|
}
|
|
stats["total"] = totalCount
|
|
var distinctCount int
|
|
if err := db.QueryRow("SELECT COUNT(DISTINCT CASE WHEN case_id IS NULL THEN 'row:'||id ELSE 'case:'||case_id END) FROM vulnerabilities "+where, args...).Scan(&distinctCount); err != nil {
|
|
return nil, err
|
|
}
|
|
stats["distinct"] = distinctCount
|
|
|
|
// 按严重程度统计
|
|
severityQuery := "SELECT severity, COUNT(*) FROM vulnerabilities " + where + " GROUP BY severity"
|
|
|
|
rows, err := db.Query(severityQuery, args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("获取严重程度统计失败: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
severityStats := make(map[string]int)
|
|
for rows.Next() {
|
|
var severity string
|
|
var count int
|
|
if err := rows.Scan(&severity, &count); err != nil {
|
|
continue
|
|
}
|
|
severityStats[severity] = count
|
|
}
|
|
stats["by_severity"] = severityStats
|
|
|
|
// 按状态统计
|
|
statusQuery := "SELECT status, COUNT(*) FROM vulnerabilities " + where + " GROUP BY status"
|
|
|
|
rows, err = db.Query(statusQuery, args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("获取状态统计失败: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
statusStats := make(map[string]int)
|
|
for rows.Next() {
|
|
var status string
|
|
var count int
|
|
if err := rows.Scan(&status, &count); err != nil {
|
|
continue
|
|
}
|
|
statusStats[status] = count
|
|
}
|
|
stats["by_status"] = statusStats
|
|
|
|
return stats, nil
|
|
}
|
|
|
|
// 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 {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
items := make([]string, 0)
|
|
for rows.Next() {
|
|
var val string
|
|
if err := rows.Scan(&val); err != nil {
|
|
continue
|
|
}
|
|
if val == "" {
|
|
continue
|
|
}
|
|
items = append(items, val)
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
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+` 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 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 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+` 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+` 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+` 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)
|
|
}
|
|
|
|
return map[string][]string{
|
|
"vulnerability_ids": vulnIDs,
|
|
"conversation_ids": conversationIDs,
|
|
"project_ids": projectIDs,
|
|
"task_ids": taskIDs,
|
|
"queue_ids": queueIDs,
|
|
"conversation_tags": conversationTags,
|
|
"task_tags": taskTags,
|
|
}, nil
|
|
}
|