From a53cd037032fac08c0f1894b9f1ae65c54b9d55f Mon Sep 17 00:00:00 2001
From: Erik Geiser <erik.geiser@redteam-pentesting.de>
Date: Wed, 28 Oct 2020 07:54:46 +0100
Subject: [PATCH] Colorize output

---
 cli/log_terminal.go  | 13 ++++++-
 cmd/fuzz/main.go     |  4 +-
 reporter/reporter.go |  6 +--
 reporter/response.go | 91 ++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 107 insertions(+), 7 deletions(-)
 create mode 100644 reporter/response.go

diff --git a/cli/log_terminal.go b/cli/log_terminal.go
index e671e8b..f3cc9ca 100644
--- a/cli/log_terminal.go
+++ b/cli/log_terminal.go
@@ -3,18 +3,25 @@ package cli
 import (
 	"fmt"
 	"io"
+	"regexp"
 	"strings"
 
 	"github.com/fd0/termstatus"
 )
 
+// taken from here (MIT licensed)
+// /~https://github.com/acarl005/stripansi/blob/5a71ef0e047df0427e87a79f27009029921f1f9b/stripansi.go#L7
+var ansiEscapeSequenceRegEx = regexp.MustCompile(
+	"[\u001B\u009B][[\\]()#;?]*(?:(?:(?:[a-zA-Z\\d]*(?:;[a-zA-Z\\d]*)*)?\u0007)|" +
+		"(?:(?:\\d{1,4}(?:;\\d{0,4})*)?[\\dA-PRZcf-ntqry=><~]))")
+
 // LogTerminal writes data to a second writer in addition to the terminal.
 type LogTerminal struct {
 	*termstatus.Terminal
 	io.Writer
 }
 
-// Printf prints a messsage with formatting.
+// Printf prints a message with formatting.
 func (lt *LogTerminal) Printf(msg string, data ...interface{}) {
 	lt.Print(fmt.Sprintf(msg, data...))
 }
@@ -26,5 +33,7 @@ func (lt *LogTerminal) Print(msg string) {
 	}
 
 	lt.Terminal.Print(msg)
-	_, _ = lt.Writer.Write([]byte(msg))
+
+	strippedMsg := ansiEscapeSequenceRegEx.ReplaceAllString(msg, "")
+	_, _ = lt.Writer.Write([]byte(strippedMsg))
 }
diff --git a/cmd/fuzz/main.go b/cmd/fuzz/main.go
index 781509a..40722f2 100644
--- a/cmd/fuzz/main.go
+++ b/cmd/fuzz/main.go
@@ -248,7 +248,7 @@ func setupTerminal(ctx context.Context, g *errgroup.Group, maxFrameRate uint, lo
 	term = statusTerm
 
 	if logfilePrefix != "" {
-		fmt.Printf("logfile is %s.log\n", logfilePrefix)
+		fmt.Printf(reporter.Bold("Logfile:")+" %s.log\n", logfilePrefix)
 
 		logfile, err := os.Create(logfilePrefix + ".log")
 		if err != nil {
@@ -469,7 +469,7 @@ func run(ctx context.Context, g *errgroup.Group, opts *Options, args []string) e
 	}
 
 	// run the reporter
-	term.Printf("input URL %v\n\n", inputURL)
+	term.Printf(reporter.Bold("Target URL:")+" %v\n\n", inputURL)
 	reporter := reporter.New(term)
 	return reporter.Display(responseCh, countCh)
 }
diff --git a/reporter/reporter.go b/reporter/reporter.go
index 78c3825..854945f 100644
--- a/reporter/reporter.go
+++ b/reporter/reporter.go
@@ -78,7 +78,7 @@ func (h *HTTPStats) Report(current string) (res []string) {
 	res = append(res, status)
 
 	for code, count := range h.StatusCodes {
-		res = append(res, fmt.Sprintf("%v: %v", code, count))
+		res = append(res, fmt.Sprintf("%s: %v", colorStatusCode(code, ""), count))
 	}
 
 	sort.Strings(res[2:])
@@ -88,7 +88,7 @@ func (h *HTTPStats) Report(current string) (res []string) {
 
 // Display shows incoming Responses.
 func (r *Reporter) Display(ch <-chan response.Response, countChannel <-chan int) error {
-	r.term.Printf("%7s %8s %8s   %-8s %s\n", "status", "header", "body", "value", "extract")
+	r.term.Printf(Bold("%7s %8s %8s   %-8s %s\n"), "status", "header", "body", "value", "extract")
 
 	stats := &HTTPStats{
 		Start:       time.Now(),
@@ -111,7 +111,7 @@ func (r *Reporter) Display(ch <-chan response.Response, countChannel <-chan int)
 		}
 
 		if !response.Hide {
-			r.term.Printf("%v\n", response)
+			r.term.Printf("%v\n", FormatResponse(response))
 			stats.ShownResponses++
 		}
 
diff --git a/reporter/response.go b/reporter/response.go
new file mode 100644
index 0000000..baa9464
--- /dev/null
+++ b/reporter/response.go
@@ -0,0 +1,91 @@
+package reporter
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"strconv"
+	"strings"
+
+	"github.com/RedTeamPentesting/monsoon/response"
+)
+
+const (
+	_ int = iota + 30 // black
+	red
+	green
+	yellow
+	blue
+	_ // magenta
+	cyan
+	_ // white
+)
+
+func colorStatusCode(statusCode int, format string) string {
+	var color int
+
+	switch statusCode / 100 {
+	case 1:
+		color = blue
+	case 2:
+		color = green
+	case 3:
+		color = cyan
+	case 4:
+		color = yellow
+	case 5:
+		color = red
+	}
+
+	if format == "" {
+		format = "%d"
+	}
+
+	return fmt.Sprintf("\033[%dm"+format+"\033[0m", color, statusCode)
+}
+
+func Bold(s string) string {
+	return "\033[1m" + s + "\033[0m"
+}
+
+func Dim(s string) string {
+	return "\033[2m" + s + "\033[0m"
+}
+
+func FormatResponse(r response.Response) string {
+	if r.Error != nil {
+		// don't print anything if the request has been cancelled
+		if r.Error == context.Canceled {
+			return ""
+		}
+		if e, ok := r.Error.(*url.Error); ok && e.Err == context.Canceled {
+			return ""
+		}
+
+		return fmt.Sprintf("%7s %18s   %v", "error", r.Error, r.Item)
+	}
+
+	res := r.HTTPResponse
+	status := fmt.Sprintf("%s %8d %8d   %-8v", colorStatusCode(res.StatusCode, "%7d"),
+		r.Header.Bytes, r.Body.Bytes, Bold(r.Item))
+	if res.StatusCode >= 300 && res.StatusCode < 400 {
+		loc, ok := res.Header["Location"]
+		if ok {
+			status += ", " + Dim("Location: ") + loc[0]
+		}
+	}
+	if len(r.Extract) > 0 {
+		status += Dim(" data: ") + Bold(strings.Join(quote(r.Extract), ", "))
+	}
+	return status
+}
+
+func quote(strs []string) []string {
+	res := make([]string, 0, len(strs))
+	for _, s := range strs {
+		r := strconv.Quote(strings.TrimSpace(s))
+		r = r[1 : len(r)-1]
+		res = append(res, r)
+	}
+	return res
+}