From 225f2d7d15a7c65975a63f617d6182bbce742c87 Mon Sep 17 00:00:00 2001 From: j3ssie Date: Thu, 12 Feb 2026 19:44:39 +0700 Subject: [PATCH] feat: add SSH and DNS asset import utilities - Add ssh_exec() and ssh_rsync() functions for remote command execution and file transfer with connection pooling - Add db_import_dns_asset() to import DNS zone records and group by domain with A/AAAA records stored separately - Add db_import_custom_asset() for flexible JSONL asset import with direct field mapping, supporting tags and custom asset types - Extend Asset model with ExternalURL, Remarks (string array), Language, Size, and LOC fields for better metadata tracking - Add backward compatibility layer for legacy JSON formats (string remarks, tags array merging) - Register new functions in function registry with proper documentation and usage examples --- docs/api/assets.mdx | 10 +- internal/database/database.go | 120 ++++++ internal/database/jsonl.go | 28 +- internal/database/models.go | 10 +- internal/database/repository/asset_repo.go | 8 + internal/database/seed.go | 82 +++- internal/executor/function_executor.go | 3 + internal/functions/cdn_functions.go | 6 +- internal/functions/constants.go | 26 ++ internal/functions/db_functions.go | 321 ++++++++++++++- internal/functions/db_functions_test.go | 379 ++++++++++++++++++ internal/functions/goja_pool.go | 9 + internal/functions/goja_runtime.go | 8 + internal/functions/ssh_functions.go | 221 ++++++++++ internal/functions/ssh_functions_test.go | 142 +++++++ .../sample-jsonl-output/custom-assets.jsonl | 3 + .../sample-jsonl-output/dns-records.txt | 6 + 17 files changed, 1349 insertions(+), 33 deletions(-) create mode 100644 internal/functions/ssh_functions.go create mode 100644 internal/functions/ssh_functions_test.go create mode 100644 test/testdata/sample-jsonl-output/custom-assets.jsonl create mode 100644 test/testdata/sample-jsonl-output/dns-records.txt diff --git a/docs/api/assets.mdx b/docs/api/assets.mdx index eb41476..371d775 100644 --- a/docs/api/assets.mdx +++ b/docs/api/assets.mdx @@ -58,7 +58,7 @@ curl "http://localhost:8002/osm/api/assets?workspace=example.com&offset=50&limit "asset_type": "web", "tech": ["nginx/1.21.0", "nodejs", "express"], "time": "245ms", - "remarks": "production", + "remarks": ["production"], "source": "httpx", "created_at": "2025-01-15T10:30:00Z", "updated_at": "2025-01-15T10:30:00Z" @@ -84,7 +84,7 @@ curl "http://localhost:8002/osm/api/assets?workspace=example.com&offset=50&limit "asset_type": "web", "tech": ["nginx/1.20.0", "php/8.1", "wordpress"], "time": "312ms", - "remarks": "admin-panel", + "remarks": ["admin-panel"], "source": "httpx", "created_at": "2025-01-15T10:31:00Z", "updated_at": "2025-01-15T10:31:00Z" @@ -122,7 +122,11 @@ curl "http://localhost:8002/osm/api/assets?workspace=example.com&offset=50&limit | `asset_type` | string | Asset type classification | | `tech` | array | Detected technologies | | `time` | string | Response time | -| `remarks` | string | Custom labels/remarks | +| `remarks` | array | Custom labels/tags/remarks | +| `language` | string | Repository/source language | +| `size` | int | Repository/file size in bytes | +| `loc` | int | Lines of code (for repository assets) | +| `blob_content` | string | Large content (raw files, keys, etc.) | | `source` | string | Discovery source (httpx, nuclei, etc.) | | `created_at` | timestamp | Creation timestamp | | `updated_at` | timestamp | Last update timestamp | diff --git a/internal/database/database.go b/internal/database/database.go index 8970e7a..addcead 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -222,6 +222,36 @@ func Migrate(ctx context.Context) error { return err } + // Add remarks column to assets table if it doesn't exist (for existing databases) + if err := addAssetRemarksColumn(ctx); err != nil { + return err + } + + // Add language column to assets table if it doesn't exist (for existing databases) + if err := addAssetLanguageColumn(ctx); err != nil { + return err + } + + // Add blob_content column to assets table if it doesn't exist (for existing databases) + if err := addAssetBlobContentColumn(ctx); err != nil { + return err + } + + // Add external_url column to assets table if it doesn't exist (for existing databases) + if err := addAssetExternalURLColumn(ctx); err != nil { + return err + } + + // Add size column to assets table if it doesn't exist (for existing databases) + if err := addAssetSizeColumn(ctx); err != nil { + return err + } + + // Add loc column to assets table if it doesn't exist (for existing databases) + if err := addAssetLOCColumn(ctx); err != nil { + return err + } + return nil } @@ -279,6 +309,96 @@ func addArtifactsOptionalColumn(ctx context.Context) error { return nil } +// addAssetRemarksColumn adds the remarks column to assets table for existing databases +func addAssetRemarksColumn(ctx context.Context) error { + _, err := db.ExecContext(ctx, "ALTER TABLE assets ADD COLUMN remarks JSON") + if err != nil { + errStr := strings.ToLower(err.Error()) + if strings.Contains(errStr, "duplicate column") || + strings.Contains(errStr, "already exists") || + strings.Contains(errStr, "sqlstate 42701") { + return nil + } + return fmt.Errorf("failed to add remarks column: %w", err) + } + return nil +} + +// addAssetLanguageColumn adds the language column to assets table for existing databases +func addAssetLanguageColumn(ctx context.Context) error { + _, err := db.ExecContext(ctx, "ALTER TABLE assets ADD COLUMN language TEXT DEFAULT ''") + if err != nil { + errStr := strings.ToLower(err.Error()) + if strings.Contains(errStr, "duplicate column") || + strings.Contains(errStr, "already exists") || + strings.Contains(errStr, "sqlstate 42701") { + return nil + } + return fmt.Errorf("failed to add language column: %w", err) + } + return nil +} + +// addAssetBlobContentColumn adds the blob_content column to assets table for existing databases +func addAssetBlobContentColumn(ctx context.Context) error { + _, err := db.ExecContext(ctx, "ALTER TABLE assets ADD COLUMN blob_content TEXT DEFAULT ''") + if err != nil { + errStr := strings.ToLower(err.Error()) + if strings.Contains(errStr, "duplicate column") || + strings.Contains(errStr, "already exists") || + strings.Contains(errStr, "sqlstate 42701") { + return nil + } + return fmt.Errorf("failed to add blob_content column: %w", err) + } + return nil +} + +// addAssetSizeColumn adds the size column to assets table for existing databases +func addAssetSizeColumn(ctx context.Context) error { + _, err := db.ExecContext(ctx, "ALTER TABLE assets ADD COLUMN size INTEGER DEFAULT 0") + if err != nil { + errStr := strings.ToLower(err.Error()) + if strings.Contains(errStr, "duplicate column") || + strings.Contains(errStr, "already exists") || + strings.Contains(errStr, "sqlstate 42701") { + return nil + } + return fmt.Errorf("failed to add size column: %w", err) + } + return nil +} + +// addAssetLOCColumn adds the loc column to assets table for existing databases +func addAssetLOCColumn(ctx context.Context) error { + _, err := db.ExecContext(ctx, "ALTER TABLE assets ADD COLUMN loc INTEGER DEFAULT 0") + if err != nil { + errStr := strings.ToLower(err.Error()) + if strings.Contains(errStr, "duplicate column") || + strings.Contains(errStr, "already exists") || + strings.Contains(errStr, "sqlstate 42701") { + return nil + } + return fmt.Errorf("failed to add loc column: %w", err) + } + return nil +} + +// addAssetExternalURLColumn adds the external_url column to assets table for existing databases +func addAssetExternalURLColumn(ctx context.Context) error { + _, err := db.ExecContext(ctx, "ALTER TABLE assets ADD COLUMN external_url TEXT DEFAULT ''") + if err != nil { + errStr := strings.ToLower(err.Error()) + if strings.Contains(errStr, "duplicate column") || + strings.Contains(errStr, "already exists") || + strings.Contains(errStr, "sqlstate 42701") { + return nil + } + return fmt.Errorf("failed to add external_url column: %w", err) + } + return nil +} + // createAssetIndexes creates indexes for the assets table func createAssetIndexes(ctx context.Context) error { indexes := []string{ diff --git a/internal/database/jsonl.go b/internal/database/jsonl.go index a991a10..6d02607 100644 --- a/internal/database/jsonl.go +++ b/internal/database/jsonl.go @@ -145,6 +145,10 @@ func (i *JSONLImporter) insertAssetBatch(ctx context.Context, assets []*Asset) ( Set("words = EXCLUDED.words"). Set("lines = EXCLUDED.lines"). Set("remarks = EXCLUDED.remarks"). + Set("language = EXCLUDED.language"). + Set("size = EXCLUDED.size"). + Set("loc = EXCLUDED.loc"). + Set("blob_content = EXCLUDED.blob_content"). Set("raw_data = EXCLUDED.raw_data"). Set("asset_type = EXCLUDED.asset_type"). Set("updated_at = EXCLUDED.updated_at"). @@ -238,8 +242,24 @@ func ParseAssetLine(line []byte, defaultWorkspace, source string) (*Asset, error if v, ok := raw["time"].(string); ok { asset.ResponseTime = v } - if v, ok := raw["remarks"].(string); ok { - asset.Labels = v + if v, ok := raw["remarks"].(string); ok && v != "" { + asset.Remarks = []string{v} + } else if v, ok := raw["remarks"].([]interface{}); ok { + asset.Remarks = interfaceSliceToStringSlice(v) + } + // Merge legacy "tags" field into Remarks + if v, ok := raw["tags"].([]interface{}); ok { + asset.Remarks = append(asset.Remarks, interfaceSliceToStringSlice(v)...) + } + // Language field + if v, ok := raw["language"].(string); ok { + asset.Language = v + } + if v, ok := raw["size"].(float64); ok { + asset.Size = int64(v) + } + if v, ok := raw["loc"].(float64); ok { + asset.LOC = int64(v) } if v, ok := raw["asset_type"].(string); ok { asset.AssetType = v @@ -257,6 +277,10 @@ func ParseAssetLine(line []byte, defaultWorkspace, source string) (*Asset, error if asset.AssetType == "" { asset.AssetType = ClassifyAssetType(asset.AssetValue) } + // Refine: if we have HTTP response data, it's an "http" asset + if asset.AssetType == "url" && (asset.StatusCode > 0 || asset.ContentLength > 0) { + asset.AssetType = "http" + } return asset, nil } diff --git a/internal/database/models.go b/internal/database/models.go index d3e7fbd..ae7dff6 100644 --- a/internal/database/models.go +++ b/internal/database/models.go @@ -224,15 +224,23 @@ type Asset struct { TLS string `bun:"tls" json:"tls,omitempty"` // Metadata + // asset type can be "http", "dns", "url", "subdomain", "ip", "repository", "file", "link", "content", "archive" AssetType string `bun:"asset_type" json:"asset_type,omitempty"` Technologies []string `bun:"technologies,type:json" json:"tech,omitempty"` ResponseTime string `bun:"response_time" json:"time,omitempty"` - Labels string `bun:"labels" json:"remarks,omitempty"` + Remarks []string `bun:"remarks,type:json" json:"remarks,omitempty"` Source string `bun:"source" json:"source,omitempty"` // e.g., "httpx", "nuclei" RawJsonData string `bun:"raw_json_data" json:"raw_json_data,omitempty"` // Original JSON RawResponse string `bun:"raw_response" json:"raw_response,omitempty"` ScreenshotBase64Data string `bun:"screenshot_base64_data" json:"screenshot_base64_data,omitempty"` + // for repository/file assets + Language string `bun:"language" json:"language,omitempty"` + Size int64 `bun:"size" json:"size,omitempty"` + LOC int64 `bun:"loc" json:"loc,omitempty"` + ExternalURL string `bun:"external_url" json:"external_url,omitempty"` + BlobContent string `bun:"blob_content" json:"blob_content,omitempty"` + // Timestamps CreatedAt time.Time `bun:"created_at,notnull,default:current_timestamp" json:"created_at"` UpdatedAt time.Time `bun:"updated_at,notnull,default:current_timestamp" json:"updated_at"` diff --git a/internal/database/repository/asset_repo.go b/internal/database/repository/asset_repo.go index 871da0d..cd075bc 100644 --- a/internal/database/repository/asset_repo.go +++ b/internal/database/repository/asset_repo.go @@ -294,6 +294,10 @@ func (r *AssetRepository) Upsert(ctx context.Context, asset *database.Asset) err Set("words = EXCLUDED.words"). Set("lines = EXCLUDED.lines"). Set("remarks = EXCLUDED.remarks"). + Set("language = EXCLUDED.language"). + Set("size = EXCLUDED.size"). + Set("loc = EXCLUDED.loc"). + Set("blob_content = EXCLUDED.blob_content"). Set("raw_data = EXCLUDED.raw_data"). Set("updated_at = EXCLUDED.updated_at"). Exec(ctx) @@ -321,6 +325,10 @@ func (r *AssetRepository) BulkUpsert(ctx context.Context, assets []*database.Ass Set("words = EXCLUDED.words"). Set("lines = EXCLUDED.lines"). Set("remarks = EXCLUDED.remarks"). + Set("language = EXCLUDED.language"). + Set("size = EXCLUDED.size"). + Set("loc = EXCLUDED.loc"). + Set("blob_content = EXCLUDED.blob_content"). Set("raw_data = EXCLUDED.raw_data"). Set("updated_at = EXCLUDED.updated_at"). Exec(ctx) diff --git a/internal/database/seed.go b/internal/database/seed.go index 0110ef5..4109a27 100644 --- a/internal/database/seed.go +++ b/internal/database/seed.go @@ -697,6 +697,7 @@ func SeedDatabase(ctx context.Context) error { HostIP: "93.184.216.34", DnsRecords: []string{"93.184.216.34"}, TLS: "TLS 1.3", + AssetType: "http", Technologies: []string{"Nginx", "CloudFlare"}, ResponseTime: "145ms", Source: "httpx", @@ -739,6 +740,7 @@ func SeedDatabase(ctx context.Context) error { HostIP: "93.184.216.35", DnsRecords: []string{"93.184.216.35"}, TLS: "TLS 1.3", + AssetType: "http", Technologies: []string{"Express", "Node.js"}, ResponseTime: "67ms", Source: "httpx", @@ -826,7 +828,7 @@ func SeedDatabase(ctx context.Context) error { DnsRecords: []string{"10.0.0.50"}, Technologies: []string{"Python", "Flask", "Gunicorn"}, ResponseTime: "23ms", - Labels: "Internal development server - no TLS", + Remarks: []string{"Internal development server - no TLS"}, Source: "httpx", LastSeenAt: threeDaysAgo, CreatedAt: oneHourAgo, @@ -982,7 +984,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.2", Technologies: []string{"Jenkins", "Java"}, ResponseTime: "67ms", - Labels: "Internal CI/CD server", + Remarks: []string{"Internal CI/CD server"}, Source: "httpx", LastSeenAt: threeDaysAgo, CreatedAt: oneHourAgo, @@ -1004,7 +1006,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"GitLab", "Ruby", "PostgreSQL"}, ResponseTime: "45ms", - Labels: "Internal Git server - redirects to login", + Remarks: []string{"Internal Git server - redirects to login"}, Source: "httpx", LastSeenAt: threeDaysAgo, CreatedAt: oneHourAgo, @@ -1050,7 +1052,7 @@ func SeedDatabase(ctx context.Context) error { DnsRecords: []string{"93.184.216.60"}, Technologies: []string{"ASP.NET", "IIS", "jQuery"}, ResponseTime: "567ms", - Labels: "Legacy system - no TLS", + Remarks: []string{"Legacy system - no TLS"}, Source: "httpx", LastSeenAt: oneWeekAgo, CreatedAt: oneHourAgo, @@ -1072,7 +1074,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"Nginx"}, ResponseTime: "1234ms", - Labels: "Beta environment - currently broken", + Remarks: []string{"Beta environment - currently broken"}, Source: "httpx", LastSeenAt: threeDaysAgo, CreatedAt: oneHourAgo, @@ -1200,7 +1202,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"Nginx"}, ResponseTime: "5000ms", - Labels: "Elasticsearch backend down", + Remarks: []string{"Elasticsearch backend down"}, Source: "httpx", LastSeenAt: oneWeekAgo, CreatedAt: oneHourAgo, @@ -1246,7 +1248,7 @@ func SeedDatabase(ctx context.Context) error { DnsRecords: []string{"192.168.1.11"}, Technologies: []string{"phpMyAdmin", "PHP", "Apache"}, ResponseTime: "45ms", - Labels: "Database admin panel", + Remarks: []string{"Database admin panel"}, Source: "httpx", LastSeenAt: oneDayAgo, CreatedAt: oneHourAgo, @@ -1289,7 +1291,7 @@ func SeedDatabase(ctx context.Context) error { DnsRecords: []string{"192.168.1.13"}, Technologies: []string{"MinIO", "Go"}, ResponseTime: "34ms", - Labels: "Object storage - access denied", + Remarks: []string{"Object storage - access denied"}, Source: "httpx", LastSeenAt: threeDaysAgo, CreatedAt: oneHourAgo, @@ -1342,7 +1344,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"FastAPI", "Python"}, ResponseTime: "45ms", - Labels: "Newly discovered API endpoint", + Remarks: []string{"Newly discovered API endpoint"}, Source: "httpx", LastSeenAt: thirtyMinsAgo, CreatedAt: thirtyMinsAgo, @@ -1366,7 +1368,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"Next.js", "React", "Vercel"}, ResponseTime: "89ms", - Labels: "New version just launched", + Remarks: []string{"New version just launched"}, Source: "httpx", LastSeenAt: oneHourAgo, CreatedAt: oneHourAgo, @@ -1393,7 +1395,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"Angular", "Node.js", "MongoDB"}, ResponseTime: "234ms", - Labels: "Portal updated with new features", + Remarks: []string{"Portal updated with new features"}, Source: "httpx", LastSeenAt: oneHourAgo, CreatedAt: oneWeekAgo, // Created a week ago @@ -1417,7 +1419,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"Vue.js", "D3.js", "PostgreSQL"}, ResponseTime: "156ms", - Labels: "Dashboard version bumped", + Remarks: []string{"Dashboard version bumped"}, Source: "httpx", LastSeenAt: twoHoursAgo, CreatedAt: twoWeeksAgo, // Created two weeks ago @@ -1441,7 +1443,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.3", Technologies: []string{"Keycloak", "Java", "Redis"}, ResponseTime: "78ms", - Labels: "Auth service patched", + Remarks: []string{"Auth service patched"}, Source: "httpx", LastSeenAt: oneDayAgo, CreatedAt: oneMonthAgo, // Created a month ago @@ -1466,7 +1468,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.2", Technologies: []string{"Nginx"}, ResponseTime: "5000ms", - Labels: "API deprecated and decommissioned", + Remarks: []string{"API deprecated and decommissioned"}, Source: "httpx", LastSeenAt: twoWeeksAgo, // Last seen 2 weeks ago - now gone CreatedAt: oneMonthAgo, @@ -1488,7 +1490,7 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.2", Technologies: []string{"Apache"}, ResponseTime: "123ms", - Labels: "Old portal removed after migration", + Remarks: []string{"Old portal removed after migration"}, Source: "httpx", LastSeenAt: oneWeekAgo, // Last seen a week ago CreatedAt: oneMonthAgo, @@ -1509,7 +1511,7 @@ func SeedDatabase(ctx context.Context) error { DnsRecords: []string{"10.0.0.200"}, Technologies: []string{}, ResponseTime: "timeout", - Labels: "Test environment taken offline", + Remarks: []string{"Test environment taken offline"}, Source: "httpx", LastSeenAt: fiveDaysAgo, // Last seen 5 days ago - offline CreatedAt: twoWeeksAgo, @@ -1535,12 +1537,52 @@ func SeedDatabase(ctx context.Context) error { TLS: "TLS 1.2", Technologies: []string{"Static HTML", "CloudFlare"}, ResponseTime: "45ms", - Labels: "Stable archive - rarely changes", + Remarks: []string{"Stable archive - rarely changes"}, Source: "httpx", LastSeenAt: threeDaysAgo, // Still accessible but rarely scanned CreatedAt: oneMonthAgo, UpdatedAt: oneMonthAgo, // Never updated since creation }, + // ============================================================ + // NON-HTTP ASSET TYPES + // DNS, repository, and external link assets + // ============================================================ + { + Workspace: "example.com", + AssetValue: "ns1.example.com", + AssetType: "dns", + HostIP: "93.184.216.90", + DnsRecords: []string{"93.184.216.90", "93.184.216.91"}, + Source: "dns-enum", + LastSeenAt: oneDayAgo, + CreatedAt: oneHourAgo, + UpdatedAt: oneHourAgo, + }, + { + Workspace: "example.com", + AssetValue: "example-corp/web-app", + AssetType: "repository", + Language: "JavaScript", + Size: 245760, + LOC: 8500, + ExternalURL: "https://github.com/example-corp/web-app", + Remarks: []string{"Main web application repository"}, + Source: "github", + LastSeenAt: oneDayAgo, + CreatedAt: oneHourAgo, + UpdatedAt: oneHourAgo, + }, + { + Workspace: "example.com", + AssetValue: "https://docs.example.com/api/v2", + URL: "https://docs.example.com/api/v2", + AssetType: "link", + Remarks: []string{"API documentation link found during crawl"}, + Source: "crawler", + LastSeenAt: oneDayAgo, + CreatedAt: oneHourAgo, + UpdatedAt: oneHourAgo, + }, } for i := range assets { @@ -2843,7 +2885,7 @@ var tableSearchColumns = map[string][]string{ "runs": {"id", "run_uuid", "run_group_id", "workflow_name", "target", "workspace", "status", "error_message"}, "step_results": {"id", "run_id", "step_name", "step_type", "status", "command", "output", "error_message"}, "artifacts": {"id", "run_id", "workspace", "name", "artifact_path", "artifact_type", "description"}, - "assets": {"workspace", "asset_value", "url", "title", "host_ip", "source", "labels"}, + "assets": {"workspace", "asset_value", "url", "title", "host_ip", "source", "remarks"}, "event_logs": {"event_id", "topic", "name", "source", "workspace", "run_id", "workflow_name", "data"}, "schedules": {"id", "name", "workflow_name", "workflow_kind", "target", "trigger_name", "schedule"}, "workspaces": {"name", "local_path", "data_source", "run_workflow"}, @@ -2880,8 +2922,8 @@ var tableAllColumns = map[string][]string{ "assets": {"id", "workspace", "asset_value", "url", "input", "scheme", "method", "path", "status_code", "content_type", "content_length", "title", "words", "lines", "host_ip", "dns_records", "tls", "asset_type", "technologies", - "response_time", "labels", "source", "raw_json_data", "raw_response", - "screenshot_base64_data", "last_seen_at", "created_at", "updated_at"}, + "response_time", "remarks", "language", "size", "loc", "blob_content", "source", "raw_json_data", "raw_response", + "screenshot_base64_data", "external_url", "last_seen_at", "created_at", "updated_at"}, "event_logs": {"id", "topic", "event_id", "name", "source_type", "source", "data_type", "data", "workspace", "run_id", "workflow_name", "processed", "processed_at", "error", "created_at"}, diff --git a/internal/executor/function_executor.go b/internal/executor/function_executor.go index 8b55981..362f375 100644 --- a/internal/executor/function_executor.go +++ b/internal/executor/function_executor.go @@ -45,6 +45,9 @@ func (e *FunctionExecutor) Execute(ctx context.Context, step *core.Step, execCtx } vars := execCtx.GetVariables() + if step.SuppressDetails { + vars["SuppressDetails"] = true + } var outputs []interface{} var err error diff --git a/internal/functions/cdn_functions.go b/internal/functions/cdn_functions.go index 05ba9d9..17e90e1 100644 --- a/internal/functions/cdn_functions.go +++ b/internal/functions/cdn_functions.go @@ -115,7 +115,7 @@ func (vf *vmFunc) cdnDelete(call goja.FunctionCall) goja.Value { if len(call.Arguments) > 1 && !goja.IsUndefined(call.Argument(1)) { mode = strings.ToLower(call.Argument(1).String()) } - jsonOnly := mode == "json" + jsonOnly := mode == "json" || vf.getContext().suppressDetails logger.Get().Debug("Calling cdnDelete", zap.String("remotePath", remotePath)) if remotePath == "undefined" || remotePath == "" { @@ -288,7 +288,7 @@ func (vf *vmFunc) cdnSyncUpload(call goja.FunctionCall) goja.Value { if len(call.Arguments) > 2 && !goja.IsUndefined(call.Argument(2)) { mode = strings.ToLower(call.Argument(2).String()) } - jsonOnly := mode == "json" + jsonOnly := mode == "json" || vf.getContext().suppressDetails logger.Get().Debug("Calling cdnSyncUpload", zap.String("localDir", localDir), zap.String("remotePrefix", remotePrefix)) result := map[string]interface{}{ @@ -404,7 +404,7 @@ func (vf *vmFunc) cdnSyncDownload(call goja.FunctionCall) goja.Value { if len(call.Arguments) > 2 && !goja.IsUndefined(call.Argument(2)) { mode = strings.ToLower(call.Argument(2).String()) } - jsonOnly := mode == "json" + jsonOnly := mode == "json" || vf.getContext().suppressDetails logger.Get().Debug("Calling cdnSyncDownload", zap.String("remotePrefix", remotePrefix), zap.String("localDir", localDir)) result := map[string]interface{}{ diff --git a/internal/functions/constants.go b/internal/functions/constants.go index 00c44ca..a633cc8 100644 --- a/internal/functions/constants.go +++ b/internal/functions/constants.go @@ -201,6 +201,12 @@ const ( FnOsSetenv = "os_setenv" // os_setenv(name, value) -> bool ) +// SSH Functions - Remote execution via SSH +const ( + FnSSHExec = "ssh_exec" // ssh_exec(host, command, user?, key_path?, password?, port?) -> string + FnSSHRsync = "ssh_rsync" // ssh_rsync(host, src, dest, user?, key_path?, password?, port?) -> bool +) + // LLM Functions - Invoke LLM from workflows const ( FnLLMInvoke = "llm_invoke" // llm_invoke(message) -> string (simple direct message) @@ -294,6 +300,10 @@ const ( FnDBImportVuln = "db_import_vuln" // db_import_vuln(workspace, json_data) -> bool FnDBImportVulnFromFile = "db_import_vuln_from_file" // db_import_vuln_from_file(workspace, file_path) -> int (count) + // DNS and custom asset import functions + FnDBImportDNSAsset = "db_import_dns_asset" // db_import_dns_asset(workspace, file_path) -> int + FnDBImportCustomAsset = "db_import_custom_asset" // db_import_custom_asset(workspace, file_path) -> map + // SARIF import functions FnDBImportSARIF = "db_import_sarif" // db_import_sarif(workspace, file_path) -> map (stats) FnConvertSARIFToMarkdown = "convert_sarif_to_markdown" // convert_sarif_to_markdown(input_path, output_path) -> bool @@ -544,6 +554,10 @@ func AllFunctions() []string { FnDBImportVuln, FnDBImportVulnFromFile, + // DNS and custom asset import functions + FnDBImportDNSAsset, + FnDBImportCustomAsset, + // SARIF import functions FnDBImportSARIF, FnConvertSARIFToMarkdown, @@ -568,6 +582,10 @@ func AllFunctions() []string { // Environment functions FnOsGetenv, FnOsSetenv, + + // SSH functions + FnSSHExec, + FnSSHRsync, } } @@ -607,6 +625,7 @@ const ( CategoryEnvironment = "environment" CategoryTypeDetection = "type_detection" CategoryLLM = "llm" + CategorySSH = "ssh" ) // CategoryInfo provides display metadata for a function category @@ -644,6 +663,7 @@ func CategoryOrder() []CategoryInfo { {CategoryInstaller, "Installer Functions", "Installer"}, {CategoryEnvironment, "Environment Functions", "Environment"}, {CategoryTypeDetection, "Type Detection Functions", "Type Detection"}, + {CategorySSH, "SSH Functions", "SSH"}, } } @@ -868,6 +888,8 @@ func FunctionRegistry() map[string][]FunctionInfo { {FnDBImportAssetFromFile, "db_import_asset_from_file(workspace, file_path)", "Import assets from JSONL file (httpx format)", "int", "db_import_asset_from_file('{{Workspace}}', '{{Output}}/httpx.jsonl')"}, {FnDBImportVuln, "db_import_vuln(workspace, json_data)", "Import single vulnerability from JSON (nuclei format)", "bool", "db_import_vuln('{{Workspace}}', '{\"template-id\":\"...\",\"info\":{\"name\":\"...\",\"severity\":\"high\"}}')"}, {FnDBImportVulnFromFile, "db_import_vuln_from_file(workspace, file_path)", "Import vulnerabilities from JSONL file (nuclei format)", "int", "db_import_vuln_from_file('{{Workspace}}', '{{Output}}/nuclei.jsonl')"}, + {FnDBImportDNSAsset, "db_import_dns_asset(workspace, file_path)", "Import DNS records from zone-style file (domain TYPE value per line), groups by domain", "int", "db_import_dns_asset('{{Workspace}}', '{{Output}}/dns-records.txt')"}, + {FnDBImportCustomAsset, "db_import_custom_asset(workspace, file_path)", "Import assets from JSONL file with direct field mapping (any Asset column can be set per line)", "map", "db_import_custom_asset('{{Workspace}}', '{{Output}}/custom-assets.jsonl')"}, {FnDBImportSARIF, "db_import_sarif(workspace, file_path)", "Import vulnerabilities from SARIF file (Semgrep, Trivy, etc.)", "map", "db_import_sarif('{{Workspace}}', '{{Output}}/semgrep.sarif')"}, {FnDBAssetDiff, "db_asset_diff(workspace)", "Get asset diff as JSONL string", "string", "db_asset_diff('{{Workspace}}')"}, {FnDBVulnDiff, "db_vuln_diff(workspace)", "Get vulnerability diff as JSONL string", "string", "db_vuln_diff('{{Workspace}}')"}, @@ -887,6 +909,10 @@ func FunctionRegistry() map[string][]FunctionInfo { {FnOsGetenv, "os_getenv(name)", "Get environment variable", "string", "os_getenv('HOME')"}, {FnOsSetenv, "os_setenv(name, value)", "Set environment variable", "bool", "os_setenv('API_KEY', 'secret')"}, }, + CategorySSH: { + {FnSSHExec, "ssh_exec(host, command, user?, key_path?, password?, port?)", "Execute command on remote host via SSH (uses connection pool)", "string", "ssh_exec('10.0.0.1', 'whoami', 'root', '~/.ssh/id_rsa')"}, + {FnSSHRsync, "ssh_rsync(host, src, dest, user?, key_path?, password?, port?)", "Copy local file/directory to remote host via rsync over SSH", "bool", "ssh_rsync('10.0.0.1', '/tmp/data.txt', '/opt/data.txt', 'root', '~/.ssh/id_rsa')"}, + }, CategoryTypeDetection: { {FnGetTypes, "get_types(input)", "Detect input type (file, folder, cidr, ip, url, domain, string)", "string", "get_types('192.168.1.0/24')"}, {FnIsFile, "is_file(path)", "Check if path is an existing regular file", "bool", "is_file('{{Output}}/results.txt')"}, diff --git a/internal/functions/db_functions.go b/internal/functions/db_functions.go index 9ce3f92..76479da 100644 --- a/internal/functions/db_functions.go +++ b/internal/functions/db_functions.go @@ -122,8 +122,14 @@ func (vf *vmFunc) updateAssetField(ctx context.Context, idStr, field string, val } switch field { - case "labels": - asset.Labels = value.String() + case "remarks": + asset.Remarks = []string{value.String()} + case "language": + asset.Language = value.String() + case "size": + asset.Size, _ = strconv.ParseInt(value.String(), 10, 64) + case "loc": + asset.LOC, _ = strconv.ParseInt(value.String(), 10, 64) case "source": asset.Source = value.String() case "asset_type": @@ -161,6 +167,50 @@ func (vf *vmFunc) updateRunField(ctx context.Context, id, field string, value go return err } +// unmarshalAssetJSON unmarshals JSON into an Asset with backward compatibility. +// Handles old format where "remarks" is a string instead of []string, +// and merges legacy "tags" array into Remarks. +func unmarshalAssetJSON(rawJSON []byte, asset *database.Asset) error { + // First try direct unmarshal + if err := json.Unmarshal(rawJSON, asset); err != nil { + // If direct unmarshal fails (e.g. "remarks" is string instead of []string), + // pre-process the JSON to fix known type mismatches + var raw map[string]interface{} + if mapErr := json.Unmarshal(rawJSON, &raw); mapErr != nil { + return err // return original error + } + // Convert "remarks" from string to []string + if remarksStr, ok := raw["remarks"].(string); ok { + raw["remarks"] = []string{remarksStr} + } + // Remove "tags" so it doesn't interfere with unmarshal + delete(raw, "tags") + fixed, marshalErr := json.Marshal(raw) + if marshalErr != nil { + return err + } + if err2 := json.Unmarshal(fixed, asset); err2 != nil { + return err2 + } + } + + // Merge legacy "tags" and handle string "remarks" into Remarks + var raw map[string]interface{} + if err := json.Unmarshal(rawJSON, &raw); err == nil { + if remarksStr, ok := raw["remarks"].(string); ok && remarksStr != "" { + asset.Remarks = []string{remarksStr} + } + if tagsArr, ok := raw["tags"].([]interface{}); ok { + for _, t := range tagsArr { + if s, ok := t.(string); ok { + asset.Remarks = append(asset.Remarks, s) + } + } + } + } + return nil +} + // dbImportAsset imports an asset from JSON data // Usage: db_import_asset('example.com', '{"host":"sub.example.com","url":"https://..."}') func (vf *vmFunc) dbImportAsset(call goja.FunctionCall) goja.Value { @@ -177,7 +227,7 @@ func (vf *vmFunc) dbImportAsset(call goja.FunctionCall) goja.Value { } var asset database.Asset - if err := json.Unmarshal([]byte(jsonData), &asset); err != nil { + if err := unmarshalAssetJSON([]byte(jsonData), &asset); err != nil { return vf.errorValue(fmt.Sprintf("invalid JSON: %v", err)) } @@ -866,7 +916,11 @@ func assetToMap(asset *database.Asset) map[string]interface{} { "asset_type": asset.AssetType, "technologies": asset.Technologies, "response_time": asset.ResponseTime, - "labels": asset.Labels, + "remarks": asset.Remarks, + "language": asset.Language, + "size": asset.Size, + "loc": asset.LOC, + "blob_content": asset.BlobContent, "source": asset.Source, "created_at": asset.CreatedAt, "updated_at": asset.UpdatedAt, @@ -2134,6 +2188,265 @@ func mapJSONToAsset(data map[string]interface{}, workspace, rawLine string) data return asset } +// domainRecords holds grouped DNS records for a single domain +type domainRecords struct { + aRecords []string + otherRecords []string // stored as "TYPE:value" +} + +// dbImportDNSAsset imports DNS records from a zone-style file and groups by domain. +// Each line: domain TYPE value (whitespace-separated, 3 fields). +// Returns count of domains imported. +// Usage: db_import_dns_asset(workspace, file_path) -> int +func (vf *vmFunc) dbImportDNSAsset(call goja.FunctionCall) goja.Value { + logger.Get().Debug("Calling " + terminal.HiGreen("dbImportDNSAsset")) + + if len(call.Arguments) < 2 { + return vf.errorValue("db_import_dns_asset requires 2 arguments: workspace, file_path") + } + + workspace := call.Argument(0).String() + filePath := call.Argument(1).String() + + if workspace == "" || workspace == "undefined" { + return vf.errorValue("workspace cannot be empty") + } + if filePath == "" || filePath == "undefined" { + return vf.errorValue("file_path cannot be empty") + } + + db := database.GetDB() + if db == nil { + return vf.errorValue("database not connected") + } + + file, err := os.Open(filePath) + if err != nil { + return vf.errorValue(fmt.Sprintf("failed to open file: %v", err)) + } + defer func() { _ = file.Close() }() + + // Group records by domain + domains := make(map[string]*domainRecords) + scanner := bufio.NewScanner(file) + + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" || strings.HasPrefix(line, ";") || strings.HasPrefix(line, "#") { + continue + } + + fields := strings.Fields(line) + if len(fields) < 3 { + continue + } + + domain := strings.TrimRight(fields[0], ".") + recordType := strings.ToUpper(fields[1]) + value := strings.TrimRight(fields[2], ".") + + if domain == "" || value == "" { + continue + } + + dr, exists := domains[domain] + if !exists { + dr = &domainRecords{} + domains[domain] = dr + } + + if recordType == "A" || recordType == "AAAA" { + dr.aRecords = append(dr.aRecords, value) + } else { + dr.otherRecords = append(dr.otherRecords, recordType+":"+value) + } + } + + if err := scanner.Err(); err != nil { + return vf.errorValue(fmt.Sprintf("error reading file: %v", err)) + } + + // Upsert each domain as an asset + ctx := context.Background() + now := time.Now() + imported := 0 + + for domain, dr := range domains { + // Build dns_records: A records as plain IPs, others as "TYPE:value" + var dnsRecords []string + dnsRecords = append(dnsRecords, dr.aRecords...) + dnsRecords = append(dnsRecords, dr.otherRecords...) + + hostIP := "" + if len(dr.aRecords) > 0 { + hostIP = dr.aRecords[0] + } + + // Check if asset already exists + var existing database.Asset + selectErr := db.NewSelect().Model(&existing). + Where("workspace = ?", workspace). + Where("asset_value = ?", domain). + Where("url = ?", ""). + Scan(ctx) + + if selectErr != nil { + // New asset - insert + asset := database.Asset{ + Workspace: workspace, + AssetValue: domain, + AssetType: "domain", + Source: "dns", + HostIP: hostIP, + DnsRecords: dnsRecords, + CreatedAt: now, + UpdatedAt: now, + LastSeenAt: now, + } + _, insertErr := db.NewInsert().Model(&asset).Exec(ctx) + if insertErr != nil { + logger.Get().Debug("failed to insert DNS asset", zap.Error(insertErr)) + continue + } + } else { + // Update existing + _, updateErr := db.NewUpdate().Model((*database.Asset)(nil)). + Set("dns_records = ?", dnsRecords). + Set("host_ip = ?", hostIP). + Set("source = ?", "dns"). + Set("updated_at = ?", now). + Set("last_seen_at = ?", now). + Where("id = ?", existing.ID). + Exec(ctx) + if updateErr != nil { + logger.Get().Debug("failed to update DNS asset", zap.Error(updateErr)) + continue + } + } + imported++ + } + + logger.Get().Debug("dbImportDNSAsset completed", + zap.String("workspace", workspace), + zap.String("file", filePath), + zap.Int("domains", imported)) + + return vf.vm.ToValue(imported) +} + +// dbImportCustomAsset imports assets from a JSONL file with direct field mapping. +// Each line is JSON that maps directly to Asset struct fields via json tags. +// Returns stats map {new, updated, errors, total}. +// Usage: db_import_custom_asset(workspace, file_path) -> map +func (vf *vmFunc) dbImportCustomAsset(call goja.FunctionCall) goja.Value { + logger.Get().Debug("Calling " + terminal.HiGreen("dbImportCustomAsset")) + + if len(call.Arguments) < 2 { + return vf.errorValue("db_import_custom_asset requires 2 arguments: workspace, file_path") + } + + workspace := call.Argument(0).String() + filePath := call.Argument(1).String() + + if workspace == "" || workspace == "undefined" { + return vf.errorValue("workspace cannot be empty") + } + if filePath == "" || filePath == "undefined" { + return vf.errorValue("file_path cannot be empty") + } + + db := database.GetDB() + if db == nil { + return vf.errorValue("database not connected") + } + + file, err := os.Open(filePath) + if err != nil { + return vf.errorValue(fmt.Sprintf("failed to open file: %v", err)) + } + defer func() { _ = file.Close() }() + + ctx := context.Background() + stats := database.ImportStats{} + + scanner := bufio.NewScanner(file) + buf := make([]byte, 0, 10*1024*1024) // 10MB buffer + scanner.Buffer(buf, 10*1024*1024) + + now := time.Now() + + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + + var asset database.Asset + if err := unmarshalAssetJSON([]byte(line), &asset); err != nil { + logger.Get().Debug("skipping invalid JSON line", zap.Error(err)) + stats.Errors++ + continue + } + + // Force workspace from function argument + asset.Workspace = workspace + asset.LastSeenAt = now + + // Check if asset already exists + var existing database.Asset + selectErr := db.NewSelect().Model(&existing). + Where("workspace = ?", workspace). + Where("asset_value = ?", asset.AssetValue). + Where("url = ?", asset.URL). + Scan(ctx) + + if selectErr != nil { + // New asset - insert + asset.CreatedAt = now + asset.UpdatedAt = now + _, insertErr := db.NewInsert().Model(&asset).Exec(ctx) + if insertErr != nil { + logger.Get().Debug("failed to insert custom asset", zap.Error(insertErr)) + stats.Errors++ + continue + } + stats.New++ + } else { + // Update existing - preserve ID and CreatedAt + asset.ID = existing.ID + asset.CreatedAt = existing.CreatedAt + asset.UpdatedAt = now + _, updateErr := db.NewUpdate().Model(&asset).WherePK().Exec(ctx) + if updateErr != nil { + logger.Get().Debug("failed to update custom asset", zap.Error(updateErr)) + stats.Errors++ + continue + } + stats.Updated++ + } + } + + if err := scanner.Err(); err != nil { + return vf.errorValue(fmt.Sprintf("error reading file: %v", err)) + } + + total := stats.New + stats.Updated + logger.Get().Debug("dbImportCustomAsset completed", + zap.String("workspace", workspace), + zap.String("file", filePath), + zap.Int("new", stats.New), + zap.Int("updated", stats.Updated), + zap.Int("errors", stats.Errors), + zap.Int("total", total)) + + return vf.vm.ToValue(map[string]interface{}{ + "new": stats.New, + "updated": stats.Updated, + "errors": stats.Errors, + "total": total, + }) +} + // dbImportVuln imports a single vulnerability from JSON data (nuclei format) // Usage: db_import_vuln(workspace, json_data) -> bool func (vf *vmFunc) dbImportVuln(call goja.FunctionCall) goja.Value { diff --git a/internal/functions/db_functions_test.go b/internal/functions/db_functions_test.go index ed52edc..43c6c75 100644 --- a/internal/functions/db_functions_test.go +++ b/internal/functions/db_functions_test.go @@ -765,3 +765,382 @@ func TestGlobToSQLLike(t *testing.T) { }) } } + +// --- 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) +} diff --git a/internal/functions/goja_pool.go b/internal/functions/goja_pool.go index 19f20d9..30faba4 100644 --- a/internal/functions/goja_pool.go +++ b/internal/functions/goja_pool.go @@ -108,6 +108,9 @@ type VMContext struct { // RuntimeVars stores variables set via set_var() for retrieval with get_var() RuntimeVars map[string]string + + // suppressDetails suppresses verbose console output (propagated from step's suppress_details) + suppressDetails bool } // VMRegistrationFunc is called to register functions on a new VM @@ -162,6 +165,7 @@ func (p *VMPool) Put(ctx *VMContext) { ctx.target = "" ctx.workspacePath = "" ctx.RuntimeVars = nil + ctx.suppressDetails = false p.pool.Put(ctx) } @@ -221,6 +225,11 @@ func (v *VMContext) SetContext(ctx map[string]interface{}) { if op, ok := ctx["Output"].(string); ok { v.workspacePath = op } + + // Extract suppress details flag (from step's suppress_details) + if sd, ok := ctx["SuppressDetails"].(bool); ok { + v.suppressDetails = sd + } } // SetVariables sets context variables on the VM diff --git a/internal/functions/goja_runtime.go b/internal/functions/goja_runtime.go index 9595dd6..fc70295 100644 --- a/internal/functions/goja_runtime.go +++ b/internal/functions/goja_runtime.go @@ -277,6 +277,10 @@ func (r *GojaRuntime) registerFunctionsOnVM(vm *goja.Runtime) { _ = vm.Set(FnDBImportVuln, vf.dbImportVuln) _ = vm.Set(FnDBImportVulnFromFile, vf.dbImportVulnFromFile) + // DNS and custom asset import functions + _ = vm.Set(FnDBImportDNSAsset, vf.dbImportDNSAsset) + _ = vm.Set(FnDBImportCustomAsset, vf.dbImportCustomAsset) + // SARIF import functions _ = vm.Set(FnDBImportSARIF, vf.dbImportSARIF) _ = vm.Set(FnConvertSARIFToMarkdown, vf.convertSARIFToMarkdown) @@ -302,6 +306,10 @@ func (r *GojaRuntime) registerFunctionsOnVM(vm *goja.Runtime) { _ = vm.Set(FnOsGetenv, vf.osGetenv) _ = vm.Set(FnOsSetenv, vf.osSetenv) + // SSH functions + _ = vm.Set(FnSSHExec, vf.sshExec) + _ = vm.Set(FnSSHRsync, vf.sshRsync) + // Console for debugging _ = vm.Set("console", map[string]interface{}{ "log": func(call goja.FunctionCall) goja.Value { diff --git a/internal/functions/ssh_functions.go b/internal/functions/ssh_functions.go new file mode 100644 index 0000000..c7c250c --- /dev/null +++ b/internal/functions/ssh_functions.go @@ -0,0 +1,221 @@ +package functions + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "time" + + "github.com/dop251/goja" + "github.com/j3ssie/osmedeus/v5/internal/core" + "github.com/j3ssie/osmedeus/v5/internal/logger" + "github.com/j3ssie/osmedeus/v5/internal/runner" + "github.com/j3ssie/osmedeus/v5/internal/terminal" + "go.uber.org/zap" +) + +// sshExpandPath expands ~ to the user's home directory +func sshExpandPath(path string) string { + if strings.HasPrefix(path, "~/") { + home, err := os.UserHomeDir() + if err == nil { + return filepath.Join(home, path[2:]) + } + } + return path +} + +// sshConfig holds parsed SSH connection parameters +type sshConfig struct { + User string + KeyFile string + Password string + Port int +} + +// parseSSHConfig extracts optional SSH parameters from goja function call arguments +// starting at paramOffset. Parameters are: user, key_path, password, port +// Returns defaults: user="root", port=22 +func parseSSHConfig(call goja.FunctionCall, paramOffset int) sshConfig { + cfg := sshConfig{ + User: "root", + Port: 22, + } + + if len(call.Arguments) > paramOffset { + v := call.Argument(paramOffset).String() + if v != "undefined" && v != "" { + cfg.User = v + } + } + if len(call.Arguments) > paramOffset+1 { + v := call.Argument(paramOffset + 1).String() + if v != "undefined" && v != "" { + cfg.KeyFile = v + } + } + if len(call.Arguments) > paramOffset+2 { + v := call.Argument(paramOffset + 2).String() + if v != "undefined" && v != "" { + cfg.Password = v + } + } + if len(call.Arguments) > paramOffset+3 { + v := call.Argument(paramOffset + 3) + if portVal, ok := v.Export().(int64); ok && portVal > 0 { + cfg.Port = int(portVal) + } else if portFloat, ok := v.Export().(float64); ok && portFloat > 0 { + cfg.Port = int(portFloat) + } + } + + return cfg +} + +// sshExec executes a command on a remote host via SSH using the connection pool +// Usage: ssh_exec(host, command, user?, key_path?, password?, port?) -> string +func (vf *vmFunc) sshExec(call goja.FunctionCall) goja.Value { + host := call.Argument(0).String() + command := call.Argument(1).String() + + logger.Get().Debug("Calling "+terminal.HiGreen(FnSSHExec), + zap.String("host", host), zap.Int("commandLength", len(command))) + + if host == "undefined" || host == "" { + logger.Get().Warn(FnSSHExec + ": host is required") + return vf.vm.ToValue("") + } + if command == "undefined" || command == "" { + logger.Get().Warn(FnSSHExec + ": command is required") + return vf.vm.ToValue("") + } + + // Parse optional SSH config (user, key_path, password, port) starting at arg index 2 + cfg := parseSSHConfig(call, 2) + + // Build RunnerConfig for SSH pool + runnerConfig := &core.RunnerConfig{ + Host: host, + Port: cfg.Port, + User: cfg.User, + KeyFile: sshExpandPath(cfg.KeyFile), + Password: cfg.Password, + } + + // Use a 5-minute timeout context + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + // Get pooled SSH connection + pool := runner.GetSSHPool() + client, poolKey, err := pool.Get(ctx, runnerConfig) + if err != nil { + logger.Get().Warn(FnSSHExec+": SSH connection failed", + zap.String("host", host), zap.Error(err)) + return vf.vm.ToValue("") + } + defer pool.Release(poolKey) + + // Create session and run command + session, err := client.NewSession() + if err != nil { + logger.Get().Warn(FnSSHExec+": failed to create SSH session", + zap.String("host", host), zap.Error(err)) + return vf.vm.ToValue("") + } + defer func() { _ = session.Close() }() + + output, err := session.CombinedOutput(command) + if err != nil { + logger.Get().Warn(FnSSHExec+": command execution failed", + zap.String("host", host), zap.String("command", command), zap.Error(err)) + // Return partial output even on error (non-zero exit code) + return vf.vm.ToValue(strings.TrimSpace(string(output))) + } + + result := strings.TrimSpace(string(output)) + logger.Get().Debug(terminal.HiGreen(FnSSHExec)+" result", + zap.String("host", host), zap.Int("outputLength", len(result))) + return vf.vm.ToValue(result) +} + +// sshRsync copies a local file/directory to a remote host via rsync over SSH +// Usage: ssh_rsync(host, src, dest, user?, key_path?, password?, port?) -> bool +func (vf *vmFunc) sshRsync(call goja.FunctionCall) goja.Value { + host := call.Argument(0).String() + src := call.Argument(1).String() + dest := call.Argument(2).String() + + logger.Get().Debug("Calling "+terminal.HiGreen(FnSSHRsync), + zap.String("host", host), zap.String("src", src), zap.String("dest", dest)) + + if host == "undefined" || host == "" { + logger.Get().Warn(FnSSHRsync + ": host is required") + return vf.vm.ToValue(false) + } + if src == "undefined" || src == "" { + logger.Get().Warn(FnSSHRsync + ": src is required") + return vf.vm.ToValue(false) + } + if dest == "undefined" || dest == "" { + logger.Get().Warn(FnSSHRsync + ": dest is required") + return vf.vm.ToValue(false) + } + + // Parse optional SSH config (user, key_path, password, port) starting at arg index 3 + cfg := parseSSHConfig(call, 3) + + // Expand local source path + src = sshExpandPath(src) + + // Build rsync destination: user@host:dest + rsyncDest := fmt.Sprintf("%s@%s:%s", cfg.User, host, dest) + + // Use a 5-minute timeout context for rsync + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + // Build rsync command with SSH options + // Add ConnectTimeout to prevent SSH from hanging on unreachable hosts + args := []string{"-avz", "-e"} + + keyPath := sshExpandPath(cfg.KeyFile) + if cfg.Password != "" && cfg.KeyFile == "" { + // Password-only auth: use sshpass wrapper + sshCmd := fmt.Sprintf("ssh -p %d -o StrictHostKeyChecking=no -o ConnectTimeout=30", cfg.Port) + args = append(args, sshCmd, src, rsyncDest) + + // Wrap with sshpass + sshpassArgs := append([]string{"-p", cfg.Password, "rsync"}, args...) + cmd := exec.CommandContext(ctx, "sshpass", sshpassArgs...) + output, err := cmd.CombinedOutput() + if err != nil { + logger.Get().Warn(FnSSHRsync+": rsync with sshpass failed", + zap.String("host", host), zap.Error(err), zap.String("output", string(output))) + return vf.vm.ToValue(false) + } + } else { + // Key-based auth (or key + password) + if keyPath != "" { + args = append(args, fmt.Sprintf("ssh -i %s -p %d -o StrictHostKeyChecking=no -o ConnectTimeout=30", keyPath, cfg.Port)) + } else { + args = append(args, fmt.Sprintf("ssh -p %d -o StrictHostKeyChecking=no -o ConnectTimeout=30", cfg.Port)) + } + args = append(args, src, rsyncDest) + + cmd := exec.CommandContext(ctx, "rsync", args...) + output, err := cmd.CombinedOutput() + if err != nil { + logger.Get().Warn(FnSSHRsync+": rsync failed", + zap.String("host", host), zap.Error(err), zap.String("output", string(output))) + return vf.vm.ToValue(false) + } + } + + logger.Get().Debug(terminal.HiGreen(FnSSHRsync)+" completed", + zap.String("host", host), zap.String("src", src), zap.String("dest", dest)) + return vf.vm.ToValue(true) +} diff --git a/internal/functions/ssh_functions_test.go b/internal/functions/ssh_functions_test.go new file mode 100644 index 0000000..e62cb5e --- /dev/null +++ b/internal/functions/ssh_functions_test.go @@ -0,0 +1,142 @@ +package functions + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- ssh_exec input validation tests --- + +func TestSSHExec_EmptyHost(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_exec("", "whoami")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, "", result) +} + +func TestSSHExec_UndefinedHost(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_exec(undefined, "whoami")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, "", result) +} + +func TestSSHExec_EmptyCommand(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_exec("10.0.0.1", "")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, "", result) +} + +func TestSSHExec_UndefinedCommand(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_exec("10.0.0.1")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, "", result) +} + +func TestSSHExec_NoArgs(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_exec()`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, "", result) +} + +// --- ssh_rsync input validation tests --- + +func TestSSHRsync_EmptyHost(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_rsync("", "/tmp/src", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSSHRsync_UndefinedHost(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_rsync(undefined, "/tmp/src", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSSHRsync_EmptySrc(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_rsync("10.0.0.1", "", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSSHRsync_EmptyDest(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_rsync("10.0.0.1", "/tmp/src", "")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSSHRsync_NoArgs(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `ssh_rsync()`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +// --- parseSSHConfig tests --- + +func TestSSHExec_DefaultUserAndPort(t *testing.T) { + // When only host and command are provided, defaults should be used (user=root, port=22) + // This will fail to connect but validates that defaults don't cause panics + registry := NewRegistry() + result, err := registry.Execute( + `ssh_exec("192.0.2.1", "echo test")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + // Connection to unreachable host will fail, returns empty string + assert.Equal(t, "", result) +} + +// NOTE: Integration tests for ssh_exec and ssh_rsync against a real SSH server +// are in the e2e test suite (test-e2e-ssh). Unit tests here only validate +// input validation since connecting to a real host would be slow/flaky. diff --git a/test/testdata/sample-jsonl-output/custom-assets.jsonl b/test/testdata/sample-jsonl-output/custom-assets.jsonl new file mode 100644 index 0000000..ff43909 --- /dev/null +++ b/test/testdata/sample-jsonl-output/custom-assets.jsonl @@ -0,0 +1,3 @@ +{"asset_value":"api.example.com","url":"https://api.example.com","asset_type":"url","source":"custom-scan","tags":["api","production"],"external_url":"https://jira.example.com/PROJ-123","remarks":"API endpoint"} +{"asset_value":"admin.example.com","asset_type":"domain","source":"dns-enum","tags":["admin"],"host_ip":"10.0.0.5"} +{"asset_value":"cdn.example.com","url":"https://cdn.example.com","asset_type":"url","source":"crawl","tech":["Cloudflare","nginx"]} diff --git a/test/testdata/sample-jsonl-output/dns-records.txt b/test/testdata/sample-jsonl-output/dns-records.txt new file mode 100644 index 0000000..9ab7d30 --- /dev/null +++ b/test/testdata/sample-jsonl-output/dns-records.txt @@ -0,0 +1,6 @@ +example.com. A 93.184.216.34 +example.com. A 93.184.216.35 +example.com. CNAME cdn.example.net. +sub.example.com. A 10.0.0.1 +sub.example.com. MX mail.example.com. +other.example.com. A 192.168.1.1