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
Manuel Bua 2020-07-09 20:57:24 +02:00
parent 6c43aab488
commit 0ff138a477
4 changed files with 48 additions and 74 deletions

View File

@ -6,30 +6,28 @@ import (
"github.com/vbauerster/mpb/v5" "github.com/vbauerster/mpb/v5"
"github.com/vbauerster/mpb/v5/cwriter" "github.com/vbauerster/mpb/v5/cwriter"
"github.com/vbauerster/mpb/v5/decor" "github.com/vbauerster/mpb/v5/decor"
"io"
"os" "os"
"regexp"
"strconv"
"strings" "strings"
"sync" "sync"
) )
type Progress struct { type Progress struct {
progress *mpb.Progress progress *mpb.Progress
progress_stdout *mpb.Progress
bar *mpb.Bar bar *mpb.Bar
total int64 total int64
initialTotal int64 initialTotal int64
captureData *captureData captureData *captureData
termWidth int termWidth int
mutex *sync.Mutex mutex *sync.Mutex
stdout *strings.Builder
stderr *strings.Builder
} }
func NewProgress(group *sync.WaitGroup) *Progress { func NewProgress(group *sync.WaitGroup) *Progress {
w := cwriter.New(os.Stdout) w := cwriter.New(os.Stderr)
tw, err := w.GetWidth() tw, err := w.GetWidth()
if err != nil { if err != nil {
panic("Couldn't determine available terminal width.") tw = 80
} }
p := &Progress{ p := &Progress{
@ -38,15 +36,10 @@ func NewProgress(group *sync.WaitGroup) *Progress {
mpb.WithOutput(os.Stderr), mpb.WithOutput(os.Stderr),
mpb.PopCompletedMode(), mpb.PopCompletedMode(),
), ),
progress_stdout: mpb.New(
mpb.WithWaitGroup(group),
mpb.WithOutput(os.Stdout),
//mpb.PopCompletedMode(),
),
termWidth: tw, termWidth: tw,
mutex: &sync.Mutex{}, mutex: &sync.Mutex{},
stdout: &strings.Builder{},
stderr: &strings.Builder{},
} }
return p return p
} }
@ -89,9 +82,7 @@ func (p *Progress) Wait() {
if p.initialTotal != p.total { if p.initialTotal != p.total {
p.bar.SetTotal(p.total, true) p.bar.SetTotal(p.total, true)
} }
p.progress.Wait() p.progress.Wait()
p.progress_stdout.Wait()
} }
// //
@ -101,36 +92,21 @@ func (p *Progress) StartStdCapture() {
p.captureData = startStdCapture() p.captureData = startStdCapture()
} }
func (p *Progress) StopStdCaptureAndShow() { func (p *Progress) StopStdCapture() {
stopStdCapture(p.captureData) stopStdCapture(p.captureData)
p.stdout.Write(p.captureData.DataStdOut.Bytes())
// stdout p.stderr.Write(p.captureData.DataStdErr.Bytes())
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.mutex.Unlock() p.mutex.Unlock()
} }
func makeLogBar(msg string) mpb.BarFiller { func (p *Progress) ShowStdOut() {
return mpb.BarFillerFunc(func(w io.Writer, _ int, st decor.Statistics) { if p.stdout.Len() > 0 {
fmt.Fprintf(w, msg) fmt.Fprint(os.Stdout, p.stdout.String())
}) }
}
func (p *Progress) ShowStdErr() {
if p.stderr.Len() > 0 {
fmt.Fprint(os.Stderr, p.stderr.String())
}
} }

View File

@ -7,9 +7,7 @@ import (
"bytes" "bytes"
"io" "io"
"os" "os"
"strings"
"sync" "sync"
"time"
) )
type captureData struct { type captureData struct {
@ -18,15 +16,11 @@ type captureData struct {
backupStderr *os.File backupStderr *os.File
writerStderr *os.File writerStderr *os.File
dataStdout string DataStdOut *bytes.Buffer
dataStderr string DataStdErr *bytes.Buffer
outStdout chan []byte outStdout chan []byte
outStderr chan []byte outStderr chan []byte
//sync sync.WaitGroup
DataStdOut []string
DataStdErr []string
} }
func startStdCapture() *captureData { func startStdCapture() *captureData {
@ -49,6 +43,9 @@ func startStdCapture() *captureData {
outStdout: make(chan []byte), outStdout: make(chan []byte),
outStderr: make(chan []byte), outStderr: make(chan []byte),
DataStdOut: &bytes.Buffer{},
DataStdErr: &bytes.Buffer{},
} }
os.Stdout = c.writerStdout os.Stdout = c.writerStdout
@ -74,23 +71,20 @@ func stopStdCapture(c *captureData) {
var wg sync.WaitGroup var wg sync.WaitGroup
stdRead := func(in <-chan []byte, outString *string, outArray *[]string) { stdRead := func(in <-chan []byte, outData *bytes.Buffer) {
defer wg.Done() defer wg.Done()
select { select {
case out := <-in: case out := <-in:
*outString = string(out) outData.Write(out)
*outArray = strings.Split(*outString, "\n") default:
if (*outArray)[len(*outArray)-1] == "" { //case <-time.After(10 * time.Millisecond):
*outArray = (*outArray)[:len(*outArray)-1]
}
case <-time.After(50 * time.Millisecond):
} }
} }
wg.Add(2) wg.Add(2)
go stdRead(c.outStdout, &c.dataStdout, &c.DataStdOut) go stdRead(c.outStdout, c.DataStdOut)
go stdRead(c.outStderr, &c.dataStderr, &c.DataStdErr) go stdRead(c.outStderr, c.DataStdErr)
wg.Wait() wg.Wait()
os.Stdout = c.backupStdout os.Stdout = c.backupStdout

View File

@ -153,6 +153,8 @@ func (r *Runner) RunEnumeration() {
} }
p.Wait() p.Wait()
p.ShowStdErr()
p.ShowStdOut()
if !results { if !results {
if r.output != nil { if r.output != nil {
@ -226,10 +228,12 @@ func (r *Runner) RunEnumeration() {
default: default:
p.StartStdCapture() p.StartStdCapture()
gologger.Errorf("Could not parse file '%s': %s\n", r.options.Templates, err) gologger.Errorf("Could not parse file '%s': %s\n", r.options.Templates, err)
p.StopStdCaptureAndShow() p.StopStdCapture()
} }
} }
p.Wait() p.Wait()
p.ShowStdErr()
p.ShowStdOut()
if !results { if !results {
if r.output != nil { if r.output != nil {
@ -239,7 +243,7 @@ func (r *Runner) RunEnumeration() {
} }
//p.StartStdCapture() //p.StartStdCapture()
gologger.Infof("No results found for the template. Happy hacking!") gologger.Infof("No results found for the template. Happy hacking!")
//p.StopStdCaptureAndShow() //p.StopStdCapture()
} }
return return
@ -254,7 +258,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
} }
p.StartStdCapture() p.StartStdCapture()
gologger.Infof("%s\n", message) gologger.Infof("%s\n", message)
p.StopStdCaptureAndShow() p.StopStdCapture()
var writer *bufio.Writer var writer *bufio.Writer
if r.output != nil { if r.output != nil {
@ -293,7 +297,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
if err != nil { if err != nil {
p.StartStdCapture() p.StartStdCapture()
gologger.Warningf("Could not create http client: %s\n", err) gologger.Warningf("Could not create http client: %s\n", err)
p.StopStdCaptureAndShow() p.StopStdCapture()
return false return false
} }
@ -322,7 +326,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
if err != nil { if err != nil {
p.StartStdCapture() p.StartStdCapture()
gologger.Warningf("Could not execute step: %s\n", err) gologger.Warningf("Could not execute step: %s\n", err)
p.StopStdCaptureAndShow() p.StopStdCapture()
} }
<-limiter <-limiter
wg.Done() wg.Done()

View File

@ -115,7 +115,7 @@ mainLoop:
if e.debug { if e.debug {
p.StartStdCapture() p.StartStdCapture()
gologger.Infof("Dumped HTTP request for %s (%s)\n\n", URL, e.template.ID) 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) dumpedRequest, err := httputil.DumpRequest(req.Request, true)
if err != nil { if err != nil {
@ -124,7 +124,7 @@ mainLoop:
} }
p.StartStdCapture() p.StartStdCapture()
fmt.Fprintf(os.Stderr, "%s", string(dumpedRequest)) fmt.Fprintf(os.Stderr, "%s", string(dumpedRequest))
p.StopStdCaptureAndShow() p.StopStdCapture()
} }
resp, err := e.httpClient.Do(req) resp, err := e.httpClient.Do(req)
@ -135,14 +135,14 @@ mainLoop:
p.Abort(1) p.Abort(1)
p.StartStdCapture() p.StartStdCapture()
gologger.Warningf("Could not do request: %s\n", err) gologger.Warningf("Could not do request: %s\n", err)
p.StopStdCaptureAndShow() p.StopStdCapture()
continue continue
} }
if e.debug { if e.debug {
p.StartStdCapture() p.StartStdCapture()
gologger.Infof("Dumped HTTP response for %s (%s)\n\n", URL, e.template.ID) gologger.Infof("Dumped HTTP response for %s (%s)\n\n", URL, e.template.ID)
p.StopStdCaptureAndShow() p.StopStdCapture()
dumpedResponse, err := httputil.DumpResponse(resp, true) dumpedResponse, err := httputil.DumpResponse(resp, true)
if err != nil { if err != nil {
@ -151,7 +151,7 @@ mainLoop:
} }
p.StartStdCapture() p.StartStdCapture()
fmt.Fprintf(os.Stderr, "%s\n", string(dumpedResponse)) fmt.Fprintf(os.Stderr, "%s\n", string(dumpedResponse))
p.StopStdCaptureAndShow() p.StopStdCapture()
} }
data, err := ioutil.ReadAll(resp.Body) data, err := ioutil.ReadAll(resp.Body)
@ -193,7 +193,7 @@ mainLoop:
// capture stdout and emit it via a mpb.BarFiller // capture stdout and emit it via a mpb.BarFiller
p.StartStdCapture() p.StartStdCapture()
e.writeOutputHTTP(compiledRequest, matcher, nil) e.writeOutputHTTP(compiledRequest, matcher, nil)
p.StopStdCaptureAndShow() p.StopStdCapture()
atomic.CompareAndSwapUint32(&e.results, 0, 1) atomic.CompareAndSwapUint32(&e.results, 0, 1)
} }
@ -216,7 +216,7 @@ mainLoop:
// capture stdout and emit it via a mpb.BarFiller // capture stdout and emit it via a mpb.BarFiller
p.StartStdCapture() p.StartStdCapture()
e.writeOutputHTTP(compiledRequest, nil, extractorResults) e.writeOutputHTTP(compiledRequest, nil, extractorResults)
p.StopStdCaptureAndShow() p.StopStdCapture()
atomic.CompareAndSwapUint32(&e.results, 0, 1) atomic.CompareAndSwapUint32(&e.results, 0, 1)
} }
@ -227,7 +227,7 @@ mainLoop:
p.StartStdCapture() p.StartStdCapture()
gologger.Verbosef("Sent HTTP request to %s\n", "http-request", URL) gologger.Verbosef("Sent HTTP request to %s\n", "http-request", URL)
p.StopStdCaptureAndShow() p.StopStdCapture()
return nil return nil
} }