snowpick/color.go

package main

import (
	"fmt"
	"os"
)

// ANSI escape sequences for terminal coloring.
// Enabled only when stderr is a terminal (not piped to a file).
var (
	colorReset  = ""
	colorBold   = ""
	colorRed    = ""
	colorGreen  = ""
	colorYellow = ""
	colorCyan   = ""
	colorDim    = ""
)

func initColors() {
	if isTerminal(os.Stderr) {
		colorReset = "\033[0m"
		colorBold = "\033[1m"
		colorRed = "\033[1;31m"
		colorGreen = "\033[1;32m"
		colorYellow = "\033[1;33m"
		colorCyan = "\033[1;36m"
		colorDim = "\033[2m"
	}
}

// isTerminal returns true if the file is connected to a terminal.
func isTerminal(f *os.File) bool {
	info, err := f.Stat()
	if err != nil {
		return false
	}
	return (info.Mode() & os.ModeCharDevice) != 0
}

func red(s string) string    { return colorRed + s + colorReset }
func green(s string) string  { return colorGreen + s + colorReset }
func yellow(s string) string { return colorYellow + s + colorReset }
func cyan(s string) string   { return colorCyan + s + colorReset }
func bold(s string) string   { return colorBold + s + colorReset }
func dim(s string) string    { return colorDim + s + colorReset }

func severityColor(count int, cats []Hint) string {
	for _, c := range cats {
		if c.Category == ImpactPII || c.Category == ImpactCredentials {
			return colorRed
		}
	}
	if count >= 1000 {
		return colorRed
	}
	if count >= 100 {
		return colorYellow
	}
	return colorCyan
}

// fmtCount formats a record count with severity coloring.
func fmtCount(count int, cats []Hint) string {
	color := severityColor(count, cats)
	return fmt.Sprintf("%s%d%s", color, count, colorReset)
}

// fmtCountSampled formats "N records, M sampled" when total > sampled.
func fmtCountSampled(total, sampled int, cats []Hint) string {
	color := severityColor(total, cats)
	if sampled > 0 && total > sampled {
		return fmt.Sprintf("%s%d%s records, %d sampled", color, total, colorReset, sampled)
	}
	return fmt.Sprintf("%s%d%s records", color, total, colorReset)
}