mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-19 02:17:38 +02:00
Add files via upload
This commit is contained in:
@@ -0,0 +1,843 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/net/idna"
|
||||
)
|
||||
|
||||
// Asset is a persistent, deduplicated target discovered manually or by recon providers.
|
||||
type Asset struct {
|
||||
ID string `json:"id"`
|
||||
ProjectID string `json:"project_id,omitempty"`
|
||||
ProjectName string `json:"project_name,omitempty"`
|
||||
Host string `json:"host"`
|
||||
IP string `json:"ip"`
|
||||
Port int `json:"port"`
|
||||
Domain string `json:"domain"`
|
||||
Protocol string `json:"protocol"`
|
||||
Title string `json:"title"`
|
||||
Server string `json:"server"`
|
||||
Country string `json:"country"`
|
||||
Province string `json:"province"`
|
||||
City string `json:"city"`
|
||||
Source string `json:"source"`
|
||||
SourceQuery string `json:"source_query"`
|
||||
Status string `json:"status"`
|
||||
Tags []string `json:"tags"`
|
||||
FirstSeenAt time.Time `json:"first_seen_at"`
|
||||
LastSeenAt time.Time `json:"last_seen_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
LastScanAt *time.Time `json:"last_scan_at,omitempty"`
|
||||
LastScanConversationID string `json:"last_scan_conversation_id,omitempty"`
|
||||
LastScanQueueID string `json:"last_scan_queue_id,omitempty"`
|
||||
LastScanTaskID string `json:"last_scan_task_id,omitempty"`
|
||||
VulnerabilityCount int `json:"vulnerability_count"`
|
||||
RiskLevel string `json:"risk_level"`
|
||||
OwnerUserID string `json:"-"`
|
||||
}
|
||||
|
||||
type AssetListFilter struct {
|
||||
Search string
|
||||
Status string
|
||||
Protocol string
|
||||
ProjectID string
|
||||
Source string
|
||||
Tag string
|
||||
Host string
|
||||
IP string
|
||||
Domain string
|
||||
Port *int
|
||||
ScanState string
|
||||
LastScanBefore *time.Time
|
||||
LastScanAfter *time.Time
|
||||
LastSeenBefore *time.Time
|
||||
LastSeenAfter *time.Time
|
||||
SortBy string
|
||||
SortOrder string
|
||||
}
|
||||
|
||||
type AssetImportResult struct {
|
||||
Created int `json:"created"`
|
||||
Updated int `json:"updated"`
|
||||
Skipped int `json:"skipped"`
|
||||
}
|
||||
|
||||
func normalizeAsset(a *Asset) {
|
||||
a.Host = strings.TrimSpace(a.Host)
|
||||
a.IP = strings.ToLower(strings.TrimSpace(a.IP))
|
||||
a.Domain = strings.ToLower(strings.TrimSpace(a.Domain))
|
||||
a.Protocol = strings.ToLower(strings.TrimSpace(a.Protocol))
|
||||
a.Title = strings.TrimSpace(a.Title)
|
||||
a.Server = strings.TrimSpace(a.Server)
|
||||
a.Country = strings.TrimSpace(a.Country)
|
||||
a.Province = strings.TrimSpace(a.Province)
|
||||
a.City = strings.TrimSpace(a.City)
|
||||
a.Source = strings.TrimSpace(a.Source)
|
||||
a.SourceQuery = strings.TrimSpace(a.SourceQuery)
|
||||
a.ProjectID = strings.TrimSpace(a.ProjectID)
|
||||
a.Status = strings.ToLower(strings.TrimSpace(a.Status))
|
||||
if a.Status == "" {
|
||||
a.Status = "active"
|
||||
}
|
||||
if a.Source == "" {
|
||||
a.Source = "manual"
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
tags := make([]string, 0, len(a.Tags))
|
||||
for _, tag := range a.Tags {
|
||||
tag = strings.TrimSpace(tag)
|
||||
if tag != "" && !seen[tag] {
|
||||
seen[tag] = true
|
||||
tags = append(tags, tag)
|
||||
}
|
||||
}
|
||||
a.Tags = tags
|
||||
// URL 型 Host 是常见输入。缺失的结构化字段在服务端同样补齐,确保
|
||||
// API、MCP 与 Web 端产生一致的去重键,而不依赖某个客户端正确解析。
|
||||
if strings.Contains(a.Host, "://") {
|
||||
if parsed, err := url.Parse(a.Host); err == nil && parsed.Hostname() != "" && parsed.User == nil {
|
||||
hostname := strings.Trim(strings.ToLower(parsed.Hostname()), "[]")
|
||||
if net.ParseIP(hostname) != nil && a.IP == "" {
|
||||
a.IP = hostname
|
||||
} else if a.Domain == "" {
|
||||
if ascii, err := idna.Lookup.ToASCII(hostname); err == nil {
|
||||
a.Domain = strings.ToLower(ascii)
|
||||
}
|
||||
}
|
||||
if a.Protocol == "" {
|
||||
a.Protocol = strings.ToLower(parsed.Scheme)
|
||||
}
|
||||
if a.Port == 0 {
|
||||
if parsed.Port() != "" {
|
||||
a.Port, _ = strconv.Atoi(parsed.Port())
|
||||
} else if a.Protocol == "https" {
|
||||
a.Port = 443
|
||||
} else if a.Protocol == "http" {
|
||||
a.Port = 80
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Recon providers occasionally return placeholders, multiple values, or
|
||||
// provider-specific identifiers in structured fields. They are optional
|
||||
// enrichment; a valid Host must not make the entire batch fail because one
|
||||
// of those fields is dirty.
|
||||
if strings.EqualFold(a.Source, "fofa") {
|
||||
if a.IP != "" && net.ParseIP(strings.Trim(a.IP, "[]")) == nil {
|
||||
a.IP = ""
|
||||
}
|
||||
if a.Domain != "" {
|
||||
ascii, err := idna.Lookup.ToASCII(strings.TrimSuffix(a.Domain, "."))
|
||||
if err != nil || !validAssetDomain(ascii) {
|
||||
a.Domain = ""
|
||||
} else {
|
||||
a.Domain = strings.ToLower(ascii)
|
||||
}
|
||||
}
|
||||
if a.Protocol != "" && !assetProtocolPattern.MatchString(a.Protocol) {
|
||||
a.Protocol = ""
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var assetProtocolPattern = regexp.MustCompile(`^[a-z][a-z0-9+.-]{0,31}$`)
|
||||
|
||||
// AssetValidationError distinguishes user-correctable asset data from storage failures.
|
||||
type AssetValidationError struct{ Message string }
|
||||
|
||||
func (e *AssetValidationError) Error() string { return e.Message }
|
||||
|
||||
func assetValidationErrorf(format string, args ...interface{}) error {
|
||||
return &AssetValidationError{Message: fmt.Sprintf(format, args...)}
|
||||
}
|
||||
|
||||
func validateAsset(a *Asset) error {
|
||||
if a == nil {
|
||||
return assetValidationErrorf("资产不能为空")
|
||||
}
|
||||
if a.Host == "" && a.IP == "" && a.Domain == "" {
|
||||
return assetValidationErrorf("资产目标不能为空")
|
||||
}
|
||||
if a.Port < 0 || a.Port > 65535 {
|
||||
return assetValidationErrorf("端口必须在 0-65535 之间")
|
||||
}
|
||||
if a.IP != "" && net.ParseIP(strings.Trim(a.IP, "[]")) == nil {
|
||||
return assetValidationErrorf("IP 地址格式无效")
|
||||
}
|
||||
if a.Domain != "" {
|
||||
ascii, err := idna.Lookup.ToASCII(strings.TrimSuffix(a.Domain, "."))
|
||||
if err != nil || !validAssetDomain(ascii) {
|
||||
return assetValidationErrorf("域名格式无效")
|
||||
}
|
||||
a.Domain = strings.ToLower(ascii)
|
||||
}
|
||||
if a.Protocol != "" && !assetProtocolPattern.MatchString(a.Protocol) {
|
||||
return assetValidationErrorf("协议格式无效")
|
||||
}
|
||||
if a.Status != "active" && a.Status != "inactive" {
|
||||
return assetValidationErrorf("资产状态必须为 active 或 inactive")
|
||||
}
|
||||
for name, value := range map[string]string{
|
||||
"Host": a.Host, "域名": a.Domain, "协议": a.Protocol, "页面标题": a.Title,
|
||||
"服务指纹": a.Server, "国家/地区": a.Country, "省份/州": a.Province, "城市": a.City,
|
||||
} {
|
||||
limit := 255
|
||||
if name == "Host" || name == "页面标题" {
|
||||
limit = 500
|
||||
}
|
||||
if utf8.RuneCountInString(value) > limit {
|
||||
return assetValidationErrorf("%s不能超过 %d 个字符", name, limit)
|
||||
}
|
||||
}
|
||||
if len(a.Tags) > 30 {
|
||||
return assetValidationErrorf("标签不能超过 30 个")
|
||||
}
|
||||
for _, tag := range a.Tags {
|
||||
if utf8.RuneCountInString(tag) > 64 {
|
||||
return assetValidationErrorf("单个标签不能超过 64 个字符")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validAssetDomain(domain string) bool {
|
||||
domain = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(domain)), ".")
|
||||
if domain == "" || len(domain) > 253 || net.ParseIP(domain) != nil {
|
||||
return false
|
||||
}
|
||||
for _, label := range strings.Split(domain, ".") {
|
||||
if len(label) == 0 || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' {
|
||||
return false
|
||||
}
|
||||
for _, r := range label {
|
||||
if (r < 'a' || r > 'z') && (r < '0' || r > '9') && r != '-' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func assetDedupKey(a *Asset) string {
|
||||
target := a.Domain
|
||||
if target == "" {
|
||||
target = a.IP
|
||||
}
|
||||
if target == "" {
|
||||
target = strings.ToLower(a.Host)
|
||||
}
|
||||
return strings.Join([]string{target, strconv.Itoa(a.Port), a.Protocol}, "|")
|
||||
}
|
||||
|
||||
func appendAssetAccess(query string, args []interface{}, access RBACListAccess, alias string) (string, []interface{}) {
|
||||
if strings.TrimSpace(access.UserID) == "" || access.Scope == RBACScopeAll {
|
||||
return query, args
|
||||
}
|
||||
prefix := ""
|
||||
if alias != "" {
|
||||
prefix = alias + "."
|
||||
}
|
||||
query += ` AND (` + prefix + `owner_user_id = ? OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'asset' AND ra.resource_id = ` + prefix + `id
|
||||
) OR (` + prefix + `project_id IS NOT NULL AND ` + prefix + `project_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM projects ap WHERE ap.id=` + prefix + `project_id AND ap.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=` + prefix + `project_id)
|
||||
)))`
|
||||
return query, append(args, access.UserID, access.UserID, access.UserID, access.UserID)
|
||||
}
|
||||
|
||||
func (db *DB) UpsertAssets(assets []*Asset, ownerUserID string, allowGlobal ...bool) (AssetImportResult, error) {
|
||||
result := AssetImportResult{}
|
||||
tx, err := db.Begin()
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
now := time.Now()
|
||||
for _, asset := range assets {
|
||||
if asset == nil {
|
||||
result.Skipped++
|
||||
continue
|
||||
}
|
||||
normalizeAsset(asset)
|
||||
if err := validateAsset(asset); err != nil {
|
||||
return result, fmt.Errorf("第 %d 个资产无效: %w", result.Created+result.Updated+result.Skipped+1, err)
|
||||
}
|
||||
key := assetDedupKey(asset)
|
||||
if key == "|0|" {
|
||||
result.Skipped++
|
||||
continue
|
||||
}
|
||||
var existingID string
|
||||
var existingOwner sql.NullString
|
||||
err := tx.QueryRow(`SELECT id,owner_user_id FROM assets WHERE dedup_key = ?`, key).Scan(&existingID, &existingOwner)
|
||||
tagsJSON, _ := json.Marshal(asset.Tags)
|
||||
if err == sql.ErrNoRows {
|
||||
asset.ID = uuid.NewString()
|
||||
asset.FirstSeenAt, asset.LastSeenAt, asset.CreatedAt, asset.UpdatedAt = now, now, now, now
|
||||
_, err = tx.Exec(`INSERT INTO assets (
|
||||
id,dedup_key,project_id,host,ip,port,domain,protocol,title,server,country,province,city,source,source_query,status,tags_json,
|
||||
first_seen_at,last_seen_at,created_at,updated_at,owner_user_id
|
||||
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
|
||||
asset.ID, key, nullIfEmpty(asset.ProjectID), asset.Host, asset.IP, asset.Port, asset.Domain, asset.Protocol, asset.Title, asset.Server,
|
||||
asset.Country, asset.Province, asset.City, asset.Source, asset.SourceQuery, asset.Status, string(tagsJSON),
|
||||
now, now, now, now, nullIfEmpty(ownerUserID))
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("创建资产失败: %w", err)
|
||||
}
|
||||
if ownerUserID != "" {
|
||||
if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_resource_assignments (id,user_id,resource_type,resource_id,created_at) SELECT ?,id,?,?,? FROM rbac_users WHERE id=?`, uuid.NewString(), "asset", asset.ID, now, ownerUserID); err != nil {
|
||||
return result, fmt.Errorf("授权新资产失败: %w", err)
|
||||
}
|
||||
}
|
||||
result.Created++
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("检查资产去重键失败: %w", err)
|
||||
}
|
||||
asset.ID = existingID
|
||||
global := len(allowGlobal) > 0 && allowGlobal[0]
|
||||
if !global && existingOwner.Valid && strings.TrimSpace(existingOwner.String) != "" && strings.TrimSpace(existingOwner.String) != strings.TrimSpace(ownerUserID) {
|
||||
result.Skipped++
|
||||
continue
|
||||
}
|
||||
_, err = tx.Exec(`UPDATE assets SET
|
||||
host=CASE WHEN ?<>'' THEN ? ELSE host END, ip=CASE WHEN ?<>'' THEN ? ELSE ip END,
|
||||
domain=CASE WHEN ?<>'' THEN ? ELSE domain END, protocol=CASE WHEN ?<>'' THEN ? ELSE protocol END,
|
||||
title=CASE WHEN ?<>'' THEN ? ELSE title END, server=CASE WHEN ?<>'' THEN ? ELSE server END,
|
||||
country=CASE WHEN ?<>'' THEN ? ELSE country END, province=CASE WHEN ?<>'' THEN ? ELSE province END,
|
||||
city=CASE WHEN ?<>'' THEN ? ELSE city END, source=CASE WHEN ?<>'' THEN ? ELSE source END,
|
||||
source_query=CASE WHEN ?<>'' THEN ? ELSE source_query END, project_id=CASE WHEN ?<>'' THEN ? ELSE project_id END,
|
||||
tags_json=CASE WHEN ?<>'[]' THEN ? ELSE tags_json END,
|
||||
last_seen_at=?, updated_at=? WHERE id=?`,
|
||||
asset.Host, asset.Host, asset.IP, asset.IP, asset.Domain, asset.Domain, asset.Protocol, asset.Protocol,
|
||||
asset.Title, asset.Title, asset.Server, asset.Server, asset.Country, asset.Country, asset.Province, asset.Province,
|
||||
asset.City, asset.City, asset.Source, asset.Source, asset.SourceQuery, asset.SourceQuery, asset.ProjectID, nullIfEmpty(asset.ProjectID), string(tagsJSON), string(tagsJSON),
|
||||
now, now, existingID)
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("更新资产失败: %w", err)
|
||||
}
|
||||
if ownerUserID != "" && (!existingOwner.Valid || strings.TrimSpace(existingOwner.String) == ownerUserID) {
|
||||
if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_resource_assignments (id,user_id,resource_type,resource_id,created_at) SELECT ?,id,?,?,? FROM rbac_users WHERE id=?`, uuid.NewString(), "asset", existingID, now, ownerUserID); err != nil {
|
||||
return result, fmt.Errorf("授权资产失败: %w", err)
|
||||
}
|
||||
}
|
||||
result.Updated++
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return result, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func assetWhere(filter AssetListFilter, access RBACListAccess) (string, []interface{}) {
|
||||
query := " WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
if q := strings.TrimSpace(filter.Search); q != "" {
|
||||
pattern := "%" + escapeAssetLike(strings.ToLower(q)) + "%"
|
||||
query += ` AND (LOWER(host) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\'
|
||||
OR LOWER(title) LIKE ? ESCAPE '\' OR LOWER(server) LIKE ? ESCAPE '\' OR LOWER(tags_json) LIKE ? ESCAPE '\')`
|
||||
for i := 0; i < 6; i++ {
|
||||
args = append(args, pattern)
|
||||
}
|
||||
}
|
||||
if filter.Status != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, filter.Status)
|
||||
}
|
||||
if filter.Protocol != "" {
|
||||
query += " AND protocol = ?"
|
||||
args = append(args, filter.Protocol)
|
||||
}
|
||||
if filter.ProjectID != "" {
|
||||
query += " AND assets.project_id = ?"
|
||||
args = append(args, filter.ProjectID)
|
||||
}
|
||||
if filter.Source != "" {
|
||||
query += " AND LOWER(assets.source) = LOWER(?)"
|
||||
args = append(args, strings.TrimSpace(filter.Source))
|
||||
}
|
||||
if tag := strings.TrimSpace(filter.Tag); tag != "" {
|
||||
pattern := "%\"" + escapeAssetLike(strings.ToLower(tag)) + "\"%"
|
||||
query += ` AND LOWER(assets.tags_json) LIKE ? ESCAPE '\'`
|
||||
args = append(args, pattern)
|
||||
}
|
||||
if filter.Host != "" {
|
||||
query += " AND LOWER(assets.host) = LOWER(?)"
|
||||
args = append(args, strings.TrimSpace(filter.Host))
|
||||
}
|
||||
if filter.IP != "" {
|
||||
query += " AND LOWER(assets.ip) = LOWER(?)"
|
||||
args = append(args, strings.TrimSpace(filter.IP))
|
||||
}
|
||||
if filter.Domain != "" {
|
||||
query += " AND LOWER(assets.domain) = LOWER(?)"
|
||||
args = append(args, strings.TrimSpace(filter.Domain))
|
||||
}
|
||||
if filter.Port != nil {
|
||||
query += " AND assets.port = ?"
|
||||
args = append(args, *filter.Port)
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(filter.ScanState)) {
|
||||
case "never":
|
||||
query += " AND " + assetEffectiveLastScanExpr + " IS NULL"
|
||||
case "scanned":
|
||||
query += " AND " + assetEffectiveLastScanExpr + " IS NOT NULL"
|
||||
}
|
||||
if filter.LastScanBefore != nil {
|
||||
query += " AND " + assetEffectiveLastScanExpr + " < ?"
|
||||
args = append(args, *filter.LastScanBefore)
|
||||
}
|
||||
if filter.LastScanAfter != nil {
|
||||
query += " AND " + assetEffectiveLastScanExpr + " > ?"
|
||||
args = append(args, *filter.LastScanAfter)
|
||||
}
|
||||
if filter.LastSeenBefore != nil {
|
||||
query += " AND assets.last_seen_at < ?"
|
||||
args = append(args, *filter.LastSeenBefore)
|
||||
}
|
||||
if filter.LastSeenAfter != nil {
|
||||
query += " AND assets.last_seen_at > ?"
|
||||
args = append(args, *filter.LastSeenAfter)
|
||||
}
|
||||
return appendAssetAccess(query, args, access, "assets")
|
||||
}
|
||||
|
||||
func escapeAssetLike(value string) string {
|
||||
value = strings.ReplaceAll(value, `\`, `\\`)
|
||||
value = strings.ReplaceAll(value, `%`, `\%`)
|
||||
return strings.ReplaceAll(value, `_`, `\_`)
|
||||
}
|
||||
|
||||
func scanAsset(scanner interface{ Scan(...interface{}) error }) (*Asset, error) {
|
||||
var a Asset
|
||||
var tags string
|
||||
var lastScanAt interface{}
|
||||
err := scanner.Scan(&a.ID, &a.ProjectID, &a.ProjectName, &a.Host, &a.IP, &a.Port, &a.Domain, &a.Protocol, &a.Title, &a.Server, &a.Country,
|
||||
&a.Province, &a.City, &a.Source, &a.SourceQuery, &a.Status, &tags, &a.FirstSeenAt, &a.LastSeenAt, &a.CreatedAt, &a.UpdatedAt,
|
||||
&lastScanAt, &a.LastScanConversationID, &a.LastScanQueueID, &a.LastScanTaskID, &a.VulnerabilityCount, &a.RiskLevel)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if parsed, ok := parseAssetScanTime(lastScanAt); ok {
|
||||
a.LastScanAt = &parsed
|
||||
}
|
||||
_ = json.Unmarshal([]byte(tags), &a.Tags)
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
func parseAssetScanTime(value interface{}) (time.Time, bool) {
|
||||
if value == nil {
|
||||
return time.Time{}, false
|
||||
}
|
||||
if parsed, ok := value.(time.Time); ok {
|
||||
return parsed, true
|
||||
}
|
||||
var raw string
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
raw = typed
|
||||
case []byte:
|
||||
raw = string(typed)
|
||||
default:
|
||||
raw = fmt.Sprint(typed)
|
||||
}
|
||||
for _, layout := range []string{
|
||||
time.RFC3339Nano,
|
||||
"2006-01-02 15:04:05.999999999-07:00",
|
||||
"2006-01-02 15:04:05.999999999Z07:00",
|
||||
"2006-01-02 15:04:05-07:00",
|
||||
"2006-01-02 15:04:05",
|
||||
} {
|
||||
if parsed, err := time.Parse(layout, strings.TrimSpace(raw)); err == nil {
|
||||
return parsed, true
|
||||
}
|
||||
}
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
const assetEffectiveLastScanExpr = `COALESCE(
|
||||
(SELECT bt.completed_at FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.completed_at IS NOT NULL LIMIT 1),
|
||||
(SELECT MAX(m.updated_at) FROM messages m WHERE m.conversation_id=assets.last_scan_conversation_id AND m.role='assistant'),
|
||||
assets.last_scan_at
|
||||
)`
|
||||
|
||||
const assetSelectColumns = `assets.id,COALESCE(assets.project_id,''),COALESCE(p.name,''),assets.host,assets.ip,assets.port,assets.domain,assets.protocol,assets.title,assets.server,assets.country,
|
||||
assets.province,assets.city,assets.source,assets.source_query,assets.status,assets.tags_json,assets.first_seen_at,assets.last_seen_at,assets.created_at,assets.updated_at,
|
||||
` + assetEffectiveLastScanExpr + `,COALESCE(assets.last_scan_conversation_id,''),COALESCE(assets.last_scan_queue_id,''),COALESCE(assets.last_scan_task_id,''),
|
||||
(SELECT COUNT(DISTINCT v.id) FROM vulnerabilities v WHERE
|
||||
(COALESCE(assets.last_scan_conversation_id,'')<>'' AND v.conversation_id=assets.last_scan_conversation_id)
|
||||
OR (COALESCE(assets.last_scan_task_id,'')<>'' AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=v.conversation_id))),
|
||||
CASE WHEN assets.last_scan_at IS NULL THEN 'unassessed' ELSE CASE COALESCE((
|
||||
SELECT MAX(CASE LOWER(COALESCE(v.severity,'')) WHEN 'critical' THEN 5 WHEN 'high' THEN 4 WHEN 'medium' THEN 3 WHEN 'low' THEN 2 WHEN 'info' THEN 1 ELSE 0 END)
|
||||
FROM vulnerabilities v WHERE
|
||||
LOWER(COALESCE(v.status,'open')) NOT IN ('fixed','false_positive','ignored') AND (
|
||||
(COALESCE(assets.last_scan_conversation_id,'')<>'' AND v.conversation_id=assets.last_scan_conversation_id)
|
||||
OR (COALESCE(assets.last_scan_task_id,'')<>'' AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=v.conversation_id))
|
||||
)
|
||||
),0) WHEN 5 THEN 'critical' WHEN 4 THEN 'high' WHEN 3 THEN 'medium' WHEN 2 THEN 'low' WHEN 1 THEN 'info' ELSE 'normal' END END`
|
||||
|
||||
// MarkAssetScanned links an asset to the conversation or batch subtask created from it.
|
||||
// The link lets the asset list show the latest scan time and vulnerabilities produced by that scan.
|
||||
func (db *DB) MarkAssetScanned(id, conversationID, queueID, taskID string, access RBACListAccess) error {
|
||||
where, args := appendAssetAccess(" WHERE id = ?", []interface{}{strings.TrimSpace(id)}, access, "assets")
|
||||
res, err := db.Exec(`UPDATE assets SET last_scan_at=?,last_scan_conversation_id=?,last_scan_queue_id=?,last_scan_task_id=?,updated_at=?`+where,
|
||||
append([]interface{}{time.Now(), strings.TrimSpace(conversationID), strings.TrimSpace(queueID), strings.TrimSpace(taskID), time.Now()}, args...)...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
if n == 0 {
|
||||
return sql.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CompleteAssetScan records completion from inside an Agent conversation. If
|
||||
// the asset was launched as a batch task, keep its task/queue link only when
|
||||
// that task belongs to the current conversation; a later ad-hoc chat scan must
|
||||
// not retain stale task associations.
|
||||
func (db *DB) CompleteAssetScan(id, conversationID string, access RBACListAccess) error {
|
||||
id = strings.TrimSpace(id)
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
if conversationID == "" {
|
||||
return fmt.Errorf("扫描对话不能为空")
|
||||
}
|
||||
where, args := appendAssetAccess(" WHERE id = ?", []interface{}{id}, access, "assets")
|
||||
now := time.Now()
|
||||
res, err := db.Exec(`UPDATE assets SET
|
||||
last_scan_at=?,last_scan_conversation_id=?,
|
||||
last_scan_queue_id=CASE WHEN EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=?) THEN last_scan_queue_id ELSE '' END,
|
||||
last_scan_task_id=CASE WHEN EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=?) THEN last_scan_task_id ELSE '' END,
|
||||
updated_at=?`+where,
|
||||
append([]interface{}{now, conversationID, conversationID, conversationID, now}, args...)...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
if n == 0 {
|
||||
return sql.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *DB) BatchTaskBelongsToQueue(taskID, queueID string) bool {
|
||||
var count int
|
||||
err := db.QueryRow(`SELECT COUNT(*) FROM batch_tasks WHERE id=? AND queue_id=?`, strings.TrimSpace(taskID), strings.TrimSpace(queueID)).Scan(&count)
|
||||
return err == nil && count > 0
|
||||
}
|
||||
|
||||
func (db *DB) ListAssets(limit, offset int, filter AssetListFilter, access RBACListAccess) ([]*Asset, int, error) {
|
||||
if limit < 1 {
|
||||
limit = 20
|
||||
}
|
||||
if limit > 100 {
|
||||
limit = 100
|
||||
}
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
where, args := assetWhere(filter, access)
|
||||
var total int
|
||||
if err := db.QueryRow("SELECT COUNT(*) FROM assets"+where, args...).Scan(&total); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
orderBy := assetOrderBy(filter.SortBy, filter.SortOrder)
|
||||
rows, err := db.Query("SELECT "+assetSelectColumns+" FROM assets LEFT JOIN projects p ON p.id=assets.project_id"+where+" ORDER BY "+orderBy+" LIMIT ? OFFSET ?", append(args, limit, offset)...)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []*Asset{}
|
||||
for rows.Next() {
|
||||
a, err := scanAsset(rows)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
items = append(items, a)
|
||||
}
|
||||
return items, total, rows.Err()
|
||||
}
|
||||
|
||||
func assetOrderBy(sortBy, sortOrder string) string {
|
||||
direction := "DESC"
|
||||
if strings.EqualFold(strings.TrimSpace(sortOrder), "asc") {
|
||||
direction = "ASC"
|
||||
}
|
||||
var expression string
|
||||
switch strings.ToLower(strings.TrimSpace(sortBy)) {
|
||||
case "last_scan_at":
|
||||
expression = assetEffectiveLastScanExpr
|
||||
// For oldest-first queries, assets that have never been scanned are the
|
||||
// most overdue and intentionally appear first. NULLs stay last for DESC.
|
||||
if direction == "ASC" {
|
||||
return "CASE WHEN " + expression + " IS NULL THEN 0 ELSE 1 END ASC, " + expression + " ASC, assets.id ASC"
|
||||
}
|
||||
return "CASE WHEN " + expression + " IS NULL THEN 1 ELSE 0 END ASC, " + expression + " DESC, assets.id ASC"
|
||||
case "first_seen_at":
|
||||
expression = "assets.first_seen_at"
|
||||
case "created_at":
|
||||
expression = "assets.created_at"
|
||||
case "updated_at":
|
||||
expression = "assets.updated_at"
|
||||
case "host":
|
||||
expression = "LOWER(assets.host)"
|
||||
case "port":
|
||||
expression = "assets.port"
|
||||
default:
|
||||
expression = "assets.last_seen_at"
|
||||
}
|
||||
return expression + " " + direction + ", assets.id ASC"
|
||||
}
|
||||
|
||||
func (db *DB) GetAsset(id string, access RBACListAccess) (*Asset, error) {
|
||||
query, args := appendAssetAccess("SELECT "+assetSelectColumns+" FROM assets LEFT JOIN projects p ON p.id=assets.project_id WHERE assets.id = ?", []interface{}{id}, access, "assets")
|
||||
return scanAsset(db.QueryRow(query, args...))
|
||||
}
|
||||
|
||||
func (db *DB) UpdateAsset(id string, a *Asset, access RBACListAccess) error {
|
||||
normalizeAsset(a)
|
||||
if err := validateAsset(a); err != nil {
|
||||
return err
|
||||
}
|
||||
key := assetDedupKey(a)
|
||||
if key == "|0|" {
|
||||
return fmt.Errorf("资产目标不能为空")
|
||||
}
|
||||
tags, _ := json.Marshal(a.Tags)
|
||||
where, args := appendAssetAccess(" WHERE id = ?", []interface{}{id}, access, "assets")
|
||||
res, err := db.Exec(`UPDATE assets SET dedup_key=?,project_id=?,host=?,ip=?,port=?,domain=?,protocol=?,title=?,server=?,country=?,province=?,city=?,source=?,source_query=?,status=?,tags_json=?,updated_at=?`+where,
|
||||
append([]interface{}{key, nullIfEmpty(a.ProjectID), a.Host, a.IP, a.Port, a.Domain, a.Protocol, a.Title, a.Server, a.Country, a.Province, a.City, a.Source, a.SourceQuery, a.Status, string(tags), time.Now()}, args...)...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
if n == 0 {
|
||||
return sql.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateAssetsProject atomically replaces the project binding for every asset.
|
||||
// It refuses the whole update when any requested asset is missing or outside
|
||||
// the caller's access scope, so a bulk action can never partially succeed.
|
||||
func (db *DB) UpdateAssetsProject(ids []string, projectID string, access RBACListAccess) (int, error) {
|
||||
unique := make([]string, 0, len(ids))
|
||||
seen := make(map[string]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[id]; exists {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
unique = append(unique, id)
|
||||
}
|
||||
if len(unique) == 0 {
|
||||
return 0, fmt.Errorf("资产列表不能为空")
|
||||
}
|
||||
|
||||
tx, err := db.Begin()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(unique)), ",")
|
||||
idArgs := make([]interface{}, len(unique))
|
||||
for i, id := range unique {
|
||||
idArgs[i] = id
|
||||
}
|
||||
countQuery, countArgs := appendAssetAccess("SELECT COUNT(*) FROM assets WHERE id IN ("+placeholders+")", idArgs, access, "assets")
|
||||
var accessible int
|
||||
if err := tx.QueryRow(countQuery, countArgs...).Scan(&accessible); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if accessible != len(unique) {
|
||||
return 0, fmt.Errorf("部分资产不存在或无权更新")
|
||||
}
|
||||
|
||||
updateArgs := []interface{}{nullIfEmpty(strings.TrimSpace(projectID)), time.Now()}
|
||||
updateArgs = append(updateArgs, idArgs...)
|
||||
updateQuery, updateArgs := appendAssetAccess("UPDATE assets SET project_id=?,updated_at=? WHERE id IN ("+placeholders+")", updateArgs, access, "assets")
|
||||
result, err := tx.Exec(updateQuery, updateArgs...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
updated, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if int(updated) != len(unique) {
|
||||
return 0, fmt.Errorf("批量更新资产失败")
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return int(updated), nil
|
||||
}
|
||||
|
||||
func (db *DB) DeleteAsset(id string, access RBACListAccess) error {
|
||||
where, args := appendAssetAccess(" WHERE id = ?", []interface{}{id}, access, "assets")
|
||||
res, err := db.Exec("DELETE FROM assets"+where, args...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
if n == 0 {
|
||||
return sql.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *DB) GetAssetStats(access RBACListAccess, requestedDays ...int) (map[string]interface{}, error) {
|
||||
days := 30
|
||||
if len(requestedDays) > 0 && (requestedDays[0] == 7 || requestedDays[0] == 30 || requestedDays[0] == 90) {
|
||||
days = requestedDays[0]
|
||||
}
|
||||
where, args := appendAssetAccess(" WHERE 1=1", nil, access, "assets")
|
||||
stats := map[string]interface{}{}
|
||||
row := db.QueryRow(`SELECT COUNT(*),COUNT(DISTINCT NULLIF(ip,'')),COUNT(DISTINCT NULLIF(domain,'')),
|
||||
COUNT(DISTINCT CASE WHEN port>0 THEN CAST(port AS TEXT) END),
|
||||
COALESCE(SUM(CASE WHEN datetime(last_seen_at)>=datetime('now','-7 days') THEN 1 ELSE 0 END),0) FROM assets`+where, args...)
|
||||
var total, ips, domains, ports, recent int
|
||||
if err := row.Scan(&total, &ips, &domains, &ports, &recent); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stats["total"], stats["ips"], stats["domains"], stats["ports"], stats["recent"] = total, ips, domains, ports, recent
|
||||
rows, err := db.Query(`SELECT CASE WHEN protocol='' THEN 'unknown' ELSE protocol END,COUNT(*) FROM assets`+where+` GROUP BY protocol ORDER BY COUNT(*) DESC LIMIT 8`, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
dist := []map[string]interface{}{}
|
||||
for rows.Next() {
|
||||
var name string
|
||||
var count int
|
||||
if err := rows.Scan(&name, &count); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dist = append(dist, map[string]interface{}{"name": name, "count": count})
|
||||
}
|
||||
stats["protocols"] = dist
|
||||
stats["period_days"] = days
|
||||
|
||||
coverage := map[string]interface{}{}
|
||||
coverageRow := db.QueryRow(`SELECT
|
||||
COALESCE(SUM(CASE WHEN last_scan_at IS NOT NULL THEN 1 ELSE 0 END),0),
|
||||
COALESCE(SUM(CASE WHEN datetime(last_scan_at)>=datetime('now','-7 days') THEN 1 ELSE 0 END),0),
|
||||
COALESCE(SUM(CASE WHEN datetime(last_scan_at)>=datetime('now','-30 days') THEN 1 ELSE 0 END),0),
|
||||
COALESCE(SUM(CASE WHEN last_scan_at IS NULL THEN 1 ELSE 0 END),0),
|
||||
COALESCE(SUM(CASE WHEN last_scan_at IS NOT NULL AND datetime(last_scan_at)<datetime('now','-30 days') THEN 1 ELSE 0 END),0),
|
||||
COALESCE(SUM(CASE WHEN status='active' THEN 1 ELSE 0 END),0)
|
||||
FROM assets`+where, args...)
|
||||
var scanned, scanned7, scanned30, neverScanned, stale, active int
|
||||
if err := coverageRow.Scan(&scanned, &scanned7, &scanned30, &neverScanned, &stale, &active); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
coverage["scanned"], coverage["scanned_7d"], coverage["scanned_30d"] = scanned, scanned7, scanned30
|
||||
coverage["never_scanned"], coverage["stale"], coverage["active"] = neverScanned, stale, active
|
||||
if total > 0 {
|
||||
coverage["rate"] = int(float64(scanned) / float64(total) * 100)
|
||||
coverage["recent_rate"] = int(float64(scanned30) / float64(total) * 100)
|
||||
} else {
|
||||
coverage["rate"], coverage["recent_rate"] = 0, 0
|
||||
}
|
||||
stats["coverage"] = coverage
|
||||
|
||||
assetDaily := map[string]map[string]int{}
|
||||
trendWhere, trendArgs := appendAssetAccess(" WHERE datetime(first_seen_at)>=datetime('now',?)", []interface{}{fmt.Sprintf("-%d days", days-1)}, access, "assets")
|
||||
trendRows, err := db.Query(`SELECT date(first_seen_at), COUNT(*)
|
||||
FROM assets`+trendWhere+` GROUP BY date(first_seen_at) ORDER BY date(first_seen_at)`, trendArgs...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for trendRows.Next() {
|
||||
var day string
|
||||
var added int
|
||||
if err := trendRows.Scan(&day, &added); err != nil {
|
||||
trendRows.Close()
|
||||
return nil, err
|
||||
}
|
||||
assetDaily[day] = map[string]int{"added": added, "inactive": 0}
|
||||
}
|
||||
if err := trendRows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
inactiveWhere, inactiveArgs := appendAssetAccess(" WHERE status='inactive' AND datetime(updated_at)>=datetime('now',?)", []interface{}{fmt.Sprintf("-%d days", days-1)}, access, "assets")
|
||||
inactiveRows, err := db.Query(`SELECT date(updated_at), COUNT(*) FROM assets`+inactiveWhere+` GROUP BY date(updated_at) ORDER BY date(updated_at)`, inactiveArgs...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for inactiveRows.Next() {
|
||||
var day string
|
||||
var inactive int
|
||||
if err := inactiveRows.Scan(&day, &inactive); err != nil {
|
||||
inactiveRows.Close()
|
||||
return nil, err
|
||||
}
|
||||
if _, ok := assetDaily[day]; !ok {
|
||||
assetDaily[day] = map[string]int{"added": 0, "inactive": 0}
|
||||
}
|
||||
assetDaily[day]["inactive"] = inactive
|
||||
}
|
||||
if err := inactiveRows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
riskDaily := map[string]map[string]int{}
|
||||
riskWhere, riskArgs := appendVulnerabilityAccessFilter(" WHERE datetime(created_at)>=datetime('now',?)", []interface{}{fmt.Sprintf("-%d days", days-1)}, access)
|
||||
riskRows, err := db.Query(`SELECT date(created_at), COUNT(*),
|
||||
COALESCE(SUM(CASE WHEN LOWER(severity) IN ('critical','high') THEN 1 ELSE 0 END),0)
|
||||
FROM vulnerabilities`+riskWhere+` GROUP BY date(created_at) ORDER BY date(created_at)`, riskArgs...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for riskRows.Next() {
|
||||
var day string
|
||||
var discovered, highRisk int
|
||||
if err := riskRows.Scan(&day, &discovered, &highRisk); err != nil {
|
||||
riskRows.Close()
|
||||
return nil, err
|
||||
}
|
||||
riskDaily[day] = map[string]int{"discovered": discovered, "high_risk": highRisk}
|
||||
}
|
||||
if err := riskRows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
assetTrend := make([]map[string]interface{}, 0, days)
|
||||
riskTrend := make([]map[string]interface{}, 0, days)
|
||||
start := time.Now().UTC().Truncate(24*time.Hour).AddDate(0, 0, -(days - 1))
|
||||
for i := 0; i < days; i++ {
|
||||
day := start.AddDate(0, 0, i).Format("2006-01-02")
|
||||
assetPoint := map[string]interface{}{"date": day, "added": 0, "inactive": 0}
|
||||
if values, ok := assetDaily[day]; ok {
|
||||
assetPoint["added"], assetPoint["inactive"] = values["added"], values["inactive"]
|
||||
}
|
||||
assetTrend = append(assetTrend, assetPoint)
|
||||
riskPoint := map[string]interface{}{"date": day, "discovered": 0, "high_risk": 0}
|
||||
if values, ok := riskDaily[day]; ok {
|
||||
riskPoint["discovered"], riskPoint["high_risk"] = values["discovered"], values["high_risk"]
|
||||
}
|
||||
riskTrend = append(riskTrend, riskPoint)
|
||||
}
|
||||
stats["asset_trend"], stats["risk_trend"] = assetTrend, riskTrend
|
||||
return stats, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,316 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestAssetURLNormalizationAndValidation(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "asset-validation.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
asset := &Asset{Host: "https://例子.测试/path", Tags: []string{" prod ", "prod"}}
|
||||
result, err := db.UpsertAssets([]*Asset{asset}, "")
|
||||
if err != nil || result.Created != 1 {
|
||||
t.Fatalf("URL asset was not created: result=%#v err=%v", result, err)
|
||||
}
|
||||
if asset.Domain != "xn--fsqu00a.xn--0zwm56d" || asset.Protocol != "https" || asset.Port != 443 {
|
||||
t.Fatalf("URL fields were not normalized: %#v", asset)
|
||||
}
|
||||
if len(asset.Tags) != 1 || asset.Tags[0] != "prod" {
|
||||
t.Fatalf("tags were not normalized: %#v", asset.Tags)
|
||||
}
|
||||
|
||||
invalid := []*Asset{
|
||||
{IP: "999.1.1.1", Status: "active"},
|
||||
{Domain: "bad_domain.example", Status: "active"},
|
||||
{Domain: "example.com", Port: 70000, Status: "active"},
|
||||
{Domain: "example.com", Protocol: "HTTP 1.1", Status: "active"},
|
||||
{Domain: "example.com", Status: "deleted"},
|
||||
}
|
||||
for _, candidate := range invalid {
|
||||
if _, err := db.UpsertAssets([]*Asset{candidate}, ""); err == nil {
|
||||
t.Fatalf("invalid asset unexpectedly accepted: %#v", candidate)
|
||||
}
|
||||
}
|
||||
|
||||
for _, host := range []string{"123", "not a formal target", "https://", "https://user:password@example.com"} {
|
||||
result, err := db.UpsertAssets([]*Asset{{Host: host}}, "")
|
||||
if err != nil || result.Created != 1 {
|
||||
t.Fatalf("opaque asset address %q was not accepted: result=%#v err=%v", host, result, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetValidationRejectsOversizedTags(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "asset-tag-validation.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
_, err = db.UpsertAssets([]*Asset{{Domain: "example.com", Tags: []string{strings.Repeat("x", 65)}}}, "")
|
||||
if err == nil || !strings.Contains(err.Error(), "标签") {
|
||||
t.Fatalf("expected tag validation error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "fofa-asset-validation.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
asset := &Asset{
|
||||
Host: "https://203.0.113.59:8443",
|
||||
IP: "203.0.113.59",
|
||||
Domain: "provider_specific_invalid_domain_59",
|
||||
Port: 8443,
|
||||
Protocol: "https",
|
||||
Source: "fofa",
|
||||
}
|
||||
result, err := db.UpsertAssets([]*Asset{asset}, "")
|
||||
if err != nil || result.Created != 1 {
|
||||
t.Fatalf("FOFA asset with dirty optional domain was not created: result=%#v err=%v", result, err)
|
||||
}
|
||||
if asset.Domain != "" || asset.IP != "203.0.113.59" {
|
||||
t.Fatalf("FOFA structured fields were not sanitized: %#v", asset)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetUpsertDeduplicatesAndUpdates(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
first := &Asset{Host: "https://example.com", Domain: "Example.COM", Port: 443, Protocol: "HTTPS", Title: "Old", Source: "fofa"}
|
||||
result, err := db.UpsertAssets([]*Asset{first}, "user-a")
|
||||
if err != nil || result.Created != 1 || result.Updated != 0 {
|
||||
t.Fatalf("first upsert = %#v, %v", result, err)
|
||||
}
|
||||
second := &Asset{Domain: "example.com", Port: 443, Protocol: "https", Title: "New", Server: "nginx", Source: "fofa"}
|
||||
result, err = db.UpsertAssets([]*Asset{second}, "user-a")
|
||||
if err != nil || result.Created != 0 || result.Updated != 1 {
|
||||
t.Fatalf("second upsert = %#v, %v", result, err)
|
||||
}
|
||||
assets, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil || total != 1 || len(assets) != 1 {
|
||||
t.Fatalf("list assets total=%d len=%d err=%v", total, len(assets), err)
|
||||
}
|
||||
if assets[0].Title != "New" || assets[0].Server != "nginx" || assets[0].Protocol != "https" {
|
||||
t.Fatalf("asset not refreshed: %#v", assets[0])
|
||||
}
|
||||
stats, err := db.GetAssetStats(RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil || stats["total"] != 1 {
|
||||
t.Fatalf("stats=%#v err=%v", stats, err)
|
||||
}
|
||||
coverage, ok := stats["coverage"].(map[string]interface{})
|
||||
if !ok || coverage["never_scanned"] != 1 || coverage["rate"] != 0 {
|
||||
t.Fatalf("coverage=%#v", stats["coverage"])
|
||||
}
|
||||
assetTrend, ok := stats["asset_trend"].([]map[string]interface{})
|
||||
if !ok || len(assetTrend) != 30 {
|
||||
t.Fatalf("asset trend=%#v", stats["asset_trend"])
|
||||
}
|
||||
riskTrend, ok := stats["risk_trend"].([]map[string]interface{})
|
||||
if !ok || len(riskTrend) != 30 {
|
||||
t.Fatalf("risk trend=%#v", stats["risk_trend"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetAccessFiltersOwners(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "assets-access.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
now := time.Now()
|
||||
if _, err := db.Exec(`INSERT INTO rbac_users (id,username,display_name,password_hash,enabled,is_builtin,created_at,updated_at) VALUES ('user-a','user-a','User A','hash',1,0,?,?)`, now, now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.UpsertAssets([]*Asset{{IP: "10.0.0.1", Port: 80, Protocol: "http"}}, "user-a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-b", Scope: RBACScopeAssigned})
|
||||
if err != nil || total != 0 {
|
||||
t.Fatalf("unexpected cross-user assets: total=%d err=%v", total, err)
|
||||
}
|
||||
_, total, err = db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeOwn})
|
||||
if err != nil || total != 1 {
|
||||
t.Fatalf("owner cannot list asset: total=%d err=%v", total, err)
|
||||
}
|
||||
assets, _, err := db.ListAssets(1, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeAssigned})
|
||||
if err != nil || len(assets) != 1 || !db.UserCanAccessResource("user-a", RBACScopeAssigned, "asset", assets[0].ID) {
|
||||
t.Fatalf("creator assignment missing: assets=%d err=%v", len(assets), err)
|
||||
}
|
||||
options, err := db.ListAssignableRBACResources("asset", "10.0.0.1", 10)
|
||||
if err != nil || len(options) != 1 {
|
||||
t.Fatalf("asset resource picker: options=%#v err=%v", options, err)
|
||||
}
|
||||
project, err := db.CreateProject(&Project{Name: "Alpha", Status: "active"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetResourceOwner("project", project.ID, "user-b"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
asset := assets[0]
|
||||
asset.ProjectID = project.ID
|
||||
if err := db.UpdateAsset(asset.ID, asset, RBACListAccess{Scope: RBACScopeAll}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projectAssets, total, err := db.ListAssets(20, 0, AssetListFilter{ProjectID: project.ID}, RBACListAccess{UserID: "user-b", Scope: RBACScopeOwn})
|
||||
if err != nil || total != 1 || len(projectAssets) != 1 || projectAssets[0].ProjectName != "Alpha" {
|
||||
t.Fatalf("project-bound asset access failed: total=%d assets=%#v err=%v", total, projectAssets, err)
|
||||
}
|
||||
if !db.UserCanAccessResource("user-b", RBACScopeOwn, "asset", asset.ID) {
|
||||
t.Fatal("project owner cannot access bound asset")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateAssetsProjectIsAtomicAndScoped(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "asset-batch-project.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
project, err := db.CreateProject(&Project{Name: "Batch Project", Status: "active"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.UpsertAssets([]*Asset{
|
||||
{IP: "192.0.2.1", Port: 80, Protocol: "http"},
|
||||
{IP: "192.0.2.2", Port: 443, Protocol: "https"},
|
||||
}, "owner-a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil || len(assets) != 2 {
|
||||
t.Fatalf("list assets: len=%d err=%v", len(assets), err)
|
||||
}
|
||||
ids := []string{assets[0].ID, assets[1].ID}
|
||||
updated, err := db.UpdateAssetsProject(ids, project.ID, RBACListAccess{UserID: "owner-a", Scope: RBACScopeOwn})
|
||||
if err != nil || updated != 2 {
|
||||
t.Fatalf("batch bind: updated=%d err=%v", updated, err)
|
||||
}
|
||||
for _, id := range ids {
|
||||
asset, err := db.GetAsset(id, RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil || asset.ProjectID != project.ID {
|
||||
t.Fatalf("asset %s was not bound: asset=%#v err=%v", id, asset, err)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := db.UpdateAssetsProject([]string{ids[0], "missing"}, "", RBACListAccess{Scope: RBACScopeAll}); err == nil {
|
||||
t.Fatal("partial batch update unexpectedly succeeded")
|
||||
}
|
||||
asset, err := db.GetAsset(ids[0], RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil || asset.ProjectID != project.ID {
|
||||
t.Fatalf("failed batch changed an asset: asset=%#v err=%v", asset, err)
|
||||
}
|
||||
|
||||
updated, err = db.UpdateAssetsProject(ids, "", RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil || updated != 2 {
|
||||
t.Fatalf("batch unbind: updated=%d err=%v", updated, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetScanLinkReturnsTimeAndRelatedVulnerabilities(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "asset-scan.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
if _, err := db.UpsertAssets([]*Asset{{IP: "192.0.2.10", Port: 443, Protocol: "https"}}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil || len(assets) != 1 {
|
||||
t.Fatalf("list assets: len=%d err=%v", len(assets), err)
|
||||
}
|
||||
conv, err := db.CreateConversation("asset scan", ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.MarkAssetScanned(assets[0].ID, conv.ID, "", "", RBACListAccess{Scope: RBACScopeAll}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.CreateVulnerability(&Vulnerability{ConversationID: conv.ID, Title: "finding", Severity: "high", Target: "192.0.2.10"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
linked, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if linked.LastScanAt == nil || linked.LastScanConversationID != conv.ID || linked.VulnerabilityCount != 1 || linked.RiskLevel != "high" {
|
||||
t.Fatalf("unexpected scan metadata: %#v", linked)
|
||||
}
|
||||
if _, err := db.Exec(`UPDATE vulnerabilities SET status='fixed' WHERE conversation_id=?`, conv.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resolved, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resolved.VulnerabilityCount != 1 || resolved.RiskLevel != "normal" {
|
||||
t.Fatalf("resolved finding should remain in history without raising current risk: %#v", resolved)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetListFlexibleFiltersAndOldestScanPagination(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "asset-query.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
assets := []*Asset{
|
||||
{IP: "192.0.2.1", Port: 443, Protocol: "https", Source: "fofa", Tags: []string{"prod"}},
|
||||
{IP: "192.0.2.2", Port: 80, Protocol: "http", Source: "manual", Tags: []string{"prod", "legacy"}},
|
||||
{Domain: "never.example.com", Port: 443, Protocol: "https", Source: "manual", Tags: []string{"prod"}},
|
||||
}
|
||||
if _, err := db.UpsertAssets(assets, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
old := time.Now().Add(-90 * 24 * time.Hour).UTC()
|
||||
recent := time.Now().Add(-24 * time.Hour).UTC()
|
||||
if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, old, assets[0].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, recent, assets[1].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
access := RBACListAccess{Scope: RBACScopeAll}
|
||||
firstPage, total, err := db.ListAssets(2, 0, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access)
|
||||
if err != nil || total != 3 || len(firstPage) != 2 {
|
||||
t.Fatalf("oldest scan page: total=%d len=%d err=%v", total, len(firstPage), err)
|
||||
}
|
||||
if firstPage[0].ID != assets[2].ID || firstPage[0].LastScanAt != nil || firstPage[1].ID != assets[0].ID {
|
||||
t.Fatalf("expected never-scanned then oldest scanned asset, got %#v", firstPage)
|
||||
}
|
||||
secondPage, _, err := db.ListAssets(2, 2, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access)
|
||||
if err != nil || len(secondPage) != 1 || secondPage[0].ID != assets[1].ID {
|
||||
t.Fatalf("unexpected second page: %#v err=%v", secondPage, err)
|
||||
}
|
||||
|
||||
never, total, err := db.ListAssets(20, 0, AssetListFilter{ScanState: "never"}, access)
|
||||
if err != nil || total != 1 || len(never) != 1 || never[0].ID != assets[2].ID {
|
||||
t.Fatalf("never-scanned filter: total=%d assets=%#v err=%v", total, never, err)
|
||||
}
|
||||
port := 443
|
||||
filtered, total, err := db.ListAssets(20, 0, AssetListFilter{Source: "fofa", Port: &port, LastScanBefore: &recent}, access)
|
||||
if err != nil || total != 1 || len(filtered) != 1 || filtered[0].ID != assets[0].ID {
|
||||
t.Fatalf("structured filters: total=%d assets=%#v err=%v", total, filtered, err)
|
||||
}
|
||||
}
|
||||
@@ -405,6 +405,19 @@ func (db *DB) initTables() error {
|
||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
|
||||
);`
|
||||
|
||||
createAssetsTable := `
|
||||
CREATE TABLE IF NOT EXISTS assets (
|
||||
id TEXT PRIMARY KEY,
|
||||
dedup_key TEXT NOT NULL UNIQUE, project_id TEXT,
|
||||
host TEXT NOT NULL DEFAULT '', ip TEXT NOT NULL DEFAULT '', port INTEGER NOT NULL DEFAULT 0,
|
||||
domain TEXT NOT NULL DEFAULT '', protocol TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '',
|
||||
server TEXT NOT NULL DEFAULT '', country TEXT NOT NULL DEFAULT '', province TEXT NOT NULL DEFAULT '', city TEXT NOT NULL DEFAULT '',
|
||||
source TEXT NOT NULL DEFAULT 'manual', source_query TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'active',
|
||||
tags_json TEXT NOT NULL DEFAULT '[]', first_seen_at DATETIME NOT NULL, last_seen_at DATETIME NOT NULL,
|
||||
created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, owner_user_id TEXT,
|
||||
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL
|
||||
);`
|
||||
|
||||
createVulnerabilityAlertSubscriptionsTable := `
|
||||
CREATE TABLE IF NOT EXISTS vulnerability_alert_subscriptions (
|
||||
user_id TEXT PRIMARY KEY,
|
||||
@@ -715,6 +728,13 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_last_seen ON assets(last_seen_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_last_scan ON assets(last_scan_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_ip ON assets(ip);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_domain ON assets(domain);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_status ON assets(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_owner ON assets(owner_user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
|
||||
@@ -823,6 +843,12 @@ func (db *DB) initTables() error {
|
||||
if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
|
||||
return fmt.Errorf("创建vulnerabilities表失败: %w", err)
|
||||
}
|
||||
if _, err := db.Exec(createAssetsTable); err != nil {
|
||||
return fmt.Errorf("创建assets表失败: %w", err)
|
||||
}
|
||||
if err := db.migrateAssetsTable(); err != nil {
|
||||
return fmt.Errorf("迁移assets表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createBatchTaskQueuesTable); err != nil {
|
||||
return fmt.Errorf("创建batch_task_queues表失败: %w", err)
|
||||
@@ -950,6 +976,32 @@ func (db *DB) migrateRobotUserSessionsTable() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateAssetsTable keeps databases created by the first asset-management release compatible.
|
||||
func (db *DB) migrateAssetsTable() error {
|
||||
columns := []struct {
|
||||
name string
|
||||
ddl string
|
||||
}{
|
||||
{"project_id", "ALTER TABLE assets ADD COLUMN project_id TEXT"},
|
||||
{"last_scan_at", "ALTER TABLE assets ADD COLUMN last_scan_at DATETIME"},
|
||||
{"last_scan_conversation_id", "ALTER TABLE assets ADD COLUMN last_scan_conversation_id TEXT NOT NULL DEFAULT ''"},
|
||||
{"last_scan_queue_id", "ALTER TABLE assets ADD COLUMN last_scan_queue_id TEXT NOT NULL DEFAULT ''"},
|
||||
{"last_scan_task_id", "ALTER TABLE assets ADD COLUMN last_scan_task_id TEXT NOT NULL DEFAULT ''"},
|
||||
}
|
||||
for _, column := range columns {
|
||||
var count int
|
||||
if err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('assets') WHERE name=?", column.name).Scan(&count); err != nil {
|
||||
return err
|
||||
}
|
||||
if count == 0 {
|
||||
if _, err := db.Exec(column.ddl); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateMessagesTable 迁移 messages 表,补充 updated_at 字段。
|
||||
// 语义:updated_at 表示该条消息最后一次被写入/更新的时间(例如助手占位消息在任务结束时更新正文)。
|
||||
func (db *DB) migrateMessagesTable() error {
|
||||
|
||||
@@ -268,6 +268,9 @@ func (db *DB) DeleteProject(id string) error {
|
||||
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
||||
}
|
||||
if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||
return fmt.Errorf("解除资产项目关联失败: %w", err)
|
||||
}
|
||||
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除项目失败: %w", err)
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user