External/nuclei
1
0
Fork 0
mirror of https://github.com/projectdiscovery/nuclei.git synced 2025-12-18 03:25:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added parallelism to workflows

Browse Source
This commit is contained in:
Ice3man543 2021-02-22 17:49:02 +05:30
parent dfa21c1896
commit ec3f85e032
5 changed files with 50 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
v2/internal/runner/processor.go
View File

@ -44,10 +44,7 @@ func (r *Runner) processWorkflowWithList(template *templates.Template) bool {
wg.Add()
go func(URL string) {
defer wg.Done()
match, err := template.CompiledWorkflow.RunWorkflow(URL)
if err != nil {
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", r.colorizer.BrightBlue(template.ID), err)
}
match := template.CompiledWorkflow.RunWorkflow(URL)
results.CAS(false, match)
}(URL)
return nil

5
v2/pkg/workflows/compile.go
View File

@ -4,12 +4,13 @@ import (
"os"
"github.com/pkg/errors"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
"gopkg.in/yaml.v2"
)
// Parse a yaml workflow file
func Parse(file string) (*Workflow, error) {
workflow := &Workflow{}
func Parse(file string, options *protocols.ExecuterOptions) (*Workflow, error) {
workflow := &Workflow{options: options}
f, err := os.Open(file)
if err != nil {

55
v2/pkg/workflows/execute.go
View File

@ -3,25 +3,32 @@ package workflows
import (
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v2/pkg/output"
"github.com/remeh/sizedwaitgroup"
"go.uber.org/atomic"
)
// RunWorkflow runs a workflow on an input and returns true or false
func (w *Workflow) RunWorkflow(input string) (bool, error) {
func (w *Workflow) RunWorkflow(input string) bool {
results := &atomic.Bool{}
swg := sizedwaitgroup.New(w.options.Options.TemplateThreads)
for _, template := range w.Workflows {
err := w.runWorkflowStep(template, input, results)
if err != nil {
return results.Load(), err
}
swg.Add()
func(template *WorkflowTemplate) {
err := w.runWorkflowStep(template, input, results, &swg)
if err != nil {
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err)
}
swg.Done()
}(template)
}
return results.Load(), nil
swg.Wait()
return results.Load()
}
// runWorkflowStep runs a workflow step for the workflow. It executes the workflow
// in a recursive manner running all subtemplates and matchers.
func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, results *atomic.Bool) error {
func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup) error {
var firstMatched bool
var mainErr error
@ -46,9 +53,6 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
}
if len(template.Matchers) > 0 {
var executionErr error
var mainErr error
for _, executer := range template.Executers {
executer.Options.Progress.AddToTotal(int64(executer.Executer.Requests()))
@ -65,10 +69,14 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
}
for _, subtemplate := range matcher.Subtemplates {
if err := w.runWorkflowStep(subtemplate, input, results); err != nil {
executionErr = err
break
}
swg.Add()
go func(subtemplate *WorkflowTemplate) {
if err := w.runWorkflowStep(subtemplate, input, results, swg); err != nil {
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", subtemplate.Template, err)
}
swg.Done()
}(subtemplate)
}
}
})
@ -80,21 +88,20 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
}
continue
}
if executionErr != nil {
if len(template.Executers) == 1 {
mainErr = executionErr
} else {
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, executionErr)
}
}
}
return mainErr
}
if len(template.Subtemplates) > 0 && firstMatched {
for _, subtemplate := range template.Subtemplates {
if err := w.runWorkflowStep(subtemplate, input, results); err != nil {
return err
}
swg.Add()
func(template *WorkflowTemplate) {
err := w.runWorkflowStep(template, input, results, swg)
if err != nil {
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err)
}
swg.Done()
}(subtemplate)
}
}
return mainErr

31
v2/pkg/workflows/execute_test.go
View File

@ -7,20 +7,20 @@ import (
"github.com/projectdiscovery/nuclei/v2/pkg/operators"
"github.com/projectdiscovery/nuclei/v2/pkg/output"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
"github.com/projectdiscovery/nuclei/v2/pkg/types"
"github.com/stretchr/testify/require"
)
func TestWorkflowsSimple(t *testing.T) {
progress, _ := progress.NewProgress(false, false, 0)
workflow := &Workflow{Workflows: []*WorkflowTemplate{
workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
{Executers: []*ProtocolExecuterPair{{
Executer: &mockExecuter{result: true}, Options: &protocols.ExecuterOptions{Progress: progress}},
}},
}}
matched, err := workflow.RunWorkflow("https://test.com")
require.Nil(t, err, "could not run workflow")
matched := workflow.RunWorkflow("https://test.com")
require.True(t, matched, "could not get correct match value")
}
@ -28,7 +28,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) {
progress, _ := progress.NewProgress(false, false, 0)
var firstInput, secondInput string
workflow := &Workflow{Workflows: []*WorkflowTemplate{
workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
{Executers: []*ProtocolExecuterPair{{
Executer: &mockExecuter{result: true, executeHook: func(input string) {
firstInput = input
@ -41,8 +41,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) {
}},
}}
matched, err := workflow.RunWorkflow("https://test.com")
require.Nil(t, err, "could not run workflow")
matched := workflow.RunWorkflow("https://test.com")
require.True(t, matched, "could not get correct match value")
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
@ -53,7 +52,7 @@ func TestWorkflowsSubtemplates(t *testing.T) {
progress, _ := progress.NewProgress(false, false, 0)
var firstInput, secondInput string
workflow := &Workflow{Workflows: []*WorkflowTemplate{
workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
{Executers: []*ProtocolExecuterPair{{
Executer: &mockExecuter{result: true, executeHook: func(input string) {
firstInput = input
@ -65,8 +64,7 @@ func TestWorkflowsSubtemplates(t *testing.T) {
}}}},
}}
matched, err := workflow.RunWorkflow("https://test.com")
require.Nil(t, err, "could not run workflow")
matched := workflow.RunWorkflow("https://test.com")
require.True(t, matched, "could not get correct match value")
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
@ -77,7 +75,7 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
progress, _ := progress.NewProgress(false, false, 0)
var firstInput, secondInput string
workflow := &Workflow{Workflows: []*WorkflowTemplate{
workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
{Executers: []*ProtocolExecuterPair{{
Executer: &mockExecuter{result: false, executeHook: func(input string) {
firstInput = input
@ -89,8 +87,7 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
}}}},
}}
matched, err := workflow.RunWorkflow("https://test.com")
require.Nil(t, err, "could not run workflow")
matched := workflow.RunWorkflow("https://test.com")
require.False(t, matched, "could not get correct match value")
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
@ -101,7 +98,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
progress, _ := progress.NewProgress(false, false, 0)
var firstInput, secondInput string
workflow := &Workflow{Workflows: []*WorkflowTemplate{
workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
{Executers: []*ProtocolExecuterPair{{
Executer: &mockExecuter{result: true, executeHook: func(input string) {
firstInput = input
@ -118,8 +115,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
}}}}}},
}}
matched, err := workflow.RunWorkflow("https://test.com")
require.Nil(t, err, "could not run workflow")
matched := workflow.RunWorkflow("https://test.com")
require.True(t, matched, "could not get correct match value")
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
@ -130,7 +126,7 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) {
progress, _ := progress.NewProgress(false, false, 0)
var firstInput, secondInput string
workflow := &Workflow{Workflows: []*WorkflowTemplate{
workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
{Executers: []*ProtocolExecuterPair{{
Executer: &mockExecuter{result: true, executeHook: func(input string) {
firstInput = input
@ -147,8 +143,7 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) {
}}}}}},
}}
matched, err := workflow.RunWorkflow("https://test.com")
require.Nil(t, err, "could not run workflow")
matched := workflow.RunWorkflow("https://test.com")
require.False(t, matched, "could not get correct match value")
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")

2
v2/pkg/workflows/workflows.go
View File

@ -6,6 +6,8 @@ import "github.com/projectdiscovery/nuclei/v2/pkg/protocols"
type Workflow struct {
// Workflows is a yaml based workflow declaration code.
Workflows []*WorkflowTemplate `yaml:"workflows,omitempty"`
options *protocols.ExecuterOptions
}
// WorkflowTemplate is a template to be ran as part of a workflow