initial commit

internal/color/color.go

@@ -0,0 +1,60 @@
+package color
+
+import (
+	"os"
+	"strings"
+
+	"github.com/fatih/color"
+)
+
+var (
+	Header      = color.New(color.FgGreen, color.Bold)
+	Bold        = color.New(color.Bold)
+	Faint       = color.New(color.Faint)
+	TCP         = color.New(color.FgCyan)
+	UDP         = color.New(color.FgMagenta)
+	LISTEN      = color.New(color.FgYellow)
+	ESTABLISHED = color.New(color.FgGreen)
+	Default     = color.New(color.FgWhite)
+)
+
+func Init(mode string) {
+	switch mode {
+	case "always":
+		color.NoColor = false
+	case "never":
+		color.NoColor = true
+	case "auto":
+		if os.Getenv("NO_COLOR") != "" || os.Getenv("TERM") == "dumb" {
+			color.NoColor = true
+		} else {
+			color.NoColor = false
+		}
+	}
+}
+
+func IsColorDisabled() bool {
+	return color.NoColor
+}
+
+func GetProtoColor(proto string) *color.Color {
+	switch strings.ToLower(proto) {
+	case "tcp":
+		return TCP
+	case "udp":
+		return UDP
+	default:
+		return Default
+	}
+}
+
+func GetStateColor(state string) *color.Color {
+	switch strings.ToUpper(state) {
+	case "LISTEN", "LISTENING":
+		return LISTEN
+	case "ESTABLISHED":
+		return ESTABLISHED
+	default:
+		return Default
+	}
+}

internal/color/color_test.go

@@ -0,0 +1,46 @@
+package color
+
+import (
+	"os"
+	"testing"
+
+	"github.com/fatih/color"
+)
+
+func TestInit(t *testing.T) {
+	testCases := []struct {
+		name      string
+		mode      string
+		noColor   string
+		term      string
+		expected  bool
+	}{
+		{"Always", "always", "", "", false},
+		{"Never", "never", "", "", true},
+		{"Auto no env", "auto", "", "xterm-256color", false},
+		{"Auto with NO_COLOR", "auto", "1", "xterm-256color", true},
+		{"Auto with TERM=dumb", "auto", "", "dumb", true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Save original env vars
+			origNoColor := os.Getenv("NO_COLOR")
+			origTerm := os.Getenv("TERM")
+
+			// Set test env vars
+			os.Setenv("NO_COLOR", tc.noColor)
+			os.Setenv("TERM", tc.term)
+
+			Init(tc.mode)
+
+			if color.NoColor != tc.expected {
+				t.Errorf("Expected color.NoColor to be %v, but got %v", tc.expected, color.NoColor)
+			}
+
+			// Restore original env vars
+			os.Setenv("NO_COLOR", origNoColor)
+			os.Setenv("TERM", origTerm)
+		})
+	}
+}