Max limit on execution of a workflow

This commit is contained in:
Ice3man543 2020-12-13 14:17:58 +05:30
parent cf2a98b03d
commit 2ae13d845a
2 changed files with 8 additions and 4 deletions

View File

@ -14,6 +14,7 @@ import (
// the template requesting process.
// nolint // false positive, options are allocated once and are necessary as is
type Options struct {
MaxWorkflowDuration int // MaxWorkflowDuration is the maximum time a workflow can run for a URL
Sandbox bool // Sandbox mode allows users to run isolated workflows with system commands disabled
Debug bool // Debug mode allows debugging request/responses for the engine
Silent bool // Silent suppresses any extra text and only writes found URLs on screen.
@ -68,6 +69,7 @@ func ParseOptions() *Options {
options := &Options{}
flag.BoolVar(&options.Sandbox, "sandbox", false, "Run workflows in isolated sandbox mode")
flag.IntVar(&options.MaxWorkflowDuration, "workflow-duration", 10, "Max time for workflow run on single URL in minutes")
flag.StringVar(&options.Target, "target", "", "Target is a single target to scan using template")
flag.Var(&options.Templates, "t", "Template input dir/file/files to run on host. Can be used multiple times. Supports globbing.")
flag.Var(&options.ExcludedTemplates, "exclude", "Template input dir/file/files to exclude. Can be used multiple times. Supports globbing.")

View File

@ -8,6 +8,7 @@ import (
"path"
"path/filepath"
"strings"
"time"
tengo "github.com/d5/tengo/v2"
"github.com/d5/tengo/v2/stdlib"
@ -132,11 +133,9 @@ func (r *Runner) processWorkflowWithList(p *progress.Progress, workflow *workflo
gologger.Warningf("Could not preload templates for workflow %s: %s\n", workflow.ID, err)
return false
}
logicBytes := []byte(workflow.Logic)
wg := sizedwaitgroup.New(r.options.BulkSize)
r.hm.Scan(func(k, _ []byte) error {
targetURL := string(k)
wg.Add()
@ -163,7 +162,10 @@ func (r *Runner) processWorkflowWithList(p *progress.Progress, workflow *workflo
variables[name] = variable
}
_, err := script.RunContext(context.Background())
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(r.options.MaxWorkflowDuration)*time.Minute)
defer cancel()
_, err := script.RunContext(ctx)
if err != nil {
gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err)
}