mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-16 16:10:21 +02:00
Add files via upload
This commit is contained in:
公明
@@ -0,0 +1,389 @@
|
||||
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 分层）。
|
||||
func GraphNodeType(category, factKey string) string {
|
||||
key := strings.ToLower(strings.TrimSpace(factKey))
|
||||
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"
|
||||
case strings.HasPrefix(key, "vuln:"):
|
||||
return "vulnerability"
|
||||
}
|
||||
c := strings.ToLower(strings.TrimSpace(category))
|
||||
switch c {
|
||||
case FactCategoryTarget:
|
||||
return "target"
|
||||
case FactCategoryExploit:
|
||||
return "exploit"
|
||||
case FactCategoryPOC:
|
||||
return "poc"
|
||||
case FactCategoryChain:
|
||||
return "chain"
|
||||
case FactCategoryFinding, "vuln":
|
||||
return "finding"
|
||||
case "auth":
|
||||
return "auth"
|
||||
case "infra", "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:<id>）
|
||||
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
|
||||
}
|
||||
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 == "auth" || c == "infra" || c == "business" {
|
||||
return true
|
||||
}
|
||||
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()
|
||||
}
|
||||
Reference in New Issue
Block a user