daffainfo
/
nuclei
mirror of https://github.com/daffainfo/nuclei.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use io.WriteCloser in output writer

dev
Alexey Zhuchkov 2021-10-30 14:19:11 +03:00
parent cd2db280bf
commit b8ebbc27f5
2 changed files with 16 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
v2/pkg/output/file_output_writer.go
View File

@ -21,14 +21,16 @@ func newFileOutputWriter(file string) (*fileWriter, error) {
} }
// WriteString writes an output to the underlying file // WriteString writes an output to the underlying file
func (w *fileWriter) Write(data []byte) error { func (w *fileWriter) Write(data []byte) (int, error) {
w.mu.Lock() w.mu.Lock()
defer w.mu.Unlock() defer w.mu.Unlock()
if _, err := w.file.Write(data); err != nil { if _, err := w.file.Write(data); err != nil {
return err return 0, err
} }
_, err := w.file.Write([]byte("\n")) if _, err := w.file.Write([]byte("\n")); err != nil {
return err return 0, err
}
return len(data) + 1, nil
} }
// Close closes the underlying writer flushing everything to disk // Close closes the underlying writer flushing everything to disk

19
v2/pkg/output/output.go
View File

@ -1,6 +1,7 @@
package output package output
import ( import (
"io"
"os" "os"
"regexp" "regexp"
"time" "time"
@ -37,9 +38,9 @@ type StandardWriter struct {
noTimestamp bool noTimestamp bool
noMetadata bool noMetadata bool
aurora aurora.Aurora aurora aurora.Aurora
outputFile *fileWriter outputFile io.WriteCloser
traceFile *fileWriter traceFile io.WriteCloser
errorFile *fileWriter errorFile io.WriteCloser
severityColors func(severity.Severity) string severityColors func(severity.Severity) string
} }
@ -99,7 +100,7 @@ type ResultEvent struct {
func NewStandardWriter(colors, noMetadata, noTimestamp, json, jsonReqResp bool, file, traceFile string, errorFile string) (*StandardWriter, error) { func NewStandardWriter(colors, noMetadata, noTimestamp, json, jsonReqResp bool, file, traceFile string, errorFile string) (*StandardWriter, error) {
auroraColorizer := aurora.NewAurora(colors) auroraColorizer := aurora.NewAurora(colors)
var outputFile *fileWriter var outputFile io.WriteCloser
if file != "" { if file != "" {
output, err := newFileOutputWriter(file) output, err := newFileOutputWriter(file)
if err != nil { if err != nil {
@ -107,7 +108,7 @@ func NewStandardWriter(colors, noMetadata, noTimestamp, json, jsonReqResp bool,
} }
outputFile = output outputFile = output
} }
var traceOutput *fileWriter var traceOutput io.WriteCloser
if traceFile != "" { if traceFile != "" {
output, err := newFileOutputWriter(traceFile) output, err := newFileOutputWriter(traceFile)
if err != nil { if err != nil {
@ -115,7 +116,7 @@ func NewStandardWriter(colors, noMetadata, noTimestamp, json, jsonReqResp bool,
} }
traceOutput = output traceOutput = output
} }
var errorOutput *fileWriter var errorOutput io.WriteCloser
if errorFile != "" { if errorFile != "" {
output, err := newFileOutputWriter(errorFile) output, err := newFileOutputWriter(errorFile)
if err != nil { if err != nil {
@ -161,7 +162,7 @@ func (w *StandardWriter) Write(event *ResultEvent) error {
if !w.json { if !w.json {
data = decolorizerRegex.ReplaceAll(data, []byte("")) data = decolorizerRegex.ReplaceAll(data, []byte(""))
} }
if writeErr := w.outputFile.Write(data); writeErr != nil { if _, writeErr := w.outputFile.Write(data); writeErr != nil {
return errors.Wrap(err, "could not write to output") return errors.Wrap(err, "could not write to output")
} }
} }
@ -198,11 +199,11 @@ func (w *StandardWriter) Request(templatePath, input, requestType string, reques
} }
if w.traceFile != nil { if w.traceFile != nil {
_ = w.traceFile.Write(data) _, _ = w.traceFile.Write(data)
} }
if requestErr != nil && w.errorFile != nil { if requestErr != nil && w.errorFile != nil {
_ = w.errorFile.Write(data) _, _ = w.errorFile.Write(data)
} }
} }