First Commit

db/finding_assets.go

@@ -0,0 +1,676 @@
+package db
+
+import (
+	"encoding/json"
+	"sort"
+	"strconv"
+	"strings"
+	"time"
+)
+
+// ---------------------------------------------------------------------------
+// Finding asset tree — the「按资产」view of the global findings list.
+//
+// 层级与 BuildCoverageGraph 同源(company → root_domain/ip/app → subdomain →
+// service → endpoint),但那份是「某个任务的范围内资产」的力导向图,这份是
+// 「全库有发现的资产」的树:只收有发现的资产及其祖先链,节点带子树聚合计数。
+// 两者的父子优先级规则必须保持一致,改一处请对照 task_scope.go 改另一处。
+// ---------------------------------------------------------------------------
+
+// FindingUnassignedAsset 是「未关联资产」的节点 key,也是列表接口的筛选哨兵:
+// 命中 asset_ids 为空、或所指资产已被删除的发现。
+const FindingUnassignedAsset = "__none__"
+
+// findingAssetTreeMaxNodes 是返回给前端的节点上限。超出时自底向上丢弃整层
+// (endpoint 优先,其次 service):它们的计数已经累加进父节点,丢节点不丢数字。
+const findingAssetTreeMaxNodes = 3000
+
+// FindingAssetNode 是资产树的一个节点。Key 与覆盖图同构:资产行是 "a:<id>"、
+// 企业是 "c:<id>"、没有资产行的根域名是合成的 "r:<domain>"、未关联桶是 "__none__"。
+type FindingAssetNode struct {
+	Key       string `json:"key"`
+	Parent    string `json:"parent,omitempty"`
+	Kind      string `json:"kind"` // company|root_domain|subdomain|ip|service|app|endpoint|none
+	Label     string `json:"label"`
+	AssetID   int64  `json:"asset_id,omitempty"`
+	CompanyID int64  `json:"company_id,omitempty"`
+	// Self 是直接挂在该资产上的发现数;Total 含全部子孙且按 finding 去重
+	// (一个发现挂多个资产时,只在共同祖先上计一次)。
+	Self        int       `json:"self"`
+	Total       int       `json:"total"`
+	Critical    int       `json:"critical"`
+	High        int       `json:"high"`
+	Medium      int       `json:"medium"`
+	Low         int       `json:"low"`
+	LastFoundAt time.Time `json:"last_found_at"`
+}
+
+// FindingAssetTree 是整棵树的一次性快照。Nodes 已排好序:同一父节点下按发现数
+// 降序、标签升序,「未关联资产」恒在最后。
+type FindingAssetTree struct {
+	Nodes        []FindingAssetNode `json:"nodes"`
+	FindingTotal int                `json:"finding_total"`
+	// Truncated=true 表示为控制体积丢弃了 DroppedKinds 里的层级。
+	Truncated    bool     `json:"truncated"`
+	DroppedKinds []string `json:"dropped_kinds,omitempty"`
+}
+
+// assetRow 是构树需要的资产字段子集。
+type assetRow struct {
+	id          int64
+	kind        string
+	companyID   int64
+	domain      string
+	rootDomain  string
+	ip          string
+	url         string
+	port        int
+	serviceType string
+	appName     string
+}
+
+func (a *assetRow) coverageNode() CoverageGraphNode {
+	return CoverageGraphNode{
+		Kind: a.kind, Domain: a.domain, RootDomain: a.rootDomain, IP: a.ip,
+		URL: a.url, Port: a.port, ServiceType: a.serviceType, AppName: a.appName,
+	}
+}
+
+// label 复用覆盖图的标签规则(URL > domain > ip > app_name > root_domain),但没有
+// URL 的服务(SMB、非 HTTP 端口等)要补上端口:否则它的标签会和宿主 IP/域名那行
+// 一模一样,树上父子两行看起来完全重复。
+func (a *assetRow) label() string {
+	if a.kind == "service" && a.url == "" {
+		if host, port := a.hostPort(); host != "" && port > 0 {
+			return host + ":" + strconv.Itoa(port)
+		}
+	}
+	n := a.coverageNode()
+	n.Key = assetKey(a.id)
+	return coverageNodeLabel(&n)
+}
+
+// hostPort 与覆盖图一致:优先 domain,其次 URL 里的 host,最后 ip。
+func (a *assetRow) hostPort() (string, int) {
+	n := a.coverageNode()
+	return hostPortOf(&n)
+}
+
+const findingAssetSelectCols = `a.id, a.type, COALESCE(a.company_id,0),
+       COALESCE(a.domain,''), COALESCE(a.root_domain,''), COALESCE(a.ip,''),
+       COALESCE(a.url,''), COALESCE(a.port,0), COALESCE(a.service_type,''),
+       COALESCE(a.app_name,'')`
+
+func scanAssetRows(rows interface {
+	Next() bool
+	Scan(...any) error
+	Err() error
+	Close() error
+}) ([]*assetRow, error) {
+	defer rows.Close()
+	var out []*assetRow
+	for rows.Next() {
+		a := &assetRow{}
+		if err := rows.Scan(&a.id, &a.kind, &a.companyID, &a.domain, &a.rootDomain,
+			&a.ip, &a.url, &a.port, &a.serviceType, &a.appName); err != nil {
+			return nil, err
+		}
+		out = append(out, a)
+	}
+	return out, rows.Err()
+}
+
+// findingAssetHit 是一条发现在构树阶段需要的最小信息。
+type findingAssetHit struct {
+	severity string
+	ts       time.Time
+	assetIDs []int64
+}
+
+// BuildFindingAssetTree 按当前筛选构建资产树。AssetScope 自身不参与(否则树会随
+// 选中节点塌缩成一条链)。
+func (d *DB) BuildFindingAssetTree(f FindingFilter) (*FindingAssetTree, error) {
+	return d.buildFindingAssetTree(f, findingAssetTreeMaxNodes)
+}
+
+// buildFindingAssetTree 是带节点上限的内部实现。maxNodes<=0 表示不截断——解析
+// AssetScope 时必须用这个模式,否则被丢掉的 endpoint 会让子树 id 集合不全。
+func (d *DB) buildFindingAssetTree(f FindingFilter, maxNodes int) (*FindingAssetTree, error) {
+	f.AssetScope = ""
+	f.assetIDs, f.assetNone, f.assetMiss = nil, false, false
+	where, args := f.where()
+
+	rows, err := d.Query(`SELECT COALESCE(f.severity,''), f.created_at,
+       COALESCE(f.asset_ids::text,'[]')
+FROM findings f LEFT JOIN tasks t ON f.task_id = t.id`+where, args...)
+	if err != nil {
+		return nil, err
+	}
+	hits := []findingAssetHit{}
+	assetIDs := map[int64]bool{}
+	for rows.Next() {
+		var h findingAssetHit
+		var aidsJSON string
+		if err := rows.Scan(&h.severity, &h.ts, &aidsJSON); err != nil {
+			rows.Close()
+			return nil, err
+		}
+		_ = json.Unmarshal([]byte(aidsJSON), &h.assetIDs)
+		for _, id := range h.assetIDs {
+			if id > 0 {
+				assetIDs[id] = true
+			}
+		}
+		hits = append(hits, h)
+	}
+	if err := rows.Err(); err != nil {
+		rows.Close()
+		return nil, err
+	}
+	rows.Close()
+
+	tree := &FindingAssetTree{Nodes: []FindingAssetNode{}, FindingTotal: len(hits)}
+	byID, err := d.loadFindingAssetRows(assetIDs)
+	if err != nil {
+		return nil, err
+	}
+
+	nodes, parentOf := d.assembleFindingAssetNodes(byID)
+	if err := d.attachCompanyNodes(nodes, parentOf); err != nil {
+		return nil, err
+	}
+
+	// 计数:一条发现沿它每个资产的祖先链向上,收集去重后的 key 集合再逐个 +1,
+	// 所以父节点不会因为一条发现挂了多个子资产而重复计数。
+	unassigned := &FindingAssetNode{Key: FindingUnassignedAsset, Kind: "none", Label: "未关联资产"}
+	touched := map[string]bool{}
+	for _, h := range hits {
+		clear(touched)
+		var direct []*FindingAssetNode
+		for _, id := range h.assetIDs {
+			node := nodes[assetKey(id)]
+			if node == nil {
+				continue
+			}
+			direct = append(direct, node)
+			for key := node.Key; key != ""; key = parentOf[key] {
+				touched[key] = true
+			}
+		}
+		if len(direct) == 0 {
+			countFinding(unassigned, h)
+			unassigned.Self++
+			continue
+		}
+		for _, node := range direct {
+			node.Self++
+		}
+		for key := range touched {
+			countFinding(nodes[key], h)
+		}
+	}
+
+	for _, node := range nodes {
+		if node.Total > 0 {
+			tree.Nodes = append(tree.Nodes, *node)
+		}
+	}
+	if unassigned.Total > 0 {
+		tree.Nodes = append(tree.Nodes, *unassigned)
+	}
+	sortFindingAssetNodes(tree.Nodes)
+	truncateFindingAssetTree(tree, maxNodes)
+	return tree, nil
+}
+
+// countFinding 把一条发现累加到节点上(总数 / 严重度分桶 / 最近发现时间)。
+func countFinding(n *FindingAssetNode, h findingAssetHit) {
+	if n == nil {
+		return
+	}
+	n.Total++
+	switch h.severity {
+	case "critical":
+		n.Critical++
+	case "high":
+		n.High++
+	case "medium":
+		n.Medium++
+	case "low":
+		n.Low++
+	}
+	if h.ts.After(n.LastFoundAt) {
+		n.LastFoundAt = h.ts
+	}
+}
+
+// loadFindingAssetRows 读取命中的资产行,并逐轮补齐祖先(service 的宿主域名/IP、
+// 子域名的根域名)。祖先自身可能没有任何发现,但树需要它们才能成形。
+func (d *DB) loadFindingAssetRows(ids map[int64]bool) (map[int64]*assetRow, error) {
+	byID := map[int64]*assetRow{}
+	if len(ids) == 0 {
+		return byID, nil
+	}
+	idList := make([]int64, 0, len(ids))
+	for id := range ids {
+		idList = append(idList, id)
+	}
+	rows, err := d.Query(`SELECT `+findingAssetSelectCols+` FROM assets a WHERE a.id = ANY($1::bigint[])`, idList)
+	if err != nil {
+		return nil, err
+	}
+	found, err := scanAssetRows(rows)
+	if err != nil {
+		return nil, err
+	}
+	for _, a := range found {
+		byID[a.id] = a
+	}
+
+	// 每轮找出还缺父节点的宿主标识,批量补一层;层数固定(endpoint→service→
+	// subdomain/ip→root_domain),4 轮足够收敛。
+	for range 4 {
+		want := missingParents(byID)
+		if want.empty() {
+			break
+		}
+		added, err := d.loadAssetsByHost(want, byID)
+		if err != nil {
+			return nil, err
+		}
+		if added == 0 {
+			break
+		}
+	}
+	return byID, nil
+}
+
+// missingHosts 是一轮补齐里要去库里找的宿主标识,按目标资产类型分开。
+type missingHosts struct {
+	services []string // endpoint 的宿主(找 service 行)
+	domains  []string // service/endpoint 的宿主域名(找 subdomain 行)
+	ips      []string // service/endpoint 的宿主 IP(找 ip 行)
+	roots    []string // 子域名的根域名(找 root_domain 行)
+}
+
+func (m missingHosts) empty() bool {
+	return len(m.services) == 0 && len(m.domains) == 0 && len(m.ips) == 0 && len(m.roots) == 0
+}
+
+// missingParents 汇总还没被加载的宿主:service(供 endpoint 挂靠)、子域名/IP(供
+// service 与 endpoint 挂靠)与根域名(供子域名挂靠)。
+func missingParents(byID map[int64]*assetRow) missingHosts {
+	haveService := map[string]bool{}
+	haveDomain := map[string]bool{}
+	haveIP := map[string]bool{}
+	haveRoot := map[string]bool{}
+	for _, a := range byID {
+		switch a.kind {
+		case "service":
+			if host, _ := a.hostPort(); host != "" {
+				haveService[host] = true
+			}
+		case "subdomain":
+			haveDomain[a.domain] = true
+		case "ip":
+			haveIP[a.ip] = true
+		case "root_domain":
+			haveRoot[a.domain] = true
+		}
+	}
+	wantService := map[string]bool{}
+	wantDomain := map[string]bool{}
+	wantIP := map[string]bool{}
+	wantRoot := map[string]bool{}
+	for _, a := range byID {
+		switch a.kind {
+		case "service", "endpoint":
+			host, _ := a.hostPort()
+			// endpoint 先找同宿主的 service;端口对不上的 service 会以 Total=0
+			// 被最终过滤掉,不会污染树。
+			if a.kind == "endpoint" && host != "" && !haveService[host] {
+				wantService[host] = true
+			}
+			if host != "" && !haveDomain[host] && !haveIP[host] && !haveRoot[host] {
+				if isIPLiteral(host) {
+					wantIP[host] = true
+				} else {
+					wantDomain[host] = true
+				}
+			}
+			if a.ip != "" && !haveIP[a.ip] {
+				wantIP[a.ip] = true
+			}
+		case "subdomain":
+			if a.rootDomain != "" && !haveRoot[a.rootDomain] {
+				wantRoot[a.rootDomain] = true
+			}
+		}
+	}
+	return missingHosts{
+		services: keysOf(wantService),
+		domains:  keysOf(wantDomain),
+		ips:      keysOf(wantIP),
+		roots:    keysOf(wantRoot),
+	}
+}
+
+func keysOf(m map[string]bool) []string {
+	if len(m) == 0 {
+		return nil
+	}
+	out := make([]string, 0, len(m))
+	for k := range m {
+		out = append(out, k)
+	}
+	return out
+}
+
+// isIPLiteral 粗判一个 host 是不是 IP 字面量(用于决定去 ip 还是 subdomain 表找宿主)。
+func isIPLiteral(host string) bool {
+	if strings.Contains(host, ":") {
+		return true // IPv6
+	}
+	if host == "" {
+		return false
+	}
+	for _, part := range strings.Split(host, ".") {
+		if part == "" {
+			return false
+		}
+		if _, err := strconv.Atoi(part); err != nil {
+			return false
+		}
+	}
+	return strings.Count(host, ".") == 3
+}
+
+// loadAssetsByHost 按宿主标识批量补齐资产行,返回本轮新增的行数。
+func (d *DB) loadAssetsByHost(want missingHosts, byID map[int64]*assetRow) (int, error) {
+	added := 0
+	load := func(q string, arg []string) error {
+		if len(arg) == 0 {
+			return nil
+		}
+		rows, err := d.Query(q, arg)
+		if err != nil {
+			return err
+		}
+		found, err := scanAssetRows(rows)
+		if err != nil {
+			return err
+		}
+		for _, a := range found {
+			if _, ok := byID[a.id]; ok {
+				continue
+			}
+			byID[a.id] = a
+			added++
+		}
+		return nil
+	}
+	if err := load(`SELECT `+findingAssetSelectCols+` FROM assets a
+WHERE a.type='service' AND (a.domain = ANY($1::text[]) OR a.ip = ANY($1::text[]))`, want.services); err != nil {
+		return 0, err
+	}
+	if err := load(`SELECT `+findingAssetSelectCols+` FROM assets a
+WHERE a.type='subdomain' AND a.domain = ANY($1::text[])`, want.domains); err != nil {
+		return 0, err
+	}
+	if err := load(`SELECT `+findingAssetSelectCols+` FROM assets a
+WHERE a.type='ip' AND a.ip = ANY($1::text[])`, want.ips); err != nil {
+		return 0, err
+	}
+	if err := load(`SELECT `+findingAssetSelectCols+` FROM assets a
+WHERE a.type='root_domain' AND a.domain = ANY($1::text[])`, want.roots); err != nil {
+		return 0, err
+	}
+	return added, nil
+}
+
+// assembleFindingAssetNodes 把资产行变成节点并连上父子关系。父节点缺位时(库里
+// 根本没有那条根域名资产)合成 "r:<domain>" 占位节点,与覆盖图的处理一致。
+func (d *DB) assembleFindingAssetNodes(byID map[int64]*assetRow) (map[string]*FindingAssetNode, map[string]string) {
+	nodes := map[string]*FindingAssetNode{}
+	parentOf := map[string]string{}
+	rootByDomain := map[string]string{}
+	subByDomain := map[string]string{}
+	ipByAddr := map[string]string{}
+	svcByHost := map[string]string{}
+	svcByHostPort := map[string]string{}
+
+	for _, a := range byID {
+		key := assetKey(a.id)
+		nodes[key] = &FindingAssetNode{
+			Key: key, Kind: a.kind, Label: a.label(),
+			AssetID: a.id, CompanyID: a.companyID,
+		}
+		switch a.kind {
+		case "root_domain":
+			if a.domain != "" {
+				rootByDomain[a.domain] = key
+			}
+		case "subdomain":
+			if a.domain != "" {
+				subByDomain[a.domain] = key
+			}
+		case "ip":
+			if a.ip != "" {
+				ipByAddr[a.ip] = key
+			}
+		case "service":
+			if host, port := a.hostPort(); host != "" {
+				svcByHost[host] = key
+				svcByHostPort[host+"|"+strconv.Itoa(port)] = key
+			}
+		}
+	}
+
+	// 子域名的根域名在库里没有资产行时,合成一个占位根,免得子域名散成顶层。
+	for _, a := range byID {
+		if a.kind != "subdomain" || a.rootDomain == "" {
+			continue
+		}
+		if _, ok := rootByDomain[a.rootDomain]; ok {
+			continue
+		}
+		key := "r:" + a.rootDomain
+		nodes[key] = &FindingAssetNode{Key: key, Kind: "root_domain", Label: a.rootDomain}
+		rootByDomain[a.rootDomain] = key
+	}
+
+	firstOf := func(keys ...string) string {
+		for _, k := range keys {
+			if k != "" {
+				if _, ok := nodes[k]; ok {
+					return k
+				}
+			}
+		}
+		return ""
+	}
+	for _, a := range byID {
+		key := assetKey(a.id)
+		var parent string
+		switch a.kind {
+		case "subdomain":
+			parent = firstOf(rootByDomain[a.rootDomain])
+		case "service":
+			host, _ := a.hostPort()
+			parent = firstOf(subByDomain[a.domain], subByDomain[host],
+				ipByAddr[a.ip], ipByAddr[host], rootByDomain[a.rootDomain], rootByDomain[host])
+		case "endpoint":
+			host, port := a.hostPort()
+			parent = firstOf(svcByHostPort[host+"|"+strconv.Itoa(port)], svcByHost[host],
+				subByDomain[host], subByDomain[a.domain], ipByAddr[host], ipByAddr[a.ip],
+				rootByDomain[a.rootDomain], rootByDomain[host])
+		}
+		if parent != "" && parent != key {
+			parentOf[key] = parent
+			nodes[key].Parent = parent
+		}
+	}
+	return nodes, parentOf
+}
+
+// attachCompanyNodes 给顶层资产(根域名 / IP / 应用)补企业父节点——只有资产确实
+// 归属了企业才会出现企业层,没归属的资产仍然自己就是顶层。
+func (d *DB) attachCompanyNodes(nodes map[string]*FindingAssetNode, parentOf map[string]string) error {
+	want := map[int64]bool{}
+	for _, n := range nodes {
+		if n.Parent != "" || n.CompanyID <= 0 {
+			continue
+		}
+		switch n.Kind {
+		case "root_domain", "ip", "app":
+			want[n.CompanyID] = true
+		}
+	}
+	if len(want) == 0 {
+		return nil
+	}
+	ids := make([]int64, 0, len(want))
+	for id := range want {
+		ids = append(ids, id)
+	}
+	rows, err := d.Query(`SELECT id, COALESCE(name,'') FROM companies WHERE id = ANY($1::bigint[])`, ids)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+	names := map[int64]string{}
+	for rows.Next() {
+		var id int64
+		var name string
+		if err := rows.Scan(&id, &name); err != nil {
+			return err
+		}
+		names[id] = name
+	}
+	if err := rows.Err(); err != nil {
+		return err
+	}
+	for id, name := range names {
+		key := companyKey(id)
+		if _, ok := nodes[key]; ok {
+			continue
+		}
+		if name == "" {
+			name = "企业 #" + strconv.FormatInt(id, 10)
+		}
+		nodes[key] = &FindingAssetNode{Key: key, Kind: "company", Label: name, CompanyID: id}
+	}
+	for _, n := range nodes {
+		if n.Parent != "" || n.CompanyID <= 0 || n.Kind == "company" {
+			continue
+		}
+		switch n.Kind {
+		case "root_domain", "ip", "app":
+			key := companyKey(n.CompanyID)
+			if _, ok := nodes[key]; !ok {
+				continue
+			}
+			n.Parent = key
+			parentOf[n.Key] = key
+		}
+	}
+	return nil
+}
+
+// sortFindingAssetNodes 排序:发现多的在前,同数按标签;「未关联资产」恒在最后。
+// 前端按数组顺序挂子节点,所以只要同一父节点下的相对顺序正确即可。
+func sortFindingAssetNodes(nodes []FindingAssetNode) {
+	sort.SliceStable(nodes, func(i, j int) bool {
+		a, b := nodes[i], nodes[j]
+		if (a.Kind == "none") != (b.Kind == "none") {
+			return b.Kind == "none"
+		}
+		if a.Total != b.Total {
+			return a.Total > b.Total
+		}
+		return a.Label < b.Label
+	})
+}
+
+// truncateFindingAssetTree 在节点过多时整层丢弃(先 endpoint 再 service)。计数已
+// 累加到父节点,丢的只是可展开的细节层级。
+func truncateFindingAssetTree(tree *FindingAssetTree, maxNodes int) {
+	if maxNodes <= 0 || len(tree.Nodes) <= maxNodes {
+		return
+	}
+	for _, kind := range []string{"endpoint", "service"} {
+		kept := tree.Nodes[:0]
+		for _, n := range tree.Nodes {
+			if n.Kind == kind {
+				continue
+			}
+			kept = append(kept, n)
+		}
+		tree.Nodes = kept
+		tree.Truncated = true
+		tree.DroppedKinds = append(tree.DroppedKinds, kind)
+		if len(tree.Nodes) <= maxNodes {
+			return
+		}
+	}
+}
+
+// applyAssetScope 把 AssetScope(节点 key)解析成可用于 SQL 的资产 id 集合。选中
+// 一个节点等于选中它的整棵子树,所以要先把树建出来再收集子孙。
+func (d *DB) applyAssetScope(f FindingFilter) (FindingFilter, error) {
+	scope := strings.TrimSpace(f.AssetScope)
+	f.assetIDs, f.assetNone, f.assetMiss = nil, false, false
+	if scope == "" {
+		return f, nil
+	}
+	if scope == FindingUnassignedAsset {
+		f.assetNone = true
+		return f, nil
+	}
+	// 不截断:被丢掉的 endpoint 同样要参与 id 收集,否则列表会少数据。
+	tree, err := d.buildFindingAssetTree(f, 0)
+	if err != nil {
+		return f, err
+	}
+	children := map[string][]FindingAssetNode{}
+	byKey := map[string]FindingAssetNode{}
+	for _, n := range tree.Nodes {
+		byKey[n.Key] = n
+		children[n.Parent] = append(children[n.Parent], n)
+	}
+	if _, ok := byKey[scope]; !ok {
+		// 选中的节点在当前筛选下已经不存在,结果应当为空而不是退化成不过滤。
+		f.assetMiss = true
+		return f, nil
+	}
+	seen := map[string]bool{scope: true}
+	queue := []string{scope}
+	for len(queue) > 0 {
+		key := queue[0]
+		queue = queue[1:]
+		if id := byKey[key].AssetID; id > 0 {
+			f.assetIDs = append(f.assetIDs, id)
+		}
+		for _, child := range children[key] {
+			if seen[child.Key] {
+				continue
+			}
+			seen[child.Key] = true
+			queue = append(queue, child.Key)
+		}
+	}
+	if len(f.assetIDs) == 0 {
+		f.assetMiss = true
+	}
+	return f, nil
+}
+
+// assetIDContainments 把资产 id 变成 jsonb 包含判断的右操作数集合,配合
+// idx_findings_asset_ids(GIN jsonb_path_ops)使用。
+func assetIDContainments(ids []int64) []string {
+	out := make([]string, 0, len(ids))
+	for _, id := range ids {
+		out = append(out, "["+strconv.FormatInt(id, 10)+"]")
+	}
+	return out
+}