package functions import ( "context" "os" "path/filepath" "testing" "time" "github.com/j3ssie/osmedeus/v5/internal/config" "github.com/j3ssie/osmedeus/v5/internal/database" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // setupTestDB initializes a test SQLite database func setupTestDB(t *testing.T) (cleanup func()) { tmpDir := t.TempDir() dbPath := filepath.Join(tmpDir, "test.sqlite") cfg := &config.Config{ BaseFolder: tmpDir, Database: config.DatabaseConfig{ DBEngine: "sqlite", DBPath: dbPath, }, } _, err := database.Connect(cfg) require.NoError(t, err) ctx := context.Background() require.NoError(t, database.Migrate(ctx)) return func() { _ = database.Close() database.SetDB(nil) } } func TestDbImportAssetFromFile(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Use the sample httpx data file testFile := "../../test/testdata/sample-jsonl-output/http-data.jsonl" // Check if file exists _, err := os.Stat(testFile) require.NoError(t, err, "sample http-data.jsonl file must exist") result, err := registry.Execute( `db_import_asset_from_file("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) // http-data.jsonl has 3 lines - result is now a map with stats stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") assert.Equal(t, 3, stats["new"]) assert.Equal(t, 3, stats["total"]) // Verify assets were imported ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var assets []database.Asset err = db.NewSelect().Model(&assets). Where("workspace = ?", "test-workspace"). Scan(ctx) require.NoError(t, err) assert.Len(t, assets, 3) // Check specific fields were mapped correctly var asset database.Asset err = db.NewSelect().Model(&asset). Where("asset_value = ?", "api.hackerone.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "http://api.hackerone.com", asset.URL) assert.Equal(t, "HackerOne API", asset.Title) assert.Equal(t, 200, asset.StatusCode) assert.Equal(t, "text/html", asset.ContentType) assert.Contains(t, asset.Technologies, "Cloudflare") assert.NotEmpty(t, asset.RawJsonData) // Check webserver and cdn_name are in remarks assert.Contains(t, asset.Remarks, "cloudflare") // webserver assert.Contains(t, asset.Remarks, "cloudflare") // cdn_name (same value here) // Check gslink asset has both nginx and cloudfront in remarks var gsAsset database.Asset err = db.NewSelect().Model(&gsAsset). Where("asset_value = ?", "gslink.hackerone.com"). Scan(ctx) require.NoError(t, err) assert.Contains(t, gsAsset.Remarks, "nginx") assert.Contains(t, gsAsset.Remarks, "cloudfront") } func TestDbImportAssetFromFile_EmptyWorkspace(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_import_asset_from_file("", "/tmp/test.jsonl")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbImportAssetFromFile_NonExistentFile(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_import_asset_from_file("test-workspace", "/nonexistent/file.jsonl")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbImportVuln(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Single nuclei-style JSON jsonData := `{"template-id":"dns-saas-service-detection","info":{"name":"DNS SaaS Service Detection","severity":"info","tags":["dns","service","discovery"],"description":"A CNAME DNS record was discovered"},"type":"dns","host":"support.hackerone.com","matched-at":"support.hackerone.com"}` result, err := registry.Execute( "db_import_vuln(\"test-workspace\", '"+jsonData+"')", map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, true, result) // Verify vulnerability was imported ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var vuln database.Vulnerability err = db.NewSelect().Model(&vuln). Where("workspace = ?", "test-workspace"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "dns-saas-service-detection", vuln.VulnInfo) assert.Equal(t, "DNS SaaS Service Detection", vuln.VulnTitle) assert.Equal(t, "info", vuln.Severity) assert.Equal(t, "support.hackerone.com", vuln.AssetValue) assert.Equal(t, "dns", vuln.AssetType) assert.Contains(t, vuln.Tags, "dns") assert.Contains(t, vuln.Tags, "service") assert.Contains(t, vuln.Tags, "discovery") } func TestDbImportVuln_InvalidJSON(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_import_vuln("test-workspace", "not valid json")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbImportVulnFromFile(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Use the sample nuclei data file testFile := "../../test/testdata/sample-jsonl-output/vuln-data.jsonl" // Check if file exists _, err := os.Stat(testFile) require.NoError(t, err, "sample vuln-data.jsonl file must exist") result, err := registry.Execute( `db_import_vuln_from_file("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) // vuln-data.jsonl has 13 lines - result is now a map with stats stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") assert.Equal(t, 13, stats["new"]) assert.Equal(t, 13, stats["total"]) // Verify vulnerabilities were imported ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var vulns []database.Vulnerability err = db.NewSelect().Model(&vulns). Where("workspace = ?", "test-workspace"). Scan(ctx) require.NoError(t, err) assert.Len(t, vulns, 13) // Check that different vulnerability types were imported var severities []string for _, v := range vulns { severities = append(severities, v.Severity) } assert.Contains(t, severities, "info") } func TestDbImportVulnFromFile_EmptyWorkspace(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_import_vuln_from_file("", "/tmp/test.jsonl")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbImportVulnFromFile_NonExistentFile(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_import_vuln_from_file("test-workspace", "/nonexistent/file.jsonl")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbImportAssetFromFile_Upsert(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Create a temp file with duplicate entries tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "assets.jsonl") content := `{"host":"example.com","url":"http://example.com","title":"First","status_code":200} {"host":"example.com","url":"http://example.com","title":"Updated","status_code":301} ` err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) result, err := registry.Execute( `db_import_asset_from_file("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) // Result is now a map with stats - 1 new, 1 updated (same asset with different data) stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") assert.Equal(t, 2, stats["total"]) // Both lines processed // Verify only one asset exists (upsert) ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var assets []database.Asset err = db.NewSelect().Model(&assets). Where("workspace = ?", "test-workspace"). Scan(ctx) require.NoError(t, err) assert.Len(t, assets, 1) // Should have the updated values assert.Equal(t, "Updated", assets[0].Title) assert.Equal(t, 301, assets[0].StatusCode) } func TestMapJSONToAsset(t *testing.T) { data := map[string]interface{}{ "host": "example.com", "url": "http://example.com", "input": "https://example.com", "scheme": "http", "method": "GET", "path": "/", "status_code": float64(200), "content_type": "text/html", "content_length": float64(1234), "title": "Example", "words": float64(100), "lines": float64(50), "host_ip": "1.2.3.4", "a": []interface{}{"1.2.3.4", "5.6.7.8"}, "tech": []interface{}{"Nginx", "PHP"}, "time": "123ms", "webserver": "nginx", } asset := mapJSONToAsset(data, "test-workspace", `{"host":"example.com"}`) assert.Equal(t, "test-workspace", asset.Workspace) assert.Equal(t, "example.com", asset.AssetValue) assert.Equal(t, "http://example.com", asset.URL) assert.Equal(t, "https://example.com", asset.Input) assert.Equal(t, "http", asset.Scheme) assert.Equal(t, "GET", asset.Method) assert.Equal(t, "/", asset.Path) assert.Equal(t, 200, asset.StatusCode) assert.Equal(t, "text/html", asset.ContentType) assert.Equal(t, int64(1234), asset.ContentLength) assert.Equal(t, "Example", asset.Title) assert.Equal(t, 100, asset.Words) assert.Equal(t, 50, asset.Lines) assert.Equal(t, "1.2.3.4", asset.HostIP) assert.Equal(t, []string{"1.2.3.4", "5.6.7.8"}, asset.DnsRecords) assert.Equal(t, []string{"Nginx", "PHP"}, asset.Technologies) assert.Equal(t, "123ms", asset.ResponseTime) assert.Equal(t, "", asset.Source) // webserver no longer mapped to source assert.Contains(t, asset.Remarks, "nginx") // webserver now in remarks assert.NotEmpty(t, asset.RawJsonData) } func TestMapJSONToAsset_HostIPFromIPHost(t *testing.T) { data := map[string]interface{}{ "host": "143.92.73.100", "url": "https://fms.example.com", "status_code": float64(200), } asset := mapJSONToAsset(data, "test-workspace", `{"host":"143.92.73.100"}`) assert.Equal(t, "https://fms.example.com", asset.AssetValue) assert.Equal(t, "143.92.73.100", asset.HostIP) // NEW: IP should be stored } func TestMapJSONToAsset_ExplicitHostIPPreferred(t *testing.T) { data := map[string]interface{}{ "host": "143.92.73.100", "host_ip": "203.0.113.42", // Different IP "url": "https://example.com", } asset := mapJSONToAsset(data, "test-workspace", `{}`) assert.Equal(t, "203.0.113.42", asset.HostIP) // Explicit value wins assert.Equal(t, "https://example.com", asset.AssetValue) } func TestMapJSONToAsset_DomainHostNoChange(t *testing.T) { data := map[string]interface{}{ "host": "example.com", "url": "https://example.com", } asset := mapJSONToAsset(data, "test-workspace", `{}`) assert.Equal(t, "", asset.HostIP) // No IP detected assert.Equal(t, "example.com", asset.AssetValue) } func TestMapJSONToVuln(t *testing.T) { data := map[string]interface{}{ "template-id": "test-vuln", "info": map[string]interface{}{ "name": "Test Vulnerability", "description": "A test vulnerability", "severity": "high", "tags": []interface{}{"tag1", "tag2"}, }, "host": "example.com", "matched-at": "http://example.com/path", "type": "http", "request": "GET / HTTP/1.1", "response": "HTTP/1.1 200 OK", } vuln := mapJSONToVuln(data, "test-workspace", `{"template-id":"test-vuln"}`) assert.Equal(t, "test-workspace", vuln.Workspace) assert.Equal(t, "test-vuln", vuln.VulnInfo) assert.Equal(t, "Test Vulnerability", vuln.VulnTitle) assert.Equal(t, "A test vulnerability", vuln.VulnDesc) assert.Equal(t, "high", vuln.Severity) assert.Equal(t, "example.com", vuln.AssetValue) // host takes precedence assert.Equal(t, "http", vuln.AssetType) assert.Equal(t, []string{"tag1", "tag2"}, vuln.Tags) assert.Equal(t, "GET / HTTP/1.1", vuln.DetailHTTPRequest) assert.Equal(t, "HTTP/1.1 200 OK", vuln.DetailHTTPResponse) assert.NotEmpty(t, vuln.RawVulnJSON) } func TestMapJSONToVuln_MatchedAtFallback(t *testing.T) { data := map[string]interface{}{ "template-id": "test-vuln", "info": map[string]interface{}{ "name": "Test", "severity": "low", }, "matched-at": "http://example.com/path", "type": "http", } vuln := mapJSONToVuln(data, "test-workspace", `{}`) // When host is not present, matched-at is used as fallback assert.Equal(t, "http://example.com/path", vuln.AssetValue) } func TestDbPartialImportAsset(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_partial_import_asset("test-workspace", "domain", "sub.example.com")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, true, result) // Verify asset was imported with correct fields ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ? AND asset_value = ?", "test-workspace", "sub.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "domain", asset.AssetType) assert.Equal(t, "test-workspace", asset.Workspace) assert.Equal(t, "sub.example.com", asset.AssetValue) } func TestDbPartialImportAsset_Upsert(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Import first time _, err := registry.Execute( `db_partial_import_asset("test-workspace", "domain", "sub.example.com")`, map[string]interface{}{}, ) require.NoError(t, err) // Import same asset again with different type result, err := registry.Execute( `db_partial_import_asset("test-workspace", "subdomain", "sub.example.com")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, true, result) // Verify only 1 row exists ctx := context.Background() db := database.GetDB() require.NotNil(t, db) count, err := db.NewSelect().Model((*database.Asset)(nil)). Where("workspace = ? AND asset_value = ?", "test-workspace", "sub.example.com"). Count(ctx) require.NoError(t, err) assert.Equal(t, 1, count) // Verify type was updated var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ? AND asset_value = ?", "test-workspace", "sub.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "subdomain", asset.AssetType) } func TestDbPartialImportAsset_MissingArgs(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_partial_import_asset("test-workspace", "domain")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbPartialImportAssetFile(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Create a temp file with 3 lines tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "domains.txt") content := "sub1.example.com\nsub2.example.com\nsub3.example.com\n" err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) result, err := registry.Execute( `db_partial_import_asset_file("test-workspace", "domain", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, int64(3), result) // Verify all 3 assets are in the DB ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var assets []database.Asset err = db.NewSelect().Model(&assets). Where("workspace = ?", "test-workspace"). Scan(ctx) require.NoError(t, err) assert.Len(t, assets, 3) // All should have the specified asset_type for _, a := range assets { assert.Equal(t, "domain", a.AssetType) } } func TestDbPartialImportAssetFile_SkipsEmptyLines(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Create a temp file with blank lines tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "domains.txt") content := "sub1.example.com\n\n \nsub2.example.com\n\n" err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) result, err := registry.Execute( `db_partial_import_asset_file("test-workspace", "domain", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, int64(2), result) } func TestDbPartialImportAssetFile_NonExistentFile(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_partial_import_asset_file("test-workspace", "domain", "/nonexistent/file.txt")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbQuickImportAsset(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Import a domain without specifying asset_type (should auto-classify) result, err := registry.Execute( `db_quick_import_asset("test-workspace", "sub.example.com")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, true, result) // Verify asset was imported ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ? AND asset_value = ?", "test-workspace", "sub.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "domain", asset.AssetType) assert.Equal(t, "test-workspace", asset.Workspace) // Verify event log was created for new asset var eventLog database.EventLog err = db.NewSelect().Model(&eventLog). Where("topic = ? AND workspace = ?", "db.new.asset", "test-workspace"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "db.new.asset", eventLog.Topic) assert.Equal(t, "db_quick_import_asset", eventLog.Source) assert.Equal(t, "function", eventLog.SourceType) assert.Contains(t, eventLog.Data, "sub.example.com") } func TestDbQuickImportAsset_WithType(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Import an IP with explicit asset_type result, err := registry.Execute( `db_quick_import_asset("test-workspace", "192.168.1.1", "ip")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, true, result) // Verify asset was imported with specified type ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ? AND asset_value = ?", "test-workspace", "192.168.1.1"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "ip", asset.AssetType) } func TestDbQuickImportAsset_NoEventOnUpdate(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() // Import asset first time _, err := registry.Execute( `db_quick_import_asset("test-workspace", "existing.example.com")`, map[string]interface{}{}, ) require.NoError(t, err) ctx := context.Background() db := database.GetDB() require.NotNil(t, db) // Count event logs count, err := db.NewSelect().Model((*database.EventLog)(nil)). Where("topic = ? AND workspace = ?", "db.new.asset", "test-workspace"). Count(ctx) require.NoError(t, err) assert.Equal(t, 1, count) // Import same asset again (update) result, err := registry.Execute( `db_quick_import_asset("test-workspace", "existing.example.com")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, true, result) // Verify no new event log was created (still 1) count, err = db.NewSelect().Model((*database.EventLog)(nil)). Where("topic = ? AND workspace = ?", "db.new.asset", "test-workspace"). Count(ctx) require.NoError(t, err) assert.Equal(t, 1, count, "should not create event log on update") } func TestDbQuickImportAsset_MissingArgs(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_quick_import_asset("test-workspace")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbResetEventLogs(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() ctx := context.Background() db := database.GetDB() require.NotNil(t, db) // Create some test event logs now := time.Now() events := []database.EventLog{ {Topic: "db.new.asset", Workspace: "workspace1", Processed: true, ProcessedAt: &now}, {Topic: "db.updated.asset", Workspace: "workspace1", Processed: true, ProcessedAt: &now}, {Topic: "run.started", Workspace: "workspace1", Processed: true, ProcessedAt: &now}, {Topic: "db.new.asset", Workspace: "workspace2", Processed: true, ProcessedAt: &now}, {Topic: "run.completed", Workspace: "workspace2", Processed: false}, // Not processed } for _, e := range events { _, err := db.NewInsert().Model(&e).Exec(ctx) require.NoError(t, err) } registry := NewRegistry() // Test 1: Reset all event logs (no filters) result, err := registry.Execute( `db_reset_event_logs()`, map[string]interface{}{}, ) require.NoError(t, err) stats := result.(map[string]interface{}) assert.Equal(t, int64(4), stats["reset"]) // 4 were processed=true assert.Equal(t, 4, stats["total"]) // Verify all are now unprocessed count, err := db.NewSelect().Model((*database.EventLog)(nil)). Where("processed = ?", false). Count(ctx) require.NoError(t, err) assert.Equal(t, 5, count) // Reset for next test - mark some as processed again _, err = db.NewUpdate().Model((*database.EventLog)(nil)). Set("processed = ?", true). Set("processed_at = ?", now). Where("topic LIKE ?", "db.%"). Exec(ctx) require.NoError(t, err) // Test 2: Reset with workspace filter result, err = registry.Execute( `db_reset_event_logs("workspace1")`, map[string]interface{}{}, ) require.NoError(t, err) stats = result.(map[string]interface{}) assert.Equal(t, int64(2), stats["reset"]) // workspace1 has 2 db.* events that are processed // Test 3: Reset with topic pattern filter // First, mark events processed again for workspace2 _, err = db.NewUpdate().Model((*database.EventLog)(nil)). Set("processed = ?", true). Set("processed_at = ?", now). Where("workspace = ?", "workspace2"). Where("topic LIKE ?", "db.%"). Exec(ctx) require.NoError(t, err) result, err = registry.Execute( `db_reset_event_logs("", "db.*")`, map[string]interface{}{}, ) require.NoError(t, err) stats = result.(map[string]interface{}) assert.Equal(t, int64(1), stats["reset"]) // Only workspace2's db.new.asset is processed // Test 4: Reset with both filters // Mark all as processed first _, err = db.NewUpdate().Model((*database.EventLog)(nil)). Set("processed = ?", true). Set("processed_at = ?", now). Where("1 = 1"). Exec(ctx) require.NoError(t, err) result, err = registry.Execute( `db_reset_event_logs("workspace1", "run.*")`, map[string]interface{}{}, ) require.NoError(t, err) stats = result.(map[string]interface{}) assert.Equal(t, int64(1), stats["reset"]) // workspace1 has 1 run.* event assert.Equal(t, 1, stats["total"]) } func TestGlobToSQLLike(t *testing.T) { tests := []struct { pattern string expected string }{ {"db.*", "db.%"}, {"run.*", "run.%"}, {"*", "%"}, {"db.?.asset", "db._.asset"}, {"test%pattern", "test\\%pattern"}, {"test_pattern", "test\\_pattern"}, {"*.asset*", "%.asset%"}, } for _, tt := range tests { t.Run(tt.pattern, func(t *testing.T) { result := globToSQLLike(tt.pattern) assert.Equal(t, tt.expected, result) }) } } // --- DNS Asset Import Tests --- func TestDbImportDNSAsset(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() testFile := "../../test/testdata/sample-jsonl-output/dns-records.txt" _, err := os.Stat(testFile) require.NoError(t, err, "sample dns-records.txt file must exist") result, err := registry.Execute( `db_import_dns_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) // 3 unique domains: example.com, sub.example.com, other.example.com assert.Equal(t, int64(3), result) // Verify assets were imported ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var assets []database.Asset err = db.NewSelect().Model(&assets). Where("workspace = ?", "test-workspace"). OrderExpr("asset_value ASC"). Scan(ctx) require.NoError(t, err) assert.Len(t, assets, 3) // Check example.com has 2 A records + 1 CNAME var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "domain", asset.AssetType) assert.Equal(t, "dns", asset.Source) assert.Equal(t, "93.184.216.34", asset.HostIP) assert.Contains(t, asset.DnsRecords, "93.184.216.34") assert.Contains(t, asset.DnsRecords, "93.184.216.35") assert.Contains(t, asset.DnsRecords, "CNAME:cdn.example.net") // Check sub.example.com has A + MX err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "sub.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "10.0.0.1", asset.HostIP) assert.Contains(t, asset.DnsRecords, "10.0.0.1") assert.Contains(t, asset.DnsRecords, "MX:mail.example.com") } func TestDbImportDNSAsset_EmptyArgs(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_import_dns_asset("", "/tmp/test.txt")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") result, err = registry.Execute( `db_import_dns_asset("ws", "")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbImportDNSAsset_CNAME(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() // Create a file with only CNAME records tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "cname-only.txt") err := os.WriteFile(testFile, []byte("www.example.com. CNAME example.com.\n"), 0644) require.NoError(t, err) registry := NewRegistry() result, err := registry.Execute( `db_import_dns_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, int64(1), result) ctx := context.Background() db := database.GetDB() var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "www.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "", asset.HostIP) // No A record assert.Contains(t, asset.DnsRecords, "CNAME:example.com") } func TestDbImportDNSAsset_Upsert(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() testFile := "../../test/testdata/sample-jsonl-output/dns-records.txt" // Import twice result1, err := registry.Execute( `db_import_dns_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, int64(3), result1) result2, err := registry.Execute( `db_import_dns_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Equal(t, int64(3), result2) // Should still only have 3 assets (no duplicates) ctx := context.Background() db := database.GetDB() var count int count, err = db.NewSelect().Model((*database.Asset)(nil)). Where("workspace = ?", "test-workspace"). Count(ctx) require.NoError(t, err) assert.Equal(t, 3, count) } // --- Custom Asset Import Tests --- func TestDbImportCustomAsset(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() testFile := "../../test/testdata/sample-jsonl-output/custom-assets.jsonl" _, err := os.Stat(testFile) require.NoError(t, err, "sample custom-assets.jsonl file must exist") result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") assert.Equal(t, 3, stats["new"]) assert.Equal(t, 3, stats["total"]) assert.Equal(t, 0, stats["errors"]) // Verify specific asset fields ctx := context.Background() db := database.GetDB() require.NotNil(t, db) var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "api.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "https://api.example.com", asset.URL) assert.Equal(t, "url", asset.AssetType) assert.Equal(t, "custom-scan", asset.Source) assert.Contains(t, asset.Remarks, "API endpoint") assert.Equal(t, "https://jira.example.com/PROJ-123", asset.ExternalURL) assert.Contains(t, asset.Remarks, "api") assert.Contains(t, asset.Remarks, "production") // Check technologies mapped err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "cdn.example.com"). Scan(ctx) require.NoError(t, err) assert.Contains(t, asset.Technologies, "Cloudflare") assert.Contains(t, asset.Technologies, "nginx") } func TestDbImportCustomAsset_EmptyArgs(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() result, err := registry.Execute( `db_import_custom_asset("", "/tmp/test.jsonl")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") result, err = registry.Execute( `db_import_custom_asset("ws", "")`, map[string]interface{}{}, ) require.NoError(t, err) assert.Contains(t, result.(string), "error:") } func TestDbImportCustomAsset_UpdateExisting(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() testFile := "../../test/testdata/sample-jsonl-output/custom-assets.jsonl" // First import result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) stats := result.(map[string]interface{}) assert.Equal(t, 3, stats["new"]) assert.Equal(t, 0, stats["updated"]) // Second import - should update result, err = registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) stats = result.(map[string]interface{}) assert.Equal(t, 0, stats["new"]) assert.Equal(t, 3, stats["updated"]) assert.Equal(t, 3, stats["total"]) } func TestDbImportCustomAsset_TagsAndExternalURL(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() // Create a JSONL with tags and external_url tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "tagged-assets.jsonl") content := `{"asset_value":"tagged.example.com","asset_type":"domain","tags":["critical","external"],"external_url":"https://bugbounty.example.com/report/42"}` + "\n" err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) registry := NewRegistry() result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) stats := result.(map[string]interface{}) assert.Equal(t, 1, stats["new"]) // Verify tags and external_url persisted ctx := context.Background() db := database.GetDB() var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "tagged.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "https://bugbounty.example.com/report/42", asset.ExternalURL) assert.Contains(t, asset.Remarks, "critical") assert.Contains(t, asset.Remarks, "external") } // --- Model Field Tests --- func TestAssetExternalURLField(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() ctx := context.Background() db := database.GetDB() require.NotNil(t, db) now := time.Now() asset := database.Asset{ Workspace: "test-workspace", AssetValue: "ext-url-test.example.com", AssetType: "domain", ExternalURL: "https://tracker.example.com/issue/99", CreatedAt: now, UpdatedAt: now, } _, err := db.NewInsert().Model(&asset).Exec(ctx) require.NoError(t, err) var fetched database.Asset err = db.NewSelect().Model(&fetched). Where("asset_value = ?", "ext-url-test.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "https://tracker.example.com/issue/99", fetched.ExternalURL) } func TestAssetRemarksField(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() ctx := context.Background() db := database.GetDB() require.NotNil(t, db) now := time.Now() asset := database.Asset{ Workspace: "test-workspace", AssetValue: "remarks-test.example.com", AssetType: "domain", Remarks: []string{"web", "api", "production"}, CreatedAt: now, UpdatedAt: now, } _, err := db.NewInsert().Model(&asset).Exec(ctx) require.NoError(t, err) var fetched database.Asset err = db.NewSelect().Model(&fetched). Where("asset_value = ?", "remarks-test.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, []string{"web", "api", "production"}, fetched.Remarks) } func TestAssetSizeAndLOCFields(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() ctx := context.Background() db := database.GetDB() require.NotNil(t, db) now := time.Now() asset := database.Asset{ Workspace: "test-workspace", AssetValue: "example-corp/web-app", AssetType: "repository", Language: "Go", Size: 512000, LOC: 12000, ExternalURL: "https://github.com/example-corp/web-app", CreatedAt: now, UpdatedAt: now, } _, err := db.NewInsert().Model(&asset).Exec(ctx) require.NoError(t, err) var fetched database.Asset err = db.NewSelect().Model(&fetched). Where("asset_value = ?", "example-corp/web-app"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "repository", fetched.AssetType) assert.Equal(t, "Go", fetched.Language) assert.Equal(t, int64(512000), fetched.Size) assert.Equal(t, int64(12000), fetched.LOC) assert.Equal(t, "https://github.com/example-corp/web-app", fetched.ExternalURL) } // --- DNS Records Key Tests --- func TestMapJSONToAsset_DnsRecordsKey(t *testing.T) { // "dns_records" key should be accepted as the primary key data := map[string]interface{}{ "host": "example.com", "url": "http://example.com", "dns_records": []interface{}{"10.0.0.1", "10.0.0.2"}, } asset := mapJSONToAsset(data, "test-workspace", `{"host":"example.com"}`) assert.Equal(t, []string{"10.0.0.1", "10.0.0.2"}, asset.DnsRecords) // "a" key should still work as a fallback dataLegacy := map[string]interface{}{ "host": "legacy.com", "url": "http://legacy.com", "a": []interface{}{"1.2.3.4"}, } assetLegacy := mapJSONToAsset(dataLegacy, "test-workspace", `{"host":"legacy.com"}`) assert.Equal(t, []string{"1.2.3.4"}, assetLegacy.DnsRecords) // When both keys exist, "dns_records" takes precedence dataBoth := map[string]interface{}{ "host": "both.com", "url": "http://both.com", "dns_records": []interface{}{"10.0.0.1"}, "a": []interface{}{"1.2.3.4"}, } assetBoth := mapJSONToAsset(dataBoth, "test-workspace", `{"host":"both.com"}`) assert.Equal(t, []string{"10.0.0.1"}, assetBoth.DnsRecords) } // --- Optional Defaults Tests --- func TestDbImportCustomAsset_OptionalDefaults(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() // Create JSONL: first line has no asset_type/source, second has them tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "defaults.jsonl") content := `{"asset_value":"no-type.example.com","url":"http://no-type.example.com"}` + "\n" + `{"asset_value":"has-type.example.com","url":"http://has-type.example.com","asset_type":"http","source":"httpx"}` + "\n" err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) registry := NewRegistry() result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`", "subdomain", "recon")`, map[string]interface{}{}, ) require.NoError(t, err) stats := result.(map[string]interface{}) assert.Equal(t, 2, stats["new"]) ctx := context.Background() db := database.GetDB() // First asset should have defaults applied var asset1 database.Asset err = db.NewSelect().Model(&asset1). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "no-type.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "subdomain", asset1.AssetType) assert.Equal(t, "recon", asset1.Source) // Second asset should keep its own values var asset2 database.Asset err = db.NewSelect().Model(&asset2). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "has-type.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "http", asset2.AssetType) assert.Equal(t, "httpx", asset2.Source) } func TestDbImportCustomAsset_PartialDefaults(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "partial.jsonl") content := `{"asset_value":"partial.example.com","url":"http://partial.example.com"}` + "\n" err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) registry := NewRegistry() // Only 3rd arg (asset_type) provided, no 4th arg (source) result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`", "dns")`, map[string]interface{}{}, ) require.NoError(t, err) stats := result.(map[string]interface{}) assert.Equal(t, 1, stats["new"]) ctx := context.Background() db := database.GetDB() var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "partial.example.com"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "dns", asset.AssetType) assert.Equal(t, "", asset.Source) } func TestDbImportCustomAsset_ContentDiscovery(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() testFile := "../../test/testdata/sample-jsonl-output/cusom-content-discovery.jsonl" _, err := os.Stat(testFile) require.NoError(t, err, "sample cusom-content-discovery.jsonl file must exist") result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`", "url", "content-discovery")`, map[string]interface{}{}, ) require.NoError(t, err) stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") assert.Equal(t, 9, stats["new"]) assert.Equal(t, 9, stats["total"]) assert.Equal(t, 0, stats["errors"]) ctx := context.Background() db := database.GetDB() require.NotNil(t, db) // Line 1: url=http://cdn.doitac.example.com/common/, no asset_value, webserver=UploadServer var asset1 database.Asset err = db.NewSelect().Model(&asset1). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "http://cdn.doitac.example.com/common/"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "http://cdn.doitac.example.com/common/", asset1.URL) assert.Contains(t, asset1.Remarks, "UploadServer") // Line 6: url=http://e-procurement.example.com/static, remarks=["Modern-App"], webserver=SGW var asset6 database.Asset err = db.NewSelect().Model(&asset6). Where("workspace = ?", "test-workspace"). Where("asset_value = ?", "http://e-procurement.example.com/static"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "http://e-procurement.example.com/static", asset6.URL) assert.Contains(t, asset6.Remarks, "Modern-App") assert.Contains(t, asset6.Remarks, "SGW") } // --- Vigolium Envelope Import Tests --- func TestDbImportCustomAsset_VigoliumEnvelope(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() // Create a small envelope JSONL with mixed record types content := `{"type":"scan","data":{"uuid":"abc","status":"completed","target":"http://example.com"}} {"type":"module","data":{"id":"xss-light","name":"XSS Light","enabled":true}} {"type":"http_record","data":{"url":"http://example.com/","method":"GET","status_code":200,"scheme":"http","hostname":"example.com","port":80,"ip":"93.184.216.34","response_content_type":"text/html","response_content_length":1256,"response_words":45,"response_title":"Example Domain","source":"scanner"}} {"type":"http_record","data":{"url":"http://example.com/login","method":"POST","status_code":302,"scheme":"http","hostname":"example.com","port":80,"ip":"93.184.216.34","response_content_type":"text/html","response_content_length":0,"response_words":0,"source":"spidering"}} {"type":"module","data":{"id":"sqli-basic","name":"SQLi Basic","enabled":true}}` testFile := filepath.Join(t.TempDir(), "vigolium-envelope.jsonl") err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) registry := NewRegistry() result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") // Only 2 http_record lines should be imported; scan and module lines skipped assert.Equal(t, 2, stats["new"]) assert.Equal(t, 2, stats["total"]) assert.Equal(t, 0, stats["errors"]) ctx := context.Background() db := database.GetDB() require.NotNil(t, db) // Verify field remapping for first http_record var asset1 database.Asset err = db.NewSelect().Model(&asset1). Where("workspace = ?", "test-workspace"). Where("url = ?", "http://example.com/"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "http://example.com/", asset1.AssetValue) // url -> asset_value fallback assert.Equal(t, "GET", asset1.Method) assert.Equal(t, 200, asset1.StatusCode) assert.Equal(t, "http", asset1.Scheme) assert.Equal(t, "93.184.216.34", asset1.HostIP) // ip -> host_ip assert.Equal(t, "text/html", asset1.ContentType) // response_content_type -> content_type assert.Equal(t, int64(1256), asset1.ContentLength) // response_content_length -> content_length assert.Equal(t, 45, asset1.Words) // response_words -> words assert.Equal(t, "Example Domain", asset1.Title) // response_title -> title assert.Equal(t, "example.com", asset1.Input) // hostname -> input assert.Equal(t, "scanner", asset1.Source) // source preserved assert.Equal(t, "web", asset1.AssetType) // default for http_record // Verify second http_record var asset2 database.Asset err = db.NewSelect().Model(&asset2). Where("workspace = ?", "test-workspace"). Where("url = ?", "http://example.com/login"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "POST", asset2.Method) assert.Equal(t, 302, asset2.StatusCode) assert.Equal(t, "spidering", asset2.Source) assert.Equal(t, "web", asset2.AssetType) // raw_json_data should be populated assert.NotEmpty(t, asset1.RawJsonData) } func TestDbImportCustomAsset_VigoliumDiscover(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() testFile := "../../test/testdata/sample-jsonl-output/vigolium-discover.jsonl" _, err := os.Stat(testFile) require.NoError(t, err, "sample vigolium-discover.jsonl file must exist") result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") // 28 http_record lines; scan and module lines should be skipped assert.Equal(t, 28, stats["new"]) assert.Equal(t, 28, stats["total"]) assert.Equal(t, 0, stats["errors"]) ctx := context.Background() db := database.GetDB() require.NotNil(t, db) // Verify a known record: the root URL request var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("url = ?", "http://localhost:3000/"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "GET", asset.Method) assert.Equal(t, 200, asset.StatusCode) assert.Equal(t, "http", asset.Scheme) assert.Equal(t, "web", asset.AssetType) assert.Equal(t, "text/html; charset=UTF-8", asset.ContentType) assert.Equal(t, "::1", asset.HostIP) assert.Equal(t, "localhost", asset.Input) } func TestDbImportCustomAsset_VigoliumSpider(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() registry := NewRegistry() testFile := "../../test/testdata/sample-jsonl-output/vigolium-spider.jsonl" _, err := os.Stat(testFile) require.NoError(t, err, "sample vigolium-spider.jsonl file must exist") result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`")`, map[string]interface{}{}, ) require.NoError(t, err) stats, ok := result.(map[string]interface{}) require.True(t, ok, "result should be a map") // 30 http_record lines; scan and module lines should be skipped assert.Equal(t, 30, stats["new"]) assert.Equal(t, 30, stats["total"]) assert.Equal(t, 0, stats["errors"]) ctx := context.Background() db := database.GetDB() require.NotNil(t, db) // Verify a known record from spider output var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("url = ?", "https://ginandjuice.shop/catalog/product/stock"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "POST", asset.Method) assert.Equal(t, 400, asset.StatusCode) assert.Equal(t, "https", asset.Scheme) assert.Equal(t, "web", asset.AssetType) assert.Equal(t, "34.246.169.176", asset.HostIP) assert.Equal(t, "ginandjuice.shop", asset.Input) assert.Equal(t, "spidering", asset.Source) } func TestDbImportCustomAsset_VigoliumWithSourceOverride(t *testing.T) { cleanup := setupTestDB(t) defer cleanup() // Envelope with no source field — should use the default from function arg content := `{"type":"http_record","data":{"url":"http://test.com/api","method":"GET","status_code":200,"scheme":"http","hostname":"test.com","ip":"1.2.3.4"}}` testFile := filepath.Join(t.TempDir(), "vigolium-no-source.jsonl") err := os.WriteFile(testFile, []byte(content), 0644) require.NoError(t, err) registry := NewRegistry() result, err := registry.Execute( `db_import_custom_asset("test-workspace", "`+testFile+`", "web", "vigolium")`, map[string]interface{}{}, ) require.NoError(t, err) stats := result.(map[string]interface{}) assert.Equal(t, 1, stats["new"]) ctx := context.Background() db := database.GetDB() var asset database.Asset err = db.NewSelect().Model(&asset). Where("workspace = ?", "test-workspace"). Where("url = ?", "http://test.com/api"). Scan(ctx) require.NoError(t, err) assert.Equal(t, "vigolium", asset.Source) // default source from arg assert.Equal(t, "web", asset.AssetType) // from envelope default, matches arg too assert.Equal(t, "1.2.3.4", asset.HostIP) // ip -> host_ip remapped assert.Equal(t, "test.com", asset.Input) // hostname -> input remapped }