From 148a9504e7dc9eceea79254792fba60f2f095851 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Wed, 21 Mar 2018 11:18:26 -0700
Subject: [PATCH] color

---
 command/commands.go                           |  5 ++
 command/meta.go                               |  3 +-
 main.go                                       | 54 ++++++++++++++++++-
 .../docs/commands/_general_options.html.md    |  3 +-
 4 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/command/commands.go b/command/commands.go
index 40b9c0317..63a050ad3 100644
--- a/command/commands.go
+++ b/command/commands.go
@@ -9,6 +9,11 @@ import (
 	"github.com/mitchellh/cli"
 )
 
+const (
+	// EnvNomadCLINoColor is an env var that toggles colored UI output.
+	EnvNomadCLINoColor = `NOMAD_CLI_NO_COLOR`
+)
+
 // DeprecatedCommand is a command that wraps an existing command and prints a
 // deprecation notice and points the user to the new command. Deprecated
 // commands are always hidden from help output.
diff --git a/command/meta.go b/command/meta.go
index de1addd11..8dc96e933 100644
--- a/command/meta.go
+++ b/command/meta.go
@@ -181,7 +181,8 @@ func generalOptionsUsage() string {
     Defaults to the "default" namespace.
 
   -no-color
-    Disables colored command output.
+    Disables colored command output. Alternatively, NOMAD_CLI_NO_COLOR may be
+    set.
 
   -ca-cert=<path>
     Path to a PEM encoded CA cert file to use to verify the
diff --git a/main.go b/main.go
index 13a07c02b..ff8216f67 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ import (
 	"github.com/hashicorp/nomad/version"
 	"github.com/mitchellh/cli"
 	"github.com/sean-/seed"
+	"golang.org/x/crypto/ssh/terminal"
 )
 
 var (
@@ -63,10 +64,39 @@ func main() {
 }
 
 func Run(args []string) int {
-	return RunCustom(args, command.Commands(nil))
+	return RunCustom(args)
 }
 
-func RunCustom(args []string, commands map[string]cli.CommandFactory) int {
+func RunCustom(args []string) int {
+	// Parse flags into env vars for global use
+	args = setupEnv(args)
+
+	// Create the meta object
+	metaPtr := new(command.Meta)
+
+	// Don't use color if disabled
+	color := true
+	if os.Getenv(command.EnvNomadCLINoColor) != "" {
+		color = false
+	}
+
+	isTerminal := terminal.IsTerminal(int(os.Stdout.Fd()))
+	metaPtr.Ui = &cli.BasicUi{
+		Reader:      os.Stdin,
+		Writer:      os.Stdout,
+		ErrorWriter: os.Stderr,
+	}
+
+	// Only use colored UI if stdout is a tty, and not disabled
+	if isTerminal && color {
+		metaPtr.Ui = &cli.ColoredUi{
+			ErrorColor: cli.UiColorRed,
+			WarnColor:  cli.UiColorYellow,
+			Ui:         metaPtr.Ui,
+		}
+	}
+
+	commands := command.Commands(metaPtr)
 	cli := &cli.CLI{
 		Name:                       "nomad",
 		Version:                    version.GetVersion().FullVersionNumber(true),
@@ -78,6 +108,7 @@ func RunCustom(args []string, commands map[string]cli.CommandFactory) int {
 		HelpFunc: groupedHelpFunc(
 			cli.BasicHelpFunc("nomad"),
 		),
+		HelpWriter: os.Stdout,
 	}
 
 	exitCode, err := cli.Run()
@@ -177,3 +208,22 @@ func printCommand(w io.Writer, name string, cmdFn cli.CommandFactory) {
 	}
 	fmt.Fprintf(w, "    %s\t%s\n", name, cmd.Synopsis())
 }
+
+// setupEnv parses args and may replace them and sets some env vars to known
+// values based on format options
+func setupEnv(args []string) []string {
+	noColor := false
+	for _, arg := range args {
+		// Check if color is set
+		if arg == "-no-color" || arg == "--no-color" {
+			noColor = true
+		}
+	}
+
+	// Put back into the env for later
+	if noColor {
+		os.Setenv(command.EnvNomadCLINoColor, "true")
+	}
+
+	return args
+}
diff --git a/website/source/docs/commands/_general_options.html.md b/website/source/docs/commands/_general_options.html.md
index 5a1811f9b..86cba60fa 100644
--- a/website/source/docs/commands/_general_options.html.md
+++ b/website/source/docs/commands/_general_options.html.md
@@ -5,7 +5,8 @@
   Overrides the `NOMAD_REGION` environment variable if set. Defaults to the
   Agent's local region.
 
-- `-no-color`: Disables colored command output.
+- `-no-color`: Disables colored command output. Alternatively,
+  `NOMAD_CLI_NO_COLOR` may be set.
 
 - `-ca-cert=<path>`: Path to a PEM encoded CA cert file to use to verify the
   Nomad server SSL certificate. Overrides the `NOMAD_CACERT` environment