nuclei icon indicating copy to clipboard operation
nuclei copied to clipboard

[FEATURE] ... Add execution method for targets, templates, and callback in the SDK module

Open tongchengbin opened this issue 5 months ago • 1 comments

Describe your feature request

We would like to request a new method in the Nuclei SDK module that allows users to specify targets, templates, and a callback method for execution.

This method should:

Targets: Accept one or multiple targets as input for scanning. Templates: Specify a set of templates to use, allowing users to choose from predefined templates or define their own template files. Callback Method: Allow users to provide a callback method that processes scan results after execution (e.g., for logging or report generation).

Describe the use case of the feature

This is a version I implemented, but I am unable to control using a separate callback for each execution, as I need to carry task-specific information for different tasks. Therefore, I cannot use the same result_callback.

func (e *NucleiEngine) ExecuteWithProvider(ctx context.Context, target provider.InputProvider, templateIds []string, callback func(event *output.ResultEvent)) error {
	if target.Count() == 0 {
		return ErrNoTargetsAvailable
	}
	var filtered []func(event *output.ResultEvent)
	if callback != nil {
		filtered = append(filtered, callback)
	}
	templatesAndWorkflows := append(e.store.Templates(), e.store.Workflows()...)
	if len(templatesAndWorkflows) == 0 {
		return ErrNoTemplatesAvailable
	}
	// load templates
	var finalTemplates []*templates.Template
	for _, id := range templateIds {
		tmpl := e.templateMaps[FormatName(id)]
		if tmpl != nil {
			finalTemplates = append(finalTemplates, tmpl)
		}
	}
	_ = e.engine.ExecuteScanWithOpts(ctx, finalTemplates, target, false)
	defer e.engine.WorkPool().Wait()
	return nil
}

func TestMultipleTarget(t *testing.T) {
	ne, err := NewNucleiEngineCtx(
		context.Background(),
		WithVerbosity(VerbosityOptions{
			Debug: false,
		}),
		WithProxy([]string{"http://127.0.0.1:9000"}, true),
		WithTemplatesOrWorkflows(TemplateSources{
			Templates: []string{
				"http/cves/2024/CVE-2024-4040.yaml"},
		}),
		WithTemplateFilters(TemplateFilters{
			ExcludeTags:   []string{"code"},
			ProtocolTypes: "http,tcp",
		}),
	)
	assert.Nil(t, err)
	assert.Equal(t, 1, len(ne.GetTemplates()))
	err = ne.LoadAllTemplates()
	ne.resultCallbacks = append(ne.resultCallbacks, func(event *output.ResultEvent) {
		if event.TemplateID != "" {
			b := FormatEvent(event)
			_, _ = os.Stdout.Write(b)
			_, _ = os.Stdout.Write([]byte("\n"))
		}
	})
	templateIds := []string{
		"CVE-2024-4040",
	}
	wg := sync.WaitGroup{}
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err = ne.ExecuteWithProvider(context.Background(), provider.NewSimpleInputProviderWithUrls("http://ctf.lostpeach.cn:49264"), templateIds, nil)
			assert.Nil(t, err)

		}()
	}
	wg.Wait()
}

Describe alternatives you've considered

No response

Additional context

No response

tongchengbin avatar Sep 20 '24 03:09 tongchengbin