mirror of https://github.com/daffainfo/nuclei.git
Do not use mpb.Progress for logging
This will cause sync issues with very fast output and will defeat the purpose of logging. Instead, buffer both stdout/stderr and show their output at the end.dev
parent
6c43aab488
commit
0ff138a477
|
@ -6,30 +6,28 @@ import (
|
|||
"github.com/vbauerster/mpb/v5"
|
||||
"github.com/vbauerster/mpb/v5/cwriter"
|
||||
"github.com/vbauerster/mpb/v5/decor"
|
||||
"io"
|
||||
"os"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type Progress struct {
|
||||
progress *mpb.Progress
|
||||
progress_stdout *mpb.Progress
|
||||
bar *mpb.Bar
|
||||
total int64
|
||||
initialTotal int64
|
||||
captureData *captureData
|
||||
termWidth int
|
||||
mutex *sync.Mutex
|
||||
stdout *strings.Builder
|
||||
stderr *strings.Builder
|
||||
}
|
||||
|
||||
func NewProgress(group *sync.WaitGroup) *Progress {
|
||||
w := cwriter.New(os.Stdout)
|
||||
w := cwriter.New(os.Stderr)
|
||||
tw, err := w.GetWidth()
|
||||
if err != nil {
|
||||
panic("Couldn't determine available terminal width.")
|
||||
tw = 80
|
||||
}
|
||||
|
||||
p := &Progress{
|
||||
|
@ -38,15 +36,10 @@ func NewProgress(group *sync.WaitGroup) *Progress {
|
|||
mpb.WithOutput(os.Stderr),
|
||||
mpb.PopCompletedMode(),
|
||||
),
|
||||
|
||||
progress_stdout: mpb.New(
|
||||
mpb.WithWaitGroup(group),
|
||||
mpb.WithOutput(os.Stdout),
|
||||
//mpb.PopCompletedMode(),
|
||||
),
|
||||
|
||||
termWidth: tw,
|
||||
mutex: &sync.Mutex{},
|
||||
stdout: &strings.Builder{},
|
||||
stderr: &strings.Builder{},
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
@ -89,9 +82,7 @@ func (p *Progress) Wait() {
|
|||
if p.initialTotal != p.total {
|
||||
p.bar.SetTotal(p.total, true)
|
||||
}
|
||||
|
||||
p.progress.Wait()
|
||||
p.progress_stdout.Wait()
|
||||
}
|
||||
|
||||
//
|
||||
|
@ -101,36 +92,21 @@ func (p *Progress) StartStdCapture() {
|
|||
p.captureData = startStdCapture()
|
||||
}
|
||||
|
||||
func (p *Progress) StopStdCaptureAndShow() {
|
||||
func (p *Progress) StopStdCapture() {
|
||||
stopStdCapture(p.captureData)
|
||||
|
||||
// stdout
|
||||
for _, captured := range p.captureData.DataStdOut {
|
||||
var r = regexp.MustCompile("(.{" + strconv.Itoa(p.termWidth) + "})")
|
||||
multiline := r.ReplaceAllString(captured, "$1\n")
|
||||
arr := strings.Split(multiline, "\n")
|
||||
|
||||
for _, msg := range arr {
|
||||
p.progress_stdout.Add(0, makeLogBar(msg)).SetTotal(0, true)
|
||||
}
|
||||
}
|
||||
|
||||
// stderr
|
||||
for _, captured := range p.captureData.DataStdErr {
|
||||
var r = regexp.MustCompile("(.{" + strconv.Itoa(p.termWidth) + "})")
|
||||
multiline := r.ReplaceAllString(captured, "$1\n")
|
||||
arr := strings.Split(multiline, "\n")
|
||||
|
||||
for _, msg := range arr {
|
||||
p.progress.Add(0, makeLogBar(msg)).SetTotal(0, true)
|
||||
}
|
||||
}
|
||||
|
||||
p.stdout.Write(p.captureData.DataStdOut.Bytes())
|
||||
p.stderr.Write(p.captureData.DataStdErr.Bytes())
|
||||
p.mutex.Unlock()
|
||||
}
|
||||
|
||||
func makeLogBar(msg string) mpb.BarFiller {
|
||||
return mpb.BarFillerFunc(func(w io.Writer, _ int, st decor.Statistics) {
|
||||
fmt.Fprintf(w, msg)
|
||||
})
|
||||
func (p *Progress) ShowStdOut() {
|
||||
if p.stdout.Len() > 0 {
|
||||
fmt.Fprint(os.Stdout, p.stdout.String())
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Progress) ShowStdErr() {
|
||||
if p.stderr.Len() > 0 {
|
||||
fmt.Fprint(os.Stderr, p.stderr.String())
|
||||
}
|
||||
}
|
||||
|
|
|
@ -7,9 +7,7 @@ import (
|
|||
"bytes"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type captureData struct {
|
||||
|
@ -18,15 +16,11 @@ type captureData struct {
|
|||
backupStderr *os.File
|
||||
writerStderr *os.File
|
||||
|
||||
dataStdout string
|
||||
dataStderr string
|
||||
DataStdOut *bytes.Buffer
|
||||
DataStdErr *bytes.Buffer
|
||||
|
||||
outStdout chan []byte
|
||||
outStderr chan []byte
|
||||
|
||||
//sync sync.WaitGroup
|
||||
|
||||
DataStdOut []string
|
||||
DataStdErr []string
|
||||
}
|
||||
|
||||
func startStdCapture() *captureData {
|
||||
|
@ -49,6 +43,9 @@ func startStdCapture() *captureData {
|
|||
|
||||
outStdout: make(chan []byte),
|
||||
outStderr: make(chan []byte),
|
||||
|
||||
DataStdOut: &bytes.Buffer{},
|
||||
DataStdErr: &bytes.Buffer{},
|
||||
}
|
||||
|
||||
os.Stdout = c.writerStdout
|
||||
|
@ -74,23 +71,20 @@ func stopStdCapture(c *captureData) {
|
|||
|
||||
var wg sync.WaitGroup
|
||||
|
||||
stdRead := func(in <-chan []byte, outString *string, outArray *[]string) {
|
||||
stdRead := func(in <-chan []byte, outData *bytes.Buffer) {
|
||||
defer wg.Done()
|
||||
|
||||
select {
|
||||
case out := <-in:
|
||||
*outString = string(out)
|
||||
*outArray = strings.Split(*outString, "\n")
|
||||
if (*outArray)[len(*outArray)-1] == "" {
|
||||
*outArray = (*outArray)[:len(*outArray)-1]
|
||||
}
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
outData.Write(out)
|
||||
default:
|
||||
//case <-time.After(10 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
|
||||
wg.Add(2)
|
||||
go stdRead(c.outStdout, &c.dataStdout, &c.DataStdOut)
|
||||
go stdRead(c.outStderr, &c.dataStderr, &c.DataStdErr)
|
||||
go stdRead(c.outStdout, c.DataStdOut)
|
||||
go stdRead(c.outStderr, c.DataStdErr)
|
||||
wg.Wait()
|
||||
|
||||
os.Stdout = c.backupStdout
|
||||
|
|
|
@ -153,6 +153,8 @@ func (r *Runner) RunEnumeration() {
|
|||
}
|
||||
|
||||
p.Wait()
|
||||
p.ShowStdErr()
|
||||
p.ShowStdOut()
|
||||
|
||||
if !results {
|
||||
if r.output != nil {
|
||||
|
@ -226,10 +228,12 @@ func (r *Runner) RunEnumeration() {
|
|||
default:
|
||||
p.StartStdCapture()
|
||||
gologger.Errorf("Could not parse file '%s': %s\n", r.options.Templates, err)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
}
|
||||
}
|
||||
p.Wait()
|
||||
p.ShowStdErr()
|
||||
p.ShowStdOut()
|
||||
|
||||
if !results {
|
||||
if r.output != nil {
|
||||
|
@ -239,7 +243,7 @@ func (r *Runner) RunEnumeration() {
|
|||
}
|
||||
//p.StartStdCapture()
|
||||
gologger.Infof("No results found for the template. Happy hacking!")
|
||||
//p.StopStdCaptureAndShow()
|
||||
//p.StopStdCapture()
|
||||
}
|
||||
|
||||
return
|
||||
|
@ -254,7 +258,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
|
|||
}
|
||||
p.StartStdCapture()
|
||||
gologger.Infof("%s\n", message)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
|
||||
var writer *bufio.Writer
|
||||
if r.output != nil {
|
||||
|
@ -293,7 +297,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
|
|||
if err != nil {
|
||||
p.StartStdCapture()
|
||||
gologger.Warningf("Could not create http client: %s\n", err)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
return false
|
||||
}
|
||||
|
||||
|
@ -322,7 +326,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
|
|||
if err != nil {
|
||||
p.StartStdCapture()
|
||||
gologger.Warningf("Could not execute step: %s\n", err)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
}
|
||||
<-limiter
|
||||
wg.Done()
|
||||
|
|
|
@ -115,7 +115,7 @@ mainLoop:
|
|||
if e.debug {
|
||||
p.StartStdCapture()
|
||||
gologger.Infof("Dumped HTTP request for %s (%s)\n\n", URL, e.template.ID)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
|
||||
dumpedRequest, err := httputil.DumpRequest(req.Request, true)
|
||||
if err != nil {
|
||||
|
@ -124,7 +124,7 @@ mainLoop:
|
|||
}
|
||||
p.StartStdCapture()
|
||||
fmt.Fprintf(os.Stderr, "%s", string(dumpedRequest))
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
}
|
||||
|
||||
resp, err := e.httpClient.Do(req)
|
||||
|
@ -135,14 +135,14 @@ mainLoop:
|
|||
p.Abort(1)
|
||||
p.StartStdCapture()
|
||||
gologger.Warningf("Could not do request: %s\n", err)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
continue
|
||||
}
|
||||
|
||||
if e.debug {
|
||||
p.StartStdCapture()
|
||||
gologger.Infof("Dumped HTTP response for %s (%s)\n\n", URL, e.template.ID)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
|
||||
dumpedResponse, err := httputil.DumpResponse(resp, true)
|
||||
if err != nil {
|
||||
|
@ -151,7 +151,7 @@ mainLoop:
|
|||
}
|
||||
p.StartStdCapture()
|
||||
fmt.Fprintf(os.Stderr, "%s\n", string(dumpedResponse))
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
}
|
||||
|
||||
data, err := ioutil.ReadAll(resp.Body)
|
||||
|
@ -193,7 +193,7 @@ mainLoop:
|
|||
// capture stdout and emit it via a mpb.BarFiller
|
||||
p.StartStdCapture()
|
||||
e.writeOutputHTTP(compiledRequest, matcher, nil)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
|
||||
atomic.CompareAndSwapUint32(&e.results, 0, 1)
|
||||
}
|
||||
|
@ -216,7 +216,7 @@ mainLoop:
|
|||
// capture stdout and emit it via a mpb.BarFiller
|
||||
p.StartStdCapture()
|
||||
e.writeOutputHTTP(compiledRequest, nil, extractorResults)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
|
||||
atomic.CompareAndSwapUint32(&e.results, 0, 1)
|
||||
}
|
||||
|
@ -227,7 +227,7 @@ mainLoop:
|
|||
|
||||
p.StartStdCapture()
|
||||
gologger.Verbosef("Sent HTTP request to %s\n", "http-request", URL)
|
||||
p.StopStdCaptureAndShow()
|
||||
p.StopStdCapture()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue