2 files changed, 133 insertions, 0 deletions
diff --git a/color/color.go b/color/color.go
new file mode 100644
index 0000000..64e0d7f
--- /dev/null
+++ b/color/color.go
@@ -0,0 +1,75 @@
+// Package color is used to prettify console output via ANSII terminal colors.
+package color
+
+import (
+	"fmt"
+)
+
+// Color name.
+type Color string
+
+// Attribute of a color.
+type Attribute string
+
+// The possible color variations.
+const (
+	escape        = "\x1b"
+	reset         = escape + "[0m"
+	seq    string = "%s%s%s"
+
+	Gray      Color = escape + "[30m"
+	Red       Color = escape + "[31m"
+	Green     Color = escape + "[32m"
+	Orange    Color = escape + "[33m"
+	Blue      Color = escape + "[34m"
+	Magenta   Color = escape + "[35m"
+	Yellow    Color = escape + "[36m"
+	LightGray Color = escape + "[37m"
+
+	BgGray      Color = escape + "[40m"
+	BgRed       Color = escape + "[41m"
+	BgGreen     Color = escape + "[42m"
+	BgOrange    Color = escape + "[43m"
+	BgBlue      Color = escape + "[44m"
+	BgMagenta   Color = escape + "[45m"
+	BgYellow    Color = escape + "[46m"
+	BgLightGray Color = escape + "[47m"
+
+	Bold         Attribute = escape + "[1m"
+	Italic       Attribute = escape + "[3m"
+	Underline    Attribute = escape + "[4m"
+	ReverseColor Attribute = escape + "[7m"
+
+	resetBold      = escape + "[22m"
+	resetItalic    = escape + "[23m"
+	resetUnderline = escape + "[24m"
+
+	Test     Color     = BgYellow
+	TestAttr Attribute = Bold
+)
+
+// Colored DTail client output enabled.
+var Colored bool
+
+// Init whether we want colored output or not.
+func Init(colored bool) {
+	Colored = colored
+}
+
+// Paint a given string in a given color.
+func Paint(c Color, s string) string {
+	return fmt.Sprintf(seq, c, s, reset)
+}
+
+// Attr adds a given attribute to a given string, such as "bold" or "italic".
+func Attr(c Attribute, s string) string {
+	switch c {
+	case Bold:
+		return fmt.Sprintf(seq, Bold, s, resetBold)
+	case Italic:
+		return fmt.Sprintf(seq, Italic, s, resetItalic)
+	case Underline:
+		return fmt.Sprintf(seq, Underline, s, resetUnderline)
+	}
+	panic("Unknown attribute")
+}
diff --git a/color/colorfy.go b/color/colorfy.go
new file mode 100644
index 0000000..9ae46f5
--- /dev/null
+++ b/color/colorfy.go
@@ -0,0 +1,58 @@
+package color
+
+import (
+	"fmt"
+	"strings"
+)
+
+// Add some color to log lines received from remote servers.
+func paintRemote(line string) string {
+	splitted := strings.Split(line, "|")
+	if splitted[2] == "100" {
+		splitted[2] = Paint(BgGreen, splitted[2])
+	} else {
+		splitted[2] = Paint(BgRed, splitted[2])
+	}
+	info := strings.Join(splitted[0:5], "|")
+	log := strings.Join(splitted[5:], "|")
+
+	if strings.HasPrefix(log, "WARN") {
+		log = Paint(BgYellow, log)
+	} else if strings.HasPrefix(log, "ERROR") {
+		log = Paint(BgRed, log)
+	} else if strings.HasPrefix(log, "FATAL") {
+		log = Attr(Bold, Paint(BgRed, log))
+	} else {
+		log = Paint(Blue, log)
+	}
+
+	return fmt.Sprintf("%s|%s", info, log)
+}
+
+// Add some color to stats generated by the client.
+func paintClientStats(line string) string {
+	splitted := strings.Split(line, "|")
+	first := strings.Join(splitted[0:4], "|")
+	connected := Paint(BgBlue, splitted[4])
+	last := strings.Join(splitted[5:], "|")
+
+	return fmt.Sprintf("%s|%s|%s", first, connected, last)
+}
+
+// Colorfy a given line based on the line's content.
+func Colorfy(line string) string {
+	if strings.HasPrefix(line, "REMOTE") {
+		return paintRemote(line)
+	}
+	if strings.HasPrefix(line, "CLIENT") && strings.Contains(line, "|stats|") {
+		return paintClientStats(line)
+	}
+	if strings.Contains(line, "ERROR") {
+		return Paint(Magenta, line)
+	}
+	if strings.Contains(line, "WARN") {
+		return Paint(Magenta, line)
+	}
+
+	return line
+}