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