From 113ccb1e0ed00c4f8ddbffe2ed5b140d22c174bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Zamanillo?= Date: Wed, 26 Aug 2020 20:05:31 +0200 Subject: [PATCH] Preload workflow templates once Fixes memory leak reported on #242 --- internal/runner/config.go | 4 ++ internal/runner/runner.go | 86 ++++++++++++++++++++++++--------------- pkg/matchers/matchers.go | 1 + pkg/templates/compile.go | 3 +- 4 files changed, 59 insertions(+), 35 deletions(-) diff --git a/internal/runner/config.go b/internal/runner/config.go index 152de9c8..93a0d2c3 100644 --- a/internal/runner/config.go +++ b/internal/runner/config.go @@ -106,6 +106,7 @@ func (r *Runner) readNucleiIgnoreFile() { if text == "" { continue } + r.templatesConfig.IgnorePaths = append(r.templatesConfig.IgnorePaths, text) } } @@ -115,12 +116,14 @@ func (r *Runner) checkIfInNucleiIgnore(item string) bool { if r.templatesConfig == nil { return false } + for _, paths := range r.templatesConfig.IgnorePaths { // If we have a path to ignore, check if it's in the item. if paths[len(paths)-1] == '/' { if strings.Contains(item, paths) { return true } + continue } // Check for file based extension in ignores @@ -128,6 +131,7 @@ func (r *Runner) checkIfInNucleiIgnore(item string) bool { return true } } + return false } diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 928ff285..56bfb0b0 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -51,6 +51,12 @@ type Runner struct { decolorizer *regexp.Regexp } +// WorkflowTemplates contains the initialized workflow templates per template group +type WorkflowTemplates struct { + Name string + Templates []*workflows.Template +} + // New creates a new client for running enumeration process. func New(options *Options) (*Runner, error) { runner := &Runner{ @@ -240,6 +246,7 @@ func (r *Runner) getParsedTemplatesFor(templatePaths []string, severities string filterBySeverity := len(severities) > 0 gologger.Infof("Loading templates...") + for _, match := range templatePaths { t, err := r.parse(match) switch tp := t.(type) { @@ -315,6 +322,7 @@ func (r *Runner) getTemplatesFor(definitions []string) []string { for _, match := range matches { if !r.checkIfInNucleiIgnore(match) { processed[match] = true + allTemplates = append(allTemplates, match) } } @@ -427,7 +435,7 @@ func (r *Runner) RunEnumeration() { case *workflows.Workflow: // workflows will dynamically adjust the totals while running, as // it can't be know in advance which requests will be called - } + } // nolint:wsl // comment } var ( @@ -575,34 +583,51 @@ func (r *Runner) processTemplateWithList(ctx context.Context, p progress.IProgre // ProcessWorkflowWithList coming from stdin or list of targets func (r *Runner) ProcessWorkflowWithList(p progress.IProgress, workflow *workflows.Workflow) { + workflowTemplatesList, err := r.PreloadTemplates(p, workflow) + if err != nil { + gologger.Warningf("Could not preload templates for workflow %s: %s\n", workflow.ID, err) + + return + } + var wg sync.WaitGroup scanner := bufio.NewScanner(strings.NewReader(r.input)) for scanner.Scan() { - text := scanner.Text() + targetURL := scanner.Text() r.limiter <- struct{}{} wg.Add(1) - go func(text string) { + go func(targetURL string) { defer wg.Done() - if err := r.ProcessWorkflow(p, workflow, text); err != nil { - gologger.Warningf("Could not run workflow for %s: %s\n", text, err) + script := tengo.NewScript([]byte(workflow.Logic)) + script.SetImports(stdlib.GetModuleMap(stdlib.AllModuleNames()...)) + + for _, workflowTemplate := range *workflowTemplatesList { + err := script.Add(workflowTemplate.Name, &workflows.NucleiVar{Templates: workflowTemplate.Templates, URL: targetURL}) + if err != nil { + gologger.Errorf("Could not initialize script for workflow '%s': %s\n", workflow.ID, err) + + continue + } + } + + _, err := script.RunContext(context.Background()) + if err != nil { + gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err) } <-r.limiter - }(text) + }(targetURL) } wg.Wait() } -// ProcessWorkflow towards an URL -func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workflow, towardURL string) error { - script := tengo.NewScript([]byte(workflow.Logic)) - script.SetImports(stdlib.GetModuleMap(stdlib.AllModuleNames()...)) - +// PreloadTemplates preload the workflow templates once +func (r *Runner) PreloadTemplates(p progress.IProgress, workflow *workflows.Workflow) (*[]WorkflowTemplates, error) { var jar *cookiejar.Jar if workflow.CookieReuse { @@ -610,10 +635,13 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf jar, err = cookiejar.New(nil) if err != nil { - return err + return nil, err } } + // Single yaml provided + var wflTemplatesList []WorkflowTemplates + for name, value := range workflow.Variables { var writer *bufio.Writer if r.output != nil { @@ -628,20 +656,19 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf if err != nil { newPath, err = r.resolvePathWithBaseFolder(filepath.Dir(workflow.GetPath()), value) if err != nil { - return err + return nil, err } } value = newPath } - // Single yaml provided - var templatesList []*workflows.Template + var wtlst []*workflows.Template if strings.HasSuffix(value, ".yaml") { t, err := templates.Parse(value) if err != nil { - return err + return nil, err } template := &workflows.Template{Progress: p} @@ -672,7 +699,7 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf } if template.DNSOptions != nil || template.HTTPOptions != nil { - templatesList = append(templatesList, template) + wtlst = append(wtlst, template) } } else { matches := []string{} @@ -682,6 +709,7 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf if !d.IsDir() && strings.HasSuffix(path, ".yaml") { matches = append(matches, path) } + return nil }, ErrorCallback: func(path string, err error) godirwalk.ErrorAction { @@ -689,18 +717,20 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf }, Unsorted: true, }) + if err != nil { - return err + return nil, err } + // 0 matches means no templates were found in directory if len(matches) == 0 { - return errors.New("no match found in the directory") + return nil, fmt.Errorf("no match found in the directory %s", value) } for _, match := range matches { t, err := templates.Parse(match) if err != nil { - return err + return nil, err } template := &workflows.Template{Progress: p} if len(t.BulkRequestsHTTP) > 0 { @@ -723,25 +753,15 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf } } if template.DNSOptions != nil || template.HTTPOptions != nil { - templatesList = append(templatesList, template) + wtlst = append(wtlst, template) } } } - err := script.Add(name, &workflows.NucleiVar{Templates: templatesList, URL: towardURL}) - if err != nil { - gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err) - return err - } + wflTemplatesList = append(wflTemplatesList, WorkflowTemplates{Name: name, Templates: wtlst}) } - _, err := script.RunContext(context.Background()) - if err != nil { - gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err) - return err - } - - return nil + return &wflTemplatesList, nil } func (r *Runner) parse(file string) (interface{}, error) { diff --git a/pkg/matchers/matchers.go b/pkg/matchers/matchers.go index 4358c009..be4b4af3 100644 --- a/pkg/matchers/matchers.go +++ b/pkg/matchers/matchers.go @@ -125,5 +125,6 @@ func (m *Matcher) isNegative(data bool) bool { if m.Negative { return !data } + return data } diff --git a/pkg/templates/compile.go b/pkg/templates/compile.go index d3622339..f15d42b2 100644 --- a/pkg/templates/compile.go +++ b/pkg/templates/compile.go @@ -1,7 +1,6 @@ package templates import ( - "errors" "fmt" "os" "path" @@ -31,7 +30,7 @@ func Parse(file string) (*Template, error) { // If no requests, and it is also not a workflow, return error. if len(template.BulkRequestsHTTP)+len(template.RequestsDNS) <= 0 { - return nil, errors.New("no requests defined") + return nil, fmt.Errorf("no requests defined for %s", template.ID) } // Compile the matchers and the extractors for http requests