CyberStrikeAI/internal/database/vulnerability_workspace.go

package database

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net/netip"
	"net/url"
	"regexp"
	"strings"
	"time"

	"github.com/google/uuid"
	"golang.org/x/net/publicsuffix"
)

type VulnerabilityTarget struct {
	Raw       string `json:"raw"`
	Host      string `json:"host"`
	Domain    string `json:"domain"`
	Scheme    string `json:"scheme,omitempty"`
	Port      string `json:"port,omitempty"`
	Path      string `json:"path,omitempty"`
	Kind      string `json:"kind"`
	GroupHost bool   `json:"group_host"`
}

const vulnerabilityTargetsVersion = 2

var vulnerabilityTargetPattern = regexp.MustCompile(`(?i)(?:[a-z][a-z0-9+.-]*://|//)[^\s<>"',;，；、（）]+|(?:\[[0-9a-f:]+\]|(?:[\p{L}\p{N}][\p{L}\p{N}.-]*\.)+[\p{L}\p{N}-]+|localhost)(?::[0-9]+)?(?:/[^\s<>"',;，；、（）]*)?`)
var vulnerabilityVersionContext = regexp.MustCompile(`(?i)(?:版本(?:号)?|version|ver|Citrix Gateway|NetScaler(?: ADC)?|Keycloak|Ivanti Connect Secure|nginx|Apache|OpenSSL|Tomcat)\s*[:=]?\s*$`)
var vulnerabilityVersionToken = regexp.MustCompile(`(?i)^\d+(?:\.\d+)*\.(?:x|final|release|snapshot|alpha\d*|beta\d*|rc\d*)$`)

func vulnerabilityTargetsJSON(raw string) string {
	targets := []VulnerabilityTarget{}
	seen := map[string]bool{}
	matches := vulnerabilityTargetPattern.FindAllStringIndex(raw, -1)
	if len(matches) == 0 {
		matches = [][]int{{0, len(raw)}}
	}
	for _, match := range matches {
		candidate := strings.TrimSpace(raw[match[0]:match[1]])
		if match[0] > 0 && strings.ContainsRune("/@", rune(raw[match[0]-1])) && !strings.Contains(candidate, "://") {
			continue
		}
		candidate = strings.TrimRight(candidate, ".,;:!?)\x60")
		host := vulnerabilityTargetHost(candidate)
		if host == "" {
			continue
		}
		_, ipErr := netip.ParseAddr(host)
		explicitHost := strings.Contains(candidate, "://") || strings.HasPrefix(candidate, "//")
		if !explicitHost && vulnerabilityVersionToken.MatchString(host) {
			continue
		}
		if ipErr == nil && !explicitHost && vulnerabilityVersionContext.MatchString(raw[:match[0]]) {
			continue
		}
		if ipErr != nil && !explicitHost {
			_, icann := publicsuffix.PublicSuffix(host)
			localHost := host == "localhost" || strings.HasSuffix(host, ".test") || strings.HasSuffix(host, ".local") || strings.HasSuffix(host, ".internal") || strings.HasSuffix(host, ".lan")
			if !icann && !localHost {
				continue
			}
		}
		target := VulnerabilityTarget{Raw: candidate, Host: host, Domain: host, Kind: "domain"}
		if ipErr == nil {
			target.Kind = "ip"
		} else {
			if domain, err := publicsuffix.EffectiveTLDPlusOne(host); err == nil {
				target.Domain = domain
			}
		}
		value := candidate
		if !strings.Contains(value, "://") && !strings.HasPrefix(value, "//") {
			value = "//" + value
		}
		if parsed, err := url.Parse(value); err == nil {
			target.Scheme = strings.ToLower(parsed.Scheme)
			target.Port = parsed.Port()
			target.Path = parsed.EscapedPath()
		}
		key := target.Host + "|" + target.Scheme + "|" + target.Port + "|" + target.Path
		if seen[key] {
			continue
		}
		seen[key] = true
		targets = append(targets, target)
	}
	hasDomain := false
	for _, target := range targets {
		if target.Kind == "domain" {
			hasDomain = true
			break
		}
	}
	for i := range targets {
		targets[i].GroupHost = targets[i].Kind == "domain" || !hasDomain
	}
	data, _ := json.Marshal(targets)
	return string(data)
}

func (db *DB) backfillVulnerabilityTargets() error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	rows, err := tx.Query("SELECT id, COALESCE(target,'') FROM vulnerabilities WHERE targets_version<?", vulnerabilityTargetsVersion)
	if err != nil {
		return err
	}
	targets := map[string]string{}
	for rows.Next() {
		var id, raw string
		if err := rows.Scan(&id, &raw); err != nil {
			rows.Close()
			return err
		}
		targets[id] = vulnerabilityTargetsJSON(raw)
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return err
	}
	for id, raw := range targets {
		if _, err := tx.Exec("UPDATE vulnerabilities SET targets_json=?,targets_version=? WHERE id=?", raw, vulnerabilityTargetsVersion, id); err != nil {
			return err
		}
	}
	return tx.Commit()
}

type VulnerabilityLinkedAsset struct {
	ID   string `json:"id"`
	Host string `json:"host"`
	Port int    `json:"port"`
}

func (db *DB) VulnerabilityLinkedAssets(v *Vulnerability, access RBACListAccess) ([]VulnerabilityLinkedAsset, error) {
	query, args := appendAssetAccess(`SELECT assets.id, assets.host, COALESCE(assets.port,0)
		FROM assets WHERE COALESCE(assets.project_id,'')=? AND EXISTS (
			SELECT 1 FROM json_each(?) jt WHERE
			json_extract(jt.value,'$.host')=LOWER(COALESCE(NULLIF(assets.domain,''),NULLIF(assets.ip,''),assets.host))
			AND (COALESCE(assets.port,0)=0 OR CAST(json_extract(jt.value,'$.port') AS INTEGER)=assets.port
			OR (COALESCE(json_extract(jt.value,'$.port'),'')='' AND
				((json_extract(jt.value,'$.scheme')='https' AND assets.port=443)
				OR (json_extract(jt.value,'$.scheme')='http' AND assets.port=80))))
		)`, []interface{}{v.ProjectID, vulnerabilityTargetsJSON(v.Target)}, access, "assets")
	rows, err := db.Query(query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	assets := []VulnerabilityLinkedAsset{}
	for rows.Next() {
		var a VulnerabilityLinkedAsset
		if err := rows.Scan(&a.ID, &a.Host, &a.Port); err != nil {
			return nil, err
		}
		assets = append(assets, a)
	}
	return assets, rows.Err()
}

// MergeVulnerabilities preserves original reports. The transaction reserves the
// writer before checking membership, ownership and project boundaries.
func (db *DB) MergeVulnerabilities(ids []string, title, reason, actor string, access RBACListAccess) (string, error) {
	if len(ids) < 2 || len(ids) > 100 || strings.TrimSpace(title) == "" || strings.TrimSpace(reason) == "" {
		return "", errors.New("请选择 2–100 条记录，并填写漏洞名称和归并理由")
	}
	tx, err := db.Begin()
	if err != nil {
		return "", err
	}
	defer tx.Rollback()
	seen := map[string]bool{}
	project := ""
	for i, id := range ids {
		if seen[id] {
			return "", errors.New("重复的漏洞 ID")
		}
		seen[id] = true
		query, args := appendVulnerabilityAccessFilter("SELECT COALESCE(project_id,''), COALESCE(case_id,'') FROM vulnerabilities WHERE id=?", []interface{}{id}, access)
		var pid, cid string
		if err := tx.QueryRow(query, args...).Scan(&pid, &cid); err != nil {
			return "", errors.New("漏洞不存在或无权归并")
		}
		if cid != "" {
			return "", errors.New("记录已归并，请先拆分原归并组")
		}
		if i == 0 {
			project = pid
		} else if project != pid {
			return "", errors.New("归并记录必须属于同一项目")
		}
	}
	id := uuid.NewString()
	if _, err := tx.Exec("INSERT INTO vulnerability_cases(id,project_id,title,reason) VALUES(?,?,?,?)", id, project, strings.TrimSpace(title), strings.TrimSpace(reason)); err != nil {
		return "", err
	}
	for _, member := range ids {
		if _, err := tx.Exec("UPDATE vulnerabilities SET case_id=? WHERE id=?", id, member); err != nil {
			return "", err
		}
	}
	rawIDs, _ := json.Marshal(ids)
	if _, err := tx.Exec("INSERT INTO vulnerability_case_events(case_id,action,actor,reason,finding_ids) VALUES(?,?,?,?,?)",
		id, "merge", actor, strings.TrimSpace(reason), string(rawIDs)); err != nil {
		return "", err
	}
	return id, tx.Commit()
}

func (db *DB) DissolveVulnerabilityCase(id, reason, actor string, access RBACListAccess) error {
	if strings.TrimSpace(reason) == "" {
		return errors.New("请填写拆分理由")
	}
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	var all, allowed int
	if err = tx.QueryRow("SELECT COUNT(*) FROM vulnerabilities WHERE case_id=?", id).Scan(&all); err != nil {
		return err
	}
	query, args := appendVulnerabilityAccessFilter("SELECT COUNT(*) FROM vulnerabilities WHERE case_id=?", []interface{}{id}, access)
	if err = tx.QueryRow(query, args...).Scan(&allowed); err != nil {
		return err
	}
	if all == 0 || all != allowed {
		return errors.New("归并组不存在或无权修改全部成员")
	}
	if _, err = tx.Exec(`INSERT INTO vulnerability_case_events(case_id,action,actor,reason,finding_ids)
		SELECT ?, 'split', ?, ?, json_group_array(id) FROM vulnerabilities WHERE case_id=?`, id, actor, reason, id); err != nil {
		return err
	}
	if _, err = tx.Exec("UPDATE vulnerabilities SET case_id=NULL WHERE case_id=?", id); err != nil {
		return err
	}
	if _, err = tx.Exec("DELETE FROM vulnerability_cases WHERE id=?", id); err != nil {
		return err
	}
	return tx.Commit()
}

type VulnerabilityHistoryItem struct {
	Kind           string         `json:"kind"`
	ConversationID string         `json:"conversation_id,omitempty"`
	Status         string         `json:"status,omitempty"`
	Evidence       string         `json:"evidence,omitempty"`
	Actor          string         `json:"actor,omitempty"`
	CreatedAt      time.Time      `json:"created_at"`
	Snapshot       *Vulnerability `json:"snapshot,omitempty"`
}

func (db *DB) VulnerabilityHistory(id string) ([]VulnerabilityHistoryItem, error) {
	rows, err := db.Query(`
		SELECT 'retest',conversation_id,status,evidence,'',created_at,
			COALESCE((SELECT snapshot_json FROM vulnerability_retests vr WHERE vr.conversation_id=vulnerability_retest_results.conversation_id),'')
		FROM vulnerability_retest_results WHERE vulnerability_id=?
		UNION ALL
		SELECT action,'','',reason,actor,created_at,'' FROM vulnerability_case_events
		WHERE EXISTS (SELECT 1 FROM json_each(finding_ids) WHERE value=?)
		UNION ALL
		SELECT 'retest_started',vr.conversation_id,'','','',c.created_at,vr.snapshot_json
		FROM vulnerability_retests vr JOIN conversations c ON c.id=vr.conversation_id
		WHERE vr.vulnerability_id=? AND NOT EXISTS (
			SELECT 1 FROM vulnerability_retest_results r WHERE r.vulnerability_id=vr.vulnerability_id AND r.conversation_id=vr.conversation_id)
		ORDER BY created_at DESC`, id, id, id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []VulnerabilityHistoryItem{}
	for rows.Next() {
		var item VulnerabilityHistoryItem
		var snapshot string
		if err := rows.Scan(&item.Kind, &item.ConversationID, &item.Status, &item.Evidence, &item.Actor, &item.CreatedAt, &snapshot); err != nil {
			return nil, err
		}
		if snapshot != "" {
			if err := json.Unmarshal([]byte(snapshot), &item.Snapshot); err != nil {
				return nil, err
			}
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func scanVulnerabilityTargets(raw string, v *Vulnerability) error {
	if err := json.Unmarshal([]byte(raw), &v.Targets); err != nil {
		return fmt.Errorf("读取漏洞目标失败: %w", err)
	}
	return nil
}

// Keep project changes from silently carrying a cross-project merge relation.
func detachVulnerabilityCaseOnProjectChange(tx *sql.Tx, id, projectID string) error {
	var oldProject, caseID string
	if err := tx.QueryRow("SELECT COALESCE(project_id,''),COALESCE(case_id,'') FROM vulnerabilities WHERE id=?", id).Scan(&oldProject, &caseID); err != nil {
		return err
	}
	if caseID == "" || oldProject == projectID {
		return nil
	}
	raw, _ := json.Marshal([]string{id})
	if _, err := tx.Exec("INSERT INTO vulnerability_case_events(case_id,action,actor,reason,finding_ids) VALUES(?,'detach','system','项目归属变更',?)", caseID, string(raw)); err != nil {
		return err
	}
	_, err := tx.Exec("UPDATE vulnerabilities SET case_id=NULL WHERE id=?", id)
	return err
}