mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-01 08:37:41 +02:00
214 lines
8.8 KiB
Go
214 lines
8.8 KiB
Go
package workflow
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"cyberstrike-ai/internal/config"
|
|
)
|
|
|
|
func TestGenerateDraftFromNaturalLanguageHighRiskAddsHITLAndValidGraph(t *testing.T) {
|
|
result, err := GenerateDraftFromNaturalLanguage(context.Background(), DraftRequest{
|
|
Prompt: "对目标资产做端口扫描，如果发现高危端口就执行加固脚本，最后输出报告",
|
|
Options: DraftOptions{
|
|
IncludeObjective: true,
|
|
AllowSchedule: false,
|
|
AllowHighRisk: false,
|
|
},
|
|
AvailableTools: []DraftTool{{Key: "nmap", Name: "nmap", Enabled: true}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GenerateDraftFromNaturalLanguage: %v", err)
|
|
}
|
|
raw, _ := json.Marshal(result.Graph)
|
|
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
|
|
t.Fatalf("generated graph should validate: %v\n%s", err, raw)
|
|
}
|
|
if !result.Audit.HighRisk || !result.Audit.NeedsHITL || len(result.Audit.RiskWarnings) == 0 {
|
|
t.Fatalf("audit did not flag high-risk HITL path: %#v", result.Audit)
|
|
}
|
|
var hasTool, hasHITL, hasConfirmation bool
|
|
for _, node := range result.Graph.Nodes {
|
|
if node.Type == "tool" && cfgString(node.Config, "tool_name") == "nmap" {
|
|
hasTool = true
|
|
}
|
|
if node.Type == "hitl" {
|
|
hasHITL = true
|
|
}
|
|
if cfgString(node.Config, "requires_human_confirmation") == "true" {
|
|
hasConfirmation = true
|
|
}
|
|
}
|
|
if !hasTool || !hasHITL || !hasConfirmation {
|
|
t.Fatalf("expected nmap tool, HITL, and confirmation marker; tool=%v hitl=%v confirmation=%v", hasTool, hasHITL, hasConfirmation)
|
|
}
|
|
}
|
|
|
|
func TestGenerateDraftAllowHighRiskStillLabelsConditionBranch(t *testing.T) {
|
|
result, err := GenerateDraftFromNaturalLanguage(context.Background(), DraftRequest{
|
|
Prompt: "如果漏洞扫描发现高危漏洞，允许生成执行修复脚本的草稿并输出报告",
|
|
Options: DraftOptions{
|
|
AllowHighRisk: true,
|
|
},
|
|
AvailableTools: []DraftTool{{Key: "nuclei", Name: "nuclei", Enabled: true}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GenerateDraftFromNaturalLanguage: %v", err)
|
|
}
|
|
raw, _ := json.Marshal(result.Graph)
|
|
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
|
|
t.Fatalf("generated graph should validate: %v\n%s", err, raw)
|
|
}
|
|
branches := map[string]bool{}
|
|
for _, edge := range result.Graph.Edges {
|
|
if branch := cfgString(edge.Config, "branch"); branch != "" {
|
|
branches[branch] = true
|
|
}
|
|
}
|
|
if !branches["true"] || !branches["false"] {
|
|
t.Fatalf("condition branches = %#v, want true and false", branches)
|
|
}
|
|
}
|
|
|
|
func TestGenerateDraftFromLLMUsesOpenAICompatibleEndpoint(t *testing.T) {
|
|
called := false
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
called = true
|
|
if r.URL.Path != "/chat/completions" {
|
|
t.Fatalf("path = %s, want /chat/completions", r.URL.Path)
|
|
}
|
|
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
|
|
t.Fatalf("authorization = %q", got)
|
|
}
|
|
var payload struct {
|
|
Temperature float64 `json:"temperature"`
|
|
ResponseFormat struct {
|
|
Type string `json:"type"`
|
|
} `json:"response_format"`
|
|
Messages []struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
} `json:"messages"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
|
t.Fatalf("decode request: %v", err)
|
|
}
|
|
if payload.Temperature != 0 || payload.ResponseFormat.Type != "json_object" {
|
|
t.Fatalf("unexpected structured output controls: temperature=%v response_format=%#v", payload.Temperature, payload.ResponseFormat)
|
|
}
|
|
if len(payload.Messages) == 0 || strings.Contains(payload.Messages[0].Content, "start|tool|agent") {
|
|
t.Fatalf("system prompt still contains pipe enum: %q", payload.Messages[0].Content)
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"meta\":{\"id\":\"llm-port-scan\",\"name\":\"端口扫描\",\"description\":\"端口扫描\",\"enabled\":true},\"graph\":{\"nodes\":[{\"id\":\"start-1\",\"type\":\"start\",\"label\":\"开始\",\"position\":{\"x\":120,\"y\":150},\"config\":{\"input_keys\":\"message, target\"}},{\"id\":\"tool-2\",\"type\":\"tool\",\"label\":\"端口扫描\",\"position\":{\"x\":330,\"y\":150},\"config\":{\"tool_name\":\"nmap\",\"arguments\":\"{\\\"target\\\":\\\"{{inputs.target}}\\\"}\",\"timeout_seconds\":\"120\",\"join_strategy\":\"all_merge\"}},{\"id\":\"output-3\",\"type\":\"output\",\"label\":\"输出报告\",\"position\":{\"x\":540,\"y\":150},\"config\":{\"source_binding\":{\"from\":\"previous\",\"field\":\"output\"},\"join_strategy\":\"all_merge\"}}],\"edges\":[{\"id\":\"edge-1\",\"source\":\"start-1\",\"target\":\"tool-2\"},{\"id\":\"edge-2\",\"source\":\"tool-2\",\"target\":\"output-3\"}],\"config\":{\"schema_version\":1}},\"capabilities\":[{\"label\":\"端口扫描\",\"tool_name\":\"nmap\",\"tool_candidates\":[\"nmap\"]}],\"audit\":{\"assumptions\":[]}}"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
result, err := GenerateDraftFromLLM(context.Background(), DraftRequest{
|
|
Prompt: "对目标做端口扫描并输出报告",
|
|
AvailableTools: []DraftTool{{Key: "nmap", Name: "nmap", Enabled: true}},
|
|
}, config.OpenAIConfig{APIKey: "test-key", BaseURL: srv.URL, Model: "test-model"}, nil)
|
|
if err != nil {
|
|
t.Fatalf("GenerateDraftFromLLM: %v", err)
|
|
}
|
|
if !called {
|
|
t.Fatal("expected LLM endpoint to be called")
|
|
}
|
|
if result.Generator != "llm" || !result.Audit.Savable || result.Meta.ID != "llm-port-scan" {
|
|
t.Fatalf("unexpected result: %#v", result)
|
|
}
|
|
for _, node := range result.Graph.Nodes {
|
|
if node.Type == "output" && cfgString(node.Config, "output_key") != "result" {
|
|
t.Fatalf("output_key = %q, want result", cfgString(node.Config, "output_key"))
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGenerateDraftFromLLMReturnsErrorOnMalformedJSON(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"meta\":{\"id\":\"bad\"},|\"graph\":{\"nodes\":[]}}"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
_, err := GenerateDraftFromLLM(context.Background(), DraftRequest{
|
|
Prompt: "随便生成一个工作流，要求所有节点都用到输出变量",
|
|
Options: DraftOptions{
|
|
IncludeObjective: true,
|
|
},
|
|
}, config.OpenAIConfig{APIKey: "test-key", BaseURL: srv.URL, Model: "test-model"}, nil)
|
|
if err == nil {
|
|
t.Fatal("expected malformed JSON error")
|
|
}
|
|
if !strings.Contains(err.Error(), "解析大模型工作流 JSON 失败") {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeLLMDraftRepairsMissingRequiredConfig(t *testing.T) {
|
|
result := normalizeLLMDraft("随便生成一个工作流", DraftRequest{}, llmDraftEnvelope{
|
|
Graph: graphDef{
|
|
Nodes: []graphNode{
|
|
{ID: "start-1", Type: "start", Label: "开始", Config: map[string]any{}},
|
|
{ID: "agent-1", Type: "agent", Label: "分析", Config: map[string]any{}},
|
|
{ID: "out-1", Type: "output", Label: "输出结果", Config: map[string]any{}},
|
|
},
|
|
Edges: []graphEdge{
|
|
{ID: "e1", Source: "start-1", Target: "agent-1"},
|
|
{ID: "e2", Source: "agent-1", Target: "out-1"},
|
|
},
|
|
},
|
|
})
|
|
raw, _ := json.Marshal(result.Graph)
|
|
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
|
|
t.Fatalf("normalized graph should validate: %v\n%s", err, raw)
|
|
}
|
|
var agentKey, outputKey string
|
|
for _, node := range result.Graph.Nodes {
|
|
switch node.Type {
|
|
case "agent":
|
|
agentKey = cfgString(node.Config, "output_key")
|
|
case "output":
|
|
outputKey = cfgString(node.Config, "output_key")
|
|
}
|
|
}
|
|
if agentKey == "" || outputKey != "result" {
|
|
t.Fatalf("agentKey=%q outputKey=%q", agentKey, outputKey)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeLLMDraftRepairsConditionBranches(t *testing.T) {
|
|
result := normalizeLLMDraft("如果发现异常则输出详情，否则输出正常", DraftRequest{}, llmDraftEnvelope{
|
|
Graph: graphDef{
|
|
Nodes: []graphNode{
|
|
{ID: "start-1", Type: "start", Label: "开始", Config: map[string]any{}},
|
|
{ID: "cond-1", Type: "condition", Label: "判断", Config: map[string]any{"expression": `{{inputs.message}} != ""`}},
|
|
{ID: "out-yes", Type: "output", Label: "异常", Config: map[string]any{}},
|
|
{ID: "out-no", Type: "output", Label: "正常", Config: map[string]any{}},
|
|
},
|
|
Edges: []graphEdge{
|
|
{ID: "e1", Source: "start-1", Target: "cond-1"},
|
|
{ID: "e2", Source: "cond-1", Target: "out-yes"},
|
|
{ID: "e3", Source: "cond-1", Target: "out-no"},
|
|
},
|
|
},
|
|
})
|
|
raw, _ := json.Marshal(result.Graph)
|
|
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
|
|
t.Fatalf("normalized graph should validate: %v\n%s", err, raw)
|
|
}
|
|
branches := map[string]bool{}
|
|
for _, edge := range result.Graph.Edges {
|
|
if edge.Source == "cond-1" {
|
|
branches[cfgString(edge.Config, "branch")] = true
|
|
}
|
|
}
|
|
if !branches["true"] || !branches["false"] {
|
|
t.Fatalf("branches = %#v, want true and false", branches)
|
|
}
|
|
}
|