Preload workflow templates once

Fixes memory leak reported on #242
dev
Víctor Zamanillo 2020-08-26 20:05:31 +02:00
parent 60005290b1
commit 113ccb1e0e
4 changed files with 59 additions and 35 deletions

View File

@ -106,6 +106,7 @@ func (r *Runner) readNucleiIgnoreFile() {
if text == "" { if text == "" {
continue continue
} }
r.templatesConfig.IgnorePaths = append(r.templatesConfig.IgnorePaths, text) r.templatesConfig.IgnorePaths = append(r.templatesConfig.IgnorePaths, text)
} }
} }
@ -115,12 +116,14 @@ func (r *Runner) checkIfInNucleiIgnore(item string) bool {
if r.templatesConfig == nil { if r.templatesConfig == nil {
return false return false
} }
for _, paths := range r.templatesConfig.IgnorePaths { for _, paths := range r.templatesConfig.IgnorePaths {
// If we have a path to ignore, check if it's in the item. // If we have a path to ignore, check if it's in the item.
if paths[len(paths)-1] == '/' { if paths[len(paths)-1] == '/' {
if strings.Contains(item, paths) { if strings.Contains(item, paths) {
return true return true
} }
continue continue
} }
// Check for file based extension in ignores // Check for file based extension in ignores
@ -128,6 +131,7 @@ func (r *Runner) checkIfInNucleiIgnore(item string) bool {
return true return true
} }
} }
return false return false
} }

View File

@ -51,6 +51,12 @@ type Runner struct {
decolorizer *regexp.Regexp 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. // New creates a new client for running enumeration process.
func New(options *Options) (*Runner, error) { func New(options *Options) (*Runner, error) {
runner := &Runner{ runner := &Runner{
@ -240,6 +246,7 @@ func (r *Runner) getParsedTemplatesFor(templatePaths []string, severities string
filterBySeverity := len(severities) > 0 filterBySeverity := len(severities) > 0
gologger.Infof("Loading templates...") gologger.Infof("Loading templates...")
for _, match := range templatePaths { for _, match := range templatePaths {
t, err := r.parse(match) t, err := r.parse(match)
switch tp := t.(type) { switch tp := t.(type) {
@ -315,6 +322,7 @@ func (r *Runner) getTemplatesFor(definitions []string) []string {
for _, match := range matches { for _, match := range matches {
if !r.checkIfInNucleiIgnore(match) { if !r.checkIfInNucleiIgnore(match) {
processed[match] = true processed[match] = true
allTemplates = append(allTemplates, match) allTemplates = append(allTemplates, match)
} }
} }
@ -427,7 +435,7 @@ func (r *Runner) RunEnumeration() {
case *workflows.Workflow: case *workflows.Workflow:
// workflows will dynamically adjust the totals while running, as // workflows will dynamically adjust the totals while running, as
// it can't be know in advance which requests will be called // it can't be know in advance which requests will be called
} } // nolint:wsl // comment
} }
var ( var (
@ -575,34 +583,51 @@ func (r *Runner) processTemplateWithList(ctx context.Context, p progress.IProgre
// ProcessWorkflowWithList coming from stdin or list of targets // ProcessWorkflowWithList coming from stdin or list of targets
func (r *Runner) ProcessWorkflowWithList(p progress.IProgress, workflow *workflows.Workflow) { 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 var wg sync.WaitGroup
scanner := bufio.NewScanner(strings.NewReader(r.input)) scanner := bufio.NewScanner(strings.NewReader(r.input))
for scanner.Scan() { for scanner.Scan() {
text := scanner.Text() targetURL := scanner.Text()
r.limiter <- struct{}{} r.limiter <- struct{}{}
wg.Add(1) wg.Add(1)
go func(text string) { go func(targetURL string) {
defer wg.Done() defer wg.Done()
if err := r.ProcessWorkflow(p, workflow, text); err != nil { script := tengo.NewScript([]byte(workflow.Logic))
gologger.Warningf("Could not run workflow for %s: %s\n", text, err) 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 <-r.limiter
}(text) }(targetURL)
} }
wg.Wait() wg.Wait()
} }
// ProcessWorkflow towards an URL // PreloadTemplates preload the workflow templates once
func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workflow, towardURL string) error { func (r *Runner) PreloadTemplates(p progress.IProgress, workflow *workflows.Workflow) (*[]WorkflowTemplates, error) {
script := tengo.NewScript([]byte(workflow.Logic))
script.SetImports(stdlib.GetModuleMap(stdlib.AllModuleNames()...))
var jar *cookiejar.Jar var jar *cookiejar.Jar
if workflow.CookieReuse { if workflow.CookieReuse {
@ -610,10 +635,13 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf
jar, err = cookiejar.New(nil) jar, err = cookiejar.New(nil)
if err != nil { if err != nil {
return err return nil, err
} }
} }
// Single yaml provided
var wflTemplatesList []WorkflowTemplates
for name, value := range workflow.Variables { for name, value := range workflow.Variables {
var writer *bufio.Writer var writer *bufio.Writer
if r.output != nil { if r.output != nil {
@ -628,20 +656,19 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf
if err != nil { if err != nil {
newPath, err = r.resolvePathWithBaseFolder(filepath.Dir(workflow.GetPath()), value) newPath, err = r.resolvePathWithBaseFolder(filepath.Dir(workflow.GetPath()), value)
if err != nil { if err != nil {
return err return nil, err
} }
} }
value = newPath value = newPath
} }
// Single yaml provided var wtlst []*workflows.Template
var templatesList []*workflows.Template
if strings.HasSuffix(value, ".yaml") { if strings.HasSuffix(value, ".yaml") {
t, err := templates.Parse(value) t, err := templates.Parse(value)
if err != nil { if err != nil {
return err return nil, err
} }
template := &workflows.Template{Progress: p} 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 { if template.DNSOptions != nil || template.HTTPOptions != nil {
templatesList = append(templatesList, template) wtlst = append(wtlst, template)
} }
} else { } else {
matches := []string{} matches := []string{}
@ -682,6 +709,7 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf
if !d.IsDir() && strings.HasSuffix(path, ".yaml") { if !d.IsDir() && strings.HasSuffix(path, ".yaml") {
matches = append(matches, path) matches = append(matches, path)
} }
return nil return nil
}, },
ErrorCallback: func(path string, err error) godirwalk.ErrorAction { ErrorCallback: func(path string, err error) godirwalk.ErrorAction {
@ -689,18 +717,20 @@ func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workf
}, },
Unsorted: true, Unsorted: true,
}) })
if err != nil { if err != nil {
return err return nil, err
} }
// 0 matches means no templates were found in directory // 0 matches means no templates were found in directory
if len(matches) == 0 { 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 { for _, match := range matches {
t, err := templates.Parse(match) t, err := templates.Parse(match)
if err != nil { if err != nil {
return err return nil, err
} }
template := &workflows.Template{Progress: p} template := &workflows.Template{Progress: p}
if len(t.BulkRequestsHTTP) > 0 { 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 { 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}) wflTemplatesList = append(wflTemplatesList, WorkflowTemplates{Name: name, Templates: wtlst})
if err != nil {
gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err)
return err
}
} }
_, err := script.RunContext(context.Background()) return &wflTemplatesList, nil
if err != nil {
gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err)
return err
}
return nil
} }
func (r *Runner) parse(file string) (interface{}, error) { func (r *Runner) parse(file string) (interface{}, error) {

View File

@ -125,5 +125,6 @@ func (m *Matcher) isNegative(data bool) bool {
if m.Negative { if m.Negative {
return !data return !data
} }
return data return data
} }

View File

@ -1,7 +1,6 @@
package templates package templates
import ( import (
"errors"
"fmt" "fmt"
"os" "os"
"path" "path"
@ -31,7 +30,7 @@ func Parse(file string) (*Template, error) {
// If no requests, and it is also not a workflow, return error. // If no requests, and it is also not a workflow, return error.
if len(template.BulkRequestsHTTP)+len(template.RequestsDNS) <= 0 { 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 // Compile the matchers and the extractors for http requests