package project import ( "fmt" "strings" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/projectprompt" ) // PathGraphCategories 攻击路径视图包含的事实分类。 var PathGraphCategories = map[string]struct{}{ FactCategoryTarget: {}, FactCategoryFinding: {}, FactCategoryChain: {}, FactCategoryExploit: {}, FactCategoryPOC: {}, "vuln": {}, } // GraphNodeType 将 fact category 映射为图节点类型(供前端样式与 ELK 分层)。 // 优先使用 category;仅 synthetic 节点(vuln:)或无 category 时才回退到 fact_key 前缀。 func GraphNodeType(category, factKey string) string { key := strings.ToLower(strings.TrimSpace(factKey)) if strings.HasPrefix(key, "vuln:") { return "vulnerability" } c := strings.ToLower(strings.TrimSpace(category)) if c != "" { switch c { case FactCategoryTarget: return "target" case FactCategoryExploit: return "exploit" case FactCategoryPOC: return "poc" case FactCategoryChain: return "chain" case FactCategoryFinding: return "finding" case "vuln": return "vulnerability" case FactCategoryAuth: return "auth" case FactCategoryInfra, FactCategoryBusiness: return "infra" case FactCategoryNote: return "note" case "missing": return "missing" default: return c } } switch { case strings.HasPrefix(key, "target/"): return "target" case strings.HasPrefix(key, "exploit/"), strings.HasPrefix(key, "evidence/"): return "exploit" case strings.HasPrefix(key, "poc/"): return "poc" case strings.HasPrefix(key, "chain/"): return "chain" case strings.HasPrefix(key, "finding/"): return "finding" case strings.HasPrefix(key, "auth/"): return "auth" case strings.HasPrefix(key, "infra/"), strings.HasPrefix(key, "business/"): return "infra" default: return "note" } } func truncateGraphLabel(summary string, maxRunes int) string { summary = strings.TrimSpace(summary) if summary == "" { return "—" } r := []rune(summary) if len(r) <= maxRunes { return summary } return string(r[:maxRunes]) + "…" } // BuildProjectFactGraph 构建项目事实图(nodes + edges)。 func BuildProjectFactGraph(db *database.DB, projectID string, view string, excludeDeprecated bool) (*database.ProjectFactGraph, error) { if db == nil { return nil, fmt.Errorf("database 未初始化") } projectID = strings.TrimSpace(projectID) if projectID == "" { return nil, fmt.Errorf("project_id 不能为空") } view = strings.TrimSpace(strings.ToLower(view)) if view == "" { view = "path" } filter := database.ProjectFactListFilter{} if excludeDeprecated { filter.ExcludeDeprecated = true } facts, err := db.ListProjectFacts(projectID, filter, 1000, 0) if err != nil { return nil, err } edges, err := db.ListProjectFactEdgesByProject(projectID) if err != nil { return nil, err } if excludeDeprecated { edges = filterDeprecatedEdges(edges) } factByKey := make(map[string]*database.ProjectFact, len(facts)) for _, f := range facts { factByKey[f.FactKey] = f } pathMode := view == "path" nodeKeys := make(map[string]struct{}) if pathMode { for _, f := range facts { if isPathGraphFact(f.Category, f.FactKey) { nodeKeys[f.FactKey] = struct{}{} } } // 路径视图中保留作为依赖目标的 auth/infra 节点 for _, e := range edges { if _, ok := nodeKeys[e.SourceFactKey]; !ok { continue } if f, ok := factByKey[e.TargetFactKey]; ok && isDependencyGraphFact(f.Category, f.FactKey) { nodeKeys[e.TargetFactKey] = struct{}{} } } } else { for _, f := range facts { nodeKeys[f.FactKey] = struct{}{} } } // 边上引用的 endpoint 纳入节点集 for _, e := range edges { if pathMode { if _, ok := nodeKeys[e.SourceFactKey]; !ok { continue } if _, ok := nodeKeys[e.TargetFactKey]; ok { // already included } else if f, ok := factByKey[e.TargetFactKey]; !ok { nodeKeys[e.TargetFactKey] = struct{}{} // 占位节点 } else if isPathGraphFact(f.Category, f.FactKey) || isDependencyGraphFact(f.Category, f.FactKey) { nodeKeys[e.TargetFactKey] = struct{}{} } else { continue } } else { nodeKeys[e.SourceFactKey] = struct{}{} nodeKeys[e.TargetFactKey] = struct{}{} } } nodes := make([]database.ProjectFactGraphNode, 0, len(nodeKeys)) for key := range nodeKeys { if f, ok := factByKey[key]; ok { nodes = append(nodes, database.ProjectFactGraphNode{ ID: f.FactKey, FactKey: f.FactKey, Category: f.Category, Label: truncateGraphLabel(f.Summary, 48), Summary: strings.TrimSpace(f.Summary), Confidence: f.Confidence, Type: GraphNodeType(f.Category, f.FactKey), Pinned: f.Pinned, }) continue } nodes = append(nodes, database.ProjectFactGraphNode{ ID: key, FactKey: key, Category: "missing", Label: key, Confidence: "tentative", Type: "missing", Pinned: false, }) } graphEdges := make([]database.ProjectFactGraphEdge, 0, len(edges)) for _, e := range edges { if pathMode { if _, ok := nodeKeys[e.SourceFactKey]; !ok { continue } if _, ok := nodeKeys[e.TargetFactKey]; !ok { continue } } else { if _, ok := nodeKeys[e.SourceFactKey]; !ok { continue } if _, ok := nodeKeys[e.TargetFactKey]; !ok { continue } } graphEdges = append(graphEdges, database.ProjectFactGraphEdge{ ID: e.ID, Source: e.SourceFactKey, Target: e.TargetFactKey, Type: e.EdgeType, Confidence: e.Confidence, }) } // related_vulnerability_id 合成边(source=fact → target=vuln:) for _, f := range facts { if _, ok := nodeKeys[f.FactKey]; !ok { continue } vid := strings.TrimSpace(f.RelatedVulnerabilityID) if vid == "" { continue } vulnNodeID := "vuln:" + vid if _, exists := nodeKeys[vulnNodeID]; !exists { nodeKeys[vulnNodeID] = struct{}{} label := "漏洞" if len(vid) >= 8 { label += " " + vid[:8] + "…" } else { label += " " + vid } nodes = append(nodes, database.ProjectFactGraphNode{ ID: vulnNodeID, FactKey: vulnNodeID, Category: "vuln", Label: label, Confidence: f.Confidence, Type: "vulnerability", Pinned: false, }) } graphEdges = append(graphEdges, database.ProjectFactGraphEdge{ ID: "vuln-link:" + f.FactKey + ":" + vid, Source: f.FactKey, Target: vulnNodeID, Type: "links_vuln", Confidence: f.Confidence, }) } return &database.ProjectFactGraph{Nodes: nodes, Edges: graphEdges}, nil } func min(a, b int) int { if a < b { return a } return b } func isPathGraphFact(category, factKey string) bool { c := strings.ToLower(strings.TrimSpace(category)) if _, ok := PathGraphCategories[c]; ok { return true } if c != "" { return false } key := strings.ToLower(strings.TrimSpace(factKey)) for _, p := range []string{"target/", "finding/", "chain/", "exploit/", "poc/", "evidence/"} { if strings.HasPrefix(key, p) { return true } } return false } func isDependencyGraphFact(category, factKey string) bool { c := strings.ToLower(strings.TrimSpace(category)) if c == FactCategoryAuth || c == FactCategoryInfra || c == FactCategoryBusiness { return true } if c != "" { return false } key := strings.ToLower(strings.TrimSpace(factKey)) return strings.HasPrefix(key, "auth/") || strings.HasPrefix(key, "infra/") || strings.HasPrefix(key, "business/") } func filterDeprecatedEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge { out := make([]*database.ProjectFactEdge, 0, len(edges)) for _, e := range edges { if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") { continue } out = append(out, e) } return out } // ParsedFactLinks 解析 links 参数(from → 当前 fact)。 type ParsedFactLinks struct { Incoming []database.ProjectFactEdgeFromInput } // ParseFactLinkInputs 从 MCP links 参数解析;空数组表示清空全部入边。 func ParseFactLinkInputs(raw interface{}) (*ParsedFactLinks, error) { if raw == nil { return nil, nil } items, ok := raw.([]interface{}) if !ok { return nil, fmt.Errorf("links 须为数组") } if len(items) == 0 { return &ParsedFactLinks{ Incoming: []database.ProjectFactEdgeFromInput{}, }, nil } parsed := &ParsedFactLinks{} for i, item := range items { m, ok := item.(map[string]interface{}) if !ok { return nil, fmt.Errorf("links[%d] 格式无效", i) } from, _ := m["from"].(string) edgeType, _ := m["type"].(string) from = strings.TrimSpace(from) edgeType = strings.TrimSpace(edgeType) if from == "" { return nil, fmt.Errorf("links[%d] 须含 from", i) } if edgeType == "" { return nil, fmt.Errorf("links[%d] 须含 type", i) } conf, _ := m["confidence"].(string) parsed.Incoming = append(parsed.Incoming, database.ProjectFactEdgeFromInput{ From: from, Type: edgeType, Confidence: strings.TrimSpace(conf), }) } return parsed, nil } // ParseFactLinksText 解析 UI 文本:`type: source_fact_key` 每行一条(from 语义)。 func ParseFactLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) { return ParseFactIncomingLinksText(text) } // FormatFactLinksText 将入边格式化为 UI 文本。 func FormatFactLinksText(edges []*database.ProjectFactEdge) string { return FormatFactIncomingLinksText(edges) } // ParseFactIncomingLinksText 解析 UI 入边文本:`type: source_fact_key` 每行一条。 func ParseFactIncomingLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) { text = strings.TrimSpace(text) if text == "" { return nil, nil } var out []database.ProjectFactEdgeFromInput for i, line := range strings.Split(text, "\n") { line = strings.TrimSpace(line) if line == "" || strings.HasPrefix(line, "#") { continue } edgeType, source, ok := strings.Cut(line, ":") if !ok { return nil, fmt.Errorf("第 %d 行格式无效,应为 type: fact_key", i+1) } edgeType = strings.TrimSpace(edgeType) source = strings.TrimSpace(source) if edgeType == "" || source == "" { return nil, fmt.Errorf("第 %d 行 type 或 fact_key 为空", i+1) } out = append(out, database.ProjectFactEdgeFromInput{From: source, Type: edgeType}) } return out, nil } // FormatFactIncomingLinksText 将入边格式化为 UI 文本。 func FormatFactIncomingLinksText(edges []*database.ProjectFactEdge) string { if len(edges) == 0 { return "" } var b strings.Builder for i, e := range edges { if i > 0 { b.WriteByte('\n') } b.WriteString(e.EdgeType) b.WriteString(": ") b.WriteString(e.SourceFactKey) } return b.String() } // FactEdgeRecordingGuidance 写入边时的 Agent 规范。 func FactEdgeRecordingGuidance() string { return projectprompt.FactEdgeRecordingGuidance() }