package cli import ( "fmt" "io" "io/ioutil" "os" "regexp" "sort" "strings" "sync" "text/template" "github.com/Masterminds/sprig" "github.com/armon/go-radix" "github.com/posener/complete" ) // CLI contains the state necessary to run subcommands and parse the // command line arguments. // // CLI also supports nested subcommands, such as "cli foo bar". To use // nested subcommands, the key in the Commands mapping below contains the // full subcommand. In this example, it would be "foo bar". // // If you use a CLI with nested subcommands, some semantics change due to // ambiguities: // // * We use longest prefix matching to find a matching subcommand. This // means if you register "foo bar" and the user executes "cli foo qux", // the "foo" command will be executed with the arg "qux". It is up to // you to handle these args. One option is to just return the special // help return code `RunResultHelp` to display help and exit. // // * The help flag "-h" or "-help" will look at all args to determine // the help function. For example: "otto apps list -h" will show the // help for "apps list" but "otto apps -h" will show it for "apps". // In the normal CLI, only the first subcommand is used. // // * The help flag will list any subcommands that a command takes // as well as the command's help itself. If there are no subcommands, // it will note this. If the CLI itself has no subcommands, this entire // section is omitted. // // * Any parent commands that don't exist are automatically created as // no-op commands that just show help for other subcommands. For example, // if you only register "foo bar", then "foo" is automatically created. // type CLI struct { // Args is the list of command-line arguments received excluding // the name of the app. For example, if the command "./cli foo bar" // was invoked, then Args should be []string{"foo", "bar"}. Args []string // Commands is a mapping of subcommand names to a factory function // for creating that Command implementation. If there is a command // with a blank string "", then it will be used as the default command // if no subcommand is specified. // // If the key has a space in it, this will create a nested subcommand. // For example, if the key is "foo bar", then to access it our CLI // must be accessed with "./cli foo bar". See the docs for CLI for // notes on how this changes some other behavior of the CLI as well. // // The factory should be as cheap as possible, ideally only allocating // a struct. The factory may be called multiple times in the course // of a command execution and certain events such as help require the // instantiation of all commands. Expensive initialization should be // deferred to function calls within the interface implementation. Commands map[string]CommandFactory // HiddenCommands is a list of commands that are "hidden". Hidden // commands are not given to the help function callback and do not // show up in autocomplete. The values in the slice should be equivalent // to the keys in the command map. HiddenCommands []string // Name defines the name of the CLI. Name string // Version of the CLI. Version string // Autocomplete enables or disables subcommand auto-completion support. // This is enabled by default when NewCLI is called. Otherwise, this // must enabled explicitly. // // Autocomplete requires the "Name" option to be set on CLI. This name // should be set exactly to the binary name that is autocompleted. // // Autocompletion is supported via the github.com/posener/complete // library. This library supports bash, zsh and fish. To add support // for other shells, please see that library. // // AutocompleteInstall and AutocompleteUninstall are the global flag // names for installing and uninstalling the autocompletion handlers // for the user's shell. The flag should omit the hyphen(s) in front of // the value. Both single and double hyphens will automatically be supported // for the flag name. These default to `autocomplete-install` and // `autocomplete-uninstall` respectively. // // AutocompleteNoDefaultFlags is a boolean which controls if the default auto- // complete flags like -help and -version are added to the output. // // AutocompleteGlobalFlags are a mapping of global flags for // autocompletion. The help and version flags are automatically added. Autocomplete bool AutocompleteInstall string AutocompleteUninstall string AutocompleteNoDefaultFlags bool AutocompleteGlobalFlags complete.Flags autocompleteInstaller autocompleteInstaller // For tests // HelpFunc is the function called to generate the generic help // text that is shown if help must be shown for the CLI that doesn't // pertain to a specific command. HelpFunc HelpFunc // HelpWriter is used to print help text and version when requested. // Defaults to os.Stderr for backwards compatibility. // It is recommended that you set HelpWriter to os.Stdout, and // ErrorWriter to os.Stderr. HelpWriter io.Writer // ErrorWriter used to output errors when a command can not be run. // Defaults to the value of HelpWriter for backwards compatibility. // It is recommended that you set HelpWriter to os.Stdout, and // ErrorWriter to os.Stderr. ErrorWriter io.Writer //--------------------------------------------------------------- // Internal fields set automatically once sync.Once autocomplete *complete.Complete commandTree *radix.Tree commandNested bool commandHidden map[string]struct{} subcommand string subcommandArgs []string topFlags []string // These are true when special global flags are set. We can/should // probably use a bitset for this one day. isHelp bool isVersion bool isAutocompleteInstall bool isAutocompleteUninstall bool } // NewClI returns a new CLI instance with sensible defaults. func NewCLI(app, version string) *CLI { return &CLI{ Name: app, Version: version, HelpFunc: BasicHelpFunc(app), Autocomplete: true, } } // IsHelp returns whether or not the help flag is present within the // arguments. func (c *CLI) IsHelp() bool { c.once.Do(c.init) return c.isHelp } // IsVersion returns whether or not the version flag is present within the // arguments. func (c *CLI) IsVersion() bool { c.once.Do(c.init) return c.isVersion } // Run runs the actual CLI based on the arguments given. func (c *CLI) Run() (int, error) { c.once.Do(c.init) // If this is a autocompletion request, satisfy it. This must be called // first before anything else since its possible to be autocompleting // -help or -version or other flags and we want to show completions // and not actually write the help or version. if c.Autocomplete && c.autocomplete.Complete() { return 0, nil } // Just show the version and exit if instructed. if c.IsVersion() && c.Version != "" { c.HelpWriter.Write([]byte(c.Version + "\n")) return 0, nil } // Just print the help when only '-h' or '--help' is passed. if c.IsHelp() && c.Subcommand() == "" { c.HelpWriter.Write([]byte(c.HelpFunc(c.helpCommands(c.Subcommand())) + "\n")) return 0, nil } // If we're attempting to install or uninstall autocomplete then handle if c.Autocomplete { // Autocomplete requires the "Name" to be set so that we know what // command to setup the autocomplete on. if c.Name == "" { return 1, fmt.Errorf( "internal error: CLI.Name must be specified for autocomplete to work") } // If both install and uninstall flags are specified, then error if c.isAutocompleteInstall && c.isAutocompleteUninstall { return 1, fmt.Errorf( "Either the autocomplete install or uninstall flag may " + "be specified, but not both.") } // If the install flag is specified, perform the install or uninstall if c.isAutocompleteInstall { if err := c.autocompleteInstaller.Install(c.Name); err != nil { return 1, err } return 0, nil } if c.isAutocompleteUninstall { if err := c.autocompleteInstaller.Uninstall(c.Name); err != nil { return 1, err } return 0, nil } } // Attempt to get the factory function for creating the command // implementation. If the command is invalid or blank, it is an error. raw, ok := c.commandTree.Get(c.Subcommand()) if !ok { c.ErrorWriter.Write([]byte(c.HelpFunc(c.helpCommands(c.subcommandParent())) + "\n")) return 127, nil } command, err := raw.(CommandFactory)() if err != nil { return 1, err } // If we've been instructed to just print the help, then print it if c.IsHelp() { c.commandHelp(c.HelpWriter, command) return 0, nil } // If there is an invalid flag, then error if len(c.topFlags) > 0 { c.ErrorWriter.Write([]byte( "Invalid flags before the subcommand. If these flags are for\n" + "the subcommand, please put them after the subcommand.\n\n")) c.commandHelp(c.ErrorWriter, command) return 1, nil } code := command.Run(c.SubcommandArgs()) if code == RunResultHelp { // Requesting help c.commandHelp(c.ErrorWriter, command) return 1, nil } return code, nil } // Subcommand returns the subcommand that the CLI would execute. For // example, a CLI from "--version version --help" would return a Subcommand // of "version" func (c *CLI) Subcommand() string { c.once.Do(c.init) return c.subcommand } // SubcommandArgs returns the arguments that will be passed to the // subcommand. func (c *CLI) SubcommandArgs() []string { c.once.Do(c.init) return c.subcommandArgs } // subcommandParent returns the parent of this subcommand, if there is one. // If there isn't on, "" is returned. func (c *CLI) subcommandParent() string { // Get the subcommand, if it is "" alread just return sub := c.Subcommand() if sub == "" { return sub } // Clear any trailing spaces and find the last space sub = strings.TrimRight(sub, " ") idx := strings.LastIndex(sub, " ") if idx == -1 { // No space means our parent is root return "" } return sub[:idx] } func (c *CLI) init() { if c.HelpFunc == nil { c.HelpFunc = BasicHelpFunc("app") if c.Name != "" { c.HelpFunc = BasicHelpFunc(c.Name) } } if c.HelpWriter == nil { c.HelpWriter = os.Stderr } if c.ErrorWriter == nil { c.ErrorWriter = c.HelpWriter } // Build our hidden commands if len(c.HiddenCommands) > 0 { c.commandHidden = make(map[string]struct{}) for _, h := range c.HiddenCommands { c.commandHidden[h] = struct{}{} } } // Build our command tree c.commandTree = radix.New() c.commandNested = false for k, v := range c.Commands { k = strings.TrimSpace(k) c.commandTree.Insert(k, v) if strings.ContainsRune(k, ' ') { c.commandNested = true } } // Go through the key and fill in any missing parent commands if c.commandNested { var walkFn radix.WalkFn toInsert := make(map[string]struct{}) walkFn = func(k string, raw interface{}) bool { idx := strings.LastIndex(k, " ") if idx == -1 { // If there is no space, just ignore top level commands return false } // Trim up to that space so we can get the expected parent k = k[:idx] if _, ok := c.commandTree.Get(k); ok { // Yay we have the parent! return false } // We're missing the parent, so let's insert this toInsert[k] = struct{}{} // Call the walk function recursively so we check this one too return walkFn(k, nil) } // Walk! c.commandTree.Walk(walkFn) // Insert any that we're missing for k := range toInsert { var f CommandFactory = func() (Command, error) { return &MockCommand{ HelpText: "This command is accessed by using one of the subcommands below.", RunResult: RunResultHelp, }, nil } c.commandTree.Insert(k, f) } } // Setup autocomplete if we have it enabled. We have to do this after // the command tree is setup so we can use the radix tree to easily find // all subcommands. if c.Autocomplete { c.initAutocomplete() } // Process the args c.processArgs() } func (c *CLI) initAutocomplete() { if c.AutocompleteInstall == "" { c.AutocompleteInstall = defaultAutocompleteInstall } if c.AutocompleteUninstall == "" { c.AutocompleteUninstall = defaultAutocompleteUninstall } if c.autocompleteInstaller == nil { c.autocompleteInstaller = &realAutocompleteInstaller{} } // We first set c.autocomplete to a noop autocompleter that outputs // to nul so that we can detect if we're autocompleting or not. If we're // not, then we do nothing. This saves a LOT of compute cycles since // initAutoCompleteSub has to walk every command. c.autocomplete = complete.New(c.Name, complete.Command{}) c.autocomplete.Out = ioutil.Discard if !c.autocomplete.Complete() { return } // Build the root command cmd := c.initAutocompleteSub("") // For the root, we add the global flags to the "Flags". This way // they don't show up on every command. if !c.AutocompleteNoDefaultFlags { cmd.Flags = map[string]complete.Predictor{ "-" + c.AutocompleteInstall: complete.PredictNothing, "-" + c.AutocompleteUninstall: complete.PredictNothing, "-help": complete.PredictNothing, "-version": complete.PredictNothing, } } cmd.GlobalFlags = c.AutocompleteGlobalFlags c.autocomplete = complete.New(c.Name, cmd) } // initAutocompleteSub creates the complete.Command for a subcommand with // the given prefix. This will continue recursively for all subcommands. // The prefix "" (empty string) can be used for the root command. func (c *CLI) initAutocompleteSub(prefix string) complete.Command { var cmd complete.Command walkFn := func(k string, raw interface{}) bool { // Ignore the empty key which can be present for default commands. if k == "" { return false } // Keep track of the full key so that we can nest further if necessary fullKey := k if len(prefix) > 0 { // If we have a prefix, trim the prefix + 1 (for the space) // Example: turns "sub one" to "one" with prefix "sub" k = k[len(prefix)+1:] } if idx := strings.Index(k, " "); idx >= 0 { // If there is a space, we trim up to the space. This turns // "sub sub2 sub3" into "sub". The prefix trim above will // trim our current depth properly. k = k[:idx] } if _, ok := cmd.Sub[k]; ok { // If we already tracked this subcommand then ignore return false } // If the command is hidden, don't record it at all if _, ok := c.commandHidden[fullKey]; ok { return false } if cmd.Sub == nil { cmd.Sub = complete.Commands(make(map[string]complete.Command)) } subCmd := c.initAutocompleteSub(fullKey) // Instantiate the command so that we can check if the command is // a CommandAutocomplete implementation. If there is an error // creating the command, we just ignore it since that will be caught // later. impl, err := raw.(CommandFactory)() if err != nil { impl = nil } // Check if it implements ComandAutocomplete. If so, setup the autocomplete if c, ok := impl.(CommandAutocomplete); ok { subCmd.Args = c.AutocompleteArgs() subCmd.Flags = c.AutocompleteFlags() } cmd.Sub[k] = subCmd return false } walkPrefix := prefix if walkPrefix != "" { walkPrefix += " " } c.commandTree.WalkPrefix(walkPrefix, walkFn) return cmd } func (c *CLI) commandHelp(out io.Writer, command Command) { // Get the template to use tpl := strings.TrimSpace(defaultHelpTemplate) if t, ok := command.(CommandHelpTemplate); ok { tpl = t.HelpTemplate() } if !strings.HasSuffix(tpl, "\n") { tpl += "\n" } // Parse it t, err := template.New("root").Funcs(sprig.TxtFuncMap()).Parse(tpl) if err != nil { t = template.Must(template.New("root").Parse(fmt.Sprintf( "Internal error! Failed to parse command help template: %s\n", err))) } // Template data data := map[string]interface{}{ "Name": c.Name, "SubcommandName": c.Subcommand(), "Help": command.Help(), } // Build subcommand list if we have it var subcommandsTpl []map[string]interface{} if c.commandNested { // Get the matching keys subcommands := c.helpCommands(c.Subcommand()) keys := make([]string, 0, len(subcommands)) for k := range subcommands { keys = append(keys, k) } // Sort the keys sort.Strings(keys) // Figure out the padding length var longest int for _, k := range keys { if v := len(k); v > longest { longest = v } } // Go through and create their structures subcommandsTpl = make([]map[string]interface{}, 0, len(subcommands)) for _, k := range keys { // Get the command raw, ok := subcommands[k] if !ok { c.ErrorWriter.Write([]byte(fmt.Sprintf( "Error getting subcommand %q", k))) } sub, err := raw() if err != nil { c.ErrorWriter.Write([]byte(fmt.Sprintf( "Error instantiating %q: %s", k, err))) } // Find the last space and make sure we only include that last part name := k if idx := strings.LastIndex(k, " "); idx > -1 { name = name[idx+1:] } subcommandsTpl = append(subcommandsTpl, map[string]interface{}{ "Name": name, "NameAligned": name + strings.Repeat(" ", longest-len(k)), "Help": sub.Help(), "Synopsis": sub.Synopsis(), }) } } data["Subcommands"] = subcommandsTpl // Write err = t.Execute(out, data) if err == nil { return } // An error, just output... c.ErrorWriter.Write([]byte(fmt.Sprintf( "Internal error rendering help: %s", err))) } // helpCommands returns the subcommands for the HelpFunc argument. // This will only contain immediate subcommands. func (c *CLI) helpCommands(prefix string) map[string]CommandFactory { // If our prefix isn't empty, make sure it ends in ' ' if prefix != "" && prefix[len(prefix)-1] != ' ' { prefix += " " } // Get all the subkeys of this command var keys []string c.commandTree.WalkPrefix(prefix, func(k string, raw interface{}) bool { // Ignore any sub-sub keys, i.e. "foo bar baz" when we want "foo bar" if !strings.Contains(k[len(prefix):], " ") { keys = append(keys, k) } return false }) // For each of the keys return that in the map result := make(map[string]CommandFactory, len(keys)) for _, k := range keys { raw, ok := c.commandTree.Get(k) if !ok { // We just got it via WalkPrefix above, so we just panic panic("not found: " + k) } // If this is a hidden command, don't show it if _, ok := c.commandHidden[k]; ok { continue } result[k] = raw.(CommandFactory) } return result } func (c *CLI) processArgs() { for i, arg := range c.Args { if arg == "--" { break } // Check for help flags. if arg == "-h" || arg == "-help" || arg == "--help" { c.isHelp = true continue } // Check for autocomplete flags if c.Autocomplete { if arg == "-"+c.AutocompleteInstall || arg == "--"+c.AutocompleteInstall { c.isAutocompleteInstall = true continue } if arg == "-"+c.AutocompleteUninstall || arg == "--"+c.AutocompleteUninstall { c.isAutocompleteUninstall = true continue } } if c.subcommand == "" { // Check for version flags if not in a subcommand. if arg == "-v" || arg == "-version" || arg == "--version" { c.isVersion = true continue } if arg != "" && arg[0] == '-' { // Record the arg... c.topFlags = append(c.topFlags, arg) } } // If we didn't find a subcommand yet and this is the first non-flag // argument, then this is our subcommand. if c.subcommand == "" && arg != "" && arg[0] != '-' { c.subcommand = arg if c.commandNested { // If the command has a space in it, then it is invalid. // Set a blank command so that it fails. if strings.ContainsRune(arg, ' ') { c.subcommand = "" return } // Determine the argument we look to to end subcommands. // We look at all arguments until one has a space. This // disallows commands like: ./cli foo "bar baz". An argument // with a space is always an argument. j := 0 for k, v := range c.Args[i:] { if strings.ContainsRune(v, ' ') { break } j = i + k + 1 } // Nested CLI, the subcommand is actually the entire // arg list up to a flag that is still a valid subcommand. searchKey := strings.Join(c.Args[i:j], " ") k, _, ok := c.commandTree.LongestPrefix(searchKey) if ok { // k could be a prefix that doesn't contain the full // command such as "foo" instead of "foobar", so we // need to verify that we have an entire key. To do that, // we look for an ending in a space or an end of string. reVerify := regexp.MustCompile(regexp.QuoteMeta(k) + `( |$)`) if reVerify.MatchString(searchKey) { c.subcommand = k i += strings.Count(k, " ") } } } // The remaining args the subcommand arguments c.subcommandArgs = c.Args[i+1:] } } // If we never found a subcommand and support a default command, then // switch to using that. if c.subcommand == "" { if _, ok := c.Commands[""]; ok { args := c.topFlags args = append(args, c.subcommandArgs...) c.topFlags = nil c.subcommandArgs = args } } } // defaultAutocompleteInstall and defaultAutocompleteUninstall are the // default values for the autocomplete install and uninstall flags. const defaultAutocompleteInstall = "autocomplete-install" const defaultAutocompleteUninstall = "autocomplete-uninstall" const defaultHelpTemplate = ` {{.Help}}{{if gt (len .Subcommands) 0}} Subcommands: {{- range $value := .Subcommands }} {{ $value.NameAligned }} {{ $value.Synopsis }}{{ end }} {{- end }} `