package db import ( "testing" ) // testSetup opens a DB and returns both stores. Skips if no PG. func testSetup(t *testing.T) (*DB, *AssetStore, *CompanyStore) { t.Helper() d, err := Open(testDSN(t)) if err != nil { t.Skipf("postgres unavailable (%v)", err) } return d, d.Assets(), d.Companies() } // deleteAsset removes one v2 asset by id. func deleteAsset(d *DB, id int64) { d.Exec(`DELETE FROM assets WHERE id = $1`, id) } // ===================================================================== // RootDomain // ===================================================================== func TestUpsertRootDomain(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id1, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "roottest.io", ICP: "A12345", TaskID: 1}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id1) if id1 == 0 { t.Fatal("expected non-zero id") } // upsert again — same id (dedup), ICP preserved id2, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "roottest.io", TaskID: 2}) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("dedup failed: %d != %d", id2, id1) } // task_ids should now contain both 1 and 2 var taskIDs []byte d.QueryRow(`SELECT task_ids FROM assets WHERE id = $1`, id1).Scan(&taskIDs) // ICP should still be set (COALESCE keeps existing) var icp *string d.QueryRow(`SELECT icp FROM assets WHERE id = $1`, id1).Scan(&icp) if icp == nil || *icp != "A12345" { t.Errorf("ICP not preserved: %v", icp) } } func TestUpsertRootDomainEmpty(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() _, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: ""}) if err == nil { t.Error("expected error for empty domain") } } // ===================================================================== // IP // ===================================================================== func TestUpsertIP(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id1, err := av2.UpsertIP(UpsertIPReq{ IP: "192.168.10.5", OpenPorts: []PortService{{Port: 22, Service: "ssh"}, {Port: 80, Service: "http"}}, TaskID: 1, }) if err != nil { t.Fatal(err) } defer deleteAsset(d, id1) // c_segment should be 192.168.10.0/24 var cseg *string d.QueryRow(`SELECT c_segment::text FROM assets WHERE id = $1`, id1).Scan(&cseg) if cseg == nil || *cseg != "192.168.10.0/24" { t.Errorf("c_segment: want 192.168.10.0/24, got %v", cseg) } // append another port via UpsertIP (merge) id2, err := av2.UpsertIP(UpsertIPReq{ IP: "192.168.10.5", OpenPorts: []PortService{{Port: 443, Service: "https"}}, TaskID: 2, }) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("dedup failed: %d != %d", id2, id1) } // verify all 3 ports present var cnt int d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE id = $1`, id1).Scan(&cnt) if cnt != 3 { t.Errorf("open_ports merge: want 3 ports, got %d", cnt) } } func TestUpsertIPBoundDomains(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id, err := av2.UpsertIP(UpsertIPReq{ IP: "10.1.2.3", BoundDomains: []string{"a.example.com"}, }) if err != nil { t.Fatal(err) } defer deleteAsset(d, id) // append a second domain id2, err := av2.UpsertIP(UpsertIPReq{ IP: "10.1.2.3", BoundDomains: []string{"b.example.com"}, }) if err != nil { t.Fatal(err) } if id2 != id { t.Errorf("dedup failed: %d != %d", id2, id) } var cnt int d.QueryRow(`SELECT array_length(bound_domains, 1) FROM assets WHERE id = $1`, id).Scan(&cnt) if cnt != 2 { t.Errorf("bound_domains merge: want 2, got %d", cnt) } } func TestAppendIPPort(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id, err := av2.UpsertIP(UpsertIPReq{IP: "10.9.8.7"}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id) if err := av2.AppendIPPort("10.9.8.7", 3306, "mysql"); err != nil { t.Fatal(err) } if err := av2.AppendIPPort("10.9.8.7", 22, "ssh"); err != nil { t.Fatal(err) } var cnt int d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE id = $1`, id).Scan(&cnt) if cnt != 2 { t.Errorf("AppendIPPort: want 2 ports, got %d", cnt) } } // ===================================================================== // Subdomain // ===================================================================== func TestUpsertSubdomain(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id, err := av2.UpsertSubdomain(UpsertSubdomainReq{ Domain: "sub.subdtest.com", RecordType: "A", RecordValue: []string{"1.2.3.4"}, ICP: "B99", TaskID: 1, }) if err != nil { t.Fatal(err) } defer d.Exec(`DELETE FROM assets WHERE domain IN ('sub.subdtest.com', 'subdtest.com') OR ip = '1.2.3.4'`) if id == 0 { t.Fatal("expected non-zero id") } // verify root_domain auto-populated var rootDomain string d.QueryRow(`SELECT COALESCE(root_domain,'') FROM assets WHERE id = $1`, id).Scan(&rootDomain) if rootDomain != "subdtest.com" { t.Errorf("root_domain: want subdtest.com, got %q", rootDomain) } // side effect: root domain asset should exist var rootCnt int d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type = 'root_domain' AND domain = 'subdtest.com'`).Scan(&rootCnt) if rootCnt != 1 { t.Error("side-effect root domain not created") } // side effect: IP asset with bound_domain var ipCnt int d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type = 'ip' AND ip = '1.2.3.4'`).Scan(&ipCnt) if ipCnt != 1 { t.Error("side-effect IP asset not created") } } func TestUpsertSubdomainDedup(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id1, err := av2.UpsertSubdomain(UpsertSubdomainReq{Domain: "dedup.subdtest2.net", RecordType: "A", RecordValue: []string{"5.5.5.5"}}) if err != nil { t.Fatal(err) } defer d.Exec(`DELETE FROM assets WHERE domain IN ('dedup.subdtest2.net', 'subdtest2.net') OR ip = '5.5.5.5'`) id2, err := av2.UpsertSubdomain(UpsertSubdomainReq{Domain: "dedup.subdtest2.net", RecordType: "A", RecordValue: []string{"5.5.5.5"}}) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("subdomain dedup failed: %d vs %d", id2, id1) } } // ===================================================================== // App // ===================================================================== func TestUpsertAppByBundle(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id1, err := av2.UpsertApp(UpsertAppReq{Name: "MyApp", BundleID: "com.example.myapp", Category: "tools"}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id1) // dedup by bundle id2, err := av2.UpsertApp(UpsertAppReq{Name: "MyApp Updated", BundleID: "com.example.myapp"}) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("app bundle dedup failed: %d vs %d", id2, id1) } } func TestUpsertAppByName(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id1, err := av2.UpsertApp(UpsertAppReq{Name: "UniqueAppNoBundle", Category: "utility"}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id1) id2, err := av2.UpsertApp(UpsertAppReq{Name: "UniqueAppNoBundle"}) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("app name dedup failed: %d vs %d", id2, id1) } } // ===================================================================== // HTTPService // ===================================================================== func TestUpsertHTTPService(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() sc := 200 cl := int64(1024) id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{ URL: "https://www.httptest.example.com/", Technologies: []string{"nginx", "vue"}, StatusCode: &sc, ContentLength: &cl, PageTitle: "Test Site", TaskID: 1, }) if err != nil { t.Fatal(err) } defer d.Exec(`DELETE FROM assets WHERE url = 'https://www.httptest.example.com' OR domain = 'www.httptest.example.com' OR domain = 'httptest.example.com'`) if id1 == 0 { t.Fatal("expected non-zero id") } // verify technology was stored var techCnt int d.QueryRow(`SELECT array_length(technologies, 1) FROM assets WHERE id = $1`, id1).Scan(&techCnt) if techCnt != 2 { t.Errorf("technologies: want 2, got %d", techCnt) } // upsert again with additional tech → should merge (append) id2, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{ URL: "https://www.httptest.example.com/", Technologies: []string{"react", "webpack"}, TaskID: 2, }) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("http service dedup failed: %d vs %d", id2, id1) } d.QueryRow(`SELECT array_length(technologies, 1) FROM assets WHERE id = $1`, id1).Scan(&techCnt) if techCnt != 4 { t.Errorf("technologies merge: want 4, got %d", techCnt) } } func TestUpsertHTTPServiceAuthAppend(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() auth1 := []map[string]any{{"type": "basic", "username": "admin", "password": "pass"}} id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{ URL: "http://authtest.example.org/", Auth: auth1, TaskID: 1, }) if err != nil { t.Fatal(err) } defer d.Exec(`DELETE FROM assets WHERE url = 'http://authtest.example.org' OR domain IN ('authtest.example.org', 'example.org')`) auth2 := []map[string]any{{"type": "bearer", "token": "tok123"}} id2, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{ URL: "http://authtest.example.org/", Auth: auth2, TaskID: 2, }) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("dedup failed: %d vs %d", id2, id1) } // auth should have 2 items now var authCnt int d.QueryRow(`SELECT cardinality(auth) FROM assets WHERE id = $1`, id1).Scan(&authCnt) if authCnt != 2 { t.Errorf("auth append: want 2, got %d", authCnt) } } // ===================================================================== // OtherService // ===================================================================== func TestUpsertOtherService(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id1, err := av2.UpsertOtherService(UpsertOtherServiceReq{ IP: "172.16.0.1", Port: 22, ServiceName: "ssh", TaskID: 1, }) if err != nil { t.Fatal(err) } defer d.Exec(`DELETE FROM assets WHERE (type = 'service' AND service_name = 'ssh' AND ip = '172.16.0.1') OR (type = 'ip' AND ip = '172.16.0.1')`) if id1 == 0 { t.Fatal("expected non-zero id") } // dedup id2, err := av2.UpsertOtherService(UpsertOtherServiceReq{ IP: "172.16.0.1", Port: 22, ServiceName: "ssh", }) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("other service dedup failed: %d vs %d", id2, id1) } // side-effect: IP asset should have port 22 in open_ports var cnt int d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE type='ip' AND ip='172.16.0.1'`).Scan(&cnt) if cnt == 0 { t.Error("side-effect: IP open_ports not populated") } } func TestUpsertOtherServiceAuthAppend(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() auth1 := []map[string]any{{"username": "admin", "password": "secret"}} id1, err := av2.UpsertOtherService(UpsertOtherServiceReq{ IP: "10.5.5.5", Port: 3306, ServiceName: "mysql", Auth: auth1, }) if err != nil { t.Fatal(err) } defer d.Exec(`DELETE FROM assets WHERE (type = 'service' AND ip = '10.5.5.5') OR (type = 'ip' AND ip = '10.5.5.5')`) auth2 := []map[string]any{{"username": "root", "password": "root"}} id2, err := av2.UpsertOtherService(UpsertOtherServiceReq{ IP: "10.5.5.5", Port: 3306, ServiceName: "mysql", Auth: auth2, }) if err != nil || id2 != id1 { t.Fatalf("dedup or error: %v, ids %d vs %d", err, id2, id1) } var authCnt int d.QueryRow(`SELECT cardinality(auth) FROM assets WHERE id = $1`, id1).Scan(&authCnt) if authCnt != 2 { t.Errorf("auth append: want 2, got %d", authCnt) } } // ===================================================================== // Endpoint // ===================================================================== func TestUpsertEndpoint(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() params1 := []map[string]any{{"location": "query", "name": "id", "value": "1"}} id1, err := av2.UpsertEndpoint(UpsertEndpointReq{ URL: "https://api.eptest.com/users?id=1", Method: "GET", Params: params1, TaskID: 1, }) if err != nil { t.Fatal(err) } defer d.Exec(`DELETE FROM assets WHERE url LIKE '%eptest.com%' OR domain IN ('api.eptest.com', 'eptest.com')`) if id1 == 0 { t.Fatal("expected non-zero id") } // dedup (same URL + method = same endpoint) id2, err := av2.UpsertEndpoint(UpsertEndpointReq{ URL: "https://api.eptest.com/users?id=1", Method: "GET", Params: []map[string]any{{"location": "query", "name": "filter", "value": "active"}}, TaskID: 2, }) if err != nil { t.Fatal(err) } if id2 != id1 { t.Errorf("endpoint dedup failed: %d vs %d", id2, id1) } // params should be merged (2 distinct params) var paramCnt int d.QueryRow(`SELECT cardinality(params) FROM assets WHERE id = $1`, id1).Scan(¶mCnt) if paramCnt != 2 { t.Errorf("params merge: want 2, got %d", paramCnt) } } func TestUpsertEndpointRequiredFields(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() _, err := av2.UpsertEndpoint(UpsertEndpointReq{URL: "", Method: "GET"}) if err == nil { t.Error("expected error for empty URL") } _, err = av2.UpsertEndpoint(UpsertEndpointReq{URL: "https://example.com/", Method: ""}) if err == nil { t.Error("expected error for empty method") } } // ===================================================================== // Query helpers // ===================================================================== func TestQueryByType(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "querytest.example.net"}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id) assets, err := av2.QueryByType("root_domain", 10, 0) if err != nil { t.Fatal(err) } found := false for _, a := range assets { if a.ID == id { found = true } } if !found { t.Error("QueryByType: inserted asset not found") } } func TestQueryByTask(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() const taskID = int64(99999) id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "taskquery.net", TaskID: taskID}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id) assets, err := av2.QueryByTask(taskID, "", 10) if err != nil { t.Fatal(err) } found := false for _, a := range assets { if a.ID == id { found = true } } if !found { t.Error("QueryByTask: inserted asset not found") } } func TestQueryByCompany(t *testing.T) { d, av2, cs := testSetup(t) defer d.Close() companyID, _, err := cs.UpsertCompany("QueryByCompanyCorp", "") if err != nil { t.Fatal(err) } defer cleanupCompany(d, companyID) cs.AddScope(companyID, []string{"qbc-test.io"}, "test") id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "qbc-test.io"}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id) if err := cs.RecomputeAttribution(); err != nil { t.Fatal(err) } assets, err := av2.QueryByCompany(companyID, "root_domain", 10) if err != nil { t.Fatal(err) } found := false for _, a := range assets { if a.ID == id { found = true } } if !found { t.Error("QueryByCompany: attributed asset not found") } } func TestCountsByType(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "counttest.example.biz"}) if err != nil { t.Fatal(err) } defer deleteAsset(d, id) counts, err := av2.CountsByType() if err != nil { t.Fatal(err) } if counts["root_domain"] == 0 { t.Error("CountsByType: root_domain count should be > 0") } } // ===================================================================== // Company attribution at insert time // ===================================================================== func TestCompanyAttributionAtInsertTime(t *testing.T) { d, av2, cs := testSetup(t) defer d.Close() companyID, _, err := cs.UpsertCompany("InsertAttrCorp", "") if err != nil { t.Fatal(err) } defer cleanupCompany(d, companyID) cs.AddScope(companyID, []string{"insertattr.com"}, "test") // Now insert a root domain — should be auto-attributed assetID, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "insertattr.com"}) if err != nil { t.Fatal(err) } defer deleteAsset(d, assetID) var cid *int64 d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, assetID).Scan(&cid) if cid == nil || *cid != companyID { t.Errorf("attribution at insert time: want company %d, got %v", companyID, cid) } } // ===================================================================== // calcCSegment // ===================================================================== func TestCalcCSegment(t *testing.T) { tests := []struct { ip string want string }{ {"192.168.1.5", "192.168.1.0/24"}, {"10.0.0.255", "10.0.0.0/24"}, {"", ""}, {"notanip", ""}, } for _, tc := range tests { got := calcCSegment(tc.ip) if got != tc.want { t.Errorf("calcCSegment(%q): want %q, got %q", tc.ip, tc.want, got) } } } // ===================================================================== // marshalStringArray // ===================================================================== func TestMarshalStringArray(t *testing.T) { got := marshalStringArray([]string{"a", "b", "c"}) if got != `{"a","b","c"}` { t.Errorf("unexpected: %q", got) } got2 := marshalStringArray(nil) if got2 != "{}" { t.Errorf("empty: %q", got2) } } // ===================================================================== // normalizeURL // ===================================================================== func TestNormalizeURL(t *testing.T) { tests := []struct { raw string want string }{ {"HTTPS://Example.COM/path", "https://example.com/path"}, {"http://example.com/", "http://example.com"}, {"http://example.com/page", "http://example.com/page"}, } for _, tc := range tests { got := normalizeURL(tc.raw) if got != tc.want { t.Errorf("normalizeURL(%q): want %q, got %q", tc.raw, tc.want, got) } } }