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
406 lines
14 KiB
Go
406 lines
14 KiB
Go
package db
|
||
|
||
import (
|
||
"fmt"
|
||
"testing"
|
||
"time"
|
||
)
|
||
|
||
// TestParseScopeLine covers classification + guardrails without a DB.
|
||
func TestParseScopeLine(t *testing.T) {
|
||
cases := []struct {
|
||
in string
|
||
kind string
|
||
wantErr bool
|
||
}{
|
||
{"example.com", "domain", false},
|
||
{"https://sub.example.com/path", "domain", false},
|
||
{"1.2.3.4", "ip", false},
|
||
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
|
||
{"198.51.100.0/24", "cidr", false},
|
||
{"co.uk", "", true}, // bare public suffix
|
||
{"not a host", "", true},
|
||
{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
|
||
}
|
||
for _, c := range cases {
|
||
r, err := ParseScopeLine(c.in)
|
||
if c.wantErr {
|
||
if err == nil {
|
||
t.Errorf("ParseScopeLine(%q) want error, got %+v", c.in, r)
|
||
}
|
||
continue
|
||
}
|
||
if err != nil {
|
||
t.Errorf("ParseScopeLine(%q) unexpected error: %v", c.in, err)
|
||
continue
|
||
}
|
||
if r.Kind != c.kind {
|
||
t.Errorf("ParseScopeLine(%q) kind=%q want %q", c.in, r.Kind, c.kind)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestParseAutoScopeLine(t *testing.T) {
|
||
cases := []struct {
|
||
name string
|
||
input string
|
||
kind string
|
||
normalized string
|
||
wantErr bool
|
||
}{
|
||
{name: "domain", input: "example.com", kind: "domain", normalized: "example.com"},
|
||
{name: "url", input: "https://sub.example.com/path", kind: "domain", normalized: "sub.example.com"},
|
||
{name: "url query", input: "https://example.com/path?source=x", kind: "domain", normalized: "example.com"},
|
||
{name: "url credentials", input: "https://user:pass@example.com/path", kind: "domain", normalized: "example.com"},
|
||
{name: "url ip", input: "http://203.0.113.10/path", kind: "ip", normalized: "203.0.113.10/32"},
|
||
{name: "ipv4", input: "203.0.113.10", kind: "ip", normalized: "203.0.113.10/32"},
|
||
{name: "ipv6", input: "2001:db8::10", kind: "ip", normalized: "2001:db8::10/128"},
|
||
{name: "cidr", input: "198.51.100.0/24", kind: "cidr", normalized: "198.51.100.0/24"},
|
||
{name: "icp latin", input: "京 ICP备 123号", kind: "icp", normalized: "京icp备123号"},
|
||
{name: "icp chinese", input: "沪网备案 9988", kind: "icp", normalized: "沪网备案9988"},
|
||
{name: "icp domain", input: "icp.example.com", kind: "domain", normalized: "icp.example.com"},
|
||
{name: "icp url query", input: "https://example.com/path?icp=1", kind: "domain", normalized: "example.com"},
|
||
// 备案号不含点号：掺了域名/版本号的描述性文字归关键词，否则会存成一条
|
||
// 永远匹配不上的死 ICP 规则。
|
||
{name: "icp with domain text", input: "备案 www.example.com", kind: "keyword", normalized: "备案 www.example.com"},
|
||
{name: "icp with version text", input: "某公司 ICP v1.0", kind: "keyword", normalized: "某公司 icp v1.0"},
|
||
{name: "icp fullwidth dot", input: "备案 例．com", kind: "keyword", normalized: "备案 例．com"},
|
||
{name: "keyword", input: " ACME Security ", kind: "keyword", normalized: "acme security"},
|
||
{name: "colon keyword", input: "ACME: Cloud: Security", kind: "keyword", normalized: "acme: cloud: security"},
|
||
{name: "empty", input: " ", wantErr: true},
|
||
{name: "invalid cidr", input: "10.0.0.0/not-a-prefix", wantErr: true},
|
||
{name: "invalid ipv4", input: "999.0.0.1", wantErr: true},
|
||
{name: "invalid ipv6", input: "2001:db8::zz", wantErr: true},
|
||
{name: "invalid ipv6 cidr", input: "2001:db8::zz/64", wantErr: true},
|
||
{name: "overbroad cidr", input: "10.0.0.0/8", wantErr: true},
|
||
{name: "bare suffix", input: "co.uk", wantErr: true},
|
||
{name: "empty domain label", input: "foo..example.com", wantErr: true},
|
||
}
|
||
for _, tc := range cases {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
rule, err := ParseAutoScopeLine(tc.input)
|
||
if tc.wantErr {
|
||
if err == nil {
|
||
t.Fatalf("ParseAutoScopeLine(%q) = %+v, want error", tc.input, rule)
|
||
}
|
||
return
|
||
}
|
||
if err != nil {
|
||
t.Fatalf("ParseAutoScopeLine(%q): %v", tc.input, err)
|
||
}
|
||
if rule.Kind != tc.kind {
|
||
t.Fatalf("kind=%q want %q", rule.Kind, tc.kind)
|
||
}
|
||
got := rule.Value
|
||
if rule.Kind == "domain" {
|
||
got = rule.Domain
|
||
} else if rule.Kind == "ip" || rule.Kind == "cidr" {
|
||
got = rule.Net
|
||
}
|
||
if got != tc.normalized {
|
||
t.Fatalf("normalized=%q want %q", got, tc.normalized)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestExplicitCompanyAttributionSurvivesScopeRebuild(t *testing.T) {
|
||
d, as, cs := testSetup(t)
|
||
defer d.Close()
|
||
|
||
stamp := time.Now().UnixNano()
|
||
explicitCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Explicit Attribution %d", stamp), "")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
autoCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Automatic Attribution %d", stamp), "")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer cleanupCompany(d, explicitCompany)
|
||
defer cleanupCompany(d, autoCompany)
|
||
|
||
domain := fmt.Sprintf("explicit-%d.invalid", stamp)
|
||
icp := fmt.Sprintf("ICP-%d", stamp)
|
||
network := fmt.Sprintf("2001:db8:%x::/64", uint64(stamp)&0xffff)
|
||
ip := fmt.Sprintf("2001:db8:%x::10", uint64(stamp)&0xffff)
|
||
|
||
// A pre-existing row with company_id and no provenance value represents old
|
||
// installations. The schema default conservatively treats it as explicit.
|
||
var assetID int64
|
||
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,company_id)
|
||
VALUES ('root_domain',$1,$1,$2) RETURNING id`, domain, explicitCompany).Scan(&assetID); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer deleteAsset(d, assetID)
|
||
appID, err := as.UpsertApp(UpsertAppReq{
|
||
Name: fmt.Sprintf("explicit-app-%d", stamp), ICP: icp, CompanyID: &explicitCompany,
|
||
})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer deleteAsset(d, appID)
|
||
autoAssetID, err := as.UpsertIP(UpsertIPReq{IP: ip})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer deleteAsset(d, autoAssetID)
|
||
|
||
rules := []ScopeInput{
|
||
{Kind: "domain", Value: domain},
|
||
{Kind: "icp", Value: icp},
|
||
{Kind: "cidr", Value: network},
|
||
}
|
||
if added, _, invalid, errs := cs.AddScopeInputs(autoCompany, rules, "test"); added != len(rules) || invalid != 0 {
|
||
t.Fatalf("AddScopeInputs: added=%d invalid=%d errors=%v", added, invalid, errs)
|
||
}
|
||
|
||
assertCompany := func(id, want int64, source string) {
|
||
t.Helper()
|
||
var got *int64
|
||
var gotSource string
|
||
if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, id).Scan(&got, &gotSource); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if got == nil || *got != want || gotSource != source {
|
||
t.Fatalf("asset %d company=%v source=%q, want %d/%q", id, got, gotSource, want, source)
|
||
}
|
||
}
|
||
assertCompany(assetID, explicitCompany, "explicit")
|
||
assertCompany(appID, explicitCompany, "explicit")
|
||
assertCompany(autoAssetID, autoCompany, "scope")
|
||
|
||
// Replacing every matching rule detaches only the automatically-owned row.
|
||
if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
|
||
{Kind: "domain", Value: fmt.Sprintf("replacement-%d.invalid", stamp)},
|
||
}, "test"); invalid != 0 || len(errs) != 0 {
|
||
t.Fatalf("UpdateScopeInputs: invalid=%d errors=%v", invalid, errs)
|
||
}
|
||
assertCompany(assetID, explicitCompany, "explicit")
|
||
assertCompany(appID, explicitCompany, "explicit")
|
||
var autoCompanyID *int64
|
||
var autoSource string
|
||
if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, autoAssetID).Scan(&autoCompanyID, &autoSource); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if autoCompanyID != nil || autoSource != "scope" {
|
||
t.Fatalf("automatic asset was not detached: company=%v source=%q", autoCompanyID, autoSource)
|
||
}
|
||
|
||
if err := cs.RecomputeAttribution(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
assertCompany(assetID, explicitCompany, "explicit")
|
||
assertCompany(appID, explicitCompany, "explicit")
|
||
|
||
// Deleting the explicitly selected company detaches through the FK, then the
|
||
// transactional rebuild may adopt the assets into a still-valid scope.
|
||
if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
|
||
{Kind: "domain", Value: domain},
|
||
{Kind: "icp", Value: icp},
|
||
}, "test"); invalid != 0 || len(errs) != 0 {
|
||
t.Fatalf("restore fallback scope: invalid=%d errors=%v", invalid, errs)
|
||
}
|
||
if err := cs.DeleteCompany(explicitCompany); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
assertCompany(assetID, autoCompany, "scope")
|
||
assertCompany(appID, autoCompany, "scope")
|
||
}
|
||
|
||
func TestParseStructuredCompanyScope(t *testing.T) {
|
||
icp, err := ParseScopeInput(ScopeInput{Kind: "ICP", Value: " 京ICP 备 123号-1\t"})
|
||
if err != nil {
|
||
t.Fatalf("parse ICP: %v", err)
|
||
}
|
||
if icp.Kind != "icp" || icp.Value != "京icp备123号-1" {
|
||
t.Fatalf("unexpected normalized ICP: %+v", icp)
|
||
}
|
||
keyword, err := ParseScopeInput(ScopeInput{Kind: "keyword", Value: " ACME Security "})
|
||
if err != nil {
|
||
t.Fatalf("parse keyword: %v", err)
|
||
}
|
||
if keyword.Value != "acme security" {
|
||
t.Fatalf("unexpected normalized keyword: %+v", keyword)
|
||
}
|
||
if _, err := ParseScopeInput(ScopeInput{Kind: "ip", Value: "example.com"}); err == nil {
|
||
t.Fatal("typed IP accepted a domain")
|
||
}
|
||
}
|
||
|
||
func TestCompanyICPAttribution(t *testing.T) {
|
||
d, err := Open(testDSN(t))
|
||
if err != nil {
|
||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||
}
|
||
// 关连接必须走 t.Cleanup 且**注册在清理之前**：t.Cleanup 是后进先出，
|
||
// 先注册关闭 → 关闭最后执行，下面的数据清理才连得上库。
|
||
// 原先这里是 `defer d.Close()`：defer 在函数返回时先跑，t.Cleanup 在那之后
|
||
// 才执行，于是清理语句全落在**已关闭的连接**上、错误又被 `_, _ =` 丢弃，
|
||
// 资产与公司就永久残留在库里。残留本身不会立刻报错，但本用例用
|
||
// `MAX(companies.id)+1` 当假 TaskID 给资产打标（见下方 suffix），
|
||
// 一旦这个数字与别的用例的任务 id 撞上，那个用例按「恰好 N 个资产」的断言
|
||
// 就会莫名失败——排查成本极高。
|
||
t.Cleanup(func() { d.Close() })
|
||
|
||
var suffix int64
|
||
if err := d.QueryRow(`SELECT COALESCE(MAX(id),0)+1 FROM companies`).Scan(&suffix); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
cs := d.Companies()
|
||
as := d.Assets()
|
||
companyID, _, err := cs.UpsertCompany(fmt.Sprintf("ICP Scope Co %d", suffix), "")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
t.Cleanup(func() {
|
||
// 不吞错误：清理失败会污染后续用例，必须让它在本次运行里显形。
|
||
if _, err := d.Exec(`DELETE FROM assets WHERE task_ids @> ARRAY[$1]::bigint[]`, suffix); err != nil {
|
||
t.Errorf("清理测试资产失败: %v", err)
|
||
}
|
||
if _, err := d.Exec(`DELETE FROM companies WHERE id=$1`, companyID); err != nil {
|
||
t.Errorf("清理测试公司失败: %v", err)
|
||
}
|
||
})
|
||
|
||
// A keyword can guide an Agent, but must never claim an asset by its name.
|
||
added, _, invalid, errs := cs.AddScopeInputs(companyID, []ScopeInput{
|
||
{Kind: "icp", Value: "京 ICP备 998877号"},
|
||
{Kind: "keyword", Value: "ICP Scope"},
|
||
}, "unit test")
|
||
if added != 2 || invalid != 0 || len(errs) != 0 {
|
||
t.Fatalf("add structured scope: added=%d invalid=%d errors=%v", added, invalid, errs)
|
||
}
|
||
|
||
rootID, err := as.UpsertRootDomain(UpsertRootDomainReq{
|
||
Domain: fmt.Sprintf("icp-scope-%d.example", suffix), ICP: "京icp备998877号", TaskID: suffix,
|
||
})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
appID, err := as.UpsertApp(UpsertAppReq{
|
||
Name: fmt.Sprintf("ICP Scope Keyword Only %d", suffix), TaskID: suffix,
|
||
})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
icpAppID, err := as.UpsertApp(UpsertAppReq{
|
||
Name: fmt.Sprintf("ICP Matched App %d", suffix), ICP: " 京 ICP备 998877号 ", TaskID: suffix,
|
||
})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
assertCompany := func(assetID int64, want *int64) {
|
||
t.Helper()
|
||
var got *int64
|
||
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, assetID).Scan(&got); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if want == nil && got != nil {
|
||
t.Fatalf("asset %d attributed by keyword: %d", assetID, *got)
|
||
}
|
||
if want != nil && (got == nil || *got != *want) {
|
||
t.Fatalf("asset %d company=%v want %d", assetID, got, *want)
|
||
}
|
||
}
|
||
assertCompany(rootID, &companyID)
|
||
assertCompany(appID, nil)
|
||
assertCompany(icpAppID, &companyID)
|
||
}
|
||
|
||
// TestCompanyScopeAttribution exercises the full loop against dev PG: create
|
||
// company (unique name), add scope, and verify auto-attribution at insert time,
|
||
// backfill of a pre-existing asset, CIDR + domain-suffix matching, and that
|
||
// out-of-scope assets stay unattributed.
|
||
func TestCompanyScopeAttribution(t *testing.T) {
|
||
d, err := Open(testDSN(t))
|
||
if err != nil {
|
||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||
}
|
||
as := d.Assets()
|
||
cs := d.Companies()
|
||
|
||
var startMax int64
|
||
if err := d.QueryRow(`SELECT COALESCE(MAX(id),0) FROM assets`).Scan(&startMax); err != nil {
|
||
d.Close()
|
||
t.Fatalf("startMax: %v", err)
|
||
}
|
||
t.Cleanup(func() {
|
||
_, _ = d.Exec(`DELETE FROM assets WHERE id > $1`, startMax)
|
||
d.Close()
|
||
})
|
||
|
||
uniq := startMax + 1
|
||
root := fmt.Sprintf("scopetest%d.com", uniq)
|
||
sub := "api." + root
|
||
ipIn := "198.51.100.9"
|
||
ipOut := "203.0.113.9"
|
||
outDomain := fmt.Sprintf("other%d.net", uniq)
|
||
|
||
cid, _, err := cs.UpsertCompany(fmt.Sprintf("ScopeCo %d", uniq), "")
|
||
if err != nil {
|
||
t.Fatalf("UpsertCompany: %v", err)
|
||
}
|
||
|
||
// a pre-existing asset (inserted BEFORE any scope) — must be back-filled.
|
||
preID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: sub})
|
||
if err != nil {
|
||
t.Fatalf("pre upsert: %v", err)
|
||
}
|
||
var preCompanyID *int64
|
||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, preID).Scan(&preCompanyID)
|
||
if preCompanyID != nil {
|
||
t.Fatalf("pre-scope asset should be unattributed, got %v", *preCompanyID)
|
||
}
|
||
|
||
cs.AddScope(cid, []string{root, "198.51.100.0/24"}, "unit test")
|
||
|
||
mustCid := func(id int64, want int64, label string) {
|
||
var cID *int64
|
||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
|
||
if cID == nil {
|
||
t.Fatalf("%s company_id = nil, want %d", label, want)
|
||
}
|
||
if *cID != want {
|
||
t.Fatalf("%s company_id = %d, want %d", label, *cID, want)
|
||
}
|
||
}
|
||
mustNil := func(id int64, label string) {
|
||
var cID *int64
|
||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
|
||
if cID != nil {
|
||
t.Fatalf("%s should be unattributed, got %d", label, *cID)
|
||
}
|
||
}
|
||
|
||
// backfill attributed the pre-existing subdomain (domain suffix match).
|
||
mustCid(preID, cid, "pre-existing subdomain (backfill)")
|
||
|
||
// insert-time attribution: ip in CIDR, another subdomain.
|
||
ipInID, err := as.UpsertIP(UpsertIPReq{IP: ipIn})
|
||
if err != nil {
|
||
t.Fatalf("UpsertIP in: %v", err)
|
||
}
|
||
mustCid(ipInID, cid, "in-CIDR ip (insert-time)")
|
||
|
||
sub2ID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: "www." + root})
|
||
if err != nil {
|
||
t.Fatalf("UpsertSubdomain: %v", err)
|
||
}
|
||
mustCid(sub2ID, cid, "new subdomain (insert-time)")
|
||
|
||
// out of scope stays unattributed.
|
||
outID, err := as.UpsertRootDomain(UpsertRootDomainReq{Domain: outDomain})
|
||
if err != nil {
|
||
t.Fatalf("UpsertRootDomain out: %v", err)
|
||
}
|
||
mustNil(outID, "out-of-scope domain")
|
||
|
||
ipOutID, err := as.UpsertIP(UpsertIPReq{IP: ipOut})
|
||
if err != nil {
|
||
t.Fatalf("UpsertIP out: %v", err)
|
||
}
|
||
mustNil(ipOutID, "out-of-CIDR ip")
|
||
}
|