nuclei/v2/internal/progress/stdcapture.go
Manuel Bua 0ff138a477 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.
2020-07-09 20:57:24 +02:00

93 lines
1.6 KiB
Go

package progress
/**
Inspired by the https://github.com/PumpkinSeed/cage module
*/
import (
"bytes"
"io"
"os"
"sync"
)
type captureData struct {
backupStdout *os.File
writerStdout *os.File
backupStderr *os.File
writerStderr *os.File
DataStdOut *bytes.Buffer
DataStdErr *bytes.Buffer
outStdout chan []byte
outStderr chan []byte
}
func startStdCapture() *captureData {
rStdout, wStdout, errStdout := os.Pipe()
if errStdout != nil {
panic(errStdout)
}
rStderr, wStderr, errStderr := os.Pipe()
if errStderr != nil {
panic(errStderr)
}
c := &captureData{
backupStdout: os.Stdout,
writerStdout: wStdout,
backupStderr: os.Stderr,
writerStderr: wStderr,
outStdout: make(chan []byte),
outStderr: make(chan []byte),
DataStdOut: &bytes.Buffer{},
DataStdErr: &bytes.Buffer{},
}
os.Stdout = c.writerStdout
os.Stderr = c.writerStderr
stdCopy := func(out chan<- []byte, reader *os.File) {
var buffer bytes.Buffer
_, _ = io.Copy(&buffer, reader)
if buffer.Len() > 0 {
out <- buffer.Bytes()
}
}
go stdCopy(c.outStdout, rStdout)
go stdCopy(c.outStderr, rStderr)
return c
}
func stopStdCapture(c *captureData) {
_ = c.writerStdout.Close()
_ = c.writerStderr.Close()
var wg sync.WaitGroup
stdRead := func(in <-chan []byte, outData *bytes.Buffer) {
defer wg.Done()
select {
case out := <-in:
outData.Write(out)
default:
//case <-time.After(10 * time.Millisecond):
}
}
wg.Add(2)
go stdRead(c.outStdout, c.DataStdOut)
go stdRead(c.outStderr, c.DataStdErr)
wg.Wait()
os.Stdout = c.backupStdout
os.Stderr = c.backupStderr
}