How to add some extra data to the custom Writer?

[Caratpine]

Hi, I'm implementing a scanner with Nuclei integrated in Go code, the code looks like this.

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path"
	"time"

	"scanner/conf"
	"scanner/internal/model"

	"github.com/logrusorgru/aurora"
	"github.com/projectdiscovery/goflags"
	"github.com/projectdiscovery/nuclei/v2/pkg/catalog/config"
	"github.com/projectdiscovery/nuclei/v2/pkg/catalog/disk"
	"github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader"
	"github.com/projectdiscovery/nuclei/v2/pkg/core"
	"github.com/projectdiscovery/nuclei/v2/pkg/core/inputs"
	"github.com/projectdiscovery/nuclei/v2/pkg/output"
	"github.com/projectdiscovery/nuclei/v2/pkg/progress"
	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/contextargs"
	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/hosterrorscache"
	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/interactsh"
	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolinit"
	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
	"github.com/projectdiscovery/nuclei/v2/pkg/types"
	"github.com/projectdiscovery/ratelimit"
	logger "github.com/sirupsen/logrus"
)

type NucleiWorker struct {
	c   *conf.Nuclei
	svc *Service

	engine          *core.Engine
	hostErrorCache  *hosterrorscache.Cache
	interactClient  *interactsh.Client
	httpTemplate    *loader.Store
	networkTemplate *loader.Store
	options         *types.Options
	configObject    *config.Config
	catalog         *disk.DiskCatalog
	executerOpts    protocols.ExecutorOptions
	writer          *NucleiWriter
}

func NewNuclei(c *conf.Config, s *Service) *NucleiWorker {
	var (
		err            error
		ps             progress.Progress
		interactshOpts *interactsh.Options
		pwd, _         = os.Getwd()
	)
	n := &NucleiWorker{
		c:   c.Nuclei,
		svc: s,

		catalog: disk.NewCatalog(path.Join(pwd, c.Nuclei.PoCPath)),
		configObject: &config.Config{
			TemplatesDirectory: path.Join(pwd, c.Nuclei.PoCPath),
		},
		options: types.DefaultOptions(),
	}
	n.options.RateLimit = int(n.c.Rate)
	n.options.Timeout = int(n.c.Timeout)

	n.hostErrorCache = hosterrorscache.New(30, hosterrorscache.DefaultMaxHostsCount, nil)
	protocolstate.Init(n.options)
	protocolinit.Init(n.options)
	ps, err = progress.NewStatsTicker(
		n.options.StatsInterval,
		n.options.EnableProgressBar,
		n.options.StatsJSON,
		n.options.Metrics,
		n.options.Cloud,
		n.options.MetricsPort)
	if err != nil {
		panic(err)
	}
	n.writer = &NucleiWriter{
		a: aurora.NewAurora(false),
	}
	if err != nil {
		panic(err)
	}
	interactshOpts = interactsh.DefaultOptions(n.writer, nil, ps)
	n.interactClient, err = interactsh.New(interactshOpts)
	if err != nil {
		panic(err)
	}
	n.executerOpts = protocols.ExecutorOptions{
		Output:          n.writer,
		Options:         n.options,
		Progress:        ps,
		Catalog:         n.catalog, // template directory
		IssuesClient:    nil,
		RateLimiter:     ratelimit.New(context.Background(), uint(n.options.RateLimit), time.Second),
		Interactsh:      n.interactClient,
		HostErrorsCache: n.hostErrorCache,
		Colorizer:       aurora.NewAurora(true),
		ResumeCfg:       types.NewResumeCfg(), // required
	}
	n.options.Tags = goflags.StringSlice{model.NucleiTagHTTP}
	n.httpTemplate, err = loader.New(loader.NewConfig(n.options, n.catalog, n.executerOpts))
	if err != nil {
		panic(err)
	}
	n.httpTemplate.Load()
	n.options.Tags = goflags.StringSlice{model.NucleiTagNetwork}
	n.networkTemplate, err = loader.New(loader.NewConfig(n.options, n.catalog, n.executerOpts))
	if err != nil {
		panic(err)
	}
	n.networkTemplate.Load()
	n.engine = core.New(n.options)
	n.engine.SetExecuterOptions(n.executerOpts)
	go n.syncTemplate()
	return n
}

func (nw *NucleiWorker) Close() {
	nw.hostErrorCache.Close()
	nw.interactClient.Close()
}

func (nw *NucleiWorker) Exec(msg *model.ScanMsg) (err error) {
	var (
		tpls  []*templates.Template
		input = &inputs.SimpleInputProvider{
			Inputs: []*contextargs.MetaInput{{Input: msg.Target}},
		}
	)
	if msg.NucleiTag == model.NucleiTagHTTP {
		tpls = nw.httpTemplate.Templates()
	}
	if msg.NucleiTag == model.NucleiTagNetwork {
		tpls = nw.networkTemplate.Templates()
	}
	_ = nw.engine.Execute(tpls, input)
	nw.engine.WorkPool().Wait()
	return
}

func (nw *NucleiWorker) syncTemplate() {
	var err error
	for {
		time.Sleep(time.Hour * 12)
		nw.options.Tags = goflags.StringSlice{model.NucleiTagHTTP}
		nw.httpTemplate, err = loader.New(loader.NewConfig(nw.options, nw.catalog, nw.executerOpts))
		if err != nil {
			logger.Errorf("syncTemplate loader.New() error(%v)", err)
		}
		nw.httpTemplate.Load()
		nw.options.Tags = goflags.StringSlice{model.NucleiTagNetwork}
		nw.networkTemplate, err = loader.New(loader.NewConfig(nw.options, nw.catalog, nw.executerOpts))
		if err != nil {
			logger.Errorf("syncTemplate loader.New() error(%v)", err)
		}
		nw.networkTemplate.Load()
	}
}

type NucleiWriter struct {
	a aurora.Aurora
}

func (w *NucleiWriter) Close() {

}

func (w *NucleiWriter) Colorizer() aurora.Aurora {
	return w.a
}

func (w *NucleiWriter) Request(templateID, url, requestType string, err error) {

}

func (w *NucleiWriter) WriteStoreDebugData(host, templateID, eventType string, data string) {

}

func (w *NucleiWriter) Write(result *output.ResultEvent) (err error) {
	logger.Infof("NucleiWriter output result(%+v)", result)
	var (
		hostname, _ = os.Hostname()
		nr          = &model.NucleiResult{
			Timestamp:   result.Timestamp,
			TemplateID:  result.TemplateID,
			Name:        result.Info.Name,
			Tags:        result.Info.Tags.ToSlice(),
			Description: result.Info.Description,
			Severity:    result.Info.SeverityHolder.Severity.String(),
			IP:          result.IP,
			Host:        result.Host,
			MatchedAt:   result.Matched,
			Type:        result.Type,
			Agent:       hostname,
			Original:    "",
			Request:     result.Request,
			Response:    result.Response,
		}
	)
	data, err := json.Marshal(nr)
	if err != nil {
		return
	}
	fmt.Println(data)
	logger.Infof("NucleiWriter push msg(%+v)", nr)
	return
}

func (w *NucleiWriter) WriteFailure(event output.InternalEvent) (err error) {
	return
}

I will only initialise the NewNuclei method once, and there are many scanning tasks. Each scanning task will execute the Exec method concurrently. I want to bring the task id in the output of the Writer. I have thought about it for a long time and have no good ideas. , any suggestions? Thank you.

[randolphcyg]

I am working on this now

[dogancanbakir]

This discussion closed automatically due to inactivity. Feel free to reopen or start new if still relevant.