projectdiscovery / nuclei

Nuclei is a fast, customizable vulnerability scanner powered by the global security community and built on a simple YAML-based DSL, enabling collaboration to tackle trending vulnerabilities on the internet. It helps you find vulnerabilities in your applications, APIs, networks, DNS, and cloud configurations.
https://docs.projectdiscovery.io/tools/nuclei
MIT License
20.85k stars 2.52k forks source link

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

Open tongchengbin opened 2 months ago

tongchengbin commented 2 months ago

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

GeorginaReeder commented 2 months ago

Thanks so much for your feature request @tongchengbin , we'll take a look into this! :)