CoCalc -- workflow_execute

GitHub Repository: projectdiscovery/nuclei
Path: blob/dev/pkg/core/workflow_execute_test.go
²⁰⁷⁰ views
1
package core
2

3
import (
4
	"context"
5
	"testing"
6

7
	"github.com/projectdiscovery/nuclei/v3/pkg/model/types/stringslice"
8
	"github.com/projectdiscovery/nuclei/v3/pkg/operators"
9
	"github.com/projectdiscovery/nuclei/v3/pkg/output"
10
	"github.com/projectdiscovery/nuclei/v3/pkg/progress"
11
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
12
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
13
	"github.com/projectdiscovery/nuclei/v3/pkg/scan"
14
	"github.com/projectdiscovery/nuclei/v3/pkg/types"
15
	"github.com/projectdiscovery/nuclei/v3/pkg/workflows"
16
	"github.com/stretchr/testify/require"
17
)
18

19
func TestWorkflowsSimple(t *testing.T) {
20
	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
21

22
	workflow := &workflows.Workflow{Options: &protocols.ExecutorOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
23
		{Executers: []*workflows.ProtocolExecuterPair{{
24
			Executer: &mockExecuter{result: true}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
25
		}},
26
	}}
27

28
	engine := &Engine{}
29
	input := contextargs.NewWithInput(context.Background(), "https://test.com")
30
	ctx := scan.NewScanContext(context.Background(), input)
31
	matched := engine.executeWorkflow(ctx, workflow)
32
	require.True(t, matched, "could not get correct match value")
33
}
34

35
func TestWorkflowsSimpleMultiple(t *testing.T) {
36
	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
37

38
	var firstInput, secondInput string
39
	workflow := &workflows.Workflow{Options: &protocols.ExecutorOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
40
		{Executers: []*workflows.ProtocolExecuterPair{{
41
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
42
				firstInput = input.Input
43
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
44
		}},
45
		{Executers: []*workflows.ProtocolExecuterPair{{
46
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
47
				secondInput = input.Input
48
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
49
		}},
50
	}}
51

52
	engine := &Engine{}
53
	input := contextargs.NewWithInput(context.Background(), "https://test.com")
54
	ctx := scan.NewScanContext(context.Background(), input)
55
	matched := engine.executeWorkflow(ctx, workflow)
56
	require.True(t, matched, "could not get correct match value")
57

58
	require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
59
	require.Equal(t, "https://test.com", secondInput, "could not get correct second input")
60
}
61

62
func TestWorkflowsSubtemplates(t *testing.T) {
63
	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
64

65
	var firstInput, secondInput string
66
	workflow := &workflows.Workflow{Options: &protocols.ExecutorOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
67
		{Executers: []*workflows.ProtocolExecuterPair{{
68
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
69
				firstInput = input.Input
70
			}, outputs: []*output.InternalWrappedEvent{
71
				{OperatorsResult: &operators.Result{}, Results: []*output.ResultEvent{{}}},
72
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
73
		}, Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{
74
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
75
				secondInput = input.Input
76
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
77
		}}}},
78
	}}
79

80
	engine := &Engine{}
81
	input := contextargs.NewWithInput(context.Background(), "https://test.com")
82
	ctx := scan.NewScanContext(context.Background(), input)
83
	matched := engine.executeWorkflow(ctx, workflow)
84
	require.True(t, matched, "could not get correct match value")
85

86
	require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
87
	require.Equal(t, "https://test.com", secondInput, "could not get correct second input")
88
}
89

90
func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
91
	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
92

93
	var firstInput, secondInput string
94
	workflow := &workflows.Workflow{Options: &protocols.ExecutorOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
95
		{Executers: []*workflows.ProtocolExecuterPair{{
96
			Executer: &mockExecuter{result: false, executeHook: func(input *contextargs.MetaInput) {
97
				firstInput = input.Input
98
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
99
		}, Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{
100
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
101
				secondInput = input.Input
102
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
103
		}}}},
104
	}}
105

106
	engine := &Engine{}
107
	input := contextargs.NewWithInput(context.Background(), "https://test.com")
108
	ctx := scan.NewScanContext(context.Background(), input)
109
	matched := engine.executeWorkflow(ctx, workflow)
110
	require.False(t, matched, "could not get correct match value")
111

112
	require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
113
	require.Equal(t, "", secondInput, "could not get correct second input")
114
}
115

116
func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
117
	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
118

119
	var firstInput, secondInput string
120
	workflow := &workflows.Workflow{Options: &protocols.ExecutorOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
121
		{Executers: []*workflows.ProtocolExecuterPair{{
122
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
123
				firstInput = input.Input
124
			}, outputs: []*output.InternalWrappedEvent{
125
				{OperatorsResult: &operators.Result{
126
					Matches:  map[string][]string{"tomcat": {}},
127
					Extracts: map[string][]string{},
128
				}},
129
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
130
		}, Matchers: []*workflows.Matcher{{Name: stringslice.StringSlice{Value: "tomcat"}, Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{
131
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
132
				secondInput = input.Input
133
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
134
		}}}}}},
135
	}}
136

137
	engine := &Engine{}
138
	input := contextargs.NewWithInput(context.Background(), "https://test.com")
139
	ctx := scan.NewScanContext(context.Background(), input)
140
	matched := engine.executeWorkflow(ctx, workflow)
141
	require.True(t, matched, "could not get correct match value")
142

143
	require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
144
	require.Equal(t, "https://test.com", secondInput, "could not get correct second input")
145
}
146

147
func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) {
148
	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
149

150
	var firstInput, secondInput string
151
	workflow := &workflows.Workflow{Options: &protocols.ExecutorOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
152
		{Executers: []*workflows.ProtocolExecuterPair{{
153
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
154
				firstInput = input.Input
155
			}, outputs: []*output.InternalWrappedEvent{
156
				{OperatorsResult: &operators.Result{
157
					Matches:  map[string][]string{"tomcat": {}},
158
					Extracts: map[string][]string{},
159
				}},
160
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
161
		}, Matchers: []*workflows.Matcher{{Name: stringslice.StringSlice{Value: "apache"}, Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{
162
			Executer: &mockExecuter{result: true, executeHook: func(input *contextargs.MetaInput) {
163
				secondInput = input.Input
164
			}}, Options: &protocols.ExecutorOptions{Progress: progressBar}},
165
		}}}}}},
166
	}}
167

168
	engine := &Engine{}
169
	input := contextargs.NewWithInput(context.Background(), "https://test.com")
170
	ctx := scan.NewScanContext(context.Background(), input)
171
	matched := engine.executeWorkflow(ctx, workflow)
172
	require.False(t, matched, "could not get correct match value")
173

174
	require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
175
	require.Equal(t, "", secondInput, "could not get correct second input")
176
}
177

178
type mockExecuter struct {
179
	result      bool
180
	executeHook func(input *contextargs.MetaInput)
181
	outputs     []*output.InternalWrappedEvent
182
}
183

184
// Compile compiles the execution generators preparing any requests possible.
185
func (m *mockExecuter) Compile() error {
186
	return nil
187
}
188

189
// Requests returns the total number of requests the rule will perform
190
func (m *mockExecuter) Requests() int {
191
	return 1
192
}
193

194
// Execute executes the protocol group and  returns true or false if results were found.
195
func (m *mockExecuter) Execute(ctx *scan.ScanContext) (bool, error) {
196
	if m.executeHook != nil {
197
		m.executeHook(ctx.Input.MetaInput)
198
	}
199
	return m.result, nil
200
}
201

202
// ExecuteWithResults executes the protocol requests and returns results instead of writing them.
203
func (m *mockExecuter) ExecuteWithResults(ctx *scan.ScanContext) ([]*output.ResultEvent, error) {
204
	if m.executeHook != nil {
205
		m.executeHook(ctx.Input.MetaInput)
206
	}
207
	for _, output := range m.outputs {
208
		ctx.LogEvent(output)
209
	}
210
	return ctx.GenerateResult(), nil
211
}
212

213
Product

Resources

Company
