2 files changed, 273 insertions, 0 deletions

diff --git a/cli/internal/logstreamer/logstreamer.go b/cli/internal/logstreamer/logstreamer.go
new file mode 100644
index 0000000..4379c25
--- /dev/null
+++ b/cli/internal/logstreamer/logstreamer.go
@@ -0,0 +1,159 @@
+// Copyright (c) 2013 Kevin van Zonneveld <kevin@vanzonneveld.net>. All rights reserved.
+// Source: https://github.com/kvz/logstreamer
+// SPDX-License-Identifier: MIT
+package logstreamer
+
+import (
+	"bytes"
+	"io"
+	"log"
+	"os"
+	"strings"
+)
+
+type Logstreamer struct {
+	Logger *log.Logger
+	buf    *bytes.Buffer
+	// If prefix == stdout, colors green
+	// If prefix == stderr, colors red
+	// Else, prefix is taken as-is, and prepended to anything
+	// you throw at Write()
+	prefix string
+	// if true, saves output in memory
+	record  bool
+	persist string
+
+	// Adds color to stdout & stderr if terminal supports it
+	colorOkay  string
+	colorFail  string
+	colorReset string
+}
+
+func NewLogstreamer(logger *log.Logger, prefix string, record bool) *Logstreamer {
+	streamer := &Logstreamer{
+		Logger:     logger,
+		buf:        bytes.NewBuffer([]byte("")),
+		prefix:     prefix,
+		record:     record,
+		persist:    "",
+		colorOkay:  "",
+		colorFail:  "",
+		colorReset: "",
+	}
+
+	if strings.HasPrefix(os.Getenv("TERM"), "xterm") {
+		streamer.colorOkay = "\x1b[32m"
+		streamer.colorFail = "\x1b[31m"
+		streamer.colorReset = "\x1b[0m"
+	}
+
+	return streamer
+}
+
+func (l *Logstreamer) Write(p []byte) (n int, err error) {
+	if n, err = l.buf.Write(p); err != nil {
+		return
+	}
+
+	err = l.OutputLines()
+	return
+}
+
+func (l *Logstreamer) Close() error {
+	if err := l.Flush(); err != nil {
+		return err
+	}
+	l.buf = bytes.NewBuffer([]byte(""))
+	return nil
+}
+
+func (l *Logstreamer) Flush() error {
+	p := make([]byte, l.buf.Len())
+	if _, err := l.buf.Read(p); err != nil {
+		return err
+	}
+
+	l.out(string(p))
+	return nil
+}
+
+func (l *Logstreamer) OutputLines() error {
+	for {
+		line, err := l.buf.ReadString('\n')
+
+		if len(line) > 0 {
+			if strings.HasSuffix(line, "\n") {
+				l.out(line)
+			} else {
+				// put back into buffer, it's not a complete line yet
+				//  Close() or Flush() have to be used to flush out
+				//  the last remaining line if it does not end with a newline
+				if _, err := l.buf.WriteString(line); err != nil {
+					return err
+				}
+			}
+		}
+
+		if err == io.EOF {
+			break
+		}
+
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (l *Logstreamer) FlushRecord() string {
+	buffer := l.persist
+	l.persist = ""
+	return buffer
+}
+
+func (l *Logstreamer) out(str string) {
+	if len(str) < 1 {
+		return
+	}
+
+	if l.record {
+		l.persist = l.persist + str
+	}
+
+	if l.prefix == "stdout" {
+		str = l.colorOkay + l.prefix + l.colorReset + " " + str
+	} else if l.prefix == "stderr" {
+		str = l.colorFail + l.prefix + l.colorReset + " " + str
+	}
+
+	l.Logger.Print(str)
+}
+
+// PrettyStdoutWriter wraps an ioWriter so it can add string
+// prefixes to every message it writes to stdout.
+type PrettyStdoutWriter struct {
+	w      io.Writer
+	Prefix string
+}
+
+var _ io.Writer = (*PrettyStdoutWriter)(nil)
+
+// NewPrettyStdoutWriter returns an instance of PrettyStdoutWriter
+func NewPrettyStdoutWriter(prefix string) *PrettyStdoutWriter {
+	return &PrettyStdoutWriter{
+		w:      os.Stdout,
+		Prefix: prefix,
+	}
+}
+
+func (psw *PrettyStdoutWriter) Write(p []byte) (int, error) {
+	str := psw.Prefix + string(p)
+	n, err := psw.w.Write([]byte(str))
+
+	if err != nil {
+		return n, err
+	}
+
+	return len(p), nil
+}
diff --git a/cli/internal/logstreamer/logstreamer_test.go b/cli/internal/logstreamer/logstreamer_test.go
new file mode 100644
index 0000000..94d8a82
--- /dev/null
+++ b/cli/internal/logstreamer/logstreamer_test.go
@@ -0,0 +1,114 @@
+// Copyright (c) 2013 Kevin van Zonneveld <kevin@vanzonneveld.net>. All rights reserved.
+// Source: https://github.com/kvz/logstreamer
+// SPDX-License-Identifier: MIT
+package logstreamer
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"log"
+	"os"
+	"os/exec"
+	"strings"
+	"testing"
+)
+
+func TestLogstreamerOk(t *testing.T) {
+	// Create a logger (your app probably already has one)
+	logger := log.New(os.Stdout, "--> ", log.Ldate|log.Ltime)
+
+	// Setup a streamer that we'll pipe cmd.Stdout to
+	logStreamerOut := NewLogstreamer(logger, "stdout", false)
+	defer logStreamerOut.Close()
+	// Setup a streamer that we'll pipe cmd.Stderr to.
+	// We want to record/buffer anything that's written to this (3rd argument true)
+	logStreamerErr := NewLogstreamer(logger, "stderr", true)
+	defer logStreamerErr.Close()
+
+	// Execute something that succeeds
+	cmd := exec.Command(
+		"ls",
+		"-al",
+	)
+	cmd.Stderr = logStreamerErr
+	cmd.Stdout = logStreamerOut
+
+	// Reset any error we recorded
+	logStreamerErr.FlushRecord()
+
+	// Execute command
+	err := cmd.Start()
+
+	// Failed to spawn?
+	if err != nil {
+		t.Fatal("ERROR could not spawn command.", err.Error())
+	}
+
+	// Failed to execute?
+	err = cmd.Wait()
+	if err != nil {
+		t.Fatal("ERROR command finished with error. ", err.Error(), logStreamerErr.FlushRecord())
+	}
+}
+
+func TestLogstreamerErr(t *testing.T) {
+	// Create a logger (your app probably already has one)
+	logger := log.New(os.Stdout, "--> ", log.Ldate|log.Ltime)
+
+	// Setup a streamer that we'll pipe cmd.Stdout to
+	logStreamerOut := NewLogstreamer(logger, "stdout", false)
+	defer logStreamerOut.Close()
+	// Setup a streamer that we'll pipe cmd.Stderr to.
+	// We want to record/buffer anything that's written to this (3rd argument true)
+	logStreamerErr := NewLogstreamer(logger, "stderr", true)
+	defer logStreamerErr.Close()
+
+	// Execute something that succeeds
+	cmd := exec.Command(
+		"ls",
+		"nonexisting",
+	)
+	cmd.Stderr = logStreamerErr
+	cmd.Stdout = logStreamerOut
+
+	// Reset any error we recorded
+	logStreamerErr.FlushRecord()
+
+	// Execute command
+	err := cmd.Start()
+
+	// Failed to spawn?
+	if err != nil {
+		logger.Print("ERROR could not spawn command. ")
+	}
+
+	// Failed to execute?
+	err = cmd.Wait()
+	if err != nil {
+		fmt.Printf("Good. command finished with %s. %s. \n", err.Error(), logStreamerErr.FlushRecord())
+	} else {
+		t.Fatal("This command should have failed")
+	}
+}
+
+func TestLogstreamerFlush(t *testing.T) {
+	const text = "Text without newline"
+
+	var buffer bytes.Buffer
+	byteWriter := bufio.NewWriter(&buffer)
+
+	logger := log.New(byteWriter, "", 0)
+	logStreamerOut := NewLogstreamer(logger, "", false)
+	defer logStreamerOut.Close()
+
+	logStreamerOut.Write([]byte(text))
+	logStreamerOut.Flush()
+	byteWriter.Flush()
+
+	s := strings.TrimSpace(buffer.String())
+
+	if s != text {
+		t.Fatalf("Expected '%s', got '%s'.", text, s)
+	}
+}