add callback support to StandardWriter (#4839)

* add callback support to StandardWriter

* minor

* use multiwriter

* minor
This commit is contained in:
Dogan Can Bakir 2024-03-11 14:37:07 +03:00 committed by GitHub
parent fe4d5f05a1
commit e9625d621f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -34,31 +34,35 @@ import (
// applyRequiredDefaults to options // applyRequiredDefaults to options
func (e *NucleiEngine) applyRequiredDefaults() { func (e *NucleiEngine) applyRequiredDefaults() {
if e.customWriter == nil { mockoutput := testutils.NewMockOutputWriter(e.opts.OmitTemplate)
mockoutput := testutils.NewMockOutputWriter(e.opts.OmitTemplate) mockoutput.WriteCallback = func(event *output.ResultEvent) {
mockoutput.WriteCallback = func(event *output.ResultEvent) { if len(e.resultCallbacks) > 0 {
if len(e.resultCallbacks) > 0 { for _, callback := range e.resultCallbacks {
for _, callback := range e.resultCallbacks { if callback != nil {
if callback != nil { callback(event)
callback(event)
}
} }
return
} }
sb := strings.Builder{} return
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.onFailureCallback != nil { sb := strings.Builder{}
mockoutput.FailureCallback = e.onFailureCallback 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.onFailureCallback != nil {
mockoutput.FailureCallback = e.onFailureCallback
}
if e.customWriter != nil {
e.customWriter = output.NewMultiWriter(e.customWriter, mockoutput)
} else {
e.customWriter = mockoutput e.customWriter = mockoutput
} }
if e.customProgress == nil { if e.customProgress == nil {
e.customProgress = &testutils.MockProgressClient{} e.customProgress = &testutils.MockProgressClient{}
} }