From dff4534807bf3431211fcff8a20009fda3c3bbc2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Wed, 6 Mar 2024 15:21:50 +0300 Subject: [PATCH 1/4] add callback support to StandardWriter --- lib/sdk_private.go | 21 +++++++++++++++++++++ pkg/output/output.go | 5 +++++ 2 files changed, 26 insertions(+) diff --git a/lib/sdk_private.go b/lib/sdk_private.go index 3281ebdbef..1a63302b79 100644 --- a/lib/sdk_private.go +++ b/lib/sdk_private.go @@ -58,6 +58,27 @@ func (e *NucleiEngine) applyRequiredDefaults() { mockoutput.FailureCallback = e.onFailureCallback } e.customWriter = mockoutput + } else { + if standardWriter, ok := e.customWriter.(*output.StandardWriter); ok { + standardWriter.WriteCallback = func(event *output.ResultEvent) { + if len(e.resultCallbacks) > 0 { + for _, callback := range e.resultCallbacks { + if callback != nil { + callback(event) + } + } + return + } + sb := strings.Builder{} + sb.WriteString(fmt.Sprintf("[%v] ", event.TemplateID)) + if event.Matched != "" { + sb.WriteString(event.Matched) + } else { + sb.WriteString(event.Host) + } + fmt.Println(sb.String()) + } + } } if e.customProgress == nil { e.customProgress = &testutils.MockProgressClient{} diff --git a/pkg/output/output.go b/pkg/output/output.go index d2893cc38e..9f4d1f8017 100644 --- a/pkg/output/output.go +++ b/pkg/output/output.go @@ -66,6 +66,7 @@ type StandardWriter struct { omitTemplate bool DisableStdout bool AddNewLinesOutputFile bool // by default this is only done for stdout + WriteCallback func(o *ResultEvent) } var decolorizerRegex = regexp.MustCompile(`\x1B\[[0-9;]*[a-zA-Z]`) @@ -229,6 +230,10 @@ func NewStandardWriter(options *types.Options) (*StandardWriter, error) { // Write writes the event to file and/or screen. func (w *StandardWriter) Write(event *ResultEvent) error { + if w.WriteCallback != nil { + w.WriteCallback(event) + } + // Enrich the result event with extra metadata on the template-path and url. if event.TemplatePath != "" { event.Template, event.TemplateURL = utils.TemplatePathURL(types.ToString(event.TemplatePath), types.ToString(event.TemplateID)) From 6faf699a92e535d60df7b64ccb39f55d25a3b0ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Wed, 6 Mar 2024 15:25:25 +0300 Subject: [PATCH 2/4] minor --- lib/sdk_private.go | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/sdk_private.go b/lib/sdk_private.go index 1a63302b79..87481346a9 100644 --- a/lib/sdk_private.go +++ b/lib/sdk_private.go @@ -59,6 +59,7 @@ func (e *NucleiEngine) applyRequiredDefaults() { } e.customWriter = mockoutput } else { + //FIXME: this is a hack to support callbacks in custom output writer if standardWriter, ok := e.customWriter.(*output.StandardWriter); ok { standardWriter.WriteCallback = func(event *output.ResultEvent) { if len(e.resultCallbacks) > 0 { From e1c5b895938db5a2311ab078165316f269d68b72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Mon, 11 Mar 2024 12:34:34 +0300 Subject: [PATCH 3/4] use multiwriter --- lib/sdk_private.go | 64 +++++++++++++++++----------------------------- 1 file changed, 23 insertions(+), 41 deletions(-) diff --git a/lib/sdk_private.go b/lib/sdk_private.go index 87481346a9..8fe69540b9 100644 --- a/lib/sdk_private.go +++ b/lib/sdk_private.go @@ -34,53 +34,35 @@ import ( // applyRequiredDefaults to options func (e *NucleiEngine) applyRequiredDefaults() { - if e.customWriter == nil { - mockoutput := testutils.NewMockOutputWriter(e.opts.OmitTemplate) - mockoutput.WriteCallback = func(event *output.ResultEvent) { - if len(e.resultCallbacks) > 0 { - for _, callback := range e.resultCallbacks { - if callback != nil { - callback(event) - } + mockoutput := testutils.NewMockOutputWriter(e.opts.OmitTemplate) + mockoutput.WriteCallback = func(event *output.ResultEvent) { + if len(e.resultCallbacks) > 0 { + for _, callback := range e.resultCallbacks { + if callback != nil { + callback(event) } - return } - sb := strings.Builder{} - sb.WriteString(fmt.Sprintf("[%v] ", event.TemplateID)) - if event.Matched != "" { - sb.WriteString(event.Matched) - } else { - sb.WriteString(event.Host) - } - fmt.Println(sb.String()) + return } - if e.onFailureCallback != nil { - mockoutput.FailureCallback = e.onFailureCallback + sb := strings.Builder{} + sb.WriteString(fmt.Sprintf("[%v] ", event.TemplateID)) + if event.Matched != "" { + sb.WriteString(event.Matched) + } else { + sb.WriteString(event.Host) } - e.customWriter = mockoutput + fmt.Println(sb.String()) + } + if e.onFailureCallback != nil { + mockoutput.FailureCallback = e.onFailureCallback + } + + if e.customWriter != nil { + e.customWriter = output.NewMultiWriter(e.customWriter, mockoutput) } else { - //FIXME: this is a hack to support callbacks in custom output writer - if standardWriter, ok := e.customWriter.(*output.StandardWriter); ok { - standardWriter.WriteCallback = func(event *output.ResultEvent) { - if len(e.resultCallbacks) > 0 { - for _, callback := range e.resultCallbacks { - if callback != nil { - callback(event) - } - } - return - } - sb := strings.Builder{} - sb.WriteString(fmt.Sprintf("[%v] ", event.TemplateID)) - if event.Matched != "" { - sb.WriteString(event.Matched) - } else { - sb.WriteString(event.Host) - } - fmt.Println(sb.String()) - } - } + e.customWriter = mockoutput } + if e.customProgress == nil { e.customProgress = &testutils.MockProgressClient{} } From 93bf4b1d7f99746573e12a98141f5a2338ad6bbc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Mon, 11 Mar 2024 12:35:49 +0300 Subject: [PATCH 4/4] minor --- pkg/output/output.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/pkg/output/output.go b/pkg/output/output.go index 9f4d1f8017..d2893cc38e 100644 --- a/pkg/output/output.go +++ b/pkg/output/output.go @@ -66,7 +66,6 @@ type StandardWriter struct { omitTemplate bool DisableStdout bool AddNewLinesOutputFile bool // by default this is only done for stdout - WriteCallback func(o *ResultEvent) } var decolorizerRegex = regexp.MustCompile(`\x1B\[[0-9;]*[a-zA-Z]`) @@ -230,10 +229,6 @@ func NewStandardWriter(options *types.Options) (*StandardWriter, error) { // Write writes the event to file and/or screen. func (w *StandardWriter) Write(event *ResultEvent) error { - if w.WriteCallback != nil { - w.WriteCallback(event) - } - // Enrich the result event with extra metadata on the template-path and url. if event.TemplatePath != "" { event.Template, event.TemplateURL = utils.TemplatePathURL(types.ToString(event.TemplatePath), types.ToString(event.TemplateID))