Add files via upload

2026-07-19 18:38:00 +02:00 · 2026-07-16 19:37:01 +08:00
parent 15cc08a95b
commit d03cc47ac3
10 changed files with 1979 additions and 2 deletions
@@ -0,0 +1,316 @@
+package database
+
+import (
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+func TestAssetURLNormalizationAndValidation(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "asset-validation.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	asset := &Asset{Host: "https://例子.测试/path", Tags: []string{" prod ", "prod"}}
+	result, err := db.UpsertAssets([]*Asset{asset}, "")
+	if err != nil || result.Created != 1 {
+		t.Fatalf("URL asset was not created: result=%#v err=%v", result, err)
+	}
+	if asset.Domain != "xn--fsqu00a.xn--0zwm56d" || asset.Protocol != "https" || asset.Port != 443 {
+		t.Fatalf("URL fields were not normalized: %#v", asset)
+	}
+	if len(asset.Tags) != 1 || asset.Tags[0] != "prod" {
+		t.Fatalf("tags were not normalized: %#v", asset.Tags)
+	}
+
+	invalid := []*Asset{
+		{IP: "999.1.1.1", Status: "active"},
+		{Domain: "bad_domain.example", Status: "active"},
+		{Domain: "example.com", Port: 70000, Status: "active"},
+		{Domain: "example.com", Protocol: "HTTP 1.1", Status: "active"},
+		{Domain: "example.com", Status: "deleted"},
+	}
+	for _, candidate := range invalid {
+		if _, err := db.UpsertAssets([]*Asset{candidate}, ""); err == nil {
+			t.Fatalf("invalid asset unexpectedly accepted: %#v", candidate)
+		}
+	}
+
+	for _, host := range []string{"123", "not a formal target", "https://", "https://user:password@example.com"} {
+		result, err := db.UpsertAssets([]*Asset{{Host: host}}, "")
+		if err != nil || result.Created != 1 {
+			t.Fatalf("opaque asset address %q was not accepted: result=%#v err=%v", host, result, err)
+		}
+	}
+}
+
+func TestAssetValidationRejectsOversizedTags(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "asset-tag-validation.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+	_, err = db.UpsertAssets([]*Asset{{Domain: "example.com", Tags: []string{strings.Repeat("x", 65)}}}, "")
+	if err == nil || !strings.Contains(err.Error(), "标签") {
+		t.Fatalf("expected tag validation error, got %v", err)
+	}
+}
+
+func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "fofa-asset-validation.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	asset := &Asset{
+		Host:     "https://203.0.113.59:8443",
+		IP:       "203.0.113.59",
+		Domain:   "provider_specific_invalid_domain_59",
+		Port:     8443,
+		Protocol: "https",
+		Source:   "fofa",
+	}
+	result, err := db.UpsertAssets([]*Asset{asset}, "")
+	if err != nil || result.Created != 1 {
+		t.Fatalf("FOFA asset with dirty optional domain was not created: result=%#v err=%v", result, err)
+	}
+	if asset.Domain != "" || asset.IP != "203.0.113.59" {
+		t.Fatalf("FOFA structured fields were not sanitized: %#v", asset)
+	}
+}
+
+func TestAssetUpsertDeduplicatesAndUpdates(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	first := &Asset{Host: "https://example.com", Domain: "Example.COM", Port: 443, Protocol: "HTTPS", Title: "Old", Source: "fofa"}
+	result, err := db.UpsertAssets([]*Asset{first}, "user-a")
+	if err != nil || result.Created != 1 || result.Updated != 0 {
+		t.Fatalf("first upsert = %#v, %v", result, err)
+	}
+	second := &Asset{Domain: "example.com", Port: 443, Protocol: "https", Title: "New", Server: "nginx", Source: "fofa"}
+	result, err = db.UpsertAssets([]*Asset{second}, "user-a")
+	if err != nil || result.Created != 0 || result.Updated != 1 {
+		t.Fatalf("second upsert = %#v, %v", result, err)
+	}
+	assets, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
+	if err != nil || total != 1 || len(assets) != 1 {
+		t.Fatalf("list assets total=%d len=%d err=%v", total, len(assets), err)
+	}
+	if assets[0].Title != "New" || assets[0].Server != "nginx" || assets[0].Protocol != "https" {
+		t.Fatalf("asset not refreshed: %#v", assets[0])
+	}
+	stats, err := db.GetAssetStats(RBACListAccess{Scope: RBACScopeAll})
+	if err != nil || stats["total"] != 1 {
+		t.Fatalf("stats=%#v err=%v", stats, err)
+	}
+	coverage, ok := stats["coverage"].(map[string]interface{})
+	if !ok || coverage["never_scanned"] != 1 || coverage["rate"] != 0 {
+		t.Fatalf("coverage=%#v", stats["coverage"])
+	}
+	assetTrend, ok := stats["asset_trend"].([]map[string]interface{})
+	if !ok || len(assetTrend) != 30 {
+		t.Fatalf("asset trend=%#v", stats["asset_trend"])
+	}
+	riskTrend, ok := stats["risk_trend"].([]map[string]interface{})
+	if !ok || len(riskTrend) != 30 {
+		t.Fatalf("risk trend=%#v", stats["risk_trend"])
+	}
+}
+
+func TestAssetAccessFiltersOwners(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "assets-access.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+	now := time.Now()
+	if _, err := db.Exec(`INSERT INTO rbac_users (id,username,display_name,password_hash,enabled,is_builtin,created_at,updated_at) VALUES ('user-a','user-a','User A','hash',1,0,?,?)`, now, now); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.UpsertAssets([]*Asset{{IP: "10.0.0.1", Port: 80, Protocol: "http"}}, "user-a"); err != nil {
+		t.Fatal(err)
+	}
+	_, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-b", Scope: RBACScopeAssigned})
+	if err != nil || total != 0 {
+		t.Fatalf("unexpected cross-user assets: total=%d err=%v", total, err)
+	}
+	_, total, err = db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeOwn})
+	if err != nil || total != 1 {
+		t.Fatalf("owner cannot list asset: total=%d err=%v", total, err)
+	}
+	assets, _, err := db.ListAssets(1, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeAssigned})
+	if err != nil || len(assets) != 1 || !db.UserCanAccessResource("user-a", RBACScopeAssigned, "asset", assets[0].ID) {
+		t.Fatalf("creator assignment missing: assets=%d err=%v", len(assets), err)
+	}
+	options, err := db.ListAssignableRBACResources("asset", "10.0.0.1", 10)
+	if err != nil || len(options) != 1 {
+		t.Fatalf("asset resource picker: options=%#v err=%v", options, err)
+	}
+	project, err := db.CreateProject(&Project{Name: "Alpha", Status: "active"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := db.SetResourceOwner("project", project.ID, "user-b"); err != nil {
+		t.Fatal(err)
+	}
+	asset := assets[0]
+	asset.ProjectID = project.ID
+	if err := db.UpdateAsset(asset.ID, asset, RBACListAccess{Scope: RBACScopeAll}); err != nil {
+		t.Fatal(err)
+	}
+	projectAssets, total, err := db.ListAssets(20, 0, AssetListFilter{ProjectID: project.ID}, RBACListAccess{UserID: "user-b", Scope: RBACScopeOwn})
+	if err != nil || total != 1 || len(projectAssets) != 1 || projectAssets[0].ProjectName != "Alpha" {
+		t.Fatalf("project-bound asset access failed: total=%d assets=%#v err=%v", total, projectAssets, err)
+	}
+	if !db.UserCanAccessResource("user-b", RBACScopeOwn, "asset", asset.ID) {
+		t.Fatal("project owner cannot access bound asset")
+	}
+}
+
+func TestUpdateAssetsProjectIsAtomicAndScoped(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "asset-batch-project.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	project, err := db.CreateProject(&Project{Name: "Batch Project", Status: "active"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.UpsertAssets([]*Asset{
+		{IP: "192.0.2.1", Port: 80, Protocol: "http"},
+		{IP: "192.0.2.2", Port: 443, Protocol: "https"},
+	}, "owner-a"); err != nil {
+		t.Fatal(err)
+	}
+	assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
+	if err != nil || len(assets) != 2 {
+		t.Fatalf("list assets: len=%d err=%v", len(assets), err)
+	}
+	ids := []string{assets[0].ID, assets[1].ID}
+	updated, err := db.UpdateAssetsProject(ids, project.ID, RBACListAccess{UserID: "owner-a", Scope: RBACScopeOwn})
+	if err != nil || updated != 2 {
+		t.Fatalf("batch bind: updated=%d err=%v", updated, err)
+	}
+	for _, id := range ids {
+		asset, err := db.GetAsset(id, RBACListAccess{Scope: RBACScopeAll})
+		if err != nil || asset.ProjectID != project.ID {
+			t.Fatalf("asset %s was not bound: asset=%#v err=%v", id, asset, err)
+		}
+	}
+
+	if _, err := db.UpdateAssetsProject([]string{ids[0], "missing"}, "", RBACListAccess{Scope: RBACScopeAll}); err == nil {
+		t.Fatal("partial batch update unexpectedly succeeded")
+	}
+	asset, err := db.GetAsset(ids[0], RBACListAccess{Scope: RBACScopeAll})
+	if err != nil || asset.ProjectID != project.ID {
+		t.Fatalf("failed batch changed an asset: asset=%#v err=%v", asset, err)
+	}
+
+	updated, err = db.UpdateAssetsProject(ids, "", RBACListAccess{Scope: RBACScopeAll})
+	if err != nil || updated != 2 {
+		t.Fatalf("batch unbind: updated=%d err=%v", updated, err)
+	}
+}
+
+func TestAssetScanLinkReturnsTimeAndRelatedVulnerabilities(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "asset-scan.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	if _, err := db.UpsertAssets([]*Asset{{IP: "192.0.2.10", Port: 443, Protocol: "https"}}, ""); err != nil {
+		t.Fatal(err)
+	}
+	assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
+	if err != nil || len(assets) != 1 {
+		t.Fatalf("list assets: len=%d err=%v", len(assets), err)
+	}
+	conv, err := db.CreateConversation("asset scan", ConversationCreateMeta{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := db.MarkAssetScanned(assets[0].ID, conv.ID, "", "", RBACListAccess{Scope: RBACScopeAll}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.CreateVulnerability(&Vulnerability{ConversationID: conv.ID, Title: "finding", Severity: "high", Target: "192.0.2.10"}); err != nil {
+		t.Fatal(err)
+	}
+	linked, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if linked.LastScanAt == nil || linked.LastScanConversationID != conv.ID || linked.VulnerabilityCount != 1 || linked.RiskLevel != "high" {
+		t.Fatalf("unexpected scan metadata: %#v", linked)
+	}
+	if _, err := db.Exec(`UPDATE vulnerabilities SET status='fixed' WHERE conversation_id=?`, conv.ID); err != nil {
+		t.Fatal(err)
+	}
+	resolved, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if resolved.VulnerabilityCount != 1 || resolved.RiskLevel != "normal" {
+		t.Fatalf("resolved finding should remain in history without raising current risk: %#v", resolved)
+	}
+}
+
+func TestAssetListFlexibleFiltersAndOldestScanPagination(t *testing.T) {
+	db, err := NewDB(filepath.Join(t.TempDir(), "asset-query.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	assets := []*Asset{
+		{IP: "192.0.2.1", Port: 443, Protocol: "https", Source: "fofa", Tags: []string{"prod"}},
+		{IP: "192.0.2.2", Port: 80, Protocol: "http", Source: "manual", Tags: []string{"prod", "legacy"}},
+		{Domain: "never.example.com", Port: 443, Protocol: "https", Source: "manual", Tags: []string{"prod"}},
+	}
+	if _, err := db.UpsertAssets(assets, ""); err != nil {
+		t.Fatal(err)
+	}
+	old := time.Now().Add(-90 * 24 * time.Hour).UTC()
+	recent := time.Now().Add(-24 * time.Hour).UTC()
+	if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, old, assets[0].ID); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, recent, assets[1].ID); err != nil {
+		t.Fatal(err)
+	}
+
+	access := RBACListAccess{Scope: RBACScopeAll}
+	firstPage, total, err := db.ListAssets(2, 0, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access)
+	if err != nil || total != 3 || len(firstPage) != 2 {
+		t.Fatalf("oldest scan page: total=%d len=%d err=%v", total, len(firstPage), err)
+	}
+	if firstPage[0].ID != assets[2].ID || firstPage[0].LastScanAt != nil || firstPage[1].ID != assets[0].ID {
+		t.Fatalf("expected never-scanned then oldest scanned asset, got %#v", firstPage)
+	}
+	secondPage, _, err := db.ListAssets(2, 2, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access)
+	if err != nil || len(secondPage) != 1 || secondPage[0].ID != assets[1].ID {
+		t.Fatalf("unexpected second page: %#v err=%v", secondPage, err)
+	}
+
+	never, total, err := db.ListAssets(20, 0, AssetListFilter{ScanState: "never"}, access)
+	if err != nil || total != 1 || len(never) != 1 || never[0].ID != assets[2].ID {
+		t.Fatalf("never-scanned filter: total=%d assets=%#v err=%v", total, never, err)
+	}
+	port := 443
+	filtered, total, err := db.ListAssets(20, 0, AssetListFilter{Source: "fofa", Port: &port, LastScanBefore: &recent}, access)
+	if err != nil || total != 1 || len(filtered) != 1 || filtered[0].ID != assets[0].ID {
+		t.Fatalf("structured filters: total=%d assets=%#v err=%v", total, filtered, err)
+	}
+}