diff --git a/.gitignore b/.gitignore index 4490b8a..66fdee0 100644 --- a/.gitignore +++ b/.gitignore @@ -50,4 +50,5 @@ OPTIMIZE.md OPTIMIZE-*.md PLANNING.md PLANNING-*.md - +INTERNAL.md +INTERNAL-*.md \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index b964096..7920517 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -188,7 +188,8 @@ osmedeus func e 'log_info("{{target}}")' # Evaluate function osmedeus --usage-example # Show all usage examples osmedeus server # Start REST API (see docs/api/ for endpoints) osmedeus server --master # Start as distributed master -osmedeus worker join # Join as distributed worker +osmedeus worker join # Join as distributed worker (ID: wosm-) +osmedeus worker join --get-public-ip # Join with public IP detection (alias: wosm-) osmedeus install binary --name # Install specific binary osmedeus install binary --all # Install all binaries osmedeus install binary --name --check # Check if binary is installed @@ -198,6 +199,7 @@ osmedeus install binary --nix-installation # Install Nix package manager osmedeus install binary --list-registry-nix-build # List Nix binaries osmedeus install binary --list-registry-direct-fetch # List direct-fetch binaries osmedeus install base --preset # Install base from preset repository +osmedeus install base --preset --keep-setting # Install base, restore previous osm-settings.yaml osmedeus install workflow --preset # Install workflows from preset repository osmedeus install validate --preset # Validate/install ready-to-use base osmedeus install env # Add binaries to PATH (auto-detects shell) @@ -270,6 +272,8 @@ REST API documentation with curl examples is in `docs/api/`. Key endpoint catego - **Decision Routing**: Uses switch/case syntax for conditional workflow branching - **Run Registry**: Tracks active runs with PID management for cancellation support - **Write Coordinator**: Batches database writes (step results, progress, artifacts) reducing I/O by ~70% +- **Install Base Backup**: `InstallBase()` automatically backs up `osm-settings.yaml` to `backup-osm-settings.yaml`; `--keep-setting` flag restores the previous settings after installation +- **Worker Identity**: Worker IDs use `wosm-` format; default alias is `wosm-` or `wosm-` when no `--alias` is provided ## SARIF Integration diff --git a/HACKING.md b/HACKING.md index 17b1db8..81a8200 100644 --- a/HACKING.md +++ b/HACKING.md @@ -34,7 +34,7 @@ osmedeus/ │ ├── console/ # Console output capture │ ├── core/ # Core types (Workflow, Step, Trigger, etc.) │ ├── database/ # SQLite/PostgreSQL via Bun ORM -│ ├── distributed/ # Distributed execution (master/worker) +│ ├── distributed/ # Distributed execution (master/worker, worker ID: wosm-) │ ├── executor/ # Workflow execution engine │ ├── fileio/ # High-performance file I/O (mmap) │ ├── functions/ # Utility functions (Goja JS runtime) @@ -1598,6 +1598,9 @@ Osmedeus supports installing base folders and workflows from curated preset repo # Install base folder from preset repository osmedeus install base --preset +# Install base and restore previous osm-settings.yaml (API keys, Redis config, etc.) +osmedeus install base --preset --keep-setting + # Install workflows from preset repository osmedeus install workflow --preset @@ -1605,6 +1608,15 @@ osmedeus install workflow --preset osmedeus install validate --preset ``` +### Settings Backup + +When `install base` runs, the entire base folder (including `osm-settings.yaml`) is deleted and replaced. To prevent losing custom settings: + +- **Automatic backup**: `osm-settings.yaml` is always backed up to `~/osmedeus-base/backup-osm-settings.yaml` before removal +- **`--keep-setting` flag**: Restores the previous `osm-settings.yaml` over the newly installed one after installation + +This is handled by the `Installer.KeepSetting` field in `internal/installer/installer.go`. + ### Environment Variables | Variable | Default | Description | diff --git a/README.md b/README.md index 6f1d805..f12f822 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ Built for both beginners and experts, it delivers powerful, composable automatio - **Template Engine** - Powerful variable interpolation with built-in and custom variables - **Utility Functions** - Rich function library with event generation, bulk processing, and JSON operations - **REST API Server** - Manage, trigger, and cancel workflows programmatically -- **Distributed Execution** - Scale with Redis-based master-worker pattern for parallel scanning +- **Distributed Execution** - Scale with Redis-based master-worker pattern for parallel scanning (workers identified as `wosm-`) - **Notifications** - Telegram bot and webhook integrations - **Cloud Storage** - S3-compatible storage for artifact management - **LLM Integration** - AI-powered workflow steps with chat completions, embeddings, and agentic tool-calling loops @@ -85,6 +85,7 @@ osmedeus func eval 'log_info("OS: " + PlatformOS + ", Arch: " + PlatformArch)' # Install from preset repositories osmedeus install base --preset +osmedeus install base --preset --keep-setting # preserve existing osm-settings.yaml osmedeus install workflow --preset # Show all usage examples diff --git a/internal/core/workflow.go b/internal/core/workflow.go index 1553d43..bd671cd 100644 --- a/internal/core/workflow.go +++ b/internal/core/workflow.go @@ -23,6 +23,12 @@ func (t *TagList) UnmarshalYAML(unmarshal func(interface{}) error) error { return nil } +// WorkflowHooks defines pre/post execution steps for workflows +type WorkflowHooks struct { + PreScanSteps []Step `yaml:"pre_scan_steps,omitempty"` + PostScanSteps []Step `yaml:"post_scan_steps,omitempty"` +} + // WorkflowHelp contains usage documentation for a workflow type WorkflowHelp struct { ExampleTargets []string `yaml:"example_targets,omitempty"` @@ -45,6 +51,9 @@ type Workflow struct { // Execution preferences (optional, can be overridden by CLI flags) Preferences *Preferences `yaml:"preferences,omitempty"` + // Hooks for pre/post scan steps + Hooks *WorkflowHooks `yaml:"hooks,omitempty"` + // Runner configuration (module-kind only) Runner RunnerType `yaml:"runner,omitempty"` RunnerConfig *RunnerConfig `yaml:"runner_config,omitempty"` @@ -209,3 +218,11 @@ func (w *Workflow) GetExampleTargets() []string { } return w.Help.ExampleTargets } + +// HookCount returns the total number of hook steps (pre + post) +func (w *Workflow) HookCount() int { + if w.Hooks == nil { + return 0 + } + return len(w.Hooks.PreScanSteps) + len(w.Hooks.PostScanSteps) +} diff --git a/internal/core/workflow_hooks_test.go b/internal/core/workflow_hooks_test.go new file mode 100644 index 0000000..9dd8528 --- /dev/null +++ b/internal/core/workflow_hooks_test.go @@ -0,0 +1,189 @@ +package core + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +func TestWorkflowHooksYAMLParsing(t *testing.T) { + t.Run("both pre and post hooks", func(t *testing.T) { + data := ` +kind: module +name: test-hooks +hooks: + pre_scan_steps: + - name: pre-1 + type: bash + command: echo pre + - name: pre-2 + type: bash + command: echo pre2 + post_scan_steps: + - name: post-1 + type: bash + command: echo post +steps: + - name: s1 + type: bash + command: echo main +` + var wf Workflow + err := yaml.Unmarshal([]byte(data), &wf) + require.NoError(t, err) + require.NotNil(t, wf.Hooks) + assert.Len(t, wf.Hooks.PreScanSteps, 2) + assert.Len(t, wf.Hooks.PostScanSteps, 1) + assert.Equal(t, "pre-1", wf.Hooks.PreScanSteps[0].Name) + assert.Equal(t, "post-1", wf.Hooks.PostScanSteps[0].Name) + }) + + t.Run("pre only", func(t *testing.T) { + data := ` +kind: module +name: test +hooks: + pre_scan_steps: + - name: pre-1 + type: bash + command: echo pre +steps: + - name: s1 + type: bash + command: echo main +` + var wf Workflow + err := yaml.Unmarshal([]byte(data), &wf) + require.NoError(t, err) + require.NotNil(t, wf.Hooks) + assert.Len(t, wf.Hooks.PreScanSteps, 1) + assert.Empty(t, wf.Hooks.PostScanSteps) + }) + + t.Run("post only", func(t *testing.T) { + data := ` +kind: module +name: test +hooks: + post_scan_steps: + - name: post-1 + type: bash + command: echo post +steps: + - name: s1 + type: bash + command: echo main +` + var wf Workflow + err := yaml.Unmarshal([]byte(data), &wf) + require.NoError(t, err) + require.NotNil(t, wf.Hooks) + assert.Empty(t, wf.Hooks.PreScanSteps) + assert.Len(t, wf.Hooks.PostScanSteps, 1) + }) + + t.Run("no hooks", func(t *testing.T) { + data := ` +kind: module +name: test +steps: + - name: s1 + type: bash + command: echo main +` + var wf Workflow + err := yaml.Unmarshal([]byte(data), &wf) + require.NoError(t, err) + assert.Nil(t, wf.Hooks) + }) + + t.Run("empty hooks block", func(t *testing.T) { + data := ` +kind: module +name: test +hooks: {} +steps: + - name: s1 + type: bash + command: echo main +` + var wf Workflow + err := yaml.Unmarshal([]byte(data), &wf) + require.NoError(t, err) + // Empty hooks block is parsed but contains no steps + if wf.Hooks != nil { + assert.Empty(t, wf.Hooks.PreScanSteps) + assert.Empty(t, wf.Hooks.PostScanSteps) + } + }) + + t.Run("flow with hooks", func(t *testing.T) { + data := ` +kind: flow +name: test-flow +hooks: + pre_scan_steps: + - name: flow-pre + type: bash + command: echo flow-pre + post_scan_steps: + - name: flow-post + type: bash + command: echo flow-post +modules: + - name: mod1 + steps: + - name: inline-step + type: bash + command: echo inline +` + var wf Workflow + err := yaml.Unmarshal([]byte(data), &wf) + require.NoError(t, err) + require.NotNil(t, wf.Hooks) + assert.Len(t, wf.Hooks.PreScanSteps, 1) + assert.Len(t, wf.Hooks.PostScanSteps, 1) + }) +} + +func TestWorkflowHookCount(t *testing.T) { + t.Run("nil hooks", func(t *testing.T) { + wf := &Workflow{} + assert.Equal(t, 0, wf.HookCount()) + }) + + t.Run("empty hooks", func(t *testing.T) { + wf := &Workflow{Hooks: &WorkflowHooks{}} + assert.Equal(t, 0, wf.HookCount()) + }) + + t.Run("pre only", func(t *testing.T) { + wf := &Workflow{ + Hooks: &WorkflowHooks{ + PreScanSteps: []Step{{Name: "pre-1"}, {Name: "pre-2"}}, + }, + } + assert.Equal(t, 2, wf.HookCount()) + }) + + t.Run("post only", func(t *testing.T) { + wf := &Workflow{ + Hooks: &WorkflowHooks{ + PostScanSteps: []Step{{Name: "post-1"}}, + }, + } + assert.Equal(t, 1, wf.HookCount()) + }) + + t.Run("both pre and post", func(t *testing.T) { + wf := &Workflow{ + Hooks: &WorkflowHooks{ + PreScanSteps: []Step{{Name: "pre-1"}, {Name: "pre-2"}}, + PostScanSteps: []Step{{Name: "post-1"}, {Name: "post-2"}, {Name: "post-3"}}, + }, + } + assert.Equal(t, 5, wf.HookCount()) + }) +} diff --git a/internal/database/models.go b/internal/database/models.go index a9a09d7..76ca8db 100644 --- a/internal/database/models.go +++ b/internal/database/models.go @@ -40,8 +40,9 @@ type Run struct { CurrentPID int `bun:"current_pid" json:"current_pid,omitempty"` // Priority and mode - RunPriority string `bun:"run_priority,notnull,default:'high'" json:"run_priority"` // low, normal, high, critical - RunMode string `bun:"run_mode,notnull,default:'local'" json:"run_mode"` // local, distributed, cloud + RunPriority string `bun:"run_priority,notnull,default:'high'" json:"run_priority"` // low, normal, high, critical + RunMode string `bun:"run_mode,notnull,default:'local'" json:"run_mode"` // local, distributed, cloud + HooksEnabled bool `bun:"hooks_enabled,default:false" json:"hooks_enabled"` // Relations Steps []*StepResult `bun:"rel:has-many,join:id=run_id" json:"steps,omitempty"` @@ -308,6 +309,7 @@ type WorkflowMeta struct { // Metadata StepCount int `bun:"step_count" json:"step_count"` ModuleCount int `bun:"module_count" json:"module_count"` + HookCount int `bun:"hook_count" json:"hook_count"` ParamsJSON string `bun:"params_json" json:"params_json"` // Serialized params // Timestamps diff --git a/internal/database/workflow_index.go b/internal/database/workflow_index.go index ef0cfc4..4738454 100644 --- a/internal/database/workflow_index.go +++ b/internal/database/workflow_index.go @@ -122,6 +122,7 @@ func upsertWorkflowMeta(ctx context.Context, w *core.Workflow, force bool) error existing.Hidden = w.Hidden existing.StepCount = len(w.Steps) existing.ModuleCount = len(w.Modules) + existing.HookCount = w.HookCount() existing.ParamsJSON = paramsJSON existing.IndexedAt = now existing.UpdatedAt = now @@ -147,6 +148,7 @@ func upsertWorkflowMeta(ctx context.Context, w *core.Workflow, force bool) error Hidden: w.Hidden, StepCount: len(w.Steps), ModuleCount: len(w.Modules), + HookCount: w.HookCount(), ParamsJSON: paramsJSON, IndexedAt: now, CreatedAt: now, diff --git a/internal/distributed/client.go b/internal/distributed/client.go index a16e892..9f43389 100644 --- a/internal/distributed/client.go +++ b/internal/distributed/client.go @@ -28,12 +28,19 @@ const ( KeyEventsPrefix = KeyPrefix + "events:" // osm:events:{topic} // Data queue keys (for worker -> master data) - KeyDataRuns = KeyPrefix + "data:runs" - KeyDataSteps = KeyPrefix + "data:steps" - KeyDataEvents = KeyPrefix + "data:events" - KeyDataArtifacts = KeyPrefix + "data:artifacts" + KeyDataRuns = KeyPrefix + "data:runs" + KeyDataSteps = KeyPrefix + "data:steps" + KeyDataEvents = KeyPrefix + "data:events" + KeyDataArtifacts = KeyPrefix + "data:artifacts" + KeyDataExecute = KeyPrefix + "data:execute" + KeyDataExecuteWorker = KeyPrefix + "data:execute:worker:" // osm:data:execute:worker:{worker_id} ) +// KeyDataExecuteForWorker returns the per-worker execute queue key. +func KeyDataExecuteForWorker(workerID string) string { + return KeyDataExecuteWorker + workerID +} + // Timeouts and intervals const ( HeartbeatInterval = 30 * time.Second @@ -302,6 +309,33 @@ func (c *Client) GetAllWorkers(ctx context.Context) ([]*WorkerInfo, error) { return workers, nil } +// GetWorker retrieves a single worker by ID +func (c *Client) GetWorker(ctx context.Context, workerID string) (*WorkerInfo, error) { + cmd := c.client.B().Hget().Key(KeyWorkers).Field(workerID).Build() + data, err := c.client.Do(ctx, cmd).ToString() + if err != nil { + if rueidis.IsRedisNil(err) { + return nil, nil + } + return nil, fmt.Errorf("failed to get worker: %w", err) + } + return UnmarshalWorkerInfo([]byte(data)) +} + +// GetWorkerByAlias retrieves a worker by its alias +func (c *Client) GetWorkerByAlias(ctx context.Context, alias string) (*WorkerInfo, error) { + workers, err := c.GetAllWorkers(ctx) + if err != nil { + return nil, err + } + for _, w := range workers { + if w.Alias == alias { + return w, nil + } + } + return nil, nil +} + // RemoveWorker removes a worker from the registry func (c *Client) RemoveWorker(ctx context.Context, workerID string) error { // Remove from both workers and heartbeat hashes diff --git a/internal/distributed/master.go b/internal/distributed/master.go index 71ca181..b784939 100644 --- a/internal/distributed/master.go +++ b/internal/distributed/master.go @@ -3,17 +3,22 @@ package distributed import ( "context" "fmt" - "github.com/j3ssie/osmedeus/v5/internal/json" "os" + "os/exec" + "strings" "sync" "time" + "github.com/j3ssie/osmedeus/v5/internal/json" + "github.com/google/uuid" "github.com/j3ssie/osmedeus/v5/internal/broker" "github.com/j3ssie/osmedeus/v5/internal/config" "github.com/j3ssie/osmedeus/v5/internal/core" "github.com/j3ssie/osmedeus/v5/internal/database" "github.com/j3ssie/osmedeus/v5/internal/database/repository" + "github.com/j3ssie/osmedeus/v5/internal/executor" + "github.com/j3ssie/osmedeus/v5/internal/functions" "github.com/j3ssie/osmedeus/v5/internal/logger" "github.com/j3ssie/osmedeus/v5/internal/terminal" "github.com/uptrace/bun" @@ -406,7 +411,7 @@ func (m *Master) persistEventLog(ctx context.Context, event *core.Event) { // dataProcessorLoop processes data from worker data queues func (m *Master) dataProcessorLoop(ctx context.Context) { - keys := []string{KeyDataRuns, KeyDataSteps, KeyDataEvents, KeyDataArtifacts} + keys := []string{KeyDataRuns, KeyDataSteps, KeyDataEvents, KeyDataArtifacts, KeyDataExecute} timeout := 1 * time.Second for { @@ -432,6 +437,18 @@ func (m *Master) dataProcessorLoop(ctx context.Context) { // processWorkerData processes data received from a worker func (m *Master) processWorkerData(ctx context.Context, key string, envelope *DataEnvelope) { + m.logger.Debug("processing worker data", + zap.String("key", key), + zap.String("type", envelope.Type), + zap.String("worker_id", envelope.WorkerID), + ) + + // Execute requests don't require a database connection + if key == KeyDataExecute { + m.processExecuteData(ctx, envelope) + return + } + if m.db == nil { m.logger.Debug("skipping data processing - no database connection", zap.String("key", key), @@ -440,12 +457,6 @@ func (m *Master) processWorkerData(ctx context.Context, key string, envelope *Da return } - m.logger.Debug("processing worker data", - zap.String("key", key), - zap.String("type", envelope.Type), - zap.String("worker_id", envelope.WorkerID), - ) - switch key { case KeyDataRuns: m.processRunData(ctx, envelope) @@ -529,3 +540,283 @@ func (m *Master) processArtifactData(ctx context.Context, envelope *DataEnvelope m.printer.Warning("Failed to create artifact %s: %s", artifact.Name, err) } } + +// ============================================================================= +// Execute Request Processing +// ============================================================================= + +// processExecuteData processes an execute request from a worker +func (m *Master) processExecuteData(ctx context.Context, envelope *DataEnvelope) { + var req ExecuteRequest + if err := json.Unmarshal(envelope.Data, &req); err != nil { + m.printer.Warning("Failed to unmarshal execute request: %s", err) + return + } + + // Resolve execute type: prefer new ExecuteType field, fall back to legacy Action + executeType := req.ExecuteType + if executeType == "" { + executeType = req.Action + } + + // Resolve target role: default to "master" for backward compatibility + targetRole := req.TargetRole + if targetRole == "" { + targetRole = "master" + } + + m.logger.Info("processing execute request from worker", + zap.String("worker_id", envelope.WorkerID), + zap.String("execute_type", executeType), + zap.String("target_role", targetRole), + zap.String("data", req.Data), + ) + + // Route requests not targeted at master to workers + if targetRole != "master" { + m.routeExecuteToWorkers(ctx, &req, envelope.WorkerID) + return + } + + switch executeType { + case "func": + // Resolve data: prefer new Data field, fall back to legacy Expr + expr := req.Data + if expr == "" { + expr = req.Expr + } + m.executeFunc(ctx, expr, envelope.WorkerID) + case "run": + // Resolve data: prefer new Data field, fall back to legacy Workflow + workflow := req.Data + if workflow == "" { + workflow = req.Workflow + } + m.executeRun(ctx, workflow, req.Target, req.Params, envelope.WorkerID) + case "bash": + command := req.Data + if command == "" { + command = req.Expr // legacy fallback + } + m.executeBash(ctx, command, envelope.WorkerID) + case "sync_to_worker": + // Data=src path on master, Target=dest path on worker + m.executeSyncToWorker(ctx, req.Data, req.Target, envelope.WorkerID) + default: + m.printer.Warning("Unknown execute type from worker %s: %s", envelope.WorkerID, executeType) + } +} + +// executeFunc runs a utility function expression on the master +func (m *Master) executeFunc(ctx context.Context, expr, workerID string) { + if expr == "" { + m.printer.Warning("Empty expression from worker %s", workerID) + return + } + + m.logger.Debug("executing function from worker", + zap.String("worker_id", workerID), + zap.String("expr", expr), + ) + + // Build context with built-in variables so functions have access to paths + execCtx := executor.BuildBuiltinVariables(m.config, nil) + + registry := functions.NewRegistry() + _, err := registry.Execute(expr, execCtx) + if err != nil { + m.printer.Warning("Failed to execute function from worker %s: %s (expr: %s)", + workerID, err, expr) + return + } + + m.logger.Info("executed function from worker", + zap.String("worker_id", workerID), + zap.String("expr", expr), + ) +} + +// executeRun submits a workflow task to the pending queue on behalf of a worker +func (m *Master) executeRun(ctx context.Context, workflow, target, params, workerID string) { + if workflow == "" || target == "" { + m.printer.Warning("Missing workflow or target from worker %s", workerID) + return + } + + m.logger.Debug("submitting run request from worker", + zap.String("worker_id", workerID), + zap.String("workflow", workflow), + zap.String("target", target), + zap.String("params", params), + ) + + // Parse params into map + taskParams := make(map[string]interface{}) + if params != "" { + for _, pair := range strings.Split(params, ",") { + pair = strings.TrimSpace(pair) + if pair == "" { + continue + } + parts := strings.SplitN(pair, "=", 2) + if len(parts) == 2 { + taskParams[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1]) + } + } + } + + task := NewTask( + uuid.NewString()[:8], + workflow, + "module", // default to module; master can detect flow if needed + target, + taskParams, + ) + + if err := m.SubmitTask(ctx, task); err != nil { + m.printer.Warning("Failed to submit task from worker %s: %s", workerID, err) + return + } + + m.logger.Info("submitted task from worker", + zap.String("worker_id", workerID), + zap.String("task_id", task.ID), + zap.String("workflow", workflow), + zap.String("target", target), + ) +} + +// executeBash runs a shell command on the master on behalf of a worker +func (m *Master) executeBash(ctx context.Context, command, workerID string) { + if command == "" { + m.printer.Warning("Empty bash command from worker %s", workerID) + return + } + + m.logger.Debug("executing bash from worker", + zap.String("worker_id", workerID), + zap.String("command", command)) + + // @NOTE: This is intentional - bash commands come from trusted workflow YAML files + // via the distributed system. Workers send requests through run_on_master('bash', cmd). + cmd := exec.CommandContext(ctx, "sh", "-c", command) + output, err := cmd.CombinedOutput() + if err != nil { + m.printer.Warning("Bash execution failed from worker %s: %s (output: %s)", + workerID, err, string(output)) + return + } + + m.logger.Info("executed bash from worker", + zap.String("worker_id", workerID), + zap.String("command", command), + zap.Int("output_len", len(output))) +} + +// executeSyncToWorker rsyncs a file/folder from the master to the requesting worker. +// src is the path on the master, dest is the path on the worker. +func (m *Master) executeSyncToWorker(ctx context.Context, src, dest, workerID string) { + if src == "" || dest == "" { + m.printer.Warning("sync_to_worker: missing src or dest from worker %s", workerID) + return + } + + // Look up the requesting worker's SSH info + worker, err := m.client.GetWorker(ctx, workerID) + if err != nil || worker == nil { + m.printer.Warning("sync_to_worker: failed to get worker %s: %v", workerID, err) + return + } + if !worker.SSHEnabled { + m.printer.Warning("sync_to_worker: worker %s does not have SSH enabled", workerID) + return + } + + host := worker.PublicIP + if host == "" { + host = worker.IPAddress + } + if host == "" { + m.printer.Warning("sync_to_worker: worker %s has no IP address", workerID) + return + } + + m.logger.Debug("syncing to worker", + zap.String("worker_id", workerID), + zap.String("host", host), + zap.String("src", src), + zap.String("dest", dest), + ) + + // Build rsync command: rsync -avz -e "ssh -i -p 22 -o StrictHostKeyChecking=no" src root@host:dest + args := []string{"-avz", "-e"} + keyPath := worker.SSHKeysPath + if keyPath != "" { + args = append(args, fmt.Sprintf("ssh -i %s -p 22 -o StrictHostKeyChecking=no -o ConnectTimeout=30", keyPath)) + } else { + args = append(args, "ssh -p 22 -o StrictHostKeyChecking=no -o ConnectTimeout=30") + } + rsyncDest := fmt.Sprintf("root@%s:%s", host, dest) + args = append(args, src, rsyncDest) + + rsyncCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) + defer cancel() + + cmd := exec.CommandContext(rsyncCtx, "rsync", args...) + output, err := cmd.CombinedOutput() + if err != nil { + m.printer.Warning("sync_to_worker: rsync to worker %s failed: %s (output: %s)", + workerID, err, string(output)) + return + } + + m.logger.Info("synced to worker", + zap.String("worker_id", workerID), + zap.String("host", host), + zap.String("src", src), + zap.String("dest", dest), + ) +} + +// routeExecuteToWorkers routes an execute request to target worker(s) via per-worker queues. +func (m *Master) routeExecuteToWorkers(ctx context.Context, req *ExecuteRequest, senderWorkerID string) { + workers, err := m.client.GetAllWorkers(ctx) + if err != nil { + m.printer.Warning("Failed to get workers for routing: %s", err) + return + } + + scope := req.TargetScope + if scope == "" { + scope = "all" + } + + var targets []*WorkerInfo + if scope == "all" { + targets = workers + } else { + for _, w := range workers { + if w.ID == scope || w.Alias == scope || w.PublicIP == scope { + targets = append(targets, w) + break + } + } + } + + if len(targets) == 0 { + m.printer.Warning("No workers matched scope %q for execute request from %s", scope, senderWorkerID) + return + } + + for _, w := range targets { + key := KeyDataExecuteForWorker(w.ID) + if err := m.client.PushData(ctx, key, "execute", req, senderWorkerID); err != nil { + m.printer.Warning("Failed to route execute to worker %s: %s", w.ID, err) + } + } + + m.logger.Info("routed execute request to workers", + zap.String("scope", scope), + zap.Int("target_count", len(targets)), + zap.String("sender", senderWorkerID)) +} diff --git a/internal/distributed/task.go b/internal/distributed/task.go index 8f2d2d8..5f692ea 100644 --- a/internal/distributed/task.go +++ b/internal/distributed/task.go @@ -51,6 +51,11 @@ type WorkerInfo struct { LastHeartbeat time.Time `json:"last_heartbeat"` TasksComplete int `json:"tasks_complete"` TasksFailed int `json:"tasks_failed"` + IPAddress string `json:"ip_address,omitempty"` + PublicIP string `json:"public_ip,omitempty"` + SSHEnabled bool `json:"ssh_enabled,omitempty"` + SSHKeysPath string `json:"ssh_keys_path,omitempty"` + Alias string `json:"alias,omitempty"` } // NewTask creates a new task with the given parameters @@ -138,6 +143,22 @@ func (t *Task) MarkCompleted() { t.CompletedAt = &now } +// ExecuteRequest represents a request from a worker to execute something on the master or another worker. +// ExecuteType "func" executes a utility function expression. +// ExecuteType "run" submits a workflow task to the pending queue. +// TargetRole controls where the request is executed: "master" or "worker". +type ExecuteRequest struct { + ExecuteType string `json:"execute_type"` // "func", "run", or "bash" + TargetRole string `json:"target_role"` // "master" or "worker" + Data string `json:"data,omitempty"` // function expression (for "func"), workflow name (for "run"), or command (for "bash") + Target string `json:"target,omitempty"` // target (for "run") + Params string `json:"params,omitempty"` // comma-separated key=value (for "run") + TargetScope string `json:"target_scope,omitempty"` // For worker-targeted: "all", alias, worker ID, or public IP + Action string `json:"action,omitempty"` // deprecated: use ExecuteType instead + Expr string `json:"expr,omitempty"` // deprecated: use Data instead + Workflow string `json:"workflow,omitempty"` // deprecated: use Data instead (for "run") +} + // MarkFailed marks the task as failed with an error message func (t *Task) MarkFailed(err string) { t.Status = TaskStatusFailed diff --git a/internal/distributed/worker.go b/internal/distributed/worker.go index 580b2e1..4e87ce7 100644 --- a/internal/distributed/worker.go +++ b/internal/distributed/worker.go @@ -3,7 +3,12 @@ package distributed import ( "context" "fmt" + "io" + "net" + "net/http" "os" + "os/exec" + "strings" "time" "github.com/google/uuid" @@ -11,11 +16,21 @@ import ( "github.com/j3ssie/osmedeus/v5/internal/core" "github.com/j3ssie/osmedeus/v5/internal/database" "github.com/j3ssie/osmedeus/v5/internal/executor" + "github.com/j3ssie/osmedeus/v5/internal/functions" "github.com/j3ssie/osmedeus/v5/internal/heuristics" + "github.com/j3ssie/osmedeus/v5/internal/json" "github.com/j3ssie/osmedeus/v5/internal/parser" "github.com/j3ssie/osmedeus/v5/internal/terminal" ) +// WorkerOptions holds optional configuration for creating a new Worker. +type WorkerOptions struct { + GetPublicIP bool + Alias string + SSHEnabled bool + SSHKeysPath string +} + // Worker represents a worker node that processes tasks type Worker struct { ID string @@ -26,34 +41,106 @@ type Worker struct { loader *parser.Loader printer *terminal.Printer + // Cleanup function for distributed hooks + unregisterHooks func() + + // Metadata + ipAddress string + publicIP string + sshEnabled bool + sshKeysPath string + alias string + // Stats tasksComplete int tasksFailed int } // NewWorker creates a new worker node -func NewWorker(cfg *config.Config) (*Worker, error) { +func NewWorker(cfg *config.Config, opts *WorkerOptions) (*Worker, error) { + if opts == nil { + opts = &WorkerOptions{} + } + client, err := NewClientFromConfig(cfg) if err != nil { return nil, fmt.Errorf("failed to create redis client: %w", err) } hostname, _ := os.Hostname() - workerID := fmt.Sprintf("%s-%s", hostname, uuid.NewString()[:8]) + workerID := fmt.Sprintf("wosm-%s", uuid.NewString()[:8]) exec := executor.NewExecutor() loader := parser.NewLoader(cfg.WorkflowsPath) exec.SetLoader(loader) - return &Worker{ - ID: workerID, - Hostname: hostname, - client: client, - config: cfg, - executor: exec, - loader: loader, - printer: terminal.NewPrinter(), - }, nil + p := terminal.NewPrinter() + + w := &Worker{ + ID: workerID, + Hostname: hostname, + client: client, + config: cfg, + executor: exec, + loader: loader, + printer: p, + ipAddress: getOutboundIP(), + sshEnabled: opts.SSHEnabled, + sshKeysPath: opts.SSHKeysPath, + alias: opts.Alias, + } + + if opts.GetPublicIP { + w.publicIP = fetchPublicIP() + if w.publicIP != "" { + p.Info("Detected public IP: %s", terminal.Cyan(w.publicIP)) + } else { + p.Warning("Could not detect public IP") + } + } + + // Default alias: wosm- or wosm- + if w.alias == "" { + if w.publicIP != "" { + w.alias = fmt.Sprintf("wosm-%s", w.publicIP) + } else if w.ipAddress != "" { + w.alias = fmt.Sprintf("wosm-%s", w.ipAddress) + } + } + + return w, nil +} + +// getOutboundIP returns the preferred outbound IP address of the machine. +// It uses a UDP dial to 8.8.8.8:80 (no actual packet is sent) to determine the source address. +func getOutboundIP() string { + conn, err := net.Dial("udp", "8.8.8.8:80") + if err != nil { + return "" + } + defer conn.Close() + addr := conn.LocalAddr().(*net.UDPAddr) + return addr.IP.String() +} + +// fetchPublicIP fetches the public IP from ipinfo.io. +func fetchPublicIP() string { + client := &http.Client{Timeout: 10 * time.Second} + req, err := http.NewRequest("GET", "https://ipinfo.io/ip", nil) + if err != nil { + return "" + } + req.Header.Set("User-Agent", core.DefaultUA) + resp, err := client.Do(req) + if err != nil { + return "" + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return "" + } + return strings.TrimSpace(string(body)) } // Run starts the worker loop @@ -75,7 +162,7 @@ func (w *Worker) Run(ctx context.Context) error { w.registerDistributedHooks() defer w.unregisterDistributedHooks() - w.printer.Success("Worker %s joined successfully", w.ID) + w.printer.Success("Worker %s joined successfully", terminal.Cyan(w.ID)) w.printer.Info("Waiting for tasks...") // Start heartbeat goroutine @@ -83,6 +170,11 @@ func (w *Worker) Run(ctx context.Context) error { defer cancelHeartbeat() go w.heartbeatLoop(heartbeatCtx) + // Start execute listener goroutine for per-worker execute requests + executeCtx, cancelExecute := context.WithCancel(ctx) + defer cancelExecute() + go w.executeListenerLoop(executeCtx) + // Main task loop for { select { @@ -111,6 +203,11 @@ func (w *Worker) register(ctx context.Context) error { Status: "idle", JoinedAt: time.Now(), LastHeartbeat: time.Now(), + IPAddress: w.ipAddress, + PublicIP: w.publicIP, + SSHEnabled: w.sshEnabled, + SSHKeysPath: w.sshKeysPath, + Alias: w.alias, } if err := w.client.RegisterWorker(ctx, info); err != nil { @@ -232,6 +329,7 @@ func (w *Worker) executeTask(ctx context.Context, task *Task) *TaskResult { Workspace: computeWorkspace(task.Target), RunPriority: "high", RunMode: "distributed", + HooksEnabled: workflow.HookCount() > 0, } // Goes through distributed hooks → Redis → master DB _ = database.CreateRun(ctx, run) @@ -297,6 +395,11 @@ func (w *Worker) updateStatus(ctx context.Context, status string, taskID string) LastHeartbeat: time.Now(), TasksComplete: w.tasksComplete, TasksFailed: w.tasksFailed, + IPAddress: w.ipAddress, + PublicIP: w.publicIP, + SSHEnabled: w.sshEnabled, + SSHKeysPath: w.sshKeysPath, + Alias: w.alias, } if err := w.client.RegisterWorker(ctx, info); err != nil { w.printer.Warning("Failed to update worker status: %s", err) @@ -322,6 +425,91 @@ func (w *Worker) GetClient() *Client { return w.client } +// ============================================================================= +// Execute Listener (per-worker execute queue) +// ============================================================================= + +// executeListenerLoop polls the per-worker execute queue for requests routed by the master. +func (w *Worker) executeListenerLoop(ctx context.Context) { + key := KeyDataExecuteForWorker(w.ID) + for { + select { + case <-ctx.Done(): + return + default: + envelope, err := w.client.PopData(ctx, key, TaskPollTimeout) + if err != nil { + if ctx.Err() != nil { + return + } + w.printer.Warning("Execute listener error: %s", err) + time.Sleep(time.Second) + continue + } + if envelope == nil { + continue + } + w.processExecuteRequest(ctx, envelope) + } + } +} + +// processExecuteRequest handles an execute request received on the worker's execute queue. +func (w *Worker) processExecuteRequest(ctx context.Context, envelope *DataEnvelope) { + var req ExecuteRequest + if err := json.Unmarshal(envelope.Data, &req); err != nil { + w.printer.Warning("Failed to unmarshal execute request: %s", err) + return + } + + executeType := req.ExecuteType + if executeType == "" { + executeType = req.Action + } + + w.printer.Info("Processing execute request: type=%s from=%s", terminal.Yellow(executeType), terminal.Cyan(envelope.WorkerID)) + + switch executeType { + case "func": + expr := req.Data + if expr == "" { + expr = req.Expr + } + execCtx := executor.BuildBuiltinVariables(w.config, nil) + registry := functions.NewRegistry() + if _, err := registry.Execute(expr, execCtx); err != nil { + w.printer.Warning("Execute func failed: %s (expr: %s)", err, expr) + } + + case "run": + workflow := req.Data + if workflow == "" { + workflow = req.Workflow + } + task := NewTask(uuid.NewString()[:8], workflow, "module", req.Target, nil) + result := w.executeTask(ctx, task) + if result.Status == TaskStatusFailed { + w.printer.Warning("Execute run failed: %s", result.Error) + } + + case "bash": + command := req.Data + if command == "" { + command = req.Expr + } + // @NOTE: This is intentional - execute requests come from trusted workflow YAML files + // via the distributed system. The master routes requests from run_on_worker() calls. + cmd := exec.CommandContext(ctx, "sh", "-c", command) + output, err := cmd.CombinedOutput() + if err != nil { + w.printer.Warning("Execute bash failed: %s (output: %s)", err, string(output)) + } + + default: + w.printer.Warning("Unknown execute type: %s", executeType) + } +} + // ============================================================================= // Helpers // ============================================================================= @@ -375,36 +563,138 @@ func (w *Worker) SendArtifact(ctx context.Context, artifact *database.Artifact) return w.client.PushData(ctx, KeyDataArtifacts, "artifact", artifact, w.ID) } +// SendExecuteRequest sends an execute request to the master via Redis queue +func (w *Worker) SendExecuteRequest(ctx context.Context, action, expr, workflow, target, params, targetRole, targetScope string) error { + req := buildExecuteRequest(action, expr, workflow, target, params, targetRole, targetScope) + return w.client.PushData(ctx, KeyDataExecute, "execute", req, w.ID) +} + +// buildExecuteRequest creates an ExecuteRequest with both new and legacy fields populated. +func buildExecuteRequest(action, expr, workflow, target, params, targetRole, targetScope string) *ExecuteRequest { + if targetRole == "" { + targetRole = "master" + } + data := expr + if action == "run" { + data = workflow + } + return &ExecuteRequest{ + ExecuteType: action, + TargetRole: targetRole, + Data: data, + Target: target, + Params: params, + TargetScope: targetScope, + // Legacy fields for backward compatibility + Action: action, + Expr: expr, + Workflow: workflow, + } +} + // ============================================================================= // Distributed Hooks Registration // ============================================================================= // registerDistributedHooks registers callbacks for database writes to use Redis queues func (w *Worker) registerDistributedHooks() { + w.unregisterHooks = RegisterDistributedHooksFromClient(w.client, w.ID) + w.printer.Info("Registered distributed hooks for database writes") +} + +// unregisterDistributedHooks removes the distributed hooks +func (w *Worker) unregisterDistributedHooks() { + if w.unregisterHooks != nil { + w.unregisterHooks() + } + w.printer.Info("Unregistered distributed hooks") +} + +// RegisterDistributedHooksFromClient registers distributed hooks using a bare +// Client and workerID, without requiring the full Worker struct. This is useful +// for one-shot operations (e.g., worker eval) that need run_on_master() routing +// without the full worker lifecycle (heartbeat, task loop, master registration). +// Returns a cleanup function that unregisters all hooks. +func RegisterDistributedHooksFromClient(client *Client, workerID string) func() { hooks := &database.DistributedHooks{ SendRun: func(ctx context.Context, run *database.Run) error { - return w.SendRunData(ctx, run) + return client.PushData(ctx, KeyDataRuns, "run", run, workerID) }, SendStepResult: func(ctx context.Context, step *database.StepResult) error { - return w.SendStepResult(ctx, step) + return client.PushData(ctx, KeyDataSteps, "step", step, workerID) }, SendEventLog: func(ctx context.Context, event *database.EventLog) error { - return w.SendEventLog(ctx, event) + return client.PushData(ctx, KeyDataEvents, "event", event, workerID) }, SendArtifact: func(ctx context.Context, artifact *database.Artifact) error { - return w.SendArtifact(ctx, artifact) + return client.PushData(ctx, KeyDataArtifacts, "artifact", artifact, workerID) }, ShouldUseRedis: func() bool { return config.ShouldUseRedisDataQueues() }, } database.RegisterDistributedHooks(hooks) - w.printer.Info("Registered distributed hooks for database writes") -} -// unregisterDistributedHooks removes the distributed hooks -func (w *Worker) unregisterDistributedHooks() { - database.UnregisterDistributedHooks() - config.SetWorkerMode(false, "") - w.printer.Info("Unregistered distributed hooks") + // Register execute hooks for run_on_master() and run_on_worker() functions + execHooks := &functions.ExecuteHooks{ + SendExecuteRequest: func(ctx context.Context, action, expr, workflow, target, params, targetRole, targetScope string) error { + req := buildExecuteRequest(action, expr, workflow, target, params, targetRole, targetScope) + return client.PushData(ctx, KeyDataExecute, "execute", req, workerID) + }, + ShouldUseRedis: func() bool { + return config.ShouldUseRedisDataQueues() + }, + ResolveWorkerSSH: func(ctx context.Context, identifier string) (*functions.WorkerSSHInfo, error) { + // Try by ID first + w, err := client.GetWorker(ctx, identifier) + if err != nil { + return nil, fmt.Errorf("failed to look up worker %q: %w", identifier, err) + } + // Try by alias if not found by ID + if w == nil { + w, err = client.GetWorkerByAlias(ctx, identifier) + if err != nil { + return nil, fmt.Errorf("failed to look up worker by alias %q: %w", identifier, err) + } + } + // Try by PublicIP if still not found + if w == nil { + workers, err := client.GetAllWorkers(ctx) + if err != nil { + return nil, fmt.Errorf("failed to list workers: %w", err) + } + for _, cand := range workers { + if cand.PublicIP == identifier || cand.IPAddress == identifier { + w = cand + break + } + } + } + if w == nil { + return nil, fmt.Errorf("worker %q not found", identifier) + } + if !w.SSHEnabled { + return nil, fmt.Errorf("worker %q does not have SSH enabled", identifier) + } + host := w.PublicIP + if host == "" { + host = w.IPAddress + } + return &functions.WorkerSSHInfo{ + ID: w.ID, + Host: host, + User: "root", + KeyPath: w.SSHKeysPath, + Alias: w.Alias, + Port: 22, + }, nil + }, + } + functions.RegisterExecuteHooks(execHooks) + + return func() { + database.UnregisterDistributedHooks() + functions.UnregisterExecuteHooks() + config.SetWorkerMode(false, "") + } } diff --git a/internal/executor/bash_executor.go b/internal/executor/bash_executor.go index 78fe5dd..e0fc938 100644 --- a/internal/executor/bash_executor.go +++ b/internal/executor/bash_executor.go @@ -2,6 +2,7 @@ package executor import ( "context" + "errors" "fmt" "os" "os/exec" @@ -103,17 +104,25 @@ func (e *BashExecutor) Execute(ctx context.Context, step *core.Step, execCtx *co return result, err } + // Extract binaries path for fallback resolution + binariesPath := "" + if bp, ok := execCtx.GetVariable("Binaries"); ok { + if bpStr, ok := bp.(string); ok { + binariesPath = bpStr + } + } + var output string // Determine execution mode if len(step.ParallelCommands) > 0 { - output, err = e.executeParallel(ctx, step.ParallelCommands, timeout) + output, err = e.executeParallel(ctx, step.ParallelCommands, timeout, binariesPath) } else if len(step.Commands) > 0 { - output, err = e.executeSequential(ctx, step.Commands, timeout) + output, err = e.executeSequential(ctx, step.Commands, timeout, binariesPath) } else if step.Command != "" { // Assemble command with structured args if present finalCmd := assembleCommand(step.Command, step.SpeedArgs, step.ConfigArgs, step.InputArgs, step.OutputArgs) - output, err = e.executeCommand(ctx, finalCmd, timeout) + output, err = e.executeCommandWithFallback(ctx, finalCmd, timeout, binariesPath) } else { err = fmt.Errorf("no command specified") } @@ -179,7 +188,7 @@ func (e *BashExecutor) executeCommand(ctx context.Context, command string, timeo } if result.ExitCode != 0 { metrics.RecordToolExecution(toolName, "failed", duration) - return result.Output, fmt.Errorf("command exited with code %d", result.ExitCode) + return result.Output, newExitCodeErrorf(result.ExitCode, "command exited with code %d", result.ExitCode) } metrics.RecordToolExecution(toolName, "success", duration) return strings.TrimSpace(result.Output), nil @@ -216,6 +225,13 @@ func (e *BashExecutor) executeCommand(ctx context.Context, command string, timeo } if err != nil { + // Extract exit code if available for fallback detection + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + exitCode := exitErr.ExitCode() + metrics.RecordToolExecution(toolName, "failed", duration) + return output, newExitCodeErrorf(exitCode, "command failed with exit code %d\nstderr: %s", exitCode, string(stderr.Bytes())) + } metrics.RecordToolExecution(toolName, "error", duration) return output, fmt.Errorf("command failed: %w\nstderr: %s", err, string(stderr.Bytes())) } @@ -224,12 +240,60 @@ func (e *BashExecutor) executeCommand(ctx context.Context, command string, timeo return strings.TrimSpace(output), nil } +// executeCommandWithFallback wraps executeCommand with automatic retry on exit code 127. +// Fallback 1: strip timeout prefix. Fallback 2: prepend binariesPath to the binary. +func (e *BashExecutor) executeCommandWithFallback(ctx context.Context, command string, timeout time.Duration, binariesPath string) (string, error) { + output, err := e.executeCommand(ctx, command, timeout) + if err == nil { + return output, nil + } + + // Only attempt fallback on exit code 127 (command not found) + var ecErr *exitCodeError + if !errors.As(err, &ecErr) || ecErr.code != 127 { + return output, err + } + + // Don't retry if context is already cancelled + if ctx.Err() != nil { + return output, err + } + + currentCmd := command + + // Fallback 1: strip timeout prefix + if result := stripTimeoutPrefix(currentCmd); result.stripped && result.command != "" { + // Use parsed duration from timeout prefix as fallback if step timeout is not set + retryTimeout := timeout + if retryTimeout == 0 && result.duration > 0 { + retryTimeout = result.duration + } + output, err = e.executeCommand(ctx, result.command, retryTimeout) + if err == nil { + return output, nil + } + // Check if still 127 for next fallback + if !errors.As(err, &ecErr) || ecErr.code != 127 { + return output, err + } + currentCmd = result.command + } + + // Fallback 2: prepend binaries path + if prepended, ok := prependBinariesPath(currentCmd, binariesPath); ok { + output, err = e.executeCommand(ctx, prepended, timeout) + return output, err + } + + return output, err +} + // executeSequential executes commands sequentially -func (e *BashExecutor) executeSequential(ctx context.Context, commands []string, timeout time.Duration) (string, error) { +func (e *BashExecutor) executeSequential(ctx context.Context, commands []string, timeout time.Duration, binariesPath string) (string, error) { var outputs []string for _, cmd := range commands { - output, err := e.executeCommand(ctx, cmd, timeout) + output, err := e.executeCommandWithFallback(ctx, cmd, timeout, binariesPath) outputs = append(outputs, output) if err != nil { return strings.Join(outputs, "\n"), err @@ -241,7 +305,7 @@ func (e *BashExecutor) executeSequential(ctx context.Context, commands []string, // executeParallel executes commands in parallel with bounded concurrency. // Uses a worker pool capped at runtime.NumCPU()*2 to prevent unbounded goroutine/memory growth. -func (e *BashExecutor) executeParallel(ctx context.Context, commands []string, timeout time.Duration) (string, error) { +func (e *BashExecutor) executeParallel(ctx context.Context, commands []string, timeout time.Duration, binariesPath string) (string, error) { type result struct { index int output string @@ -268,7 +332,7 @@ func (e *BashExecutor) executeParallel(ctx context.Context, commands []string, t go func() { defer wg.Done() for work := range workQueue { - output, err := e.executeCommand(ctx, work.command, timeout) + output, err := e.executeCommandWithFallback(ctx, work.command, timeout, binariesPath) results <- result{index: work.index, output: output, err: err} } }() diff --git a/internal/executor/cmd_fallback.go b/internal/executor/cmd_fallback.go new file mode 100644 index 0000000..ea72419 --- /dev/null +++ b/internal/executor/cmd_fallback.go @@ -0,0 +1,176 @@ +package executor + +import ( + "fmt" + "strconv" + "strings" + "time" +) + +// exitCodeError wraps an error with its numeric exit code so callers +// can check for specific codes (e.g., 127 = command not found). +type exitCodeError struct { + code int + msg string +} + +func (e *exitCodeError) Error() string { + return e.msg +} + +// newExitCodeError creates an exitCodeError with the given code and message. +func newExitCodeError(code int, msg string) *exitCodeError { + return &exitCodeError{code: code, msg: msg} +} + +// newExitCodeErrorf creates an exitCodeError with formatted message. +func newExitCodeErrorf(code int, format string, args ...any) *exitCodeError { + return &exitCodeError{code: code, msg: fmt.Sprintf(format, args...)} +} + +// stripTimeoutResult holds the result of parsing a timeout prefix. +type stripTimeoutResult struct { + command string // remaining command after stripping prefix + duration time.Duration // parsed duration from the timeout prefix (0 if unparseable) + stripped bool // true if a timeout prefix was found and stripped +} + +// stripTimeoutPrefix removes a "timeout" command prefix from a command string. +// It handles various flag forms: -k VAL, --kill-after=VAL, -s SIG, --signal=SIG, +// --foreground, --preserve-status, -v, --verbose. +// Returns a stripTimeoutResult with the remaining command, parsed duration, and whether stripping occurred. +func stripTimeoutPrefix(command string) stripTimeoutResult { + fields := strings.Fields(command) + if len(fields) == 0 { + return stripTimeoutResult{command: command} + } + + // First token must be "timeout" (or a path ending in /timeout) + base := fields[0] + if idx := strings.LastIndex(base, "/"); idx >= 0 { + base = base[idx+1:] + } + if base != "timeout" { + return stripTimeoutResult{command: command} + } + + i := 1 // skip "timeout" + + // Parse optional flags before the DURATION argument + for i < len(fields) { + arg := fields[i] + + // Flags that take a separate value: -k VAL, -s SIG + if arg == "-k" || arg == "--kill-after" || arg == "-s" || arg == "--signal" { + i += 2 // skip flag + value + continue + } + + // Flags with = form: --kill-after=VAL, --signal=SIG + if strings.HasPrefix(arg, "--kill-after=") || strings.HasPrefix(arg, "--signal=") { + i++ + continue + } + + // Short form -k5s (value attached) + if len(arg) > 2 && arg[0] == '-' && arg[1] == 'k' { + i++ + continue + } + + // Boolean flags + if arg == "--foreground" || arg == "--preserve-status" || arg == "-v" || arg == "--verbose" { + i++ + continue + } + + // Not a recognized flag — this should be the DURATION + break + } + + // Skip the DURATION argument + if i >= len(fields) { + // No duration found — malformed, don't strip + return stripTimeoutResult{command: command} + } + durationStr := fields[i] + parsedDuration := parseTimeoutDuration(durationStr) + i++ // skip duration + + // Everything after duration is the actual command + if i >= len(fields) { + // Nothing after duration — no command to run + return stripTimeoutResult{command: "", duration: parsedDuration, stripped: true} + } + + return stripTimeoutResult{ + command: strings.Join(fields[i:], " "), + duration: parsedDuration, + stripped: true, + } +} + +// parseTimeoutDuration parses a GNU coreutils timeout duration string. +// Supports formats: plain number (seconds), or number with suffix s/m/h/d. +// Returns 0 if the string cannot be parsed. +func parseTimeoutDuration(s string) time.Duration { + if s == "" { + return 0 + } + + // Check for suffix + last := s[len(s)-1] + switch last { + case 's': + return parseDurationNumber(s[:len(s)-1], time.Second) + case 'm': + return parseDurationNumber(s[:len(s)-1], time.Minute) + case 'h': + return parseDurationNumber(s[:len(s)-1], time.Hour) + case 'd': + return parseDurationNumber(s[:len(s)-1], 24*time.Hour) + default: + // No suffix — default is seconds + return parseDurationNumber(s, time.Second) + } +} + +// parseDurationNumber parses a numeric string and multiplies by the given unit. +// Supports both integer and floating-point values. Returns 0 on parse error. +func parseDurationNumber(s string, unit time.Duration) time.Duration { + if s == "" { + return 0 + } + f, err := strconv.ParseFloat(s, 64) + if err != nil { + return 0 + } + return time.Duration(f * float64(unit)) +} + +// prependBinariesPath prepends binariesPath to the first token (binary name) +// of a command if it doesn't already contain a path separator. +// Returns the modified command and true if the path was prepended. +func prependBinariesPath(command, binariesPath string) (string, bool) { + if binariesPath == "" || command == "" { + return command, false + } + + fields := strings.Fields(command) + if len(fields) == 0 { + return command, false + } + + binary := fields[0] + + // Don't prepend if binary already has a path + if strings.Contains(binary, "/") { + return command, false + } + + // Ensure binariesPath doesn't have trailing slash + binariesPath = strings.TrimRight(binariesPath, "/") + + fields[0] = binariesPath + "/" + binary + return strings.Join(fields, " "), true +} diff --git a/internal/executor/cmd_fallback_test.go b/internal/executor/cmd_fallback_test.go new file mode 100644 index 0000000..c9a6510 --- /dev/null +++ b/internal/executor/cmd_fallback_test.go @@ -0,0 +1,512 @@ +package executor + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/j3ssie/osmedeus/v5/internal/core" + "github.com/j3ssie/osmedeus/v5/internal/runner" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- stripTimeoutPrefix tests --- + +func TestStripTimeoutPrefix_BasicTimeout(t *testing.T) { + r := stripTimeoutPrefix("timeout 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) + assert.Equal(t, 30*time.Second, r.duration) +} + +func TestStripTimeoutPrefix_WithKillAfterFlag(t *testing.T) { + r := stripTimeoutPrefix("timeout -k 10 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) + assert.Equal(t, 30*time.Second, r.duration) +} + +func TestStripTimeoutPrefix_WithKillAfterEquals(t *testing.T) { + r := stripTimeoutPrefix("timeout --kill-after=10s 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_WithSignalFlag(t *testing.T) { + r := stripTimeoutPrefix("timeout -s SIGKILL 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_WithSignalEquals(t *testing.T) { + r := stripTimeoutPrefix("timeout --signal=TERM 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_WithForeground(t *testing.T) { + r := stripTimeoutPrefix("timeout --foreground 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_WithPreserveStatus(t *testing.T) { + r := stripTimeoutPrefix("timeout --preserve-status 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_WithVerbose(t *testing.T) { + r := stripTimeoutPrefix("timeout -v 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_MultipleFlags(t *testing.T) { + r := stripTimeoutPrefix("timeout -k 10 -s SIGKILL --foreground 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_ShortKillAttached(t *testing.T) { + r := stripTimeoutPrefix("timeout -k5s 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_NoCommandAfterDuration(t *testing.T) { + r := stripTimeoutPrefix("timeout 30") + assert.True(t, r.stripped) + assert.Equal(t, "", r.command) + assert.Equal(t, 30*time.Second, r.duration) +} + +func TestStripTimeoutPrefix_NotTimeout(t *testing.T) { + r := stripTimeoutPrefix("nuclei -t templates") + assert.False(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) +} + +func TestStripTimeoutPrefix_EmptyString(t *testing.T) { + r := stripTimeoutPrefix("") + assert.False(t, r.stripped) + assert.Equal(t, "", r.command) +} + +func TestStripTimeoutPrefix_FullPathTimeout(t *testing.T) { + r := stripTimeoutPrefix("/usr/bin/timeout 30 nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) + assert.Equal(t, 30*time.Second, r.duration) +} + +func TestStripTimeoutPrefix_DurationWithSuffix(t *testing.T) { + r := stripTimeoutPrefix("timeout 2h nuclei -t templates") + assert.True(t, r.stripped) + assert.Equal(t, "nuclei -t templates", r.command) + assert.Equal(t, 2*time.Hour, r.duration) +} + +// --- parseTimeoutDuration tests --- + +func TestParseTimeoutDuration_Seconds(t *testing.T) { + assert.Equal(t, 30*time.Second, parseTimeoutDuration("30")) +} + +func TestParseTimeoutDuration_SecondsWithSuffix(t *testing.T) { + assert.Equal(t, 30*time.Second, parseTimeoutDuration("30s")) +} + +func TestParseTimeoutDuration_Minutes(t *testing.T) { + assert.Equal(t, 5*time.Minute, parseTimeoutDuration("5m")) +} + +func TestParseTimeoutDuration_Hours(t *testing.T) { + assert.Equal(t, 2*time.Hour, parseTimeoutDuration("2h")) +} + +func TestParseTimeoutDuration_Days(t *testing.T) { + assert.Equal(t, 24*time.Hour, parseTimeoutDuration("1d")) +} + +func TestParseTimeoutDuration_Float(t *testing.T) { + assert.Equal(t, time.Duration(1.5*float64(time.Second)), parseTimeoutDuration("1.5")) +} + +func TestParseTimeoutDuration_FloatWithSuffix(t *testing.T) { + assert.Equal(t, time.Duration(2.5*float64(time.Minute)), parseTimeoutDuration("2.5m")) +} + +func TestParseTimeoutDuration_Empty(t *testing.T) { + assert.Equal(t, time.Duration(0), parseTimeoutDuration("")) +} + +func TestParseTimeoutDuration_Invalid(t *testing.T) { + assert.Equal(t, time.Duration(0), parseTimeoutDuration("abc")) +} + +// --- prependBinariesPath tests --- + +func TestPrependBinariesPath_Normal(t *testing.T) { + cmd, ok := prependBinariesPath("nuclei -t templates", "/opt/tools") + assert.True(t, ok) + assert.Equal(t, "/opt/tools/nuclei -t templates", cmd) +} + +func TestPrependBinariesPath_AlreadyHasPath(t *testing.T) { + cmd, ok := prependBinariesPath("/usr/bin/nuclei -t templates", "/opt/tools") + assert.False(t, ok) + assert.Equal(t, "/usr/bin/nuclei -t templates", cmd) +} + +func TestPrependBinariesPath_RelativePath(t *testing.T) { + cmd, ok := prependBinariesPath("./nuclei -t templates", "/opt/tools") + assert.False(t, ok) + assert.Equal(t, "./nuclei -t templates", cmd) +} + +func TestPrependBinariesPath_EmptyBinariesPath(t *testing.T) { + cmd, ok := prependBinariesPath("nuclei -t templates", "") + assert.False(t, ok) + assert.Equal(t, "nuclei -t templates", cmd) +} + +func TestPrependBinariesPath_EmptyCommand(t *testing.T) { + cmd, ok := prependBinariesPath("", "/opt/tools") + assert.False(t, ok) + assert.Equal(t, "", cmd) +} + +func TestPrependBinariesPath_TrailingSlash(t *testing.T) { + cmd, ok := prependBinariesPath("nuclei -t templates", "/opt/tools/") + assert.True(t, ok) + assert.Equal(t, "/opt/tools/nuclei -t templates", cmd) +} + +func TestPrependBinariesPath_SingleWord(t *testing.T) { + cmd, ok := prependBinariesPath("nuclei", "/opt/tools") + assert.True(t, ok) + assert.Equal(t, "/opt/tools/nuclei", cmd) +} + +// --- exitCodeError tests --- + +func TestExitCodeError_ErrorMessage(t *testing.T) { + err := newExitCodeErrorf(127, "command not found: %s", "nuclei") + assert.Equal(t, "command not found: nuclei", err.Error()) + assert.Equal(t, 127, err.code) +} + +func TestExitCodeError_ErrorsAs(t *testing.T) { + err := newExitCodeErrorf(127, "command not found") + var ecErr *exitCodeError + require.True(t, errors.As(err, &ecErr)) + assert.Equal(t, 127, ecErr.code) +} + +func TestExitCodeError_NotExitCodeError(t *testing.T) { + err := errors.New("some other error") + var ecErr *exitCodeError + assert.False(t, errors.As(err, &ecErr)) +} + +func TestNewExitCodeError(t *testing.T) { + err := newExitCodeError(1, "failed") + assert.Equal(t, 1, err.code) + assert.Equal(t, "failed", err.Error()) +} + +// --- mockRunner for integration tests --- + +type mockFallbackRunner struct { + // responses maps command strings to their mock results + responses map[string]*mockResponse +} + +type mockResponse struct { + output string + exitCode int +} + +func newMockFallbackRunner() *mockFallbackRunner { + return &mockFallbackRunner{ + responses: make(map[string]*mockResponse), + } +} + +func (m *mockFallbackRunner) addResponse(command string, output string, exitCode int) { + m.responses[command] = &mockResponse{output: output, exitCode: exitCode} +} + +func (m *mockFallbackRunner) Execute(_ context.Context, command string) (*runner.CommandResult, error) { + if resp, ok := m.responses[command]; ok { + return &runner.CommandResult{ + Output: resp.output, + ExitCode: resp.exitCode, + }, nil + } + // Default: command not found + return &runner.CommandResult{ + Output: "sh: command not found", + ExitCode: 127, + }, nil +} + +func (m *mockFallbackRunner) Setup(_ context.Context) error { return nil } +func (m *mockFallbackRunner) Cleanup(_ context.Context) error { return nil } +func (m *mockFallbackRunner) Type() core.RunnerType { return core.RunnerTypeHost } +func (m *mockFallbackRunner) IsRemote() bool { return false } +func (m *mockFallbackRunner) SetPIDCallbacks(_, _ runner.PIDCallback) {} +func (m *mockFallbackRunner) CopyFromRemote(_ context.Context, _, _ string) error { + return nil +} + +// --- Integration tests: BashExecutor fallback chain --- + +func TestBashExecutor_Fallback_StripTimeout(t *testing.T) { + mock := newMockFallbackRunner() + // Original command with timeout fails (127) + // Stripped command succeeds + mock.addResponse("nuclei -t templates", "found 5 results", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeCommandWithFallback( + context.Background(), + "timeout 30 nuclei -t templates", + 0, + "", + ) + + assert.NoError(t, err) + assert.Equal(t, "found 5 results", output) +} + +func TestBashExecutor_Fallback_PrependBinariesPath(t *testing.T) { + mock := newMockFallbackRunner() + // Original command fails (127) + // Prepended command succeeds + mock.addResponse("/opt/tools/nuclei -t templates", "found 3 results", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeCommandWithFallback( + context.Background(), + "nuclei -t templates", + 0, + "/opt/tools", + ) + + assert.NoError(t, err) + assert.Equal(t, "found 3 results", output) +} + +func TestBashExecutor_Fallback_StripTimeoutThenPrependBinaries(t *testing.T) { + mock := newMockFallbackRunner() + // Original "timeout 30 nuclei -t templates" → 127 + // Stripped "nuclei -t templates" → still 127 + // Prepended "/opt/tools/nuclei -t templates" → success + mock.addResponse("/opt/tools/nuclei -t templates", "found 7 results", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeCommandWithFallback( + context.Background(), + "timeout 30 nuclei -t templates", + 0, + "/opt/tools", + ) + + assert.NoError(t, err) + assert.Equal(t, "found 7 results", output) +} + +func TestBashExecutor_Fallback_NoFallbackOnNon127(t *testing.T) { + mock := newMockFallbackRunner() + // Command fails with exit code 1 (not 127) + mock.addResponse("nuclei -t templates", "error output", 1) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + _, err := exec.executeCommandWithFallback( + context.Background(), + "nuclei -t templates", + 0, + "/opt/tools", + ) + + require.Error(t, err) + var ecErr *exitCodeError + require.True(t, errors.As(err, &ecErr)) + assert.Equal(t, 1, ecErr.code) +} + +func TestBashExecutor_Fallback_SuccessNoFallback(t *testing.T) { + mock := newMockFallbackRunner() + // Command succeeds on first try + mock.addResponse("nuclei -t templates", "all good", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeCommandWithFallback( + context.Background(), + "nuclei -t templates", + 0, + "/opt/tools", + ) + + assert.NoError(t, err) + assert.Equal(t, "all good", output) +} + +func TestBashExecutor_Fallback_AllFallbacksFail(t *testing.T) { + mock := newMockFallbackRunner() + // All commands return 127 (everything is command not found) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + _, err := exec.executeCommandWithFallback( + context.Background(), + "timeout 30 nuclei -t templates", + 0, + "/opt/tools", + ) + + require.Error(t, err) + var ecErr *exitCodeError + require.True(t, errors.As(err, &ecErr)) + assert.Equal(t, 127, ecErr.code) +} + +func TestBashExecutor_Fallback_CancelledContextNoRetry(t *testing.T) { + mock := newMockFallbackRunner() + // All responses return 127 + // But we expect no fallback because context is cancelled + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel immediately + + _, err := exec.executeCommandWithFallback( + ctx, + "timeout 30 nuclei -t templates", + 0, + "/opt/tools", + ) + + // Should fail due to context cancellation, not 127 fallback + require.Error(t, err) +} + +func TestBashExecutor_Fallback_TimeoutFlagsWithEquals(t *testing.T) { + mock := newMockFallbackRunner() + // timeout with --kill-after= and --signal= flags + mock.addResponse("nuclei -t templates", "found results", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeCommandWithFallback( + context.Background(), + "timeout --kill-after=10s --signal=TERM 30 nuclei -t templates", + 0, + "", + ) + + assert.NoError(t, err) + assert.Equal(t, "found results", output) +} + +func TestBashExecutor_Fallback_BinaryWithPathNoFallback2(t *testing.T) { + mock := newMockFallbackRunner() + // Binary already has a path prefix — fallback 2 (prepend) should not apply + // Fallback 1 (strip timeout) should work + mock.addResponse("/usr/local/bin/nuclei -t templates", "found", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeCommandWithFallback( + context.Background(), + "timeout 30 /usr/local/bin/nuclei -t templates", + 0, + "/opt/tools", + ) + + assert.NoError(t, err) + assert.Equal(t, "found", output) +} + +func TestBashExecutor_Fallback_SequentialUsesWithFallback(t *testing.T) { + mock := newMockFallbackRunner() + // First command needs fallback (strip timeout) + mock.addResponse("nuclei -t templates", "result1", 0) + // Second command succeeds directly + mock.addResponse("echo done", "done", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeSequential( + context.Background(), + []string{"timeout 30 nuclei -t templates", "echo done"}, + 0, + "", + ) + + assert.NoError(t, err) + assert.Contains(t, output, "result1") + assert.Contains(t, output, "done") +} + +func TestBashExecutor_Fallback_ParallelUsesWithFallback(t *testing.T) { + mock := newMockFallbackRunner() + mock.addResponse("nuclei -t templates", "result1", 0) + mock.addResponse("echo done", "done", 0) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + output, err := exec.executeParallel( + context.Background(), + []string{"timeout 30 nuclei -t templates", "echo done"}, + 0, + "", + ) + + assert.NoError(t, err) + assert.Contains(t, output, "result1") + assert.Contains(t, output, "done") +} + +func TestBashExecutor_Fallback_StripTimeoutNonZeroExitAfterStrip(t *testing.T) { + mock := newMockFallbackRunner() + // After stripping timeout, command fails with non-127 exit code + mock.addResponse("nuclei -t templates", "some error", 2) + + exec := NewBashExecutor(nil) + exec.SetRunner(mock) + + _, err := exec.executeCommandWithFallback( + context.Background(), + "timeout 30 nuclei -t templates", + time.Duration(0), + "/opt/tools", + ) + + // Should return the exit code 2 error, not try fallback 2 + require.Error(t, err) + var ecErr *exitCodeError + require.True(t, errors.As(err, &ecErr)) + assert.Equal(t, 2, ecErr.code) +} diff --git a/internal/executor/executor.go b/internal/executor/executor.go index 66a8e6d..82cc20c 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -1159,6 +1159,14 @@ func (e *Executor) ExecuteModule(ctx context.Context, module *core.Workflow, par } } + // Set HooksEnabled variable for state export + execCtx.SetVariable("HooksEnabled", module.HookCount() > 0) + + // Execute pre-scan hook steps + if module.Hooks != nil && len(module.Hooks.PreScanSteps) > 0 { + e.executeHookSteps(ctx, "pre_scan_steps", module.Hooks.PreScanSteps, execCtx) + } + // Execute steps e.logger.Debug("Starting step execution loop", zap.Int("total_steps", len(module.Steps)), @@ -1297,6 +1305,11 @@ func (e *Executor) ExecuteModule(ctx context.Context, module *core.Workflow, par } } + // Execute post-scan hook steps + if module.Hooks != nil && len(module.Hooks.PostScanSteps) > 0 { + e.executeHookSteps(ctx, "post_scan_steps", module.Hooks.PostScanSteps, execCtx) + } + result.Status = core.RunStatusCompleted result.EndTime = time.Now() result.Exports = execCtx.Exports @@ -1816,6 +1829,23 @@ func (e *Executor) ExecuteFlow(ctx context.Context, flow *core.Workflow, params } } + // Set HooksEnabled variable for state export + execCtx.SetVariable("HooksEnabled", flow.HookCount() > 0) + + // Execute pre-scan hook steps + if flow.Hooks != nil && len(flow.Hooks.PreScanSteps) > 0 { + // Flow needs a runner for hook steps — create a host runner + binaryPath, _ := os.Executable() + hookRunner, hookErr := runner.NewRunner(flow, binaryPath) + if hookErr == nil { + if setupErr := hookRunner.Setup(ctx); setupErr == nil { + e.stepDispatcher.SetRunner(hookRunner) + e.executeHookSteps(ctx, "pre_scan_steps", flow.Hooks.PreScanSteps, execCtx) + _ = hookRunner.Cleanup(ctx) + } + } + } + // Parse excluded modules excludeList := parseExcludeList(params["exclude_modules"]) fuzzyExcludeList := parseExcludeList(params["fuzzy_exclude_modules"]) @@ -2077,6 +2107,19 @@ func (e *Executor) ExecuteFlow(ctx context.Context, flow *core.Workflow, params return result, ctx.Err() } + // Execute post-scan hook steps + if flow.Hooks != nil && len(flow.Hooks.PostScanSteps) > 0 { + binaryPath, _ := os.Executable() + hookRunner, hookErr := runner.NewRunner(flow, binaryPath) + if hookErr == nil { + if setupErr := hookRunner.Setup(ctx); setupErr == nil { + e.stepDispatcher.SetRunner(hookRunner) + e.executeHookSteps(ctx, "post_scan_steps", flow.Hooks.PostScanSteps, execCtx) + _ = hookRunner.Cleanup(ctx) + } + } + } + result.Status = core.RunStatusCompleted result.EndTime = time.Now() result.Exports = execCtx.Exports @@ -2132,6 +2175,30 @@ func (e *Executor) ExecuteFlow(ctx context.Context, flow *core.Workflow, params return result, nil } +// executeHookSteps runs a list of hook steps (pre_scan or post_scan). +// Hook failures are non-fatal — they log a warning and continue. +func (e *Executor) executeHookSteps(ctx context.Context, hookName string, steps []core.Step, execCtx *core.ExecutionContext) { + for i := range steps { + select { + case <-ctx.Done(): + return + default: + } + + step := &steps[i] + if step.Type == "" { + step.Type = core.StepTypeBash + } + _, err := e.executeStep(ctx, step, execCtx) + if err != nil { + execCtx.Logger.Warn("Hook step failed", + zap.String("hook", hookName), + zap.String("step", step.Name), + zap.Error(err)) + } + } +} + // executeStep executes a single step func (e *Executor) executeStep(ctx context.Context, step *core.Step, execCtx *core.ExecutionContext) (*core.StepResult, error) { stepLogger := logger.WithStep(execCtx.WorkflowName, execCtx.RunUUID, step.Name) diff --git a/internal/executor/remote_bash_executor.go b/internal/executor/remote_bash_executor.go index e0cb55d..a1318b3 100644 --- a/internal/executor/remote_bash_executor.go +++ b/internal/executor/remote_bash_executor.go @@ -2,6 +2,7 @@ package executor import ( "context" + "errors" "fmt" "strings" "sync" @@ -87,16 +88,24 @@ func (e *RemoteBashExecutor) Execute(ctx context.Context, step *core.Step, execC _ = r.Cleanup(cleanupCtx) }() + // Extract binaries path for fallback resolution + binariesPath := "" + if bp, ok := execCtx.GetVariable("Binaries"); ok { + if bpStr, ok := bp.(string); ok { + binariesPath = bpStr + } + } + // Execute command(s) using the runner var output string if len(step.ParallelCommands) > 0 { - output, err = e.executeParallel(ctx, r, step.ParallelCommands, timeout) + output, err = e.executeParallel(ctx, r, step.ParallelCommands, timeout, binariesPath) } else if len(step.Commands) > 0 { - output, err = e.executeSequential(ctx, r, step.Commands, timeout) + output, err = e.executeSequential(ctx, r, step.Commands, timeout, binariesPath) } else if step.Command != "" { // Assemble command with structured args if present finalCmd := assembleCommand(step.Command, step.SpeedArgs, step.ConfigArgs, step.InputArgs, step.OutputArgs) - output, err = e.executeCommand(ctx, r, finalCmd, timeout) + output, err = e.executeCommandWithFallback(ctx, r, finalCmd, timeout, binariesPath) } else { err = fmt.Errorf("no command specified") } @@ -186,18 +195,65 @@ func (e *RemoteBashExecutor) executeCommand(ctx context.Context, r runner.Runner } if cmdResult.ExitCode != 0 { - return cmdResult.Output, fmt.Errorf("command exited with code %d", cmdResult.ExitCode) + return cmdResult.Output, newExitCodeErrorf(cmdResult.ExitCode, "command exited with code %d", cmdResult.ExitCode) } return strings.TrimSpace(cmdResult.Output), nil } +// executeCommandWithFallback wraps executeCommand with automatic retry on exit code 127. +// Fallback 1: strip timeout prefix. Fallback 2: prepend binariesPath to the binary. +func (e *RemoteBashExecutor) executeCommandWithFallback(ctx context.Context, r runner.Runner, command string, timeout time.Duration, binariesPath string) (string, error) { + output, err := e.executeCommand(ctx, r, command, timeout) + if err == nil { + return output, nil + } + + // Only attempt fallback on exit code 127 (command not found) + var ecErr *exitCodeError + if !errors.As(err, &ecErr) || ecErr.code != 127 { + return output, err + } + + // Don't retry if context is already cancelled + if ctx.Err() != nil { + return output, err + } + + currentCmd := command + + // Fallback 1: strip timeout prefix + if result := stripTimeoutPrefix(currentCmd); result.stripped && result.command != "" { + // Use parsed duration from timeout prefix as fallback if step timeout is not set + retryTimeout := timeout + if retryTimeout == 0 && result.duration > 0 { + retryTimeout = result.duration + } + output, err = e.executeCommand(ctx, r, result.command, retryTimeout) + if err == nil { + return output, nil + } + if !errors.As(err, &ecErr) || ecErr.code != 127 { + return output, err + } + currentCmd = result.command + } + + // Fallback 2: prepend binaries path + if prepended, ok := prependBinariesPath(currentCmd, binariesPath); ok { + output, err = e.executeCommand(ctx, r, prepended, timeout) + return output, err + } + + return output, err +} + // executeSequential executes commands sequentially -func (e *RemoteBashExecutor) executeSequential(ctx context.Context, r runner.Runner, commands []string, timeout time.Duration) (string, error) { +func (e *RemoteBashExecutor) executeSequential(ctx context.Context, r runner.Runner, commands []string, timeout time.Duration, binariesPath string) (string, error) { var outputs []string for _, cmd := range commands { - output, err := e.executeCommand(ctx, r, cmd, timeout) + output, err := e.executeCommandWithFallback(ctx, r, cmd, timeout, binariesPath) outputs = append(outputs, output) if err != nil { return strings.Join(outputs, "\n"), err @@ -208,7 +264,7 @@ func (e *RemoteBashExecutor) executeSequential(ctx context.Context, r runner.Run } // executeParallel executes commands in parallel -func (e *RemoteBashExecutor) executeParallel(ctx context.Context, r runner.Runner, commands []string, timeout time.Duration) (string, error) { +func (e *RemoteBashExecutor) executeParallel(ctx context.Context, r runner.Runner, commands []string, timeout time.Duration, binariesPath string) (string, error) { type cmdResult struct { index int output string @@ -222,7 +278,7 @@ func (e *RemoteBashExecutor) executeParallel(ctx context.Context, r runner.Runne wg.Add(1) go func(idx int, command string) { defer wg.Done() - output, err := e.executeCommand(ctx, r, command, timeout) + output, err := e.executeCommandWithFallback(ctx, r, command, timeout, binariesPath) results <- cmdResult{index: idx, output: output, err: err} }(i, cmd) } diff --git a/internal/executor/state_export.go b/internal/executor/state_export.go index ed0c22c..71a388c 100644 --- a/internal/executor/state_export.go +++ b/internal/executor/state_export.go @@ -24,6 +24,23 @@ func buildExportContext(result *core.WorkflowResult, execCtx *core.ExecutionCont ctx.WorkspacePath = execCtx.WorkspacePath ctx.WorkspaceName = execCtx.WorkspaceName ctx.Params = execCtx.Params + + // Read hook/run metadata from execution context variables + if v, ok := execCtx.GetVariable("HooksEnabled"); ok { + if b, ok := v.(bool); ok { + ctx.HooksEnabled = b + } + } + if v, ok := execCtx.GetVariable("RunMode"); ok { + if s, ok := v.(string); ok { + ctx.RunMode = s + } + } + if v, ok := execCtx.GetVariable("RunPriority"); ok { + if s, ok := v.(string); ok { + ctx.RunPriority = s + } + } } // Populate/override from result diff --git a/internal/functions/constants.go b/internal/functions/constants.go index 922ee34..ed5208b 100644 --- a/internal/functions/constants.go +++ b/internal/functions/constants.go @@ -111,6 +111,8 @@ const ( FnPickValid = "pick_valid" // pick_valid(v1, v2, ..., v10) -> any (first valid value) FnRunModule = "run_module" // run_module(module, target, params?) -> string (run osmedeus module) FnRunFlow = "run_flow" // run_flow(flow, target, params?) -> string (run osmedeus flow) + FnRunOnMaster = "run_on_master" // run_on_master(action, ...args) -> bool (execute on master node) + FnRunOnWorker = "run_on_worker" // run_on_worker(scope, action, ...args) -> bool (execute on worker nodes) FnExecPython = "exec_python" // exec_python(code) -> string (run inline Python, prefer python3) FnExecPythonFile = "exec_python_file" // exec_python_file(path) -> string (run Python file, prefer python3) ) @@ -228,8 +230,11 @@ const ( // 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 + 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 + FnSyncFromMaster = "sync_from_master" // sync_from_master(src, dest) -> bool + FnSyncFromWorker = "sync_from_worker" // sync_from_worker(identifier, ip, src, dest) -> bool + FnRsyncToWorker = "rsync_to_worker" // rsync_to_worker(identifier, ip, src, dest) -> bool ) // LLM Functions - Invoke LLM from workflows @@ -246,6 +251,12 @@ const ( FnExtractTo = "extract_to" // extract_to(source, dest) -> bool (auto-detect .zip, .tar.gz, .tar.bz2, .tar.xz, .tgz; removes dest first) ) +// Snapshot Functions - Workspace export/import +const ( + FnSnapshotExport = "snapshot_export" // snapshot_export(workspace, dest?) -> string (zip path on success, empty on failure) + FnSnapshotImport = "snapshot_import" // snapshot_import(source) -> string (workspace name on success, empty on failure) +) + // Diff Functions - Compare files const ( FnExtractDiff = "extract_diff" // extract_diff(file1, file2) -> string (lines only in file2) @@ -424,6 +435,8 @@ func AllFunctions() []string { FnPickValid, FnRunModule, FnRunFlow, + FnRunOnMaster, + FnRunOnWorker, FnExecPython, FnExecPythonFile, @@ -512,6 +525,10 @@ func AllFunctions() []string { FnUnzipDir, FnExtractTo, + // Snapshot Functions + FnSnapshotExport, + FnSnapshotImport, + // Diff Functions FnExtractDiff, @@ -612,6 +629,9 @@ func AllFunctions() []string { // SSH functions FnSSHExec, FnSSHRsync, + FnSyncFromMaster, + FnSyncFromWorker, + FnRsyncToWorker, } } @@ -642,6 +662,7 @@ const ( CategoryCDNStorage = "cdn_storage" CategoryUnixCommands = "unix_commands" CategoryArchive = "archive" + CategorySnapshot = "snapshot" CategoryDiff = "diff" CategoryOutput = "output" CategoryURLProcessing = "url_processing" @@ -652,6 +673,7 @@ const ( CategoryTypeDetection = "type_detection" CategoryLLM = "llm" CategorySSH = "ssh" + CategoryDistributed = "distributed" ) // CategoryInfo provides display metadata for a function category @@ -681,6 +703,7 @@ func CategoryOrder() []CategoryInfo { {CategoryCDNStorage, "CDN/Storage Functions", "CDN/Storage"}, {CategoryUnixCommands, "Unix Command Wrappers", "Unix"}, {CategoryArchive, "Archive Functions (Go)", "Archive"}, + {CategorySnapshot, "Snapshot Functions", "Snapshot"}, {CategoryDiff, "Diff Functions", "Diff"}, {CategoryOutput, "Output Functions", "Output"}, {CategoryURLProcessing, "URL Processing Functions", "URL"}, @@ -690,6 +713,7 @@ func CategoryOrder() []CategoryInfo { {CategoryEnvironment, "Environment Functions", "Environment"}, {CategoryTypeDetection, "Type Detection Functions", "Type Detection"}, {CategorySSH, "SSH Functions", "SSH"}, + {CategoryDistributed, "Distributed Functions", "Distributed"}, } } @@ -849,6 +873,10 @@ func FunctionRegistry() map[string][]FunctionInfo { {FnUnzipDir, "unzip_dir(source, dest)", "Unzip archive using Go archive/zip", "bool", "unzip_dir('/tmp/archive.zip', '/tmp/extracted')"}, {FnExtractTo, "extract_to(source, dest)", "Auto-detect archive format (.zip, .tar.gz, .tar.bz2, .tar.xz, .tgz) and extract to dest (removes dest first)", "bool", "extract_to('/tmp/repo.tar.gz', '/tmp/repo')"}, }, + CategorySnapshot: { + {FnSnapshotExport, "snapshot_export(workspace, dest?)", "Export workspace as ZIP snapshot (subprocess)", "string", "snapshot_export('example.com')"}, + {FnSnapshotImport, "snapshot_import(source)", "Import workspace from ZIP snapshot (subprocess, force mode)", "string", "snapshot_import('/path/to/example.com_snapshot.zip')"}, + }, CategoryDiff: { {FnExtractDiff, "extract_diff(file1, file2)", "Lines only in file2 (new content)", "string", "extract_diff('{{Output}}/old-subs.txt', '{{Output}}/new-subs.txt')"}, }, @@ -939,6 +967,13 @@ func FunctionRegistry() map[string][]FunctionInfo { 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')"}, + {FnSyncFromMaster, "sync_from_master(src, dest)", "Request master to rsync file/folder to this worker via SSH", "bool", "sync_from_master('/opt/osmedeus/base/wordlists', '{{BaseFolder}}/wordlists')"}, + {FnSyncFromWorker, "sync_from_worker(identifier, ip, src, dest)", "Pull file/folder from a worker via rsync over SSH", "bool", "sync_from_worker('worker-1', '10.0.0.2', '/opt/output/results.txt', '/tmp/results.txt')"}, + {FnRsyncToWorker, "rsync_to_worker(identifier, ip, src, dest)", "Push file/folder to a worker via rsync over SSH", "bool", "rsync_to_worker('worker-1', '10.0.0.2', '/tmp/data.txt', '/opt/data.txt')"}, + }, + CategoryDistributed: { + {FnRunOnMaster, "run_on_master(action, ...args)", "Execute on master: 'func' runs JS expression, 'run' submits workflow, 'bash' runs shell command", "bool", "run_on_master('bash', 'nmap -sV target.com')"}, + {FnRunOnWorker, "run_on_worker(scope, action, ...args)", "Execute on worker(s): scope='all'/alias/ID/IP; actions: 'func', 'run', 'bash'", "bool", "run_on_worker('all', 'bash', 'apt update && apt install -y nmap')"}, }, CategoryTypeDetection: { {FnGetTypes, "get_types(input)", "Detect input type (file, folder, cidr, ip, url, domain, string)", "string", "get_types('192.168.1.0/24')"}, diff --git a/internal/functions/db_functions_test.go b/internal/functions/db_functions_test.go index 91fc837..88e1342 100644 --- a/internal/functions/db_functions_test.go +++ b/internal/functions/db_functions_test.go @@ -334,7 +334,7 @@ func TestMapJSONToAsset(t *testing.T) { 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.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) } diff --git a/internal/functions/execute_hooks.go b/internal/functions/execute_hooks.go new file mode 100644 index 0000000..af2a4a4 --- /dev/null +++ b/internal/functions/execute_hooks.go @@ -0,0 +1,106 @@ +package functions + +import ( + "context" + "sync" + + "github.com/j3ssie/osmedeus/v5/internal/logger" + "go.uber.org/zap" +) + +// WorkerSSHInfo holds SSH connection details for a worker node. +// Lives in the functions package to avoid circular imports between +// functions and distributed packages. +type WorkerSSHInfo struct { + ID string + Host string // PublicIP or IPAddress + User string // default "root" + KeyPath string // SSHKeysPath + Alias string + Port int // default 22 +} + +// ExecuteHooks provides callback functions for sending execute requests +// to the master node when running in distributed worker mode. +// This follows the same pattern as database.DistributedHooks to avoid +// import cycles between functions and distributed packages. +type ExecuteHooks struct { + // SendExecuteRequest pushes an execute request via Redis. + // targetRole is "master" or "worker" to control where the request is executed. + // targetScope specifies which worker(s) to target: "all", alias, worker ID, or public IP. + SendExecuteRequest func(ctx context.Context, action, expr, workflow, target, params, targetRole, targetScope string) error + + // ShouldUseRedis returns true if we are in distributed worker mode + ShouldUseRedis func() bool + + // ResolveWorkerSSH looks up SSH connection info for a worker by ID, alias, or IP. + ResolveWorkerSSH func(ctx context.Context, identifier string) (*WorkerSSHInfo, error) +} + +var ( + executeHooks *ExecuteHooks + executeHookMu sync.RWMutex +) + +// RegisterExecuteHooks registers callbacks for distributed execute requests. +// Called by the distributed package at worker startup. +func RegisterExecuteHooks(hooks *ExecuteHooks) { + executeHookMu.Lock() + defer executeHookMu.Unlock() + executeHooks = hooks +} + +// UnregisterExecuteHooks removes the execute hooks. +func UnregisterExecuteHooks() { + executeHookMu.Lock() + defer executeHookMu.Unlock() + executeHooks = nil +} + +// trySendExecuteRequest attempts to send an execute request via Redis if in distributed mode. +// targetRole is "master" or "worker" to control where the request is executed. +// targetScope specifies which worker(s) to target: "all", alias, worker ID, or public IP (only used when targetRole="worker"). +// Returns true if sent to Redis, false if caller should fall back to local execution. +func trySendExecuteRequest(action, expr, workflow, target, params, targetRole, targetScope string) bool { + executeHookMu.RLock() + hooks := executeHooks + executeHookMu.RUnlock() + + if hooks == nil || hooks.ShouldUseRedis == nil || !hooks.ShouldUseRedis() { + return false + } + + if hooks.SendExecuteRequest == nil { + return false + } + + if err := hooks.SendExecuteRequest(context.Background(), action, expr, workflow, target, params, targetRole, targetScope); err != nil { + logger.Get().Warn("trySendExecuteRequest: failed to send to Redis, falling back to local", + zap.String("action", action), + zap.String("targetRole", targetRole), + zap.Error(err)) + return false + } + return true +} + +// tryResolveWorkerSSH attempts to resolve SSH connection info for a worker +// by ID, alias, or IP address. Returns nil if hooks are not registered or resolution fails. +func tryResolveWorkerSSH(identifier string) *WorkerSSHInfo { + executeHookMu.RLock() + hooks := executeHooks + executeHookMu.RUnlock() + + if hooks == nil || hooks.ResolveWorkerSSH == nil { + return nil + } + + info, err := hooks.ResolveWorkerSSH(context.Background(), identifier) + if err != nil { + logger.Get().Warn("tryResolveWorkerSSH failed", + zap.String("identifier", identifier), + zap.Error(err)) + return nil + } + return info +} diff --git a/internal/functions/goja_runtime.go b/internal/functions/goja_runtime.go index c165ef3..5d374c6 100644 --- a/internal/functions/goja_runtime.go +++ b/internal/functions/goja_runtime.go @@ -120,6 +120,8 @@ func (r *GojaRuntime) registerFunctionsOnVM(vm *goja.Runtime) { _ = vm.Set(FnPickValid, vf.pickValid) _ = vm.Set(FnRunModule, vf.runModule) _ = vm.Set(FnRunFlow, vf.runFlow) + _ = vm.Set(FnRunOnMaster, vf.runOnMaster) + _ = vm.Set(FnRunOnWorker, vf.runOnWorker) _ = vm.Set(FnExecPython, vf.execPython) _ = vm.Set(FnExecPythonFile, vf.execPythonFile) @@ -208,6 +210,10 @@ func (r *GojaRuntime) registerFunctionsOnVM(vm *goja.Runtime) { _ = vm.Set(FnUnzipDir, vf.unzipDir) _ = vm.Set(FnExtractTo, vf.extractTo) + // Snapshot functions + _ = vm.Set(FnSnapshotExport, vf.snapshotExport) + _ = vm.Set(FnSnapshotImport, vf.snapshotImport) + // Diff functions _ = vm.Set(FnExtractDiff, vf.extractDiff) @@ -310,6 +316,9 @@ func (r *GojaRuntime) registerFunctionsOnVM(vm *goja.Runtime) { // SSH functions _ = vm.Set(FnSSHExec, vf.sshExec) _ = vm.Set(FnSSHRsync, vf.sshRsync) + _ = vm.Set(FnSyncFromMaster, vf.syncFromMaster) + _ = vm.Set(FnSyncFromWorker, vf.syncFromWorker) + _ = vm.Set(FnRsyncToWorker, vf.rsyncToWorker) // Console for debugging _ = vm.Set("console", map[string]interface{}{ diff --git a/internal/functions/ssh_functions.go b/internal/functions/ssh_functions.go index c7c250c..0131f90 100644 --- a/internal/functions/ssh_functions.go +++ b/internal/functions/ssh_functions.go @@ -219,3 +219,191 @@ func (vf *vmFunc) sshRsync(call goja.FunctionCall) goja.Value { zap.String("host", host), zap.String("src", src), zap.String("dest", dest)) return vf.vm.ToValue(true) } + +// resolveWorkerSSHParams merges WorkerSSHInfo (from Redis lookup) with an explicit IP override. +// Returns the resolved host, user, keyPath, and port for rsync. +func resolveWorkerSSHParams(info *WorkerSSHInfo, ip string) (host, user, keyPath string, port int) { + user = "root" + port = 22 + if info != nil { + host = info.Host + user = info.User + keyPath = info.KeyPath + if info.Port > 0 { + port = info.Port + } + } + // Explicit IP overrides resolved host + if ip != "undefined" && ip != "" { + host = ip + } + return +} + +// executeRsync runs rsync with SSH options. src and dest should already include +// user@host: prefix where appropriate. +func executeRsync(src, dest, keyPath string, port int) bool { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + args := []string{"-avz", "-e"} + keyPath = sshExpandPath(keyPath) + if keyPath != "" { + args = append(args, fmt.Sprintf("ssh -i %s -p %d -o StrictHostKeyChecking=no -o ConnectTimeout=30", keyPath, port)) + } else { + args = append(args, fmt.Sprintf("ssh -p %d -o StrictHostKeyChecking=no -o ConnectTimeout=30", port)) + } + args = append(args, src, dest) + + cmd := exec.CommandContext(ctx, "rsync", args...) + output, err := cmd.CombinedOutput() + if err != nil { + logger.Get().Warn("rsync failed", + zap.String("src", src), zap.String("dest", dest), + zap.Error(err), zap.String("output", string(output))) + return false + } + return true +} + +// syncFromMaster requests the master to rsync a file/folder to this worker. +// In distributed mode, sends a "sync_to_worker" execute request to the master. +// In standalone mode, falls back to local cp -r. +// Usage: sync_from_master(src, dest) -> bool +func (vf *vmFunc) syncFromMaster(call goja.FunctionCall) goja.Value { + src := call.Argument(0).String() + dest := call.Argument(1).String() + + logger.Get().Debug("Calling "+terminal.HiGreen(FnSyncFromMaster), + zap.String("src", src), zap.String("dest", dest)) + + if src == "undefined" || src == "" { + logger.Get().Warn(FnSyncFromMaster + ": src is required") + return vf.vm.ToValue(false) + } + if dest == "undefined" || dest == "" { + logger.Get().Warn(FnSyncFromMaster + ": dest is required") + return vf.vm.ToValue(false) + } + + // Try distributed: send "sync_to_worker" request to master. + // Master will look up this worker's SSH info and rsync src to worker:dest. + // Data=src, Target=dest, TargetRole="master" + if trySendExecuteRequest("sync_to_worker", src, "", dest, "", "master", "") { + logger.Get().Debug(terminal.HiGreen(FnSyncFromMaster)+" sent to master via Redis", + zap.String("src", src), zap.String("dest", dest)) + return vf.vm.ToValue(true) + } + + // Fallback: local cp -r (standalone mode) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + cmd := exec.CommandContext(ctx, "cp", "-r", src, dest) + output, err := cmd.CombinedOutput() + if err != nil { + logger.Get().Warn(FnSyncFromMaster+": local cp fallback failed", + zap.String("src", src), zap.String("dest", dest), + zap.Error(err), zap.String("output", string(output))) + return vf.vm.ToValue(false) + } + + logger.Get().Debug(terminal.HiGreen(FnSyncFromMaster)+" completed via local cp", + zap.String("src", src), zap.String("dest", dest)) + return vf.vm.ToValue(true) +} + +// syncFromWorker pulls a file/folder from a specific worker via rsync over SSH. +// Usage: sync_from_worker(identifier, ip, src, dest) -> bool +func (vf *vmFunc) syncFromWorker(call goja.FunctionCall) goja.Value { + identifier := call.Argument(0).String() + ip := call.Argument(1).String() + src := call.Argument(2).String() + dest := call.Argument(3).String() + + logger.Get().Debug("Calling "+terminal.HiGreen(FnSyncFromWorker), + zap.String("identifier", identifier), zap.String("ip", ip), + zap.String("src", src), zap.String("dest", dest)) + + if identifier == "undefined" || identifier == "" { + logger.Get().Warn(FnSyncFromWorker + ": identifier is required") + return vf.vm.ToValue(false) + } + if src == "undefined" || src == "" { + logger.Get().Warn(FnSyncFromWorker + ": src is required") + return vf.vm.ToValue(false) + } + if dest == "undefined" || dest == "" { + logger.Get().Warn(FnSyncFromWorker + ": dest is required") + return vf.vm.ToValue(false) + } + + // Resolve worker SSH info from Redis + workerSSH := tryResolveWorkerSSH(identifier) + host, user, keyPath, port := resolveWorkerSSHParams(workerSSH, ip) + + if host == "" { + logger.Get().Warn(FnSyncFromWorker+": no host resolved", + zap.String("identifier", identifier), zap.String("ip", ip)) + return vf.vm.ToValue(false) + } + + // rsync pull: user@host:src -> dest + rsyncSrc := fmt.Sprintf("%s@%s:%s", user, host, src) + ok := executeRsync(rsyncSrc, dest, keyPath, port) + + if ok { + logger.Get().Debug(terminal.HiGreen(FnSyncFromWorker)+" completed", + zap.String("host", host), zap.String("src", src), zap.String("dest", dest)) + } + return vf.vm.ToValue(ok) +} + +// rsyncToWorker pushes a file/folder to a specific worker via rsync over SSH. +// Usage: rsync_to_worker(identifier, ip, src, dest) -> bool +func (vf *vmFunc) rsyncToWorker(call goja.FunctionCall) goja.Value { + identifier := call.Argument(0).String() + ip := call.Argument(1).String() + src := call.Argument(2).String() + dest := call.Argument(3).String() + + logger.Get().Debug("Calling "+terminal.HiGreen(FnRsyncToWorker), + zap.String("identifier", identifier), zap.String("ip", ip), + zap.String("src", src), zap.String("dest", dest)) + + if identifier == "undefined" || identifier == "" { + logger.Get().Warn(FnRsyncToWorker + ": identifier is required") + return vf.vm.ToValue(false) + } + if src == "undefined" || src == "" { + logger.Get().Warn(FnRsyncToWorker + ": src is required") + return vf.vm.ToValue(false) + } + if dest == "undefined" || dest == "" { + logger.Get().Warn(FnRsyncToWorker + ": dest is required") + return vf.vm.ToValue(false) + } + + // Resolve worker SSH info from Redis + workerSSH := tryResolveWorkerSSH(identifier) + host, user, keyPath, port := resolveWorkerSSHParams(workerSSH, ip) + + if host == "" { + logger.Get().Warn(FnRsyncToWorker+": no host resolved", + zap.String("identifier", identifier), zap.String("ip", ip)) + return vf.vm.ToValue(false) + } + + // Expand local source path + src = sshExpandPath(src) + + // rsync push: src -> user@host:dest + rsyncDest := fmt.Sprintf("%s@%s:%s", user, host, dest) + ok := executeRsync(src, rsyncDest, keyPath, port) + + if ok { + logger.Get().Debug(terminal.HiGreen(FnRsyncToWorker)+" completed", + zap.String("host", host), zap.String("src", src), zap.String("dest", dest)) + } + return vf.vm.ToValue(ok) +} diff --git a/internal/functions/ssh_functions_test.go b/internal/functions/ssh_functions_test.go index e62cb5e..74a9833 100644 --- a/internal/functions/ssh_functions_test.go +++ b/internal/functions/ssh_functions_test.go @@ -137,6 +137,157 @@ func TestSSHExec_DefaultUserAndPort(t *testing.T) { 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. +// --- sync_from_master input validation tests --- + +func TestSyncFromMaster_EmptySrc(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_master("", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSyncFromMaster_EmptyDest(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_master("/tmp/src", "")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSyncFromMaster_NoArgs(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_master()`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +// --- sync_from_worker input validation tests --- + +func TestSyncFromWorker_EmptyIdentifier(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_worker("", "10.0.0.2", "/tmp/src", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSyncFromWorker_EmptySrc(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_worker("worker-1", "10.0.0.2", "", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSyncFromWorker_EmptyDest(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_worker("worker-1", "10.0.0.2", "/tmp/src", "")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSyncFromWorker_NoArgs(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_worker()`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestSyncFromWorker_NoHost(t *testing.T) { + // No SSH hooks registered and no explicit IP -> no host resolved + registry := NewRegistry() + result, err := registry.Execute( + `sync_from_worker("worker-1", "", "/tmp/src", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +// --- rsync_to_worker input validation tests --- + +func TestRsyncToWorker_EmptyIdentifier(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `rsync_to_worker("", "10.0.0.2", "/tmp/src", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestRsyncToWorker_EmptySrc(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `rsync_to_worker("worker-1", "10.0.0.2", "", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestRsyncToWorker_EmptyDest(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `rsync_to_worker("worker-1", "10.0.0.2", "/tmp/src", "")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestRsyncToWorker_NoArgs(t *testing.T) { + registry := NewRegistry() + result, err := registry.Execute( + `rsync_to_worker()`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +func TestRsyncToWorker_NoHost(t *testing.T) { + // No SSH hooks registered and no explicit IP -> no host resolved + registry := NewRegistry() + result, err := registry.Execute( + `rsync_to_worker("worker-1", "", "/tmp/src", "/tmp/dest")`, + map[string]interface{}{}, + ) + + require.NoError(t, err) + assert.Equal(t, false, result) +} + +// NOTE: Integration tests for ssh_exec, ssh_rsync, and sync functions 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/internal/functions/util_functions.go b/internal/functions/util_functions.go index 70ce132..938bf05 100644 --- a/internal/functions/util_functions.go +++ b/internal/functions/util_functions.go @@ -1085,6 +1085,116 @@ func (vf *vmFunc) moveFile(call goja.FunctionCall) goja.Value { return vf.vm.ToValue(true) } +// snapshotExport exports a workspace as a ZIP snapshot by running osmedeus snapshot export as a subprocess. +// Usage: snapshot_export(workspace, dest?) -> string (zip path on success, empty on failure) +func (vf *vmFunc) snapshotExport(call goja.FunctionCall) goja.Value { + workspace := call.Argument(0).String() + dest := call.Argument(1).String() + + if workspace == "undefined" || workspace == "" { + logger.Get().Warn("snapshot_export: workspace is required") + return vf.vm.ToValue("") + } + if dest == "undefined" { + dest = "" + } + + // Find the current executable + exePath, err := os.Executable() + if err != nil { + logger.Get().Warn("snapshot_export: failed to find executable", zap.Error(err)) + return vf.vm.ToValue("") + } + + // Build command arguments + args := []string{"snapshot", "export", workspace} + if dest != "" { + args = append(args, "-o", dest) + } + + logger.Get().Debug("Calling "+terminal.HiGreen("snapshot_export"), + zap.String("workspace", workspace), + zap.String("dest", dest)) + + // @NOTE: This is intentional - snapshot_export() is a utility function exposed to workflow + // definitions for exporting workspaces. Input comes from trusted workflow YAML files. + cmd := exec.Command(exePath, args...) + output, err := cmd.CombinedOutput() + if err != nil { + logger.Get().Warn("snapshot_export: command failed", + zap.String("workspace", workspace), + zap.Error(err), + zap.String("output", string(output))) + return vf.vm.ToValue("") + } + + // Parse output to extract the zip file path from "File: " line + var zipPath string + for _, line := range strings.Split(string(output), "\n") { + line = strings.TrimSpace(line) + if idx := strings.Index(line, "File:"); idx >= 0 { + zipPath = strings.TrimSpace(line[idx+len("File:"):]) + break + } + } + + logger.Get().Debug(terminal.HiGreen("snapshot_export")+" result", + zap.String("workspace", workspace), + zap.String("zipPath", zipPath)) + return vf.vm.ToValue(zipPath) +} + +// snapshotImport imports a workspace from a ZIP snapshot by running osmedeus snapshot import as a subprocess. +// Usage: snapshot_import(source) -> string (workspace name on success, empty on failure) +func (vf *vmFunc) snapshotImport(call goja.FunctionCall) goja.Value { + source := call.Argument(0).String() + + if source == "undefined" || source == "" { + logger.Get().Warn("snapshot_import: source is required") + return vf.vm.ToValue("") + } + + // Find the current executable + exePath, err := os.Executable() + if err != nil { + logger.Get().Warn("snapshot_import: failed to find executable", zap.Error(err)) + return vf.vm.ToValue("") + } + + // Build command arguments with --force to skip interactive confirmation + args := []string{"snapshot", "import", source, "--force"} + + logger.Get().Debug("Calling "+terminal.HiGreen("snapshot_import"), + zap.String("source", source)) + + // @NOTE: This is intentional - snapshot_import() is a utility function exposed to workflow + // definitions for importing workspaces. Input comes from trusted workflow YAML files. + cmd := exec.Command(exePath, args...) + output, err := cmd.CombinedOutput() + if err != nil { + logger.Get().Warn("snapshot_import: command failed", + zap.String("source", source), + zap.Error(err), + zap.String("output", string(output))) + return vf.vm.ToValue("") + } + + // Parse output to extract workspace name from "Workspace: " line + var workspaceName string + for _, line := range strings.Split(string(output), "\n") { + line = strings.TrimSpace(line) + if idx := strings.Index(line, "Workspace:"); idx >= 0 { + workspaceName = strings.TrimSpace(line[idx+len("Workspace:"):]) + break + } + } + + logger.Get().Debug(terminal.HiGreen("snapshot_import")+" result", + zap.String("source", source), + zap.String("workspace", workspaceName)) + return vf.vm.ToValue(workspaceName) +} + // parseParamsToFlags parses a comma-separated "key=value" string into -p flags. // E.g. "threads=10,deep=true" -> ["-p", "threads=10", "-p", "deep=true"] func parseParamsToFlags(params string) []string { @@ -1172,6 +1282,213 @@ func (vf *vmFunc) runFlow(call goja.FunctionCall) goja.Value { return vf.runOsmedeus("-f", flow, target, params, "run_flow") } +// runOnMaster executes an action on the master node in distributed mode. +// In distributed mode, the request is sent via Redis to the master. +// In standalone mode, falls back to local execution. +// +// Usage: +// +// run_on_master('func', 'db_import_sarif("ws", "/path/file.sarif")') -> bool +// run_on_master('run', 'subdomain', 'example.com', 'threads=10') -> bool +// run_on_master('bash', 'nmap -sV target.com') -> bool +func (vf *vmFunc) runOnMaster(call goja.FunctionCall) goja.Value { + action := call.Argument(0).String() + if action == "undefined" || action == "" { + logger.Get().Warn("run_on_master: action is required (func, run, or bash)") + return vf.vm.ToValue(false) + } + + log := logger.Get() + + switch action { + case "func": + expr := call.Argument(1).String() + if expr == "undefined" || expr == "" { + log.Warn("run_on_master: expression is required for 'func' action") + return vf.vm.ToValue(false) + } + + log.Debug("Calling "+terminal.HiGreen("run_on_master"), + zap.String("action", "func"), + zap.String("expr", expr)) + + // Try distributed path first + if trySendExecuteRequest("func", expr, "", "", "", "master", "") { + log.Debug("run_on_master: sent func request to master via Redis") + return vf.vm.ToValue(true) + } + + // Fallback: execute locally using the current runtime + log.Warn("run_on_master: not in distributed mode, executing locally", + zap.String("action", action), + zap.String("hint", "ensure worker mode is active and Redis is configured")) + execCtx := make(map[string]interface{}) + _, err := vf.runtime.Execute(expr, execCtx) + if err != nil { + log.Warn("run_on_master: local execution failed", + zap.String("expr", expr), + zap.Error(err)) + return vf.vm.ToValue(false) + } + return vf.vm.ToValue(true) + + case "run": + workflow := call.Argument(1).String() + target := call.Argument(2).String() + params := call.Argument(3).String() + + if workflow == "undefined" || workflow == "" || target == "undefined" || target == "" { + log.Warn("run_on_master: workflow and target are required for 'run' action") + return vf.vm.ToValue(false) + } + if params == "undefined" { + params = "" + } + + log.Debug("Calling "+terminal.HiGreen("run_on_master"), + zap.String("action", "run"), + zap.String("workflow", workflow), + zap.String("target", target), + zap.String("params", params)) + + // Try distributed path first + if trySendExecuteRequest("run", "", workflow, target, params, "master", "") { + log.Debug("run_on_master: sent run request to master via Redis") + return vf.vm.ToValue(true) + } + + // Fallback: run locally as subprocess + log.Warn("run_on_master: not in distributed mode, running locally", + zap.String("action", action), + zap.String("hint", "ensure worker mode is active and Redis is configured")) + vf.runOsmedeus("-m", workflow, target, params, "run_on_master") + return vf.vm.ToValue(true) + + case "bash": + command := call.Argument(1).String() + if command == "undefined" || command == "" { + log.Warn("run_on_master: command is required for 'bash' action") + return vf.vm.ToValue(false) + } + + log.Debug("Calling "+terminal.HiGreen("run_on_master"), + zap.String("action", "bash"), + zap.String("command", command)) + + // Try distributed path first + if trySendExecuteRequest("bash", command, "", "", "", "master", "") { + log.Debug("run_on_master: sent bash request to master via Redis") + return vf.vm.ToValue(true) + } + + // Fallback: execute locally + log.Warn("run_on_master: not in distributed mode, executing bash locally") + // @NOTE: This is intentional - run_on_master('bash', cmd) is a utility function + // exposed to workflow definitions. Input comes from trusted workflow YAML files. + cmd := exec.Command("sh", "-c", command) + output, err := cmd.CombinedOutput() + if err != nil { + log.Warn("run_on_master: local bash execution failed", + zap.String("command", command), zap.Error(err), + zap.String("output", string(output))) + return vf.vm.ToValue(false) + } + return vf.vm.ToValue(true) + + default: + log.Warn("run_on_master: unknown action", + zap.String("action", action), + zap.String("valid", "func, run, bash")) + return vf.vm.ToValue(false) + } +} + +// runOnWorker executes an action on worker node(s) in distributed mode. +// In distributed mode, the request is routed via Redis through the master to target workers. +// In standalone mode, falls back to local execution. +// +// Usage: +// +// run_on_worker('all', 'func', 'log_info("hello")') -> bool +// run_on_worker('scanner-1', 'run', 'subdomain', 'example.com', 'threads=10') -> bool +// run_on_worker('all', 'bash', 'apt update && apt install -y nmap') -> bool +func (vf *vmFunc) runOnWorker(call goja.FunctionCall) goja.Value { + scope := call.Argument(0).String() + if scope == "undefined" || scope == "" { + scope = "all" + } + action := call.Argument(1).String() + if action == "undefined" || action == "" { + logger.Get().Warn("run_on_worker: action is required (func, run, bash)") + return vf.vm.ToValue(false) + } + + log := logger.Get() + + switch action { + case "func": + expr := call.Argument(2).String() + if expr == "undefined" || expr == "" { + log.Warn("run_on_worker: expression required for 'func' action") + return vf.vm.ToValue(false) + } + if trySendExecuteRequest("func", expr, "", "", "", "worker", scope) { + return vf.vm.ToValue(true) + } + // Fallback: execute locally + log.Warn("run_on_worker: not in distributed mode, executing locally") + execCtx := make(map[string]interface{}) + _, err := vf.runtime.Execute(expr, execCtx) + if err != nil { + log.Warn("run_on_worker: local execution failed", zap.Error(err)) + return vf.vm.ToValue(false) + } + return vf.vm.ToValue(true) + + case "run": + workflow := call.Argument(2).String() + target := call.Argument(3).String() + params := call.Argument(4).String() + if workflow == "undefined" || workflow == "" || target == "undefined" || target == "" { + log.Warn("run_on_worker: workflow and target required for 'run' action") + return vf.vm.ToValue(false) + } + if params == "undefined" { + params = "" + } + if trySendExecuteRequest("run", "", workflow, target, params, "worker", scope) { + return vf.vm.ToValue(true) + } + // Fallback: run locally + log.Warn("run_on_worker: not in distributed mode, running locally") + vf.runOsmedeus("-m", workflow, target, params, "run_on_worker") + return vf.vm.ToValue(true) + + case "bash": + command := call.Argument(2).String() + if command == "undefined" || command == "" { + log.Warn("run_on_worker: command required for 'bash' action") + return vf.vm.ToValue(false) + } + if trySendExecuteRequest("bash", command, "", "", "", "worker", scope) { + return vf.vm.ToValue(true) + } + // Fallback: execute locally + log.Warn("run_on_worker: not in distributed mode, executing bash locally") + // @NOTE: This is intentional - run_on_worker('scope', 'bash', cmd) is a utility function + // exposed to workflow definitions. Input comes from trusted workflow YAML files. + cmd := exec.Command("sh", "-c", command) + _, _ = cmd.CombinedOutput() + return vf.vm.ToValue(true) + + default: + log.Warn("run_on_worker: unknown action (expected 'func', 'run', or 'bash'). "+ + "Did you forget the scope argument? Correct usage: run_on_worker(scope, action, ...args)", + zap.String("action", action)) + return vf.vm.ToValue(false) + } +} + // copyFileBuffered copies a file using buffered I/O (memory-efficient for large files) func copyFileBuffered(source, dest string, mode os.FileMode) error { srcFile, err := os.Open(source) diff --git a/internal/installer/binary.go b/internal/installer/binary.go index c3d05b9..73062f6 100644 --- a/internal/installer/binary.go +++ b/internal/installer/binary.go @@ -317,12 +317,12 @@ func InstallBinary(name string, registry BinaryRegistry, binariesFolder string, terminal.Gray(terminal.SymbolBowtie), terminal.HiBlue(name)) return nil } - // Non-core tools: show exists, still copy to external-binaries - fmt.Printf("[%s] Binary '%s' already available in PATH, copying to external-binaries\n", + // Non-core tools: show exists, still symlink to external-binaries + fmt.Printf("[%s] Binary '%s' already available in PATH, symlinking to external-binaries\n", terminal.Gray(terminal.SymbolBowtie), terminal.HiBlue(name)) entry, ok := registry[name] if ok { - _ = CopyInstalledBinaryToFolder(name, &entry, binariesFolder) + _ = SymlinkInstalledBinaryToFolder(name, &entry, binariesFolder) } return nil } @@ -377,9 +377,9 @@ func InstallBinary(name string, registry BinaryRegistry, binariesFolder string, return err } - // After successful command execution, copy the binary to external-binaries - if err := CopyInstalledBinaryToFolder(name, &entry, binariesFolder); err != nil { - logger.Get().Warn("Failed to copy binary to external-binaries folder", + // After successful command execution, symlink the binary to external-binaries + if err := SymlinkInstalledBinaryToFolder(name, &entry, binariesFolder); err != nil { + logger.Get().Warn("Failed to symlink binary to external-binaries folder", zap.String("name", name), zap.Error(err)) // Don't return error - installation succeeded, this is just an optimization @@ -596,10 +596,27 @@ func copyFile(src, dest string) error { return os.Chmod(dest, srcInfo.Mode()) } -// CopyInstalledBinaryToFolder finds a binary using LookPath and copies it to the destination folder +// symlinkOrCopyFile creates a symlink from src to dest, falling back to copy if symlink fails +func symlinkOrCopyFile(src, dest string) error { + // Remove existing file/symlink at dest + if _, err := os.Lstat(dest); err == nil { + if err := os.Remove(dest); err != nil { + return fmt.Errorf("failed to remove existing destination: %w", err) + } + } + // Try symlink first, fall back to copy + if err := os.Symlink(src, dest); err != nil { + logger.Get().Warn("Symlink failed, falling back to copy", + zap.String("src", src), zap.String("dest", dest), zap.Error(err)) + return copyFile(src, dest) + } + return nil +} + +// SymlinkInstalledBinaryToFolder finds a binary using LookPath and symlinks it to the destination folder // Uses the validate command (valide-command) from registry to locate the binary -// Returns nil if binary not found (installation may have failed) or copy succeeds -func CopyInstalledBinaryToFolder(name string, entry *BinaryEntry, destFolder string) error { +// Returns nil if binary not found (installation may have failed) or symlink succeeds +func SymlinkInstalledBinaryToFolder(name string, entry *BinaryEntry, destFolder string) error { // Determine which command to look for lookupCmd := name if entry != nil && entry.ValidateCommand != "" { @@ -631,15 +648,15 @@ func CopyInstalledBinaryToFolder(name string, entry *BinaryEntry, destFolder str return nil } - logger.Get().Info("Copying binary to external-binaries", + logger.Get().Info("Symlinking binary to external-binaries", zap.String("from", binaryPath), zap.String("to", destPath)) - if err := copyFile(binaryPath, destPath); err != nil { - return fmt.Errorf("failed to copy binary: %w", err) + if err := symlinkOrCopyFile(binaryPath, destPath); err != nil { + return fmt.Errorf("failed to symlink binary: %w", err) } - return os.Chmod(destPath, 0755) + return nil } // ListBinaries returns all binary names in the registry diff --git a/internal/installer/go_install.go b/internal/installer/go_install.go index f71d6b5..9224d52 100644 --- a/internal/installer/go_install.go +++ b/internal/installer/go_install.go @@ -2,7 +2,6 @@ package installer import ( "fmt" - "io" "os" "os/exec" "path/filepath" @@ -86,18 +85,18 @@ func InstallBinaryViaGo(binaryName string, goPackage string, binariesFolder stri return fmt.Errorf("failed to install %s via go install: %w", pkg, err) } - // Copy binary from GOBIN to binaries folder if specified + // Symlink binary from GOBIN to binaries folder if specified if binariesFolder != "" { - if err := copyGoBinaryToFolder(binaryName, binariesFolder); err != nil { - return fmt.Errorf("failed to copy binary to folder: %w", err) + if err := symlinkGoBinaryToFolder(binaryName, binariesFolder); err != nil { + return fmt.Errorf("failed to symlink binary to folder: %w", err) } } return nil } -// copyGoBinaryToFolder finds a binary installed by go and copies it to the target folder -func copyGoBinaryToFolder(binaryName string, binariesFolder string) error { +// symlinkGoBinaryToFolder finds a binary installed by go and symlinks it to the target folder +func symlinkGoBinaryToFolder(binaryName string, binariesFolder string) error { goBinPath, err := GetGoBinPath() if err != nil { return err @@ -121,33 +120,11 @@ func copyGoBinaryToFolder(binaryName string, binariesFolder string) error { destPath := filepath.Join(binariesFolder, binaryName) - logger.Get().Info("Copying binary to external folder", + logger.Get().Info("Symlinking binary to external folder", zap.String("src", srcPath), zap.String("dest", destPath)) - // Copy the file - srcFile, err := os.Open(srcPath) - if err != nil { - return fmt.Errorf("failed to open source binary: %w", err) - } - defer func() { _ = srcFile.Close() }() - - destFile, err := os.Create(destPath) - if err != nil { - return fmt.Errorf("failed to create destination binary: %w", err) - } - defer func() { _ = destFile.Close() }() - - if _, err := io.Copy(destFile, srcFile); err != nil { - return fmt.Errorf("failed to copy binary: %w", err) - } - - // Make executable - if err := os.Chmod(destPath, 0755); err != nil { - return fmt.Errorf("failed to make binary executable: %w", err) - } - - return nil + return symlinkOrCopyFile(srcPath, destPath) } // GetGoPackageName returns the go_install package path for a binary diff --git a/internal/installer/installer.go b/internal/installer/installer.go index fc1183c..c7fb5af 100644 --- a/internal/installer/installer.go +++ b/internal/installer/installer.go @@ -11,6 +11,10 @@ import ( const ( // DatabaseFileName is the name of the database file to preserve during base installation DatabaseFileName = "database-osm.sqlite" + // SettingsFileName is the name of the settings file + SettingsFileName = "osm-settings.yaml" + // SettingsBackupName is the name of the settings backup file + SettingsBackupName = "backup-osm-settings.yaml" ) // Installer handles installation of workflows, base folder, and binaries @@ -20,6 +24,7 @@ type Installer struct { BinariesFolder string CustomHeaders map[string]string Printer *terminal.Printer + KeepSetting bool } // NewInstaller creates a new Installer with the given paths @@ -116,6 +121,20 @@ func (i *Installer) InstallBase(source string) error { } } + // Backup osm-settings.yaml before base folder is removed + settingsPath := filepath.Join(i.BaseFolder, SettingsFileName) + var settingsBackupPath string + + if _, err := os.Stat(settingsPath); err == nil { + i.Printer.Info("Backing up settings...") + settingsBackupPath, err = backupFile(settingsPath) + if err != nil { + i.Printer.Warning("Failed to backup settings: %s", err) + } else { + defer func() { _ = os.Remove(settingsBackupPath) }() + } + } + // Fetch source to temp directory i.Printer.Info("Fetching source...") tempDir, err := FetchToTemp(source, i.CustomHeaders) @@ -177,6 +196,26 @@ func (i *Installer) InstallBase(source string) error { } } + // Save settings backup to persistent path for reference + if settingsBackupPath != "" { + persistentBackupPath := filepath.Join(i.BaseFolder, SettingsBackupName) + if err := restoreFile(settingsBackupPath, persistentBackupPath); err != nil { + i.Printer.Warning("Failed to save settings backup: %s", err) + } else { + i.Printer.Info("Previous settings backed up to: %s", terminal.Cyan(persistentBackupPath)) + } + + // If --keep-setting, restore the old settings over the new one + if i.KeepSetting { + newSettingsPath := filepath.Join(i.BaseFolder, SettingsFileName) + if err := restoreFile(settingsBackupPath, newSettingsPath); err != nil { + i.Printer.Warning("Failed to restore settings: %s", err) + } else { + i.Printer.Success("Previous settings restored to: %s", terminal.Cyan(newSettingsPath)) + } + } + } + i.Printer.Success("Base folder installed successfully to %s", terminal.Gray(i.BaseFolder)) return nil } diff --git a/internal/installer/nix.go b/internal/installer/nix.go index 61b2389..ed95f62 100644 --- a/internal/installer/nix.go +++ b/internal/installer/nix.go @@ -117,18 +117,18 @@ func InstallBinaryViaNix(binaryName string, nixPackage string, binariesFolder st // Store output for caller to display NixInstallOutput = string(output) - // Copy binary from Nix profile to binaries folder + // Symlink binary from Nix profile to binaries folder if binariesFolder != "" { - if err := copyNixBinaryToFolder(binaryName, binariesFolder); err != nil { - return fmt.Errorf("failed to copy binary to folder: %w", err) + if err := symlinkNixBinaryToFolder(binaryName, binariesFolder); err != nil { + return fmt.Errorf("failed to symlink binary to folder: %w", err) } } return nil } -// copyNixBinaryToFolder finds a binary installed by Nix and copies it to the target folder -func copyNixBinaryToFolder(binaryName string, binariesFolder string) error { +// symlinkNixBinaryToFolder finds a binary installed by Nix and symlinks it to the target folder +func symlinkNixBinaryToFolder(binaryName string, binariesFolder string) error { // Find the binary path using 'which' cmd := exec.Command("which", binaryName) output, err := cmd.Output() @@ -161,33 +161,11 @@ func copyNixBinaryToFolder(binaryName string, binariesFolder string) error { destPath := filepath.Join(binariesFolder, binaryName) - logger.Get().Info("Copying binary to external folder", + logger.Get().Info("Symlinking binary to external folder", zap.String("src", srcPath), zap.String("dest", destPath)) - // Copy the file - srcFile, err := os.Open(srcPath) - if err != nil { - return fmt.Errorf("failed to open source binary: %w", err) - } - defer func() { _ = srcFile.Close() }() - - destFile, err := os.Create(destPath) - if err != nil { - return fmt.Errorf("failed to create destination binary: %w", err) - } - defer func() { _ = destFile.Close() }() - - if _, err := io.Copy(destFile, srcFile); err != nil { - return fmt.Errorf("failed to copy binary: %w", err) - } - - // Make executable - if err := os.Chmod(destPath, 0755); err != nil { - return fmt.Errorf("failed to make binary executable: %w", err) - } - - return nil + return symlinkOrCopyFile(srcPath, destPath) } // GetNixPackageName returns the Nix package name for a binary diff --git a/internal/state/export.go b/internal/state/export.go index 13f91a6..c210d9f 100644 --- a/internal/state/export.go +++ b/internal/state/export.go @@ -31,24 +31,33 @@ func Export(stateFile string, ctx *ExportContext) error { UpdatedAt: time.Now(), } - // Always use context data first for run info (has correct in-memory status) - // The DB read happens before status is updated, so it returns stale data. - // Context data comes from the in-memory result which has the correct status. - if ctx != nil { - export.Run = runInfoFromContext(ctx) - } - - // Optionally enrich with DB data for fields not available in context - // (currently all fields are available in context, so this is just for future-proofing) - if export.Run == nil && ctx != nil && ctx.RunUUID != "" && db != nil { + // Prefer DB Run record as primary data source — it has correct TotalSteps + // (calculated via calculateTotalSteps which sums module steps for flows), + // CompletedSteps (updated by WriteCoordinator), RunMode, and RunPriority. + // Context provides status and completed_at because the DB UpdateRunStatus + // call happens AFTER state export inside the executor. + if ctx != nil && ctx.RunUUID != "" && db != nil { var run database.Run err := db.NewSelect().Model(&run). Where("run_uuid = ?", ctx.RunUUID).Scan(dbCtx) if err == nil { export.Run = runInfoFromDB(&run) + // Override status and completed_at from context — DB still has + // "running" status at export time since UpdateRunStatus runs after export + if ctx.Status != "" { + export.Run.Status = ctx.Status + } + if ctx.CompletedAt != nil { + export.Run.CompletedAt = ctx.CompletedAt + } } } + // Fallback to context-only if DB is unavailable + if export.Run == nil && ctx != nil { + export.Run = runInfoFromContext(ctx) + } + // Try to load workspace from database first workspaceLoaded := false if ctx != nil && ctx.WorkspaceName != "" && db != nil { @@ -99,6 +108,9 @@ func runInfoFromDB(run *database.Run) *RunInfo { ErrorMessage: run.ErrorMessage, TotalSteps: run.TotalSteps, CompletedSteps: run.CompletedSteps, + HooksEnabled: run.HooksEnabled, + RunMode: run.RunMode, + RunPriority: run.RunPriority, } } @@ -119,6 +131,9 @@ func runInfoFromContext(ctx *ExportContext) *RunInfo { ErrorMessage: ctx.ErrorMessage, TotalSteps: ctx.TotalSteps, CompletedSteps: ctx.CompletedSteps, + HooksEnabled: ctx.HooksEnabled, + RunMode: ctx.RunMode, + RunPriority: ctx.RunPriority, } } diff --git a/internal/state/types.go b/internal/state/types.go index 8007dab..b55f4eb 100644 --- a/internal/state/types.go +++ b/internal/state/types.go @@ -24,6 +24,9 @@ type RunInfo struct { ErrorMessage string `json:"error_message,omitempty"` TotalSteps int `json:"total_steps"` CompletedSteps int `json:"completed_steps"` + HooksEnabled bool `json:"hooks_enabled"` + RunMode string `json:"run_mode,omitempty"` + RunPriority string `json:"run_priority,omitempty"` } // WorkspaceInfo contains workspace information for export @@ -61,5 +64,8 @@ type ExportContext struct { ErrorMessage string TotalSteps int CompletedSteps int + HooksEnabled bool + RunMode string + RunPriority string Artifacts []string } diff --git a/pkg/cli/db.go b/pkg/cli/db.go index b23ae3e..afae300 100644 --- a/pkg/cli/db.go +++ b/pkg/cli/db.go @@ -34,6 +34,7 @@ var ( dbClear bool dbListTables bool dbIncludeHeavy bool + dbCleanWS bool ) // defaultHiddenColumns are columns hidden by default for all tables @@ -135,6 +136,8 @@ func init() { dbCmd.PersistentFlags().BoolVar(&dbClear, "clear", false, "clear all records from the specified table (requires --table and --force)") dbCmd.PersistentFlags().BoolVar(&dbIncludeHeavy, "include-heavy", false, "include large fields (raw_response, screenshot, blob_content) in output") + dbCleanCmd.Flags().BoolVar(&dbCleanWS, "clean-ws", false, "also remove workspace data directory (e.g. ~/workspaces-osmedeus)") + dbIndexWorkflowCmd.Flags().BoolVar(&dbIndexForce, "force", false, "force re-index all workflows regardless of checksum") dbIndexCmd.AddCommand(dbIndexWorkflowCmd) @@ -256,6 +259,24 @@ func runDBClean(cmd *cobra.Command, args []string) error { printer.Info("Database: %s", getDatabaseInfo(cfg, db)) } + // Clean workspace data directory if --clean-ws is set + if dbCleanWS { + wsPath := cfg.GetWorkspacesDir() + if wsPath == "" { + printer.Warning("Workspaces path not configured, skipping workspace cleanup") + } else { + printer.Info("Removing workspace data: %s", wsPath) + if err := os.RemoveAll(wsPath); err != nil { + return fmt.Errorf("failed to remove workspaces directory: %w", err) + } + // Recreate the empty directory + if err := os.MkdirAll(wsPath, 0755); err != nil { + return fmt.Errorf("failed to recreate workspaces directory: %w", err) + } + printer.Success("Workspace data cleaned: %s", wsPath) + } + } + return nil } diff --git a/pkg/cli/function.go b/pkg/cli/function.go index 90e87b2..aadae02 100644 --- a/pkg/cli/function.go +++ b/pkg/cli/function.go @@ -2,6 +2,7 @@ package cli import ( "bufio" + "context" "fmt" "io" "os" @@ -11,8 +12,10 @@ import ( "syscall" "time" + "github.com/google/uuid" "github.com/j3ssie/osmedeus/v5/internal/config" "github.com/j3ssie/osmedeus/v5/internal/database" + "github.com/j3ssie/osmedeus/v5/internal/distributed" "github.com/j3ssie/osmedeus/v5/internal/executor" "github.com/j3ssie/osmedeus/v5/internal/functions" "github.com/j3ssie/osmedeus/v5/internal/template" @@ -105,6 +108,10 @@ func init() { evalCmd.Flags().BoolVar(&funcRepeat, "repeat", false, "repeat run after completion") evalCmd.Flags().StringVar(&funcRepeatWaitTime, "repeat-wait-time", "5s", "wait time between repeats (e.g., 30s, 20m, 10h, 1d)") + // Distributed mode flags for eval commands + evalCmd.Flags().StringVar(&redisURL, "redis-url", "", "Redis connection URL for distributed mode (enables run_on_master/run_on_worker routing)") + functionEvalCmd.Flags().StringVar(&redisURL, "redis-url", "", "Redis connection URL for distributed mode (enables run_on_master/run_on_worker routing)") + functionCmd.AddCommand(functionEvalCmd) functionCmd.AddCommand(functionListCmd) } @@ -140,56 +147,40 @@ func runFunctionEval(cmd *cobra.Command, args []string) error { } } - // Determine script source: --function-file > -f flag > positional arg > -e flag > stdin - var script string - - // Read script from function file if provided - if funcFunctionFile != "" { - data, err := os.ReadFile(funcFunctionFile) - if err != nil { - printer.Error("Failed to read function file: %s", err) - return fmt.Errorf("failed to read function file: %w", err) - } - script = strings.TrimSpace(string(data)) - } else if evalFunctionName != "" { - // Handle -f/--function flag: build script from function name + positional args - var quotedArgs []string - for _, arg := range args { - quotedArgs = append(quotedArgs, fmt.Sprintf("%q", arg)) - } - script = fmt.Sprintf("%s(%s)", evalFunctionName, strings.Join(quotedArgs, ", ")) - } else if len(args) > 0 && args[0] != "-" { - if len(args) > 1 { - // Multiple args: treat first as function name, rest as arguments - // e.g., "func_name arg1 arg2" → "func_name("arg1", "arg2")" - var quotedArgs []string - for _, arg := range args[1:] { - quotedArgs = append(quotedArgs, fmt.Sprintf("%q", arg)) + // If --redis-url is provided, register distributed hooks for master-mode eval + if redisURL != "" { + cfg := config.Get() + if cfg != nil { + redisCfg, err := distributed.ParseRedisURL(redisURL) + if err != nil { + return err } - script = fmt.Sprintf("%s(%s)", args[0], strings.Join(quotedArgs, ", ")) - } else { - // Single arg: check if it is a bare function name and add () if needed - script = normalizeScriptExpression(args[0]) + cfg.Redis = *redisCfg } - } else if evalScript != "" { - // Script provided via -e flag - script = evalScript - } else if evalStdin || (len(args) > 0 && args[0] == "-") { - // Read script from stdin - data, err := io.ReadAll(os.Stdin) + if cfg == nil || !cfg.IsRedisConfigured() { + return fmt.Errorf("redis not configured") + } + client, err := distributed.NewClientFromConfig(cfg) if err != nil { - printer.Error("Failed to read from stdin: %s", err) - return fmt.Errorf("failed to read from stdin: %w", err) + return fmt.Errorf("failed to create redis client: %w", err) } - script = strings.TrimSpace(string(data)) + defer client.Close() + ctx := context.Background() + if err := client.Ping(ctx); err != nil { + return fmt.Errorf("failed to connect to redis: %w", err) + } + hostname, _ := os.Hostname() + masterID := fmt.Sprintf("%s-master-eval-%s", hostname, uuid.NewString()[:8]) + config.SetWorkerMode(true, masterID) + cleanup := distributed.RegisterDistributedHooksFromClient(client, masterID) + defer cleanup() + printer.Info("Distributed mode enabled (id: %s)", terminal.Cyan(masterID)) } - if script != "" { - script = normalizeScriptExpression(script) - } - - if script == "" { - return fmt.Errorf("no script provided: use positional argument, -e flag, --function-file, or --stdin") + script, err := resolveEvalScript(args) + if err != nil { + printer.Error("%s", err) + return err } // Main execution loop @@ -452,6 +443,62 @@ func runFunctionList(cmd *cobra.Command, args []string) error { return nil } +// resolveEvalScript determines the script to execute from various sources: +// --function-file > -f flag > positional arg > -e flag > --stdin. +// Returns the resolved script string or an error. +func resolveEvalScript(args []string) (string, error) { + var script string + + // Read script from function file if provided + if funcFunctionFile != "" { + data, err := os.ReadFile(funcFunctionFile) + if err != nil { + return "", fmt.Errorf("failed to read function file: %w", err) + } + script = strings.TrimSpace(string(data)) + } else if evalFunctionName != "" { + // Handle -f/--function flag: build script from function name + positional args + var quotedArgs []string + for _, arg := range args { + quotedArgs = append(quotedArgs, fmt.Sprintf("%q", arg)) + } + script = fmt.Sprintf("%s(%s)", evalFunctionName, strings.Join(quotedArgs, ", ")) + } else if len(args) > 0 && args[0] != "-" { + if len(args) > 1 { + // Multiple args: treat first as function name, rest as arguments + // e.g., "func_name arg1 arg2" → "func_name("arg1", "arg2")" + var quotedArgs []string + for _, arg := range args[1:] { + quotedArgs = append(quotedArgs, fmt.Sprintf("%q", arg)) + } + script = fmt.Sprintf("%s(%s)", args[0], strings.Join(quotedArgs, ", ")) + } else { + // Single arg: check if it is a bare function name and add () if needed + script = normalizeScriptExpression(args[0]) + } + } else if evalScript != "" { + // Script provided via -e flag + script = evalScript + } else if evalStdin || (len(args) > 0 && args[0] == "-") { + // Read script from stdin + data, err := io.ReadAll(os.Stdin) + if err != nil { + return "", fmt.Errorf("failed to read from stdin: %w", err) + } + script = strings.TrimSpace(string(data)) + } + + if script != "" { + script = normalizeScriptExpression(script) + } + + if script == "" { + return "", fmt.Errorf("no script provided: use positional argument, -e flag, --function-file, or --stdin") + } + + return script, nil +} + // normalizeScriptExpression checks if the input is a bare function name (without parentheses) // and adds () to make it a valid function call. This allows "osmedeus eval cdn_ls_tree" // to work the same as "osmedeus eval cdn_ls_tree()" diff --git a/pkg/cli/install.go b/pkg/cli/install.go index 563d040..7a726cf 100644 --- a/pkg/cli/install.go +++ b/pkg/cli/install.go @@ -36,6 +36,7 @@ var ( validateSample bool validatePreset bool installEnvAll bool + keepSetting bool goGetterSources []string goGetterDest string listRegistryNixBuild bool @@ -270,6 +271,7 @@ func runInstallBase(cmd *cobra.Command, args []string) error { cfg.BinariesPath, headers, ) + inst.KeepSetting = keepSetting if err := inst.InstallBase(presetURL); err != nil { return err } @@ -333,6 +335,7 @@ func runInstallBase(cmd *cobra.Command, args []string) error { cfg.BinariesPath, headers, ) + inst.KeepSetting = keepSetting if err := inst.InstallBase(source); err != nil { return err @@ -726,7 +729,7 @@ func installBinariesViaNix(names []string, registry installer.BinaryRegistry, bi if !installer.IsCoreUnixTool(name) { entry, ok := registry[name] if ok { - _ = installer.CopyInstalledBinaryToFolder(name, &entry, binariesFolder) + _ = installer.SymlinkInstalledBinaryToFolder(name, &entry, binariesFolder) } } printer.Info("Binary '%s' already available in PATH, skipping", terminal.HiBlue(name)) @@ -1138,6 +1141,7 @@ func runInstallValidate(cmd *cobra.Command, args []string) error { cfg.BinariesPath, headers, ) + inst.KeepSetting = keepSetting if err := inst.InstallBase(presetURL); err != nil { return err } @@ -1586,6 +1590,7 @@ func init() { installBaseCmd.Flags().BoolVar(&baseSample, "sample", false, "initialize base folder from embedded sample (replaces existing base folder)") installBaseCmd.Flags().BoolVar(&basePreset, "preset", false, "install from OSM_PRESET_URL environment variable (default: DEFAULT_BASE_REPO)") + installBaseCmd.Flags().BoolVar(&keepSetting, "keep-setting", false, "restore previous osm-settings.yaml after base installation") installWorkflowCmd.Flags().BoolVar(&workflowPreset, "preset", false, "install from OSM_WORKFLOW_URL environment variable (default: DEFAULT_WORKFLOW_REPO)") // Note: --force flag is now global (defined in root.go) @@ -1621,6 +1626,7 @@ func init() { installValidateCmd.Flags().BoolVar(&validateSample, "sample", false, "initialize base folder from embedded sample (replaces existing base folder)") installValidateCmd.Flags().BoolVar(&validatePreset, "preset", false, "install ready-to-use base from default repository") + installValidateCmd.Flags().BoolVar(&keepSetting, "keep-setting", false, "restore previous osm-settings.yaml after base installation") } // parseCustomHeaders converts the string slice of "Key: Value" pairs to a map @@ -1782,7 +1788,7 @@ func installBinariesParallel(names []string, registry installer.BinaryRegistry, if !installer.IsCoreUnixTool(name) { entry, ok := registry[name] if ok { - _ = installer.CopyInstalledBinaryToFolder(name, &entry, binariesFolder) + _ = installer.SymlinkInstalledBinaryToFolder(name, &entry, binariesFolder) } } status[name] = "installed" diff --git a/pkg/cli/run.go b/pkg/cli/run.go index ee80edb..9e719df 100644 --- a/pkg/cli/run.go +++ b/pkg/cli/run.go @@ -1110,6 +1110,7 @@ func createCLIRunRecord(ctx context.Context, cfg *config.Config, workflow *core. Workspace: workspace, RunPriority: "critical", // CLI runs execute immediately RunMode: "local", + HooksEnabled: workflow.HookCount() > 0, } if err := database.CreateRun(ctx, run); err != nil { diff --git a/pkg/cli/usage.go b/pkg/cli/usage.go index ca0738d..0ed1a9d 100644 --- a/pkg/cli/usage.go +++ b/pkg/cli/usage.go @@ -324,7 +324,9 @@ func UsageWorker() string { ` + terminal.BoldCyan("▶ Subcommands") + ` • ` + terminal.Yellow("join") + ` - Join the distributed worker pool - • ` + terminal.Yellow("status") + ` - Show worker pool status + • ` + terminal.Yellow("status") + ` - Show worker pool status (alias: ls); use ` + terminal.Yellow("--json") + ` for JSON output + • ` + terminal.Yellow("set") + ` - Update a worker field (alias, public-ip, ssh-enabled, ssh-keys-path) + • ` + terminal.Yellow("eval") + ` - Evaluate a function expression with distributed hooks ` + docsFooter() } @@ -344,6 +346,9 @@ func UsageWorkerJoin() string { ` + terminal.Green("# Join using a specific Redis URL") + ` osmedeus worker join ` + terminal.Yellow("--redis-url") + ` redis://user:pass@localhost:6379/0 + ` + terminal.Green("# Join and auto-detect public IP") + ` + osmedeus worker join ` + terminal.Yellow("--get-public-ip") + ` + ` + docsFooter() } @@ -352,6 +357,69 @@ func UsageWorkerStatus() string { return terminal.BoldCyan("◆ Description") + ` Display the status of all workers connected to the Redis server. + Use ` + terminal.Yellow("--json") + ` to output worker info as JSON for scripting and automation. + +` + terminal.BoldCyan("▷ Examples") + ` + ` + terminal.Green("# Show worker status as a table") + ` + osmedeus worker status + + ` + terminal.Green("# Output worker info as JSON") + ` + osmedeus worker status ` + terminal.Yellow("--json") + ` + +` + docsFooter() +} + +// UsageWorkerEval returns the Long description for the worker eval command +func UsageWorkerEval() string { + return terminal.BoldCyan("◆ Description") + ` + Evaluate a utility function expression with distributed hooks registered. + + This connects to Redis and registers run_on_master() hooks so that + expressions can route calls to the master node. Useful for one-shot + operations from a worker context (e.g., inside Docker or CI pipelines) + without running a full worker loop. + +` + terminal.BoldCyan("▷ Examples") + ` + ` + terminal.Green("# Simple function eval with distributed hooks") + ` + osmedeus worker eval 'log_info("hello from worker eval")' ` + terminal.Yellow("--redis-url") + ` redis://localhost:6379 + + ` + terminal.Green("# Route a call to the master node") + ` + osmedeus worker eval 'run_on_master("func", "log_info(\"routed via redis\")")' ` + terminal.Yellow("--redis-url") + ` redis://localhost:6379 + + ` + terminal.Green("# With target variable") + ` + osmedeus worker eval 'log_info("hello")' ` + terminal.Yellow("-t") + ` example.com ` + terminal.Yellow("--redis-url") + ` redis://localhost:6379 + + ` + terminal.Green("# Read script from stdin") + ` + echo 'run_on_master("func", "db_import_sarif(\"ws\", \"/path/f.sarif\")")' | osmedeus worker eval ` + terminal.Yellow("--stdin") + ` ` + terminal.Yellow("--redis-url") + ` redis://localhost:6379 + +` + docsFooter() +} + +// UsageWorkerSet returns the Long description for the worker set command +func UsageWorkerSet() string { + return terminal.BoldCyan("◆ Description") + ` + Update a field on a registered worker. The worker can be identified by its + ID or alias. + +` + terminal.BoldCyan("▶ Valid Fields") + ` + • ` + terminal.Yellow("alias") + ` - Human-friendly name for the worker + • ` + terminal.Yellow("public-ip") + ` - Public IP address + • ` + terminal.Yellow("ssh-enabled") + ` - Whether SSH is enabled (true/false) + • ` + terminal.Yellow("ssh-keys-path") + ` - Path to SSH keys + +` + terminal.BoldCyan("▷ Examples") + ` + ` + terminal.Green("# Set an alias for a worker") + ` + osmedeus worker set alias scanner-1 + + ` + terminal.Green("# Set public IP") + ` + osmedeus worker set scanner-1 public-ip 203.0.113.10 + + ` + terminal.Green("# Enable SSH") + ` + osmedeus worker set scanner-1 ssh-enabled true + + ` + terminal.Green("# With custom Redis URL") + ` + osmedeus worker set alias prod-1 ` + terminal.Yellow("--redis-url") + ` redis://localhost:6379 + ` + docsFooter() } @@ -510,12 +578,14 @@ func UsageDBSeed() string { func UsageDBClean() string { return terminal.BoldCyan("◆ Description") + ` Remove all data from all database tables. + Use --clean-ws to also remove workspace data (e.g. ~/workspaces-osmedeus). ` + terminal.Yellow("WARNING:") + ` This is a destructive operation that cannot be undone. Use the --force flag to skip the confirmation prompt. -` + terminal.BoldCyan("▷ Example") + ` +` + terminal.BoldCyan("▷ Examples") + ` ` + terminal.Green("osmedeus db clean --force") + ` + ` + terminal.Green("osmedeus db clean --force --clean-ws") + ` ` + docsFooter() } @@ -744,6 +814,9 @@ func UsageAllExamples() string { ` + terminal.Green("# Check worker status") + ` osmedeus worker status + ` + terminal.Green("# Evaluate function with distributed hooks (one-shot)") + ` + osmedeus worker eval 'run_on_master("func", "log_info(\"hello\")")' ` + terminal.Yellow("--redis-url") + ` redis://localhost:6379 + ` + terminal.BoldCyan("▶ Install Examples") + ` ` + terminal.Green("# Install binary") + ` osmedeus install binary ` + terminal.Yellow("--name") + ` nuclei @@ -900,6 +973,7 @@ func UsageFullExample() string { ` + terminal.Gray("───────────────────────────────────────────────────────────────────") + ` osmedeus worker join Join the distributed worker pool osmedeus worker status Show worker pool status + osmedeus worker eval