ARTEX — AI 自主渗透测试系统(Go 后端 + Next.js 前端,单二进制内嵌)。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
94 lines
2.5 KiB
Go
94 lines
2.5 KiB
Go
package db
|
|
|
|
import "testing"
|
|
|
|
func TestExplorationFlow(t *testing.T) {
|
|
d, err := Open(testDSN(t))
|
|
if err != nil {
|
|
t.Skipf("postgres unavailable (%v) — skipping", err)
|
|
}
|
|
defer d.Close()
|
|
|
|
expID, err := d.CreateExploration("test", "拿下测试目标")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) // cascades nodes/edges/activity
|
|
es := d.Exploration(expID)
|
|
|
|
// goal node + two intents
|
|
goal, err := es.AddGoal(map[string]any{"text": "getadmin", "vulnclass": "authz"}, "human")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := es.AddIntent(map[string]any{"summary": "enumerate endpoints"}, 5, nil, "planner"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
i2, err := es.AddIntent(map[string]any{"summary": "test idor"}, 8, nil, "planner")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// frontier ordered by priority desc → i2(8) before i1(5)
|
|
fr, err := es.Frontier(10)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(fr) != 2 || fr[0].ID != i2 {
|
|
t.Fatalf("frontier order wrong: %+v", fr)
|
|
}
|
|
|
|
// atomic claim: first wins, second on same id fails
|
|
ok, err := es.ClaimIntent(i2, "worker-1")
|
|
if err != nil || !ok {
|
|
t.Fatalf("claim i2: ok=%v err=%v", ok, err)
|
|
}
|
|
ok2, _ := es.ClaimIntent(i2, "worker-2")
|
|
if ok2 {
|
|
t.Fatalf("double-claim should fail")
|
|
}
|
|
|
|
// finding yields from intent, proves goal
|
|
find, err := es.AddNode("finding", map[string]any{"vulnclass": "idor", "severity": "high"}, 9, "confirmed", "worker-1", nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := es.Link(i2, "yields", find); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := es.Link(find, "proves", goal); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := es.SetNodeState(goal, "met"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// activity poll by id cursor
|
|
id1, err := es.AppendActivity(Activity{Worker: "worker-1", Kind: "tool_use", Tool: "Bash", Summary: "ran curl", Detail: "full output"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
items, cursor, err := es.ActivityList(nil, 0, 100)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(items) != 1 || cursor != id1 {
|
|
t.Fatalf("activity list: items=%d cursor=%d", len(items), cursor)
|
|
}
|
|
det, _ := es.ActivityDetail(id1)
|
|
if det != "full output" {
|
|
t.Fatalf("detail want 'full output', got %q", det)
|
|
}
|
|
// incremental: nothing new after cursor
|
|
items2, _, _ := es.ActivityList(nil, cursor, 100)
|
|
if len(items2) != 0 {
|
|
t.Fatalf("incremental poll should be empty, got %d", len(items2))
|
|
}
|
|
|
|
// stats
|
|
st, _ := es.Stats()
|
|
if st["intent"] != 2 || st["goal"] != 1 || st["finding"] != 1 {
|
|
t.Fatalf("stats: %+v", st)
|
|
}
|
|
}
|