dela
0335d572de
ci / go (push) Waiting to run
ci / go-db (agent) (push) Waiting to run
ci / go-db (config) (push) Waiting to run
ci / go-db (db) (push) Waiting to run
ci / go-db (evidence) (push) Waiting to run
ci / go-db (llmrec) (push) Waiting to run
ci / go-db (server) (push) Waiting to run
detections / detections (push) Waiting to run
web / web (push) Waiting to run
docs / links (push) Canceled after 0s
451 lines
14 KiB
Go
451 lines
14 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/Autumn-27/artex/db"
|
|
)
|
|
|
|
// testDB opens a DB connection, skipping if PG is unavailable.
|
|
func testDB(t *testing.T) *db.DB {
|
|
t.Helper()
|
|
dsn, _, err := db.DSN()
|
|
if err != nil {
|
|
t.Skipf("no database config (%v)", err)
|
|
}
|
|
d, err := db.Open(dsn)
|
|
if err != nil {
|
|
t.Skipf("postgres unavailable (%v)", err)
|
|
}
|
|
return d
|
|
}
|
|
|
|
// callInsertAssets calls the insert_assets tool with the given payload.
|
|
func callInsertAssets(t *testing.T, ts *ToolSet, payload any) map[string]any {
|
|
t.Helper()
|
|
raw, _ := json.Marshal(payload)
|
|
tool := ts.insertAssets()
|
|
res, err := tool.Call(context.Background(), raw, nil)
|
|
if err != nil {
|
|
t.Fatalf("insertAssets Call error: %v", err)
|
|
}
|
|
text := res.Flatten()
|
|
var out map[string]any
|
|
if err := json.Unmarshal([]byte(text), &out); err != nil {
|
|
t.Fatalf("unmarshal result: %v\nraw: %s", err, text)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// =====================================================================
|
|
// TestInsertAssetsSubdomainSideEffects
|
|
// 子域名插入 → 自动创建 root_domain + IP 资产，IP 绑定域名
|
|
// =====================================================================
|
|
func TestInsertAssetsSubdomainSideEffects(t *testing.T) {
|
|
d := testDB(t)
|
|
defer d.Close()
|
|
|
|
ts := NewToolSet(nil, "")
|
|
ts.SetAssetStore(d.Assets(), d.Companies())
|
|
defer d.Exec(`DELETE FROM assets WHERE domain IN ('ia-sub.sideeffect-test.com','sideeffect-test.com') OR ip='7.8.9.10'`)
|
|
|
|
out := callInsertAssets(t, ts, map[string]any{
|
|
"assets": []any{
|
|
map[string]any{
|
|
"type": "subdomain",
|
|
"domain": "ia-sub.sideeffect-test.com",
|
|
"record_type": "A",
|
|
"record_value": []string{"7.8.9.10"},
|
|
},
|
|
},
|
|
"task_id": 999,
|
|
})
|
|
|
|
// no errors
|
|
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
|
t.Errorf("unexpected errors: %v", errs)
|
|
}
|
|
results, _ := out["results"].([]any)
|
|
if len(results) == 0 {
|
|
t.Fatal("no results returned")
|
|
}
|
|
|
|
// root_domain should exist
|
|
var rootCnt int
|
|
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type='root_domain' AND domain='sideeffect-test.com'`).Scan(&rootCnt)
|
|
if rootCnt != 1 {
|
|
t.Errorf("side-effect: root_domain not created, got %d", rootCnt)
|
|
}
|
|
|
|
// IP asset should exist with bound_domains containing our subdomain
|
|
var ipID int64
|
|
var boundDomains []byte
|
|
d.QueryRow(`SELECT id, array_to_json(bound_domains)::text FROM assets WHERE type='ip' AND ip='7.8.9.10'`).Scan(&ipID, &boundDomains)
|
|
if ipID == 0 {
|
|
t.Error("side-effect: IP asset not created")
|
|
}
|
|
var domains []string
|
|
json.Unmarshal(boundDomains, &domains)
|
|
found := false
|
|
for _, d := range domains {
|
|
if d == "ia-sub.sideeffect-test.com" {
|
|
found = true
|
|
}
|
|
}
|
|
if !found {
|
|
t.Errorf("side-effect: bound_domains should contain subdomain, got %v", domains)
|
|
}
|
|
|
|
// record_value stored as array
|
|
var rvRaw []byte
|
|
d.QueryRow(`SELECT array_to_json(record_value)::text FROM assets WHERE type='subdomain' AND domain='ia-sub.sideeffect-test.com'`).Scan(&rvRaw)
|
|
var rv []string
|
|
json.Unmarshal(rvRaw, &rv)
|
|
if len(rv) == 0 || rv[0] != "7.8.9.10" {
|
|
t.Errorf("record_value stored incorrectly: %v", rv)
|
|
}
|
|
}
|
|
|
|
// =====================================================================
|
|
// TestInsertAssetsMultiIPSubdomain
|
|
// 多个 IP 的子域名：所有 IP 都应存入 record_value[]，各自创建 IP 资产
|
|
// =====================================================================
|
|
func TestInsertAssetsMultiIPSubdomain(t *testing.T) {
|
|
d := testDB(t)
|
|
defer d.Close()
|
|
|
|
ts := NewToolSet(nil, "")
|
|
ts.SetAssetStore(d.Assets(), d.Companies())
|
|
defer d.Exec(`DELETE FROM assets WHERE domain IN ('multi.multiip-test.io','multiip-test.io') OR ip IN ('1.1.1.1','2.2.2.2')`)
|
|
|
|
out := callInsertAssets(t, ts, map[string]any{
|
|
"assets": []any{
|
|
map[string]any{
|
|
"type": "subdomain",
|
|
"domain": "multi.multiip-test.io",
|
|
"record_type": "A",
|
|
"record_value": []string{"1.1.1.1", "2.2.2.2"},
|
|
},
|
|
},
|
|
})
|
|
|
|
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
|
t.Errorf("unexpected errors: %v", errs)
|
|
}
|
|
|
|
// Both IPs should have IP assets
|
|
var ip1Cnt, ip2Cnt int
|
|
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type='ip' AND ip='1.1.1.1'`).Scan(&ip1Cnt)
|
|
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type='ip' AND ip='2.2.2.2'`).Scan(&ip2Cnt)
|
|
if ip1Cnt != 1 {
|
|
t.Error("IP 1.1.1.1 asset not created")
|
|
}
|
|
if ip2Cnt != 1 {
|
|
t.Error("IP 2.2.2.2 asset not created")
|
|
}
|
|
|
|
// record_value should contain both IPs
|
|
var rvRaw []byte
|
|
d.QueryRow(`SELECT array_to_json(record_value)::text FROM assets WHERE type='subdomain' AND domain='multi.multiip-test.io'`).Scan(&rvRaw)
|
|
var rv []string
|
|
json.Unmarshal(rvRaw, &rv)
|
|
if len(rv) != 2 {
|
|
t.Errorf("record_value: want 2 IPs, got %v", rv)
|
|
}
|
|
}
|
|
|
|
// =====================================================================
|
|
// TestInsertAssetsHTTPServiceTechnologies
|
|
// HTTP 服务插入：technologies 存储并可读回；IP 存在时域名和端口写入 IP 资产
|
|
// =====================================================================
|
|
func TestInsertAssetsHTTPServiceTechnologies(t *testing.T) {
|
|
d := testDB(t)
|
|
defer d.Close()
|
|
|
|
ts := NewToolSet(nil, "")
|
|
ts.SetAssetStore(d.Assets(), d.Companies())
|
|
defer d.Exec(`DELETE FROM assets WHERE url='https://tech-test.example.com' OR domain IN ('tech-test.example.com','example.com') OR ip='3.4.5.6'`)
|
|
|
|
out := callInsertAssets(t, ts, map[string]any{
|
|
"assets": []any{
|
|
map[string]any{
|
|
"type": "service",
|
|
"url": "https://tech-test.example.com",
|
|
"service_ip": "3.4.5.6",
|
|
"technologies": []string{"Nginx", "Vue.js", "Cloudflare"},
|
|
"status_code": 200,
|
|
"page_title": "Tech Test Site",
|
|
},
|
|
},
|
|
"task_id": 888,
|
|
})
|
|
|
|
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
|
t.Errorf("unexpected errors: %v", errs)
|
|
}
|
|
|
|
// technologies should be stored
|
|
var techCnt int
|
|
d.QueryRow(`SELECT array_length(technologies,1) FROM assets WHERE url='https://tech-test.example.com'`).Scan(&techCnt)
|
|
if techCnt != 3 {
|
|
t.Errorf("technologies: want 3, got %d", techCnt)
|
|
}
|
|
|
|
// QueryByType should return technologies correctly (verifies array_to_json scan)
|
|
assets, err := d.Assets().QueryByType("service", 50, 0)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var found *db.Asset
|
|
for _, a := range assets {
|
|
if a.URL == "https://tech-test.example.com" {
|
|
found = a
|
|
break
|
|
}
|
|
}
|
|
if found == nil {
|
|
t.Fatal("service not found via QueryByType")
|
|
}
|
|
if len(found.Technologies) != 3 {
|
|
t.Errorf("QueryByType: technologies roundtrip failed, got %v", found.Technologies)
|
|
}
|
|
|
|
// side effect: IP asset should exist with bound_domains containing the service domain
|
|
var ipID int64
|
|
var bdRaw []byte
|
|
var portCnt int
|
|
d.QueryRow(`SELECT id, array_to_json(bound_domains)::text FROM assets WHERE type='ip' AND ip='3.4.5.6'`).Scan(&ipID, &bdRaw)
|
|
if ipID == 0 {
|
|
t.Error("side-effect: IP asset not created for HTTP service IP")
|
|
}
|
|
var bd []string
|
|
json.Unmarshal(bdRaw, &bd)
|
|
hasDomain := false
|
|
for _, dom := range bd {
|
|
if dom == "tech-test.example.com" {
|
|
hasDomain = true
|
|
}
|
|
}
|
|
if !hasDomain {
|
|
t.Errorf("side-effect: IP bound_domains missing service domain, got %v", bd)
|
|
}
|
|
|
|
// side effect: IP open_ports should contain port 443
|
|
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE type='ip' AND ip='3.4.5.6'`).Scan(&portCnt)
|
|
if portCnt == 0 {
|
|
t.Error("side-effect: IP open_ports not set for HTTP service")
|
|
}
|
|
}
|
|
|
|
// =====================================================================
|
|
// TestInsertAssetsOtherService
|
|
// 非 HTTP 服务：c_segment 自动生成，IP 资产含 open_ports 和 bound_domains
|
|
// =====================================================================
|
|
func TestInsertAssetsOtherService(t *testing.T) {
|
|
d := testDB(t)
|
|
defer d.Close()
|
|
|
|
ts := NewToolSet(nil, "")
|
|
ts.SetAssetStore(d.Assets(), d.Companies())
|
|
defer d.Exec(`DELETE FROM assets WHERE
|
|
(type='service' AND ip='10.20.30.40') OR
|
|
(type='ip' AND ip='10.20.30.40') OR
|
|
domain IN ('db.othersvc-test.com','othersvc-test.com')`)
|
|
|
|
out := callInsertAssets(t, ts, map[string]any{
|
|
"assets": []any{
|
|
map[string]any{
|
|
"type": "service",
|
|
"ip": "10.20.30.40",
|
|
"domain": "db.othersvc-test.com",
|
|
"port": 3306,
|
|
"service_name": "mysql",
|
|
},
|
|
},
|
|
})
|
|
|
|
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
|
t.Errorf("unexpected errors: %v", errs)
|
|
}
|
|
|
|
// c_segment should be auto-set on the service
|
|
var cseg *string
|
|
d.QueryRow(`SELECT c_segment::text FROM assets WHERE type='service' AND ip='10.20.30.40'`).Scan(&cseg)
|
|
if cseg == nil || *cseg != "10.20.30.0/24" {
|
|
t.Errorf("c_segment: want 10.20.30.0/24, got %v", cseg)
|
|
}
|
|
|
|
// IP side-effect: port 3306 in open_ports
|
|
var portCnt int
|
|
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE type='ip' AND ip='10.20.30.40'`).Scan(&portCnt)
|
|
if portCnt == 0 {
|
|
t.Error("side-effect: IP open_ports should contain port 3306")
|
|
}
|
|
|
|
// IP side-effect: bound_domains contains the service domain
|
|
var bdRaw []byte
|
|
d.QueryRow(`SELECT array_to_json(bound_domains)::text FROM assets WHERE type='ip' AND ip='10.20.30.40'`).Scan(&bdRaw)
|
|
var bd []string
|
|
json.Unmarshal(bdRaw, &bd)
|
|
hasDomain := false
|
|
for _, dom := range bd {
|
|
if dom == "db.othersvc-test.com" {
|
|
hasDomain = true
|
|
}
|
|
}
|
|
if !hasDomain {
|
|
t.Errorf("side-effect: IP bound_domains missing service domain, got %v", bd)
|
|
}
|
|
}
|
|
|
|
// =====================================================================
|
|
// TestInsertAssetsMixedBatch
|
|
// 混合批量插入：一次调用插入多种类型
|
|
// =====================================================================
|
|
func TestInsertAssetsMixedBatch(t *testing.T) {
|
|
d := testDB(t)
|
|
defer d.Close()
|
|
|
|
ts := NewToolSet(nil, "")
|
|
ts.SetAssetStore(d.Assets(), d.Companies())
|
|
defer d.Exec(`DELETE FROM assets WHERE
|
|
domain IN ('batch-sub.batch-test.org','batch-test.org') OR
|
|
ip='55.66.77.88' OR
|
|
url='https://batch-test.org/api' OR
|
|
(type='endpoint' AND url='https://batch-test.org/api/users')`)
|
|
|
|
out := callInsertAssets(t, ts, map[string]any{
|
|
"assets": []any{
|
|
// root_domain
|
|
map[string]any{"type": "root_domain", "domain": "batch-test.org"},
|
|
// subdomain with A record
|
|
map[string]any{"type": "subdomain", "domain": "batch-sub.batch-test.org", "record_type": "A", "record_value": []string{"55.66.77.88"}},
|
|
// HTTP service
|
|
map[string]any{"type": "service", "url": "https://batch-test.org/api", "technologies": []string{"Go", "PostgreSQL"}, "status_code": 200},
|
|
// endpoint
|
|
map[string]any{"type": "endpoint", "url": "https://batch-test.org/api/users", "method": "GET"},
|
|
},
|
|
"task_id": 777,
|
|
})
|
|
|
|
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
|
t.Errorf("unexpected errors: %v", errs)
|
|
}
|
|
results, _ := out["results"].([]any)
|
|
if len(results) != 4 {
|
|
t.Errorf("mixed batch: want 4 results, got %d", len(results))
|
|
}
|
|
|
|
// verify all types exist in DB
|
|
types := []string{"root_domain", "subdomain", "service", "endpoint"}
|
|
for _, typ := range types {
|
|
var cnt int
|
|
switch typ {
|
|
case "root_domain":
|
|
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND domain='batch-test.org'`, typ).Scan(&cnt)
|
|
case "subdomain":
|
|
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND domain='batch-sub.batch-test.org'`, typ).Scan(&cnt)
|
|
case "service":
|
|
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND url='https://batch-test.org/api'`, typ).Scan(&cnt)
|
|
case "endpoint":
|
|
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND url='https://batch-test.org/api/users'`, typ).Scan(&cnt)
|
|
}
|
|
if cnt != 1 {
|
|
t.Errorf("mixed batch: %s not found in DB", typ)
|
|
}
|
|
}
|
|
}
|
|
|
|
// =====================================================================
|
|
// TestInsertAssetsDedup
|
|
// 幂等写入：同一资产插入两次，返回相同 ID
|
|
// =====================================================================
|
|
func TestInsertAssetsDedup(t *testing.T) {
|
|
d := testDB(t)
|
|
defer d.Close()
|
|
|
|
ts := NewToolSet(nil, "")
|
|
ts.SetAssetStore(d.Assets(), d.Companies())
|
|
defer d.Exec(`DELETE FROM assets WHERE domain='dedup-ia.deduptest.net' OR domain='deduptest.net'`)
|
|
|
|
payload := map[string]any{
|
|
"assets": []any{
|
|
map[string]any{"type": "root_domain", "domain": "deduptest.net"},
|
|
},
|
|
}
|
|
|
|
out1 := callInsertAssets(t, ts, payload)
|
|
out2 := callInsertAssets(t, ts, payload)
|
|
|
|
getID := func(out map[string]any) float64 {
|
|
results, _ := out["results"].([]any)
|
|
if len(results) == 0 {
|
|
return 0
|
|
}
|
|
m, _ := results[0].(map[string]any)
|
|
id, _ := m["id"].(float64)
|
|
return id
|
|
}
|
|
|
|
id1, id2 := getID(out1), getID(out2)
|
|
if id1 == 0 || id1 != id2 {
|
|
t.Errorf("dedup: want same ID on double insert, got %v vs %v", id1, id2)
|
|
}
|
|
}
|
|
|
|
// =====================================================================
|
|
// TestInsertAssetsRejectsHostnameIPPerItem
|
|
// 一批里混入 ip 填了主机名的一条 → 只有那条失败，其余照常入库，
|
|
// 且错误里带得上 index 和改正方法，Agent 下一轮能自己修好。
|
|
// =====================================================================
|
|
func TestInsertAssetsRejectsHostnameIPPerItem(t *testing.T) {
|
|
d := testDB(t)
|
|
defer d.Close()
|
|
|
|
ts := NewToolSet(nil, "")
|
|
ts.SetAssetStore(d.Assets(), d.Companies())
|
|
defer d.Exec(`DELETE FROM assets WHERE ip IN ('198.51.100.23','cdn.badip-test.com') OR domain='badip-test.com'`)
|
|
|
|
out := callInsertAssets(t, ts, map[string]any{
|
|
"assets": []any{
|
|
map[string]any{"type": "root_domain", "domain": "badip-test.com"},
|
|
map[string]any{"type": "ip", "ip": "cdn.badip-test.com"},
|
|
map[string]any{"type": "ip", "ip": "198.51.100.23"},
|
|
},
|
|
})
|
|
|
|
// The two valid entries must survive the bad one — a whole-batch failure
|
|
// would make the agent re-send assets that were already fine.
|
|
results, _ := out["results"].([]any)
|
|
if len(results) != 2 {
|
|
t.Fatalf("results=%v, want the 2 valid assets", out["results"])
|
|
}
|
|
|
|
errsRaw, _ := out["errors"].([]any)
|
|
if len(errsRaw) != 1 {
|
|
t.Fatalf("errors=%v, want exactly the invalid entry", out["errors"])
|
|
}
|
|
entry, _ := errsRaw[0].(map[string]any)
|
|
if index, _ := entry["index"].(float64); int(index) != 1 {
|
|
t.Fatalf("error index=%v, want 1", entry["index"])
|
|
}
|
|
message, _ := entry["error"].(string)
|
|
for _, want := range []string{"cdn.badip-test.com", "type=subdomain", "A/AAAA"} {
|
|
if !strings.Contains(message, want) {
|
|
t.Fatalf("error message %q lacks %q — agent cannot act on it", message, want)
|
|
}
|
|
}
|
|
|
|
// The rejected value must not have reached the table.
|
|
var stored int
|
|
if err := d.QueryRow(`SELECT count(*) FROM assets WHERE ip='cdn.badip-test.com'`).Scan(&stored); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if stored != 0 {
|
|
t.Fatalf("rejected hostname still stored in assets.ip (%d rows)", stored)
|
|
}
|
|
}
|