cmd/vulnreport: add color to logs (and refactor)

Refactor logging and add colors.

Colors are on by default but can be turned off with
`vulnreport -color=false [subcommand]...`, or by setting
the environment variable `NO_COLOR`.

Change-Id: Iaea30dfeb045b2b690683f257cf49676c785951a
Reviewed-on: https://go-review.googlesource.com/c/vulndb/+/562555
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Damien Neil <[email protected]>

Author: tatianab

cmd/vulnreport/color/color.go

@@ -0,0 +1,59 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package color
+
+const (
+	escape = "\033["
+	end    = "m"
+
+	// Reset all attributes and colors to their defaults.
+	Reset = escape + "0" + end
+
+	// Attribute codes
+	Bold      = escape + "1" + end
+	Faint     = escape + "2" + end
+	Underline = escape + "4" + end
+	Blink     = escape + "5" + end
+
+	// Foreground colors
+	Black   = escape + "30" + end
+	Red     = escape + "31" + end
+	Green   = escape + "32" + end
+	Yellow  = escape + "33" + end
+	Blue    = escape + "34" + end
+	Magenta = escape + "35" + end
+	Cyan    = escape + "36" + end
+	White   = escape + "37" + end
+
+	// Background colors
+	BlackBG   = escape + "40" + end
+	RedBG     = escape + "41" + end
+	GreenBG   = escape + "42" + end
+	YellowBG  = escape + "43" + end
+	BlueBG    = escape + "44" + end
+	MagentaBG = escape + "45" + end
+	CyanBG    = escape + "46" + end
+	WhiteBG   = escape + "47" + end
+
+	// High intensity foreground colors
+	BlackHi   = escape + "90" + end
+	RedHi     = escape + "91" + end
+	GreenHi   = escape + "92" + end
+	YellowHi  = escape + "93" + end
+	BlueHi    = escape + "94" + end
+	MagentaHi = escape + "95" + end
+	CyanHi    = escape + "96" + end
+	WhiteHi   = escape + "97" + end
+
+	// High intensity background colors
+	BlackHiBG   = escape + "100" + end
+	RedHiBG     = escape + "101" + end
+	GreenHiBG   = escape + "102" + end
+	YellowHiBG  = escape + "103" + end
+	BlueHiBG    = escape + "104" + end
+	MagentaHiBG = escape + "105" + end
+	CyanHiBG    = escape + "106" + end
+	WhiteHiBG   = escape + "107" + end
+)

cmd/vulnreport/log/log.go

@@ -5,58 +5,117 @@
 package log
 
 import (
+	"fmt"
 	"io"
 	"os"
 
 	"log"
+
+	"golang.org/x/vulndb/cmd/vulnreport/color"
 )
 
+func SetQuiet() {
+	loggers[infoLvl] = log.New(io.Discard, "", 0)
+}
+
+func RemoveColor() {
+	for lvl := range loggers {
+		loggers[lvl].SetPrefix(metas[lvl].prefix)
+	}
+	colorize = false
+}
+
+const (
+	infoLvl int = iota
+	outLvl
+	warnLvl
+	errLvl
+)
+
+func defaultLoggers() []*log.Logger {
+	ls := make([]*log.Logger, len(metas))
+	for lvl, lm := range metas {
+		ls[lvl] = log.New(lm.w, lm.color+lm.prefix, 0)
+	}
+	return ls
+}
+
 var (
-	infolog *log.Logger
-	outlog  *log.Logger
-	warnlog *log.Logger
-	errlog  *log.Logger
+	loggers = defaultLoggers()
+
+	// Whether to display colors in logs.
+	colorize bool = true
+
+	metas = []*metadata{
+		infoLvl: {
+			prefix: "info: ",
+			color:  color.Faint,
+			w:      os.Stderr,
+		},
+		outLvl: {
+			prefix: "",
+			color:  color.Reset,
+			w:      os.Stdout,
+		},
+		warnLvl: {
+			prefix: "WARNING: ",
+			color:  color.YellowHi,
+			w:      os.Stderr,
+		},
+		errLvl: {
+			prefix: "ERROR: ",
+			color:  color.RedHi,
+			w:      os.Stderr,
+		},
+	}
 )
 
-func Init(quiet bool) {
-	if quiet {
-		infolog = log.New(io.Discard, "", 0)
-	} else {
-		infolog = log.New(os.Stderr, "info: ", 0)
+type metadata struct {
+	prefix string
+	color  string
+	w      io.Writer
+}
+
+func printf(lvl int, format string, v ...any) {
+	println(lvl, fmt.Sprintf(format, v...))
+}
+
+func println(lvl int, v ...any) {
+	l := loggers[lvl]
+	l.Println(v...)
+	if colorize {
+		fmt.Fprint(l.Writer(), color.Reset)
 	}
-	outlog = log.New(os.Stdout, "", 0)
-	warnlog = log.New(os.Stderr, "WARNING: ", 0)
-	errlog = log.New(os.Stderr, "ERROR: ", 0)
 }
 
 func Infof(format string, v ...any) {
-	infolog.Printf(format, v...)
+	printf(infoLvl, format, v...)
 }
 
 func Outf(format string, v ...any) {
-	outlog.Printf(format, v...)
+	printf(outLvl, format, v...)
 }
 
 func Warnf(format string, v ...any) {
-	warnlog.Printf(format, v...)
+	printf(warnLvl, format, v...)
 }
 
 func Errf(format string, v ...any) {
-	errlog.Printf(format, v...)
+	printf(errLvl, format, v...)
 }
 
 func Info(v ...any) {
-	infolog.Println(v...)
+	println(infoLvl, v...)
 }
 
 func Out(v ...any) {
-	outlog.Println(v...)
+	println(outLvl, v...)
 }
 
 func Warn(v ...any) {
-	warnlog.Println(v...)
+	println(warnLvl, v...)
 }
 
 func Err(v ...any) {
-	errlog.Println(v...)
+	println(errLvl, v...)
 }

cmd/vulnreport/main.go

@@ -22,10 +22,10 @@ var (
 	githubToken = flag.String("ghtoken", "", "GitHub access token (default: value of VULN_GITHUB_ACCESS_TOKEN)")
 	cpuprofile  = flag.String("cpuprofile", "", "write cpuprofile to this file")
 	quiet       = flag.Bool("q", false, "quiet mode (suppress info logs)")
+	colorize    = flag.Bool("color", os.Getenv("NO_COLOR") == "", "show colors in logs")
 )
 
 func init() {
-	vlog.Init(*quiet)
 	out := flag.CommandLine.Output()
 	flag.Usage = func() {
 		fmt.Fprintf(out, "usage: vulnreport [flags] [cmd] [args]\n\n")
@@ -67,6 +67,13 @@ func main() {
 		log.Fatal("subcommand required")
 	}
 
+	if *quiet {
+		vlog.SetQuiet()
+	}
+	if !*colorize {
+		vlog.RemoveColor()
+	}
+
 	if *githubToken == "" {
 		*githubToken = os.Getenv("VULN_GITHUB_ACCESS_TOKEN")
 	}