github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/main.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"sort"
     9  	"strings"
    10  	"text/tabwriter"
    11  
    12  	// These packages have init() funcs which check os.Args and drop directly
    13  	// into their command logic. This is because they are run as separate
    14  	// processes along side of a task. By early importing them we can avoid
    15  	// additional code being imported and thus reserving memory
    16  	_ "github.com/hashicorp/nomad/client/logmon"
    17  	_ "github.com/hashicorp/nomad/drivers/docker/docklog"
    18  	_ "github.com/hashicorp/nomad/drivers/shared/executor"
    19  
    20  	"github.com/hashicorp/nomad/command"
    21  	"github.com/hashicorp/nomad/version"
    22  	colorable "github.com/mattn/go-colorable"
    23  	"github.com/mitchellh/cli"
    24  	"github.com/sean-/seed"
    25  	"golang.org/x/crypto/ssh/terminal"
    26  )
    27  
    28  var (
    29  	// Hidden hides the commands from both help and autocomplete. Commands that
    30  	// users should not be running should be placed here, versus hiding
    31  	// subcommands from the main help, which should be filtered out of the
    32  	// commands above.
    33  	hidden = []string{
    34  		"alloc-status",
    35  		"check",
    36  		"client-config",
    37  		"debug",
    38  		"eval-status",
    39  		"executor",
    40  		"keygen",
    41  		"keyring",
    42  		"logmon",
    43  		"node-drain",
    44  		"node-status",
    45  		"server-force-leave",
    46  		"server-join",
    47  		"server-members",
    48  		"syslog",
    49  		"docker_logger",
    50  		"operator raft _info",
    51  		"operator raft _logs",
    52  		"operator raft _state",
    53  	}
    54  
    55  	// aliases is the list of aliases we want users to be aware of. We hide
    56  	// these form the help output but autocomplete them.
    57  	aliases = []string{
    58  		"fs",
    59  		"init",
    60  		"inspect",
    61  		"logs",
    62  		"plan",
    63  		"validate",
    64  	}
    65  
    66  	// Common commands are grouped separately to call them out to operators.
    67  	commonCommands = []string{
    68  		"run",
    69  		"stop",
    70  		"status",
    71  		"alloc",
    72  		"job",
    73  		"node",
    74  		"agent",
    75  	}
    76  )
    77  
    78  func init() {
    79  	seed.Init()
    80  }
    81  
    82  func main() {
    83  	os.Exit(Run(os.Args[1:]))
    84  }
    85  
    86  func Run(args []string) int {
    87  	return RunCustom(args)
    88  }
    89  
    90  func RunCustom(args []string) int {
    91  	// Parse flags into env vars for global use
    92  	args = setupEnv(args)
    93  
    94  	// Create the meta object
    95  	metaPtr := new(command.Meta)
    96  
    97  	// Don't use color if disabled
    98  	color := true
    99  	if os.Getenv(command.EnvNomadCLINoColor) != "" {
   100  		color = false
   101  	}
   102  
   103  	isTerminal := terminal.IsTerminal(int(os.Stdout.Fd()))
   104  	metaPtr.Ui = &cli.BasicUi{
   105  		Reader:      os.Stdin,
   106  		Writer:      colorable.NewColorableStdout(),
   107  		ErrorWriter: colorable.NewColorableStderr(),
   108  	}
   109  
   110  	// The Nomad agent never outputs color
   111  	agentUi := &cli.BasicUi{
   112  		Reader:      os.Stdin,
   113  		Writer:      os.Stdout,
   114  		ErrorWriter: os.Stderr,
   115  	}
   116  
   117  	// Only use colored UI if stdout is a tty, and not disabled
   118  	if isTerminal && color {
   119  		metaPtr.Ui = &cli.ColoredUi{
   120  			ErrorColor: cli.UiColorRed,
   121  			WarnColor:  cli.UiColorYellow,
   122  			InfoColor:  cli.UiColorGreen,
   123  			Ui:         metaPtr.Ui,
   124  		}
   125  	}
   126  
   127  	commands := command.Commands(metaPtr, agentUi)
   128  	cli := &cli.CLI{
   129  		Name:                       "nomad",
   130  		Version:                    version.GetVersion().FullVersionNumber(true),
   131  		Args:                       args,
   132  		Commands:                   commands,
   133  		HiddenCommands:             hidden,
   134  		Autocomplete:               true,
   135  		AutocompleteNoDefaultFlags: true,
   136  		HelpFunc: groupedHelpFunc(
   137  			cli.BasicHelpFunc("nomad"),
   138  		),
   139  		HelpWriter: os.Stdout,
   140  	}
   141  
   142  	exitCode, err := cli.Run()
   143  	if err != nil {
   144  		fmt.Fprintf(os.Stderr, "Error executing CLI: %s\n", err.Error())
   145  		return 1
   146  	}
   147  
   148  	return exitCode
   149  }
   150  
   151  func groupedHelpFunc(f cli.HelpFunc) cli.HelpFunc {
   152  	return func(commands map[string]cli.CommandFactory) string {
   153  		var b bytes.Buffer
   154  		tw := tabwriter.NewWriter(&b, 0, 2, 6, ' ', 0)
   155  
   156  		fmt.Fprintf(tw, "Usage: nomad [-version] [-help] [-autocomplete-(un)install] <command> [args]\n\n")
   157  		fmt.Fprintf(tw, "Common commands:\n")
   158  		for _, v := range commonCommands {
   159  			printCommand(tw, v, commands[v])
   160  		}
   161  
   162  		// Filter out common commands and aliased commands from the other
   163  		// commands output
   164  		otherCommands := make([]string, 0, len(commands))
   165  		for k := range commands {
   166  			found := false
   167  			for _, v := range commonCommands {
   168  				if k == v {
   169  					found = true
   170  					break
   171  				}
   172  			}
   173  
   174  			for _, v := range aliases {
   175  				if k == v {
   176  					found = true
   177  					break
   178  				}
   179  			}
   180  
   181  			if !found {
   182  				otherCommands = append(otherCommands, k)
   183  			}
   184  		}
   185  		sort.Strings(otherCommands)
   186  
   187  		fmt.Fprintf(tw, "\n")
   188  		fmt.Fprintf(tw, "Other commands:\n")
   189  		for _, v := range otherCommands {
   190  			printCommand(tw, v, commands[v])
   191  		}
   192  
   193  		tw.Flush()
   194  
   195  		return strings.TrimSpace(b.String())
   196  	}
   197  }
   198  
   199  func printCommand(w io.Writer, name string, cmdFn cli.CommandFactory) {
   200  	cmd, err := cmdFn()
   201  	if err != nil {
   202  		panic(fmt.Sprintf("failed to load %q command: %s", name, err))
   203  	}
   204  	fmt.Fprintf(w, "    %s\t%s\n", name, cmd.Synopsis())
   205  }
   206  
   207  // setupEnv parses args and may replace them and sets some env vars to known
   208  // values based on format options
   209  func setupEnv(args []string) []string {
   210  	noColor := false
   211  	for _, arg := range args {
   212  		// Check if color is set
   213  		if arg == "-no-color" || arg == "--no-color" {
   214  			noColor = true
   215  		}
   216  	}
   217  
   218  	// Put back into the env for later
   219  	if noColor {
   220  		os.Setenv(command.EnvNomadCLINoColor, "true")
   221  	}
   222  
   223  	return args
   224  }