Add files via upload

This commit is contained in:
公明
2026-08-15 02:12:18 +08:00
committed by GitHub
parent b07a645d34
commit 9f38fda15f
30 changed files with 4519 additions and 0 deletions
+99
View File
@@ -0,0 +1,99 @@
package project
import (
"fmt"
"strings"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
)
// AppendSystemPromptBlock 将附加块追加到 system prompt。
func AppendSystemPromptBlock(base, block string) string {
base = strings.TrimSpace(base)
block = strings.TrimSpace(block)
if block == "" {
return base
}
if base == "" {
return block
}
return base + "\n\n" + block
}
const (
factIndexFooterGetDetail = "需要完整内容(攻击链、POC、请求响应等)时必须调用 get_project_fact(fact_key),禁止凭摘要臆造细节。"
factIndexFooterWriteHint = "写入事实 links 时用 from(来源 fact_key → 当前 fact),如 finding 上 {from:target/*, type:discovered_on}body 写可复现全流程(发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀)。"
factIndexFooterEmpty = "需要写入请使用 upsert_project_fact;需要详情请调用 get_project_fact(fact_key)。"
)
// BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引(key + summary + 关系边 + 攻击路径,不含 body)。
func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) {
if db == nil || !cfg.Enabled {
return "", nil
}
projectID = strings.TrimSpace(projectID)
if projectID == "" {
return "", nil
}
proj, err := db.GetProject(projectID)
if err != nil {
return "", err
}
facts, err := db.ListProjectFactsForIndex(projectID, cfg.DefaultInjectDeprecated)
if err != nil {
return "", err
}
allEdges, _ := db.ListProjectFactEdgesByProject(projectID)
_, incomingByTarget := indexEdgeGroupMaps(allEdges)
if len(facts) == 0 {
return wrapFactIndexBlock(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n(暂无事实)\n%s", proj.Name, proj.ID, factIndexFooterEmpty)), nil
}
sortFactsForIndex(facts)
maxRunes := cfg.FactIndexMaxRunesEffective()
pathMaxRunes := cfg.FactIndexPathMaxRunesEffective()
footer := factIndexFooterGetDetail + "\n" + factIndexFooterWriteHint
footerRunes := len([]rune(footer))
factsBudget := maxRunes - pathMaxRunes - footerRunes
if factsBudget < 800 {
factsBudget = maxRunes - footerRunes
pathMaxRunes = 0
}
indexedKeys := make(map[string]struct{}, len(facts))
var b strings.Builder
b.WriteString(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n", proj.Name, proj.ID))
used := len([]rune(b.String()))
omitted := 0
for _, f := range facts {
indexedKeys[f.FactKey] = struct{}{}
line := fmt.Sprintf("- [%s] %s — %s (%s)", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence)
line += FormatFactIndexLinksHint(f.FactKey, incomingByTarget[f.FactKey])
line += "\n"
lineRunes := len([]rune(line))
if used+lineRunes > factsBudget {
omitted++
continue
}
b.WriteString(line)
used += lineRunes
}
if omitted > 0 {
b.WriteString(fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", omitted))
}
if pathSection := BuildFactPathOverviewSection(allEdges, indexedKeys, pathMaxRunes); pathSection != "" {
b.WriteString("\n")
b.WriteString(pathSection)
}
b.WriteString(footer)
return wrapFactIndexBlock(b.String()), nil
}
+56
View File
@@ -0,0 +1,56 @@
package project
import "strings"
// FactIndexSectionHeading 黑板索引可读标题行前缀(块内保留,供 Agent 阅读)。
const FactIndexSectionHeading = "## 项目黑板索引"
// FactIndexSectionStartMarker / EndMarkerHTML 注释边界,供程序化替换;对模型无指令语义。
const (
FactIndexSectionStartMarker = "<!-- fact-index-start -->"
FactIndexSectionEndMarker = "<!-- fact-index-end -->"
)
// ReplaceFactIndexSection 用 freshIndex 替换 content 中已有的项目黑板索引段。
// freshIndex 须为 BuildFactIndexBlock 的完整输出。起止 HTML 注释缺失时返回 (_, false)。
func ReplaceFactIndexSection(content, freshIndex string) (string, bool) {
freshIndex = strings.TrimSpace(freshIndex)
if freshIndex == "" {
return content, false
}
start, ok := factIndexSectionStart(content)
if !ok {
return content, false
}
end, ok := factIndexSectionEnd(content, start)
if !ok || end <= start {
return content, false
}
return content[:start] + freshIndex + content[end:], true
}
// wrapFactIndexBlock 为 BuildFactIndexBlock 正文加上统一起止 HTML 注释边界。
func wrapFactIndexBlock(content string) string {
content = strings.TrimSpace(content)
return FactIndexSectionStartMarker + "\n" + content + "\n" + FactIndexSectionEndMarker + "\n"
}
func factIndexSectionStart(content string) (int, bool) {
idx := strings.Index(content, FactIndexSectionStartMarker)
if idx < 0 {
return 0, false
}
return idx, true
}
func factIndexSectionEnd(content string, start int) (int, bool) {
if start < 0 || start >= len(content) {
return 0, false
}
tail := content[start:]
idx := strings.LastIndex(tail, FactIndexSectionEndMarker)
if idx < 0 {
return 0, false
}
return start + idx + len(FactIndexSectionEndMarker), true
}
+154
View File
@@ -0,0 +1,154 @@
package project
import (
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func sampleFactIndexWithFacts(projectLabel, summary string) string {
return wrapFactIndexBlock("## 项目黑板索引(project: " + projectLabel + ", id: x\n" +
"- [target/a] target — " + summary + " (tentative)\n" +
factIndexFooterGetDetail + "\n" +
factIndexFooterWriteHint)
}
func TestReplaceFactIndexSection(t *testing.T) {
t.Parallel()
oldIndex := sampleFactIndexWithFacts("p1", "old summary")
newIndex := sampleFactIndexWithFacts("p1", "new summary")
t.Run("replaces index before next section", func(t *testing.T) {
content := "你是助手\n\n" + oldIndex + "\n\n## 图片分析\n看截图"
out, ok := ReplaceFactIndexSection(content, newIndex)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "old summary") {
t.Fatalf("old index should be gone: %q", out)
}
if !strings.Contains(out, "new summary") || !strings.Contains(out, "## 图片分析") {
t.Fatalf("expected new index and preserved vision section: %q", out)
}
if strings.Count(out, FactIndexSectionStartMarker) != 1 || strings.Count(out, FactIndexSectionEndMarker) != 1 {
t.Fatalf("expected exactly one start/end marker pair: %q", out)
}
})
t.Run("replaces index at end", func(t *testing.T) {
content := "## 项目测试范围\nscope\n\n" + oldIndex
out, ok := ReplaceFactIndexSection(content, newIndex)
if !ok {
t.Fatal("expected replacement")
}
if !strings.Contains(out, "## 项目测试范围") || !strings.Contains(out, "new summary") {
t.Fatalf("scope preserved, index updated: %q", out)
}
})
t.Run("summary with false markdown header does not truncate early", func(t *testing.T) {
summaryWithFakeHeader := "see\n\n## fake header in summary"
old := sampleFactIndexWithFacts("p1", summaryWithFakeHeader)
newIdx := sampleFactIndexWithFacts("p1", "new summary")
content := old + "\n\n## 图片分析\nvision"
out, ok := ReplaceFactIndexSection(content, newIdx)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "fake header in summary") {
t.Fatalf("old index tail should be fully removed: %q", out)
}
})
t.Run("summary containing end marker text does not truncate early", func(t *testing.T) {
summary := "note " + FactIndexSectionEndMarker + " in summary"
old := sampleFactIndexWithFacts("p1", summary)
newIdx := sampleFactIndexWithFacts("p1", "clean")
content := old + "\n\n## 图片分析\nvision"
out, ok := ReplaceFactIndexSection(content, newIdx)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "in summary") {
t.Fatalf("old block should be fully removed: %q", out)
}
})
t.Run("missing html markers does not replace", func(t *testing.T) {
legacy := "## 项目黑板索引(project: p1, id: x\n- [a] note — old (tentative)\n"
newIdx := sampleFactIndexWithFacts("p1", "new")
out, ok := ReplaceFactIndexSection("prefix\n\n"+legacy, newIdx)
if ok {
t.Fatalf("expected no replacement without markers: %q", out)
}
})
t.Run("empty facts block", func(t *testing.T) {
oldEmpty := wrapFactIndexBlock("## 项目黑板索引(project: p1, id: x\n(暂无事实)\n" + factIndexFooterEmpty)
newEmpty := sampleFactIndexWithFacts("p1", "first fact")
out, ok := ReplaceFactIndexSection(oldEmpty, newEmpty)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "(暂无事实)") {
t.Fatalf("old empty block should be gone: %q", out)
}
})
t.Run("no marker", func(t *testing.T) {
_, ok := ReplaceFactIndexSection("no blackboard here", newIndex)
if ok {
t.Fatal("expected false when marker missing")
}
})
t.Run("empty fresh index", func(t *testing.T) {
_, ok := ReplaceFactIndexSection(oldIndex, " ")
if ok {
t.Fatal("expected false for empty fresh index")
}
})
}
func TestFactIndexSectionBounds_useHTMLMarkers(t *testing.T) {
t.Parallel()
body := sampleFactIndexWithFacts("p", "line with\n\n## not a real section") + "TAIL_SHOULD_DROP"
start, ok := factIndexSectionStart(body)
if !ok || !strings.HasPrefix(body[start:], FactIndexSectionStartMarker) {
t.Fatalf("start should be at html start marker, got %d", start)
}
end, ok := factIndexSectionEnd(body, start)
if !ok || body[end:] != "\nTAIL_SHOULD_DROP" {
t.Fatalf("end should be after end marker, got remainder %q", body[end:])
}
}
func TestBuildFactIndexBlock_includesHTMLMarkers(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "marker-proj"})
if err != nil {
t.Fatal(err)
}
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true})
if err != nil {
t.Fatal(err)
}
if !strings.HasPrefix(strings.TrimSpace(block), FactIndexSectionStartMarker) {
t.Fatalf("block should start with start marker: %q", block)
}
if !strings.Contains(block, FactIndexSectionEndMarker) {
t.Fatalf("block should include end marker: %q", block)
}
}
+256
View File
@@ -0,0 +1,256 @@
package project
import (
"fmt"
"regexp"
"strings"
"cyberstrike-ai/internal/database"
)
var (
bodyDepFactLine = regexp.MustCompile(`(?im)^[\s\-*]*依赖事实\s*[:]\s*([a-zA-Z0-9][a-zA-Z0-9._/-]*)`)
bodyRelFactLine = regexp.MustCompile(`(?im)^[\s\-*]*相关\s*fact_key\s*[:]\s*([a-zA-Z0-9][a-zA-Z0-9._/-]*)`)
bodyAssocSection = regexp.MustCompile(`(?im)^##\s*关联\s*$`)
bodySyncLinksHead = "结构化关系边(自动同步)"
)
// ParseLinksFromBody 从 body「关联」段落解析 from 语义的关系边(无显式 links 时的兜底)。
func ParseLinksFromBody(body string) []database.ProjectFactEdgeFromInput {
body = strings.TrimSpace(body)
if body == "" {
return nil
}
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeFromInput
add := func(key, edgeType string) {
key = strings.TrimSpace(key)
if key == "" {
return
}
if err := database.ValidateFactKey(key); err != nil {
return
}
sig := edgeType + "\x00" + key
if _, ok := seen[sig]; ok {
return
}
seen[sig] = struct{}{}
out = append(out, database.ProjectFactEdgeFromInput{From: key, Type: edgeType})
}
for _, m := range bodyDepFactLine.FindAllStringSubmatch(body, -1) {
if len(m) > 1 {
add(m[1], "depends_on")
}
}
for _, m := range bodyRelFactLine.FindAllStringSubmatch(body, -1) {
if len(m) > 1 {
add(m[1], "supports")
}
}
// 自动同步块:type: key
syncBlock := extractBodySyncLinksBlock(body)
for _, line := range strings.Split(syncBlock, "\n") {
line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "-"))
if line == "" {
continue
}
edgeType, source, ok := strings.Cut(line, ":")
if !ok {
continue
}
edgeType = strings.TrimSpace(edgeType)
source = strings.TrimSpace(source)
if err := database.ValidateProjectFactEdgeType(edgeType); err != nil {
continue
}
add(source, edgeType)
}
if len(out) == 0 {
return nil
}
return out
}
func extractBodySyncLinksBlock(body string) string {
lines := strings.Split(body, "\n")
var b strings.Builder
inAssoc := false
inSync := false
for _, line := range lines {
trim := strings.TrimSpace(line)
if bodyAssocSection.MatchString(trim) {
inAssoc = true
inSync = false
continue
}
if inAssoc && strings.HasPrefix(trim, "## ") && !strings.HasPrefix(trim, "## 关联") {
break
}
if inAssoc && strings.Contains(trim, bodySyncLinksHead) {
inSync = true
continue
}
if inSync {
if trim == "" || strings.HasPrefix(trim, "-") || strings.Contains(trim, ":") {
if strings.HasPrefix(trim, "-") || (strings.Contains(trim, ":") && !strings.Contains(trim, "related_vulnerability")) {
b.WriteString(trim)
b.WriteByte('\n')
}
} else if strings.HasPrefix(trim, "##") {
break
}
}
}
return b.String()
}
// SyncBodyLinksSection 将入边镜像写入 body 的「关联」段(人读用;结构化以 links 为准)。
func SyncBodyLinksSection(body string, edges []*database.ProjectFactEdge) string {
body = strings.TrimSpace(body)
block := formatBodySyncLinksBlock(edges)
if block == "" {
return body
}
if body == "" {
return "## 关联\n" + block
}
lines := strings.Split(body, "\n")
var out []string
inAssoc := false
replaced := false
for i := 0; i < len(lines); i++ {
trim := strings.TrimSpace(lines[i])
if bodyAssocSection.MatchString(trim) {
inAssoc = true
out = append(out, lines[i])
// 跳过旧同步块
j := i + 1
for j < len(lines) {
t := strings.TrimSpace(lines[j])
if strings.HasPrefix(t, "## ") {
break
}
if strings.Contains(t, bodySyncLinksHead) {
for j < len(lines) {
t2 := strings.TrimSpace(lines[j])
if t2 != "" && !strings.HasPrefix(t2, "-") && !strings.Contains(t2, ":") && !strings.Contains(t2, bodySyncLinksHead) {
if strings.HasPrefix(t2, "##") {
break
}
}
j++
if j < len(lines) && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
break
}
if j >= len(lines) {
break
}
if j > i+1 && strings.TrimSpace(lines[j-1]) == "" && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
break
}
}
break
}
j++
}
out = append(out, block)
i = j - 1
replaced = true
continue
}
out = append(out, lines[i])
}
if !replaced {
if !inAssoc {
out = append(out, "", "## 关联", block)
} else {
out = append(out, block)
}
}
return strings.TrimSpace(strings.Join(out, "\n"))
}
func formatBodySyncLinksBlock(edges []*database.ProjectFactEdge) string {
if len(edges) == 0 {
return fmt.Sprintf("- %s:\n (暂无)", bodySyncLinksHead)
}
var b strings.Builder
b.WriteString("- ")
b.WriteString(bodySyncLinksHead)
b.WriteString(":\n")
for _, e := range edges {
b.WriteString(fmt.Sprintf(" - %s: %s\n", e.EdgeType, e.SourceFactKey))
}
return strings.TrimRight(b.String(), "\n")
}
// ResolveFactLinksForUpsert 合并显式 links、links_text 与 body 解析结果。
func ResolveFactLinksForUpsert(explicit []database.ProjectFactEdgeFromInput, linksText *string, body string, explicitSet bool) ([]database.ProjectFactEdgeFromInput, bool, error) {
if explicitSet {
if len(explicit) > 0 {
return explicit, true, nil
}
if linksText != nil {
parsed, err := ParseFactLinksText(*linksText)
if err != nil {
return nil, true, err
}
if parsed == nil {
return []database.ProjectFactEdgeFromInput{}, true, nil
}
return parsed, true, nil
}
return []database.ProjectFactEdgeFromInput{}, true, nil
}
if parsed := ParseLinksFromBody(body); len(parsed) > 0 {
return parsed, true, nil
}
return nil, false, nil
}
// MergeLinkFromInputsUnique 合并多组 from 入边输入并去重。
func MergeLinkFromInputsUnique(groups ...[]database.ProjectFactEdgeFromInput) []database.ProjectFactEdgeFromInput {
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeFromInput
for _, g := range groups {
for _, in := range g {
sig := in.Type + "\x00" + in.From
if _, ok := seen[sig]; ok {
continue
}
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
continue
}
if err := database.ValidateFactKey(in.From); err != nil {
continue
}
seen[sig] = struct{}{}
out = append(out, in)
}
}
return out
}
// MergeLinkInputsUnique 合并多组 link 输入并去重(内部出边写入用)。
func MergeLinkInputsUnique(groups ...[]database.ProjectFactEdgeInput) []database.ProjectFactEdgeInput {
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeInput
for _, g := range groups {
for _, in := range g {
sig := in.Type + "\x00" + in.To
if _, ok := seen[sig]; ok {
continue
}
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
continue
}
if err := database.ValidateFactKey(in.To); err != nil {
continue
}
seen[sig] = struct{}{}
out = append(out, in)
}
}
return out
}
+68
View File
@@ -0,0 +1,68 @@
package project
import (
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestParseLinksFromBodyDependsOn(t *testing.T) {
t.Parallel()
body := "## 关联\n- 依赖事实: target/api\n- 相关 fact_key: auth/session"
links := ParseLinksFromBody(body)
if len(links) != 2 {
t.Fatalf("want 2 links, got %d", len(links))
}
}
func TestSyncBodyLinksSection(t *testing.T) {
t.Parallel()
body := "## 结论\nx\n\n## 关联\n- 依赖事实: old/key"
edges := []*database.ProjectFactEdge{{EdgeType: "discovered_on", SourceFactKey: "target/a"}}
out := SyncBodyLinksSection(body, edges)
if !strings.Contains(out, "discovered_on: target/a") {
t.Fatalf("missing synced edge: %q", out)
}
}
func TestFactGraphIntegration(t *testing.T) {
dir := t.TempDir()
dbPath := filepath.Join(dir, "test.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "g"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat, summary string }{
{"target/root", "target", "root"},
{"finding/x", "finding", "finding x"},
} {
_, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.summary, Confidence: "confirmed",
})
if err != nil {
t.Fatal(err)
}
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/x", []database.ProjectFactEdgeFromInput{
{From: "target/root", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
if len(graph.Nodes) < 2 || len(graph.Edges) < 1 {
t.Fatalf("expected graph nodes/edges, got %d/%d", len(graph.Nodes), len(graph.Edges))
}
}
+407
View File
@@ -0,0 +1,407 @@
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:<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
}
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()
}
+96
View File
@@ -0,0 +1,96 @@
package project
import (
"cyberstrike-ai/internal/database"
)
// ApplyFactOutgoingLinks 替换某事实的出边(links 为 nil 时不修改)。
func ApplyFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput) error {
if links == nil {
return nil
}
return db.ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID, links)
}
// ResolveFactLinkInputs 合并 links 数组与 links_text 文本(数组优先)。
func ResolveFactLinkInputs(links []database.ProjectFactEdgeFromInput, linksText string) ([]database.ProjectFactEdgeFromInput, error) {
if len(links) > 0 {
return links, nil
}
return ParseFactLinksText(linksText)
}
// ApplyFactIncomingLinks 替换某事实的入边(links 为 nil 时不修改)。
func ApplyFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput) error {
if links == nil {
return nil
}
return db.ReplaceIncomingProjectFactEdges(projectID, targetFactKey, links)
}
// PersistFactIncomingLinks 写入入边并可选同步当前事实 body「关联」段。
func PersistFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput, syncBody bool) error {
if links == nil {
return nil
}
if err := ApplyFactIncomingLinks(db, projectID, targetFactKey, links); err != nil {
return err
}
if !syncBody {
return nil
}
f, err := db.GetProjectFactByKey(projectID, targetFactKey)
if err != nil {
return nil
}
in, err := db.ListIncomingProjectFactEdges(projectID, targetFactKey)
if err != nil {
return err
}
f.Body = SyncBodyLinksSection(f.Body, in)
_, err = db.UpsertProjectFact(f)
return err
}
// PersistFactLinksFromParsed 写入解析后的 linksparsed 为 nil 表示不修改)。
func PersistFactLinksFromParsed(db *database.DB, projectID, factKey, sourceConversationID string, parsed *ParsedFactLinks, syncBody bool) error {
if parsed == nil || parsed.Incoming == nil {
return nil
}
return PersistFactIncomingLinks(db, projectID, factKey, parsed.Incoming, syncBody)
}
// PersistFactOutgoingLinks 写入出边(图连线等低层 APIbody 同步请用 PersistFactIncomingLinks)。
func PersistFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput, syncBody bool) error {
if links == nil {
return nil
}
return ApplyFactOutgoingLinks(db, projectID, sourceFactKey, sourceConversationID, links)
}
// LinkCountMap 项目内各 fact 的入/出边计数。
type LinkCountMap map[string]LinkCounts
// LinkCounts 单 fact 的入/出边数。
type LinkCounts struct {
Outgoing int `json:"outgoing"`
Incoming int `json:"incoming"`
}
// LoadProjectFactLinkCounts 批量加载边计数。
func LoadProjectFactLinkCounts(db *database.DB, projectID string) (LinkCountMap, error) {
edges, err := db.ListProjectFactEdgesByProject(projectID)
if err != nil {
return nil, err
}
m := LinkCountMap{}
for _, e := range edges {
c := m[e.SourceFactKey]
c.Outgoing++
m[e.SourceFactKey] = c
c = m[e.TargetFactKey]
c.Incoming++
m[e.TargetFactKey] = c
}
return m, nil
}
+296
View File
@@ -0,0 +1,296 @@
package project
import (
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestParseFactLinksText(t *testing.T) {
t.Parallel()
inputs, err := ParseFactLinksText("discovered_on: target/api\nleads_to: finding/swagger")
if err != nil {
t.Fatal(err)
}
if len(inputs) != 2 {
t.Fatalf("want 2 links, got %d", len(inputs))
}
if inputs[0].Type != "discovered_on" || inputs[0].From != "target/api" {
t.Fatalf("unexpected first link: %+v", inputs[0])
}
}
func TestParseFactIncomingLinksText(t *testing.T) {
t.Parallel()
inputs, err := ParseFactIncomingLinksText("leads_to: finding/swagger\ndepends_on: target/api")
if err != nil {
t.Fatal(err)
}
if len(inputs) != 2 {
t.Fatalf("want 2 links, got %d", len(inputs))
}
if inputs[0].Type != "leads_to" || inputs[0].From != "finding/swagger" {
t.Fatalf("unexpected first link: %+v", inputs[0])
}
}
func TestFormatFactIncomingLinksText(t *testing.T) {
t.Parallel()
text := FormatFactIncomingLinksText([]*database.ProjectFactEdge{
{EdgeType: "leads_to", SourceFactKey: "finding/a"},
{EdgeType: "depends_on", SourceFactKey: "target/b"},
})
want := "leads_to: finding/a\ndepends_on: target/b"
if text != want {
t.Fatalf("got %q want %q", text, want)
}
}
func TestParseFactLinkInputsEmptyClears(t *testing.T) {
t.Parallel()
parsed, err := ParseFactLinkInputs([]interface{}{})
if err != nil {
t.Fatal(err)
}
if parsed == nil || parsed.Incoming == nil || len(parsed.Incoming) != 0 {
t.Fatalf("empty array should clear incoming links, got %v", parsed)
}
}
func TestParseFactLinkInputsFrom(t *testing.T) {
t.Parallel()
raw := []interface{}{
map[string]interface{}{
"from": "target/primary_domain",
"type": "discovered_on",
},
}
parsed, err := ParseFactLinkInputs(raw)
if err != nil {
t.Fatal(err)
}
if len(parsed.Incoming) != 1 || parsed.Incoming[0].From != "target/primary_domain" {
t.Fatalf("unexpected incoming: %+v", parsed.Incoming)
}
}
func TestParseFactLinkInputsRequiresFrom(t *testing.T) {
t.Parallel()
raw := []interface{}{
map[string]interface{}{
"to": "target/primary_domain",
"type": "discovered_on",
},
}
_, err := ParseFactLinkInputs(raw)
if err == nil {
t.Fatal("expected error when from is missing")
}
}
func TestGraphNodeType(t *testing.T) {
t.Parallel()
if GraphNodeType("chain", "chain/x") != "chain" {
t.Fatal("chain category")
}
if GraphNodeType("finding", "finding/x") != "finding" {
t.Fatal("finding category")
}
if GraphNodeType("exploit", "exploit/x") != "exploit" {
t.Fatal("exploit category")
}
if GraphNodeType("finding", "evidence/x") != "finding" {
t.Fatal("category should override evidence key prefix")
}
if GraphNodeType("note", "target/x") != "note" {
t.Fatal("category should override target key prefix")
}
if GraphNodeType("vuln", "finding/x") != "vulnerability" {
t.Fatal("vuln category maps to vulnerability node type")
}
if GraphNodeType("", "target/x") != "target" {
t.Fatal("empty category falls back to target key prefix")
}
}
func TestBuildProjectFactGraphPreservesStoredEdgeDirection(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "path-edges"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat string }{
{"target/primary_domain", "target"},
{"chain/full_attack_path", "chain"},
{"finding/mysql_public", "finding"},
{"exploit/mysql_creds_extract", "exploit"},
} {
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
{From: "exploit/mysql_creds_extract", Type: "exploits"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "chain/full_attack_path", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/mysql_creds_extract", []database.ProjectFactEdgeFromInput{
{From: "chain/full_attack_path", Type: "leads_to"},
}); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
want := map[string]struct{}{
"target/primary_domain|discovered_on|finding/mysql_public": {},
"exploit/mysql_creds_extract|exploits|finding/mysql_public": {},
"target/primary_domain|discovered_on|chain/full_attack_path": {},
"chain/full_attack_path|leads_to|exploit/mysql_creds_extract": {},
}
for _, e := range graph.Edges {
key := e.Source + "|" + e.Type + "|" + e.Target
delete(want, key)
}
if len(want) > 0 {
t.Fatalf("missing expected stored-direction edges: %v", want)
}
countInOut := func(factKey string) (out, in int) {
for _, e := range graph.Edges {
if e.Source == factKey {
out++
}
if e.Target == factKey {
in++
}
}
return out, in
}
if out, in := countInOut("chain/full_attack_path"); out != 1 || in != 1 {
t.Fatalf("chain/full_attack_path want out=1 in=1 got out=%d in=%d", out, in)
}
if out, in := countInOut("exploit/mysql_creds_extract"); out != 1 || in != 1 {
t.Fatalf("exploit/mysql_creds_extract want out=1 in=1 got out=%d in=%d", out, in)
}
}
func TestPersistFactLinksFromUsesFromAsIncoming(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "from-links"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat string }{
{"target/primary_domain", "target"},
{"finding/sqli", "finding"},
} {
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
}
parsed := &ParsedFactLinks{
Incoming: []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
},
}
if err := PersistFactLinksFromParsed(db, p.ID, "finding/sqli", "", parsed, false); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
want := "target/primary_domain|discovered_on|finding/sqli"
for _, e := range graph.Edges {
key := e.Source + "|" + e.Type + "|" + e.Target
if key == want {
return
}
}
t.Fatalf("expected edge %s, got %+v", want, graph.Edges)
}
func TestFormatOutgoingLinksHint(t *testing.T) {
t.Parallel()
hint := FormatOutgoingLinksHint([]*database.ProjectFactEdge{
{EdgeType: "discovered_on", TargetFactKey: "target/a"},
})
if hint == "" || hint[0] != ' ' {
t.Fatalf("unexpected hint: %q", hint)
}
}
func TestReplaceIncomingAllowsNotYetCreatedSource(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "parallel-links"})
if err != nil {
t.Fatal(err)
}
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: "exploit/sqli", Category: "exploit", Summary: "exploit", Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/sqli", []database.ProjectFactEdgeFromInput{
{From: "finding/sqli_endpoint", Type: "exploits"},
}); err != nil {
t.Fatalf("incoming edge should not require source fact to exist yet: %v", err)
}
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: "finding/sqli_endpoint", Category: "finding", Summary: "finding", Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
in, err := db.ListIncomingProjectFactEdges(p.ID, "exploit/sqli")
if err != nil || len(in) != 1 || in[0].SourceFactKey != "finding/sqli_endpoint" {
t.Fatalf("expected persisted edge from finding, got %+v err=%v", in, err)
}
}
func TestValidateProjectFactEdgeType(t *testing.T) {
t.Parallel()
if err := database.ValidateProjectFactEdgeType("leads_to"); err != nil {
t.Fatal(err)
}
if err := database.ValidateProjectFactEdgeType("invalid"); err == nil {
t.Fatal("expected error")
}
}
+231
View File
@@ -0,0 +1,231 @@
package project
import (
"fmt"
"sort"
"strings"
"cyberstrike-ai/internal/database"
)
var factIndexEdgeTypeOrder = []string{
"discovered_on", "leads_to", "enables", "depends_on", "exploits", "contains", "part_of", "supports",
}
func filterIndexEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge {
if len(edges) == 0 {
return nil
}
out := make([]*database.ProjectFactEdge, 0, len(edges))
for _, e := range edges {
if e == nil {
continue
}
if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") {
continue
}
edgeType := strings.ToLower(strings.TrimSpace(e.EdgeType))
if _, ok := database.ValidProjectFactEdgeTypes[edgeType]; !ok {
continue
}
out = append(out, e)
}
return out
}
func edgeConfidenceSuffix(confidence string) string {
c := strings.ToLower(strings.TrimSpace(confidence))
if c == "" || c == "confirmed" {
return ""
}
return " (" + c + ")"
}
func formatRelationHintPart(e *database.ProjectFactEdge) string {
return fmt.Sprintf("%s←%s%s", e.EdgeType, e.SourceFactKey, edgeConfidenceSuffix(e.Confidence))
}
func formatOutgoingHintPart(e *database.ProjectFactEdge) string {
return fmt.Sprintf("%s→%s%s", e.EdgeType, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence))
}
func formatIncomingHintPart(e *database.ProjectFactEdge) string {
return formatRelationHintPart(e)
}
func joinEdgeHintParts(edges []*database.ProjectFactEdge, formatter func(*database.ProjectFactEdge) string) string {
parts := make([]string, 0, len(edges))
for _, e := range edges {
parts = append(parts, formatter(e))
}
return strings.Join(parts, ", ")
}
// FormatOutgoingLinksHint 黑板索引用出边摘要(全部有效边类型,不截断)。
func FormatOutgoingLinksHint(edges []*database.ProjectFactEdge) string {
edges = filterIndexEdges(edges)
if len(edges) == 0 {
return ""
}
return " {出边: " + joinEdgeHintParts(edges, formatOutgoingHintPart) + "}"
}
// FormatIncomingLinksHint 黑板索引用入边摘要(全部有效边类型,不截断)。
func FormatIncomingLinksHint(edges []*database.ProjectFactEdge) string {
edges = filterIndexEdges(edges)
if len(edges) == 0 {
return ""
}
return " {入边: " + joinEdgeHintParts(edges, formatIncomingHintPart) + "}"
}
// FormatFactIndexLinksHint 黑板索引行内关系边(from → 当前 fact,与 upsert links 一致)。
func FormatFactIndexLinksHint(_ string, incoming []*database.ProjectFactEdge) string {
in := filterIndexEdges(incoming)
if len(in) == 0 {
return ""
}
return " {关系边: " + joinEdgeHintParts(in, formatRelationHintPart) + "}"
}
func indexEdgeGroupMaps(edges []*database.ProjectFactEdge) (outgoing, incoming map[string][]*database.ProjectFactEdge) {
outgoing = map[string][]*database.ProjectFactEdge{}
incoming = map[string][]*database.ProjectFactEdge{}
for _, e := range filterIndexEdges(edges) {
outgoing[e.SourceFactKey] = append(outgoing[e.SourceFactKey], e)
incoming[e.TargetFactKey] = append(incoming[e.TargetFactKey], e)
}
return outgoing, incoming
}
func relationOverviewLine(e *database.ProjectFactEdge) string {
return fmt.Sprintf("- %s → %s%s · %s", e.SourceFactKey, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence), e.EdgeType)
}
func indexEdgeSortKey(e *database.ProjectFactEdge) (int, int, string) {
confRank := 0
if strings.EqualFold(strings.TrimSpace(e.Confidence), "tentative") {
confRank = 1
}
typeRank := len(factIndexEdgeTypeOrder) + 1
for i, t := range factIndexEdgeTypeOrder {
if strings.EqualFold(e.EdgeType, t) {
typeRank = i
break
}
}
return confRank, typeRank, e.SourceFactKey + ">" + e.TargetFactKey + ">" + e.EdgeType
}
func sortIndexOverviewEdges(edges []*database.ProjectFactEdge) {
sort.SliceStable(edges, func(i, j int) bool {
ci, ti, ki := indexEdgeSortKey(edges[i])
cj, tj, kj := indexEdgeSortKey(edges[j])
if ci != cj {
return ci < cj
}
if ti != tj {
return ti < tj
}
return ki < kj
})
}
// BuildFactPathOverviewSection 生成事实关系速览(全部有效边类型,不含 body)。
func BuildFactPathOverviewSection(edges []*database.ProjectFactEdge, indexedKeys map[string]struct{}, maxRunes int) string {
if maxRunes <= 0 {
return ""
}
candidates := filterIndexEdges(edges)
if len(candidates) == 0 {
return ""
}
filtered := make([]*database.ProjectFactEdge, 0, len(candidates))
for _, e := range candidates {
if len(indexedKeys) > 0 {
if _, ok := indexedKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := indexedKeys[e.TargetFactKey]; !ok {
continue
}
}
filtered = append(filtered, e)
}
if len(filtered) == 0 {
return ""
}
sortIndexOverviewEdges(filtered)
header := "### 攻击路径(事实关系)\n"
header += "source → target · type(与攻击路径图/库中方向一致;写入时在目标 fact 的 links 用 from 声明来源)\n"
var b strings.Builder
b.WriteString(header)
used := len([]rune(header))
omitted := 0
for _, e := range filtered {
line := relationOverviewLine(e) + "\n"
lineRunes := len([]rune(line))
if used+lineRunes > maxRunes {
omitted++
continue
}
b.WriteString(line)
used += lineRunes
}
if omitted > 0 {
extra := fmt.Sprintf("(另有 %d 条关系边未列入,请 get_project_fact 查看完整关系。)\n", omitted)
if used+len([]rune(extra)) <= maxRunes {
b.WriteString(extra)
}
}
if used <= len([]rune(header)) {
return ""
}
return b.String()
}
func factIndexSortPriority(f *database.ProjectFact) int {
if f == nil {
return 0
}
score := 0
if f.Pinned {
score += 1000
}
c := strings.ToLower(strings.TrimSpace(f.Category))
switch c {
case FactCategoryTarget:
score += 400
case FactCategoryFinding, FactCategoryChain:
score += 300
case FactCategoryExploit, FactCategoryPOC:
score += 250
case "auth", "infra", "business":
score += 200
case "note":
score += 50
default:
key := strings.ToLower(strings.TrimSpace(f.FactKey))
if strings.HasPrefix(key, "target/") {
score += 400
} else if strings.HasPrefix(key, "finding/") || strings.HasPrefix(key, "chain/") {
score += 300
}
}
if strings.EqualFold(strings.TrimSpace(f.Confidence), "confirmed") {
score += 80
}
return score
}
func sortFactsForIndex(facts []*database.ProjectFact) {
sort.SliceStable(facts, func(i, j int) bool {
pi, pj := factIndexSortPriority(facts[i]), factIndexSortPriority(facts[j])
if pi != pj {
return pi > pj
}
return facts[i].UpdatedAt.After(facts[j].UpdatedAt)
})
}
+161
View File
@@ -0,0 +1,161 @@
package project
import (
"fmt"
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestFormatIncomingLinksHint(t *testing.T) {
t.Parallel()
hint := FormatIncomingLinksHint([]*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "finding/x", Confidence: "tentative"},
})
if !strings.Contains(hint, "入边:") {
t.Fatalf("expected 入边 label: %q", hint)
}
if !strings.Contains(hint, "discovered_on←finding/x") {
t.Fatalf("unexpected hint: %q", hint)
}
if !strings.Contains(hint, "tentative") {
t.Fatalf("expected tentative in hint: %q", hint)
}
}
func TestFormatIncomingLinksHint_allEdges(t *testing.T) {
t.Parallel()
edges := make([]*database.ProjectFactEdge, 0, 5)
for i := 1; i <= 5; i++ {
edges = append(edges, &database.ProjectFactEdge{
EdgeType: "discovered_on",
SourceFactKey: fmt.Sprintf("finding/f%d", i),
Confidence: "tentative",
})
}
hint := FormatIncomingLinksHint(edges)
if strings.Contains(hint, "+") {
t.Fatalf("should not truncate with +N: %q", hint)
}
for i := 1; i <= 5; i++ {
if !strings.Contains(hint, fmt.Sprintf("finding/f%d", i)) {
t.Fatalf("missing edge f%d in hint: %q", i, hint)
}
}
}
func TestFormatFactIndexLinksHint_incomingOnly(t *testing.T) {
t.Parallel()
in := []*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "target/dev", Confidence: "tentative"},
{EdgeType: "exploits", SourceFactKey: "exploit/rce", Confidence: "confirmed"},
}
hint := FormatFactIndexLinksHint("finding/sqli", in)
if !strings.Contains(hint, "关系边:") {
t.Fatalf("missing 关系边 label: %q", hint)
}
if !strings.Contains(hint, "discovered_on←target/dev") {
t.Fatalf("missing discovered_on: %q", hint)
}
if !strings.Contains(hint, "exploits←exploit/rce") {
t.Fatalf("missing exploits: %q", hint)
}
if strings.Contains(hint, "出边") || strings.Contains(hint, "入边") {
t.Fatalf("should not use legacy 出边/入边 labels: %q", hint)
}
}
func TestFormatFactIndexLinksHint_includesAuxiliaryEdgeTypes(t *testing.T) {
t.Parallel()
in := []*database.ProjectFactEdge{{EdgeType: "supports", SourceFactKey: "note/log"}}
hint := FormatFactIndexLinksHint("finding/x", in)
if !strings.Contains(hint, "supports←note/log") {
t.Fatalf("supports edge should be included: %q", hint)
}
}
func TestBuildFactPathOverviewSection(t *testing.T) {
t.Parallel()
edges := []*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "target/dev", TargetFactKey: "finding/sqli", Confidence: "tentative"},
{EdgeType: "exploits", SourceFactKey: "exploit/rce", TargetFactKey: "finding/sqli", Confidence: "confirmed"},
{EdgeType: "supports", SourceFactKey: "note/log", TargetFactKey: "finding/sqli"},
}
keys := map[string]struct{}{
"target/dev": {}, "finding/sqli": {}, "exploit/rce": {}, "note/log": {},
}
section := BuildFactPathOverviewSection(edges, keys, 800)
if !strings.Contains(section, "### 攻击路径(事实关系)") {
t.Fatalf("missing header: %q", section)
}
if !strings.Contains(section, "target/dev → finding/sqli") {
t.Fatalf("missing discovered_on line: %q", section)
}
if !strings.Contains(section, "exploit/rce → finding/sqli") {
t.Fatalf("missing exploits line: %q", section)
}
if !strings.Contains(section, "note/log → finding/sqli") {
t.Fatalf("supports edge should be included: %q", section)
}
}
func TestBuildFactIndexBlock_withLinksAndPathOverview(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "path-proj"})
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "target/dev",
Category: "target",
Summary: "dev 子域",
Confidence: "confirmed",
})
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "finding/sqli",
Category: "finding",
Summary: "时间盲注",
Confidence: "tentative",
})
if err != nil {
t.Fatal(err)
}
_, err = db.AddProjectFactEdge(proj.ID, database.ProjectFactEdgeInput{
To: "finding/sqli",
Type: "discovered_on",
}, "target/dev", "")
if err != nil {
t.Fatal(err)
}
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true, FactIndexMaxRunes: 6500, FactIndexPathMaxRunes: 1000})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(block, "关系边: discovered_on←target/dev") {
t.Fatalf("finding line should include relation hint: %q", block)
}
if !strings.Contains(block, "### 攻击路径(事实关系)") {
t.Fatalf("missing relation overview: %q", block)
}
if !strings.Contains(block, "target/dev → finding/sqli") {
t.Fatalf("missing overview edge: %q", block)
}
}
+23
View File
@@ -0,0 +1,23 @@
package project
import "cyberstrike-ai/internal/projectprompt"
// FactRecordingIncrementalRhythmMarkdown 见 projectprompt。
func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string {
return projectprompt.FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent)
}
// FactRecordingBlackboardSection 见 projectprompt。
func FactRecordingBlackboardSection(coordinatorDelegate bool) string {
return projectprompt.FactRecordingBlackboardSection(coordinatorDelegate)
}
// FactRecordingSubAgentSection 见 projectprompt。
func FactRecordingSubAgentSection() string {
return projectprompt.FactRecordingSubAgentSection()
}
// FactRecordingBlackboardSectionMarkdown 见 projectprompt。
func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string {
return projectprompt.FactRecordingBlackboardSectionMarkdown(coordinatorDelegate)
}
+135
View File
@@ -0,0 +1,135 @@
package project
import (
"fmt"
"strings"
"cyberstrike-ai/internal/projectprompt"
)
// 事实 category 常量(写入 upsert_project_fact 的 category 字段)。
const (
FactCategoryTarget = "target"
FactCategoryAuth = "auth"
FactCategoryInfra = "infra"
FactCategoryBusiness = "business"
FactCategoryFinding = "finding"
FactCategoryChain = "chain"
FactCategoryExploit = "exploit"
FactCategoryPOC = "poc"
FactCategoryNote = "note"
)
// RequiresAttackChainBody 判断该事实是否应携带可复现的攻击链 / exploit 详情(写在 body,非仅 summary)。
func RequiresAttackChainBody(category, factKey string) bool {
c := strings.ToLower(strings.TrimSpace(category))
switch c {
case FactCategoryFinding, FactCategoryChain, FactCategoryExploit, FactCategoryPOC, "vuln":
return true
}
key := strings.ToLower(strings.TrimSpace(factKey))
for _, prefix := range []string{"finding/", "chain/", "exploit/", "poc/"} {
if strings.HasPrefix(key, prefix) {
return true
}
}
return false
}
// IsSparseFactBody 攻击链类事实 body 过短或缺少关键段落时返回 true(软校验,不阻断写入)。
func IsSparseFactBody(category, factKey, body string) bool {
if !RequiresAttackChainBody(category, factKey) {
return false
}
body = strings.TrimSpace(body)
if body == "" {
return true
}
lower := strings.ToLower(body)
// 至少应包含可复现线索:步骤/请求/命令/代码块 之一
hasSteps := strings.Contains(lower, "攻击链") || strings.Contains(lower, "## 攻击") ||
strings.Contains(lower, "## exploit") || strings.Contains(lower, "## poc")
hasHTTP := strings.Contains(lower, "```http") || strings.Contains(lower, "```bash") ||
strings.Contains(lower, "curl ") || strings.Contains(lower, "get ") || strings.Contains(lower, "post ")
hasReq := strings.Contains(lower, "请求") || strings.Contains(lower, "响应") || strings.Contains(lower, "payload")
// 无攻击链/POC/请求等结构线索,视为仅结论性描述(不论长短)
return !(hasSteps || hasHTTP || hasReq)
}
// FactBodyTemplate 按 category 返回建议的 body Markdown 骨架(供 Agent 填入真实内容)。
func FactBodyTemplate(category, factKey string) string {
if RequiresAttackChainBody(category, factKey) {
return attackChainFactBodyTemplate
}
return envFactBodyTemplate
}
const attackChainFactBodyTemplate = `## 结论(可验证,一句话)
<勿仅写「存在漏洞」;写明类型 + 位置 + 触发条件>
## 目标与入口
- 目标: <URL / IP:Port / 主机名>
- 入口: <路径 / 接口 / 参数>
- 前置条件: <匿名 / 角色 / Cookie / 其他依赖>
## 攻击链(逐步可复现)
1. <侦察/发现>
2. <利用/触发>
3. <影响证明(读文件、RCE 回显、越权数据等)>
## Exploit / POC
### 请求
` + "```http\n<METHOD> <path> HTTP/1.1\nHost: ...\n...\n\n<body>\n```" + `
### 响应 / 现象
<关键响应片段、状态码、差异点>
### 命令 / 脚本(如有)
` + "```bash\n<command>\n```" + `
## 关键证据
- <工具输出摘要 / 截图路径 / 会话或消息 ID>
## 关联
- related_vulnerability_id: <可选,对应 record_vulnerability 的 id>
- linksupsert 参数): [{ "from": "<fact_key>", "type": "discovered_on|..." }]from → 当前 fact
- 依赖事实(body 可读镜像): <fact_key,如 auth/session_cookie>
## 备注与不确定性
<待验证假设、环境差异、绕过尝试记录>`
const envFactBodyTemplate = `## 摘要
<该事实的核心认知>
## 细节
<端口/版本/路径/凭据特征/业务规则等>
## 来源与证据
<命令输出、响应片段、发现时间>
## 关联
- 相关 fact_key: <可选>`
// FactRecordingGuidanceBlock 写入系统提示:要求事实沉淀攻击链上下文而非仅结论。
func FactRecordingGuidanceBlock() string {
return projectprompt.FactRecordingGuidanceBlock()
}
// SparseBodyWarning 攻击链类事实 body 不足时的工具返回提示(不阻断保存)。
func SparseBodyWarning(category, factKey string) string {
if !IsSparseFactBody(category, factKey, "") {
return ""
}
return fmt.Sprintf(
"\n\n⚠ 提示:category=%q / fact_key=%q 属于攻击链类事实,但 body 为空或过简。请补充完整攻击链与 POC(参考模板),便于后续审计复现。\n建议 body 骨架:\n%s",
category, factKey, FactBodyTemplate(category, factKey),
)
}
// SparseBodyWarningIfNeeded 根据实际 body 判断是否追加警告。
func SparseBodyWarningIfNeeded(category, factKey, body string) string {
if !IsSparseFactBody(category, factKey, body) {
return ""
}
return SparseBodyWarning(category, factKey)
}
+42
View File
@@ -0,0 +1,42 @@
package project
import (
"strings"
"testing"
)
func TestRequiresAttackChainBody(t *testing.T) {
cases := []struct {
cat, key string
want bool
}{
{"finding", "note/misc", true},
{"note", "finding/sqli-login", true},
{"target", "target/primary_domain", false},
{"auth", "auth/admin_cookie", false},
{"chain", "x", true},
{"", "exploit/rce-upload", true},
}
for _, tc := range cases {
if got := RequiresAttackChainBody(tc.cat, tc.key); got != tc.want {
t.Errorf("RequiresAttackChainBody(%q,%q)=%v want %v", tc.cat, tc.key, got, tc.want)
}
}
}
func TestIsSparseFactBody(t *testing.T) {
long := strings.Repeat("x", 150)
if !IsSparseFactBody("finding", "finding/x", "") {
t.Error("empty body should be sparse")
}
if !IsSparseFactBody("finding", "finding/x", long) {
t.Error("body without repro clues should be sparse")
}
body := "## 攻击链\n1. step\n## Exploit\n```http\nGET / HTTP/1.1\n```\n"
if IsSparseFactBody("finding", "finding/x", body) {
t.Error("structured body should not be sparse")
}
if IsSparseFactBody("target", "target/x", "") {
t.Error("env fact empty body is ok")
}
}
+99
View File
@@ -0,0 +1,99 @@
package project
import (
"encoding/json"
"fmt"
"strings"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
)
// projectScopePayload 解析 projects.scope_json(约定字段,可扩展)。
type projectScopePayload struct {
Targets []string `json:"targets"`
Exclude []string `json:"exclude"`
Notes string `json:"notes"`
}
// BuildScopeBlock 将项目 scope_json 格式化为 Agent 可读的授权范围块。
func BuildScopeBlock(proj *database.Project) string {
if proj == nil {
return ""
}
raw := strings.TrimSpace(proj.ScopeJSON)
if raw == "" {
return ""
}
var payload projectScopePayload
if err := json.Unmarshal([]byte(raw), &payload); err != nil {
return fmt.Sprintf("## 项目测试范围(project: %s\nscope_json 非合法 JSON,请人工核对配置)\n```\n%s\n```\n"+
"仅对明确授权目标执行测试;超出范围须停止并说明。\n", proj.Name, truncateRunes(raw, 800))
}
var b strings.Builder
b.WriteString(fmt.Sprintf("## 项目测试范围(project: %s, id: %s\n", proj.Name, proj.ID))
b.WriteString("以下为授权边界,**必须遵守**:仅测试列出的 targets,避开 exclude,不得擅自扩大范围。\n")
if len(payload.Targets) > 0 {
b.WriteString("\n**允许测试(targets**\n")
for _, t := range payload.Targets {
t = strings.TrimSpace(t)
if t != "" {
b.WriteString("- " + t + "\n")
}
}
}
if len(payload.Exclude) > 0 {
b.WriteString("\n**明确排除(exclude**\n")
for _, t := range payload.Exclude {
t = strings.TrimSpace(t)
if t != "" {
b.WriteString("- " + t + "\n")
}
}
}
if n := strings.TrimSpace(payload.Notes); n != "" {
b.WriteString("\n**说明(notes**\n" + n + "\n")
}
if len(payload.Targets) == 0 && len(payload.Exclude) == 0 && strings.TrimSpace(payload.Notes) == "" {
b.WriteString("\nscope_json 已配置但未识别 targets/exclude/notes 字段,原始内容供参考)\n```json\n")
b.WriteString(truncateRunes(raw, 1200))
b.WriteString("\n```\n")
}
b.WriteString("\n若目标不在 targets 内或命中 exclude,不得主动扫描/利用;需用户明确扩大授权后再继续。\n")
return b.String()
}
func truncateRunes(s string, max int) string {
r := []rune(s)
if len(r) <= max {
return s
}
return string(r[:max]) + "…"
}
// BuildProjectBlackboardBlock 组合测试范围 + 事实黑板索引。
func BuildProjectBlackboardBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) {
projectID = strings.TrimSpace(projectID)
if projectID == "" {
return "", nil
}
proj, err := db.GetProject(projectID)
if err != nil {
return "", err
}
parts := []string{}
if scope := strings.TrimSpace(BuildScopeBlock(proj)); scope != "" {
parts = append(parts, scope)
}
index, err := BuildFactIndexBlock(db, projectID, cfg)
if err != nil {
return "", err
}
if strings.TrimSpace(index) != "" {
parts = append(parts, index)
}
return strings.Join(parts, "\n\n"), nil
}
+40
View File
@@ -0,0 +1,40 @@
package project
import (
"strings"
"testing"
"cyberstrike-ai/internal/database"
)
func TestBuildScopeBlock_targetsExcludeNotes(t *testing.T) {
proj := &database.Project{
ID: "p1",
Name: "Acme",
ScopeJSON: `{"targets":["https://app.example.com"],"exclude":["*.cdn.example.com"],"notes":"仅 Web 层"}`,
}
block := BuildScopeBlock(proj)
if !strings.Contains(block, "https://app.example.com") {
t.Fatalf("missing target: %s", block)
}
if !strings.Contains(block, "cdn.example.com") {
t.Fatalf("missing exclude: %s", block)
}
if !strings.Contains(block, "仅 Web 层") {
t.Fatalf("missing notes: %s", block)
}
}
func TestBuildScopeBlock_empty(t *testing.T) {
if BuildScopeBlock(&database.Project{Name: "X"}) != "" {
t.Fatal("expected empty")
}
}
func TestBuildScopeBlock_invalidJSON(t *testing.T) {
proj := &database.Project{Name: "X", ScopeJSON: `{not json`}
block := BuildScopeBlock(proj)
if !strings.Contains(block, "非合法 JSON") {
t.Fatalf("unexpected: %s", block)
}
}
+21
View File
@@ -0,0 +1,21 @@
package project
import "cyberstrike-ai/internal/database"
// GetProjectStats 聚合项目统计(含待补全事实数)。
func GetProjectStats(db *database.DB, projectID string) (*database.ProjectStats, error) {
stats, err := db.GetProjectStatsCounts(projectID)
if err != nil {
return nil, err
}
rows, err := db.ListProjectFactsForSparseCheck(projectID)
if err != nil {
return nil, err
}
for _, r := range rows {
if IsSparseFactBody(r.Category, r.FactKey, r.Body) {
stats.SparseFactCount++
}
}
return stats, nil
}
+26
View File
@@ -0,0 +1,26 @@
package project
import "strings"
// VisionImageSectionMarker 图片分析 section 标题(与 AppendVisionImageAnalysisIfReady 注入一致)。
const VisionImageSectionMarker = "## 图片分析"
// VisionImageAnalysisSection 单/多代理共用的图片分析提示(analyze_image;上下文仅保留文字摘要)。
func VisionImageAnalysisSection() string {
var b strings.Builder
b.WriteString(VisionImageSectionMarker)
b.WriteString("\n\n")
b.WriteString("- 遇到图片文件(截图、验证码、登录页、报告配图)时,若存在工具 analyze_image,请传入服务器上的文件路径进行分析。\n")
b.WriteString("- 不要对二进制图片使用 read_file 指望理解内容;用户消息中「📎 xxx.png: /path」即为可传给 analyze_image 的路径。\n")
b.WriteString("- 验证码类:若已从页面或接口保存为本地图片(如 captcha.png),用 analyze_imagequestion 写明「只输出验证码字符」;识别失败则刷新验证码后重新保存再识;复杂滑块/行为验证码勿指望单次识图成功。\n")
b.WriteString("- 委派子代理时,若子任务含验证码/截图识读,在 task description 中写明图片路径与期望输出格式。\n")
return b.String()
}
// AppendVisionImageAnalysisIfReady 仅在 vision.enabled 且 model 已配置时追加图片分析提示。
func AppendVisionImageAnalysisIfReady(base string, visionReady bool) string {
if !visionReady {
return base
}
return AppendSystemPromptBlock(base, VisionImageAnalysisSection())
}
+69
View File
@@ -0,0 +1,69 @@
package project
import (
"fmt"
"os"
"path/filepath"
"strings"
)
func sanitizeWorkspacePathSegment(s string) string {
s = strings.TrimSpace(s)
if s == "" {
return "default"
}
s = strings.ReplaceAll(s, string(filepath.Separator), "-")
s = strings.ReplaceAll(s, "/", "-")
s = strings.ReplaceAll(s, "\\", "-")
s = strings.ReplaceAll(s, "..", "__")
if len(s) > 180 {
s = s[:180]
}
return s
}
// WorkspaceRootDir returns the relative workspace root for downloads and local analysis.
// Project-bound sessions share projects/<id>/; otherwise conversations/<id>/.
func WorkspaceRootDir(configuredBase, projectID, conversationID string) string {
base := strings.TrimSpace(configuredBase)
if base == "" {
base = filepath.Join("tmp", "workspace")
}
if pid := strings.TrimSpace(projectID); pid != "" {
return filepath.Join(base, "projects", sanitizeWorkspacePathSegment(pid))
}
conv := strings.TrimSpace(conversationID)
if conv == "" {
conv = "default"
}
return filepath.Join(base, "conversations", sanitizeWorkspacePathSegment(conv))
}
// EnsureWorkspace creates the workspace directory and returns its absolute path.
func EnsureWorkspace(root string) (string, error) {
abs, err := filepath.Abs(strings.TrimSpace(root))
if err != nil {
return "", fmt.Errorf("workspace abs: %w", err)
}
if err := os.MkdirAll(abs, 0o755); err != nil {
return "", fmt.Errorf("workspace mkdir: %w", err)
}
return abs, nil
}
// BuildWorkspaceBlock instructs the agent to use the session workspace instead of /tmp.
func BuildWorkspaceBlock(absPath string) string {
absPath = strings.TrimSpace(absPath)
if absPath == "" {
return ""
}
return fmt.Sprintf(`## 会话工作目录(下载与本地分析)
**必须使用以下目录**保存 curl/wget 下载的文件、临时 HTML/JS,以及 read_file/glob/grep 的检索范围:
`+"`%s`"+`
- **禁止**使用系统 `+"`/tmp`"+` 或其它全局临时目录(多项目/多会话会互窜遗留文件)。
- 下载示例:`+"`curl -o '%s/page.html' 'https://target/'`"+`exec 时可将 `+"`workdir`"+` 设为该目录。
- 读取下载产物或临时分析文件前,用 glob/grep/read_file **限定在该目录**下搜索,勿在 `+"`/tmp`"+` 盲目检索。
- 当用户询问“当前目录”“项目根目录”或应用自身文件时,优先按服务进程当前工作目录理解;不要把空的会话工作目录误当成项目根目录。`, absPath, absPath)
}
+58
View File
@@ -0,0 +1,58 @@
package project
import (
"os"
"path/filepath"
"strings"
"testing"
)
func TestWorkspaceRootDirProjectScoped(t *testing.T) {
got := WorkspaceRootDir("", "proj-1", "conv-1")
want := filepath.Join("tmp", "workspace", "projects", "proj-1")
if got != want {
t.Fatalf("got %q want %q", got, want)
}
}
func TestWorkspaceRootDirConversationScoped(t *testing.T) {
got := WorkspaceRootDir("/data/ws", "", "conv-abc")
want := filepath.Join("/data/ws", "conversations", "conv-abc")
if got != want {
t.Fatalf("got %q want %q", got, want)
}
}
func TestEnsureWorkspaceCreatesDir(t *testing.T) {
root := filepath.Join(t.TempDir(), "nested", "workspace")
abs, err := EnsureWorkspace(root)
if err != nil {
t.Fatalf("EnsureWorkspace: %v", err)
}
st, err := os.Stat(abs)
if err != nil {
t.Fatalf("Stat: %v", err)
}
if !st.IsDir() {
t.Fatal("expected directory")
}
}
func TestBuildWorkspaceBlockMentionsPath(t *testing.T) {
block := BuildWorkspaceBlock("/opt/csai/tmp/workspace/projects/p1")
if block == "" {
t.Fatal("expected non-empty block")
}
if !strings.Contains(block, "/opt/csai/tmp/workspace/projects/p1") {
t.Fatalf("block missing path: %s", block)
}
if !strings.Contains(block, "/tmp") {
t.Fatalf("block should warn about /tmp: %s", block)
}
if !strings.Contains(block, "当前目录") || !strings.Contains(block, "服务进程当前工作目录") {
t.Fatalf("block should distinguish current/project dir from workspace: %s", block)
}
if !strings.Contains(block, "不要把空的会话工作目录误当成项目根目录") {
t.Fatalf("block should warn about empty workspace confusion: %s", block)
}
}