shell.go

package commandline

import (
	"context"
	"fmt"
	"k8s.io/apimachinery/pkg/util/sets"
	"os"
	"os/signal"
	"os/user"
	"strings"
	"syscall"

	"github.com/c-bata/go-prompt"
	"github.com/c-bata/go-prompt/completer"
	"github.com/pkg/errors"
	"github.com/urfave/cli/v2"
)

const shellCommandName = "shell"

// GenerateShellCommand generates the shell command with a list of commands that the shell should take.
// TODO: this needs a more comprehensive test suite.
// TODO: support ctrl+c.
func GenerateShellCommand(shellCommands []*cli.Command) *cli.Command {
	// explicitly exclude shell if included
	capturedCommands := pruneShellCommands(shellCommands)

	// make sure tty is open, this will not be the case in distroless containers
	_, err := syscall.Open("/dev/tty", syscall.O_RDONLY, 0)
	shellAvailable := err == nil

	return &cli.Command{
		Name:  shellCommandName,
		Usage: "start an interactive shell.",
		Flags: []cli.Flag{
			&LogLevel,
		},
		Action: func(c *cli.Context) (err error) {
			SetLogLevel(c)

			console := cli.NewApp()
			console.Commands = capturedCommands
			console.Action = func(c *cli.Context) error {
				fmt.Printf("Command not found. Type 'help' for a list of commands or \"%s\", \"%s\" or \"%s\" to exit.\n", quitCommand, exitCommand, quitCommandShort)
				return nil
			}

			if c.Args().Len() == 0 {
				err := console.RunContext(c.Context, strings.Fields("cmd help"))
				if err != nil {
					return errors.Wrap(err, "could not show help")
				}
			}

			sigs := make(chan os.Signal)
			ctx, cancel := context.WithCancel(c.Context)
			go func() {
				for sig := range sigs {
					if sig == syscall.SIGINT || sig == syscall.SIGTERM {
						cancel()
						fmt.Printf("\n(type \"%s\", \"%s\" or \"%s\" to exit)\n\n >", quitCommand, exitCommand, quitCommandShort)
					}
				}
			}()
			signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
			defer func() {
				signal.Stop(sigs)
				close(sigs)
			}()

			if !shellAvailable {
				fmt.Println("Shell is not available in this environment because /dev/tty is not available. This is expected in containerized images")
				return nil
			}

			interactive := newInteractiveClient(ctx, capturedCommands, console)
			for {
				p := prompt.New(
					interactive.executor,
					interactive.completor,
					prompt.OptionCompletionWordSeparator(completer.FilePathCompletionSeparator),
					prompt.OptionMaxSuggestion(3),
					prompt.OptionLivePrefix(livePrefix),
				)
				p.Run()
			}
		},
	}
}

// pruneShellCommands gets a list of commands including the shell command.
func pruneShellCommands(commands []*cli.Command) (prunedCommands []*cli.Command) {
	// initialize shell commands
	nameSet := sets.NewString()
	for _, command := range commands {
		if command.Name != shellCommandName {
			prunedCommands = append(prunedCommands, command)
		}
		if !nameSet.Has(command.Name) {
			fmt.Printf("Command %s already exists, skipping\n", command.Name)
		}

		nameSet.Insert(command.Name)
	}

	return prunedCommands
}

// ShellCmd is used to launch an interactive shell. This is useful for repeatedly interacting with the cli or testing.

// livePrefix generates a prefix with the current user directory.
// this is useful for file relative commands like creating a new geth account.
func livePrefix() (prefix string, useLivePrefix bool) {
	pwd, err := os.Getwd()
	useLivePrefix = true
	if err != nil {
		prefix += " $ "
		return
	}

	prefix += " " + pwd + " $ "

	u, err := user.Current()
	if err != nil {
		return
	}

	prefix = strings.ReplaceAll(prefix, u.HomeDir, "~")
	return
}

// interactiveClient object.
type interactiveClient struct {
	// cli app
	app *cli.App
	// ctx is the parent context
	//nolint: containedctx
	ctx context.Context
	// shellCommands are all shell commands supported by the interactive client
	shellCommands []*cli.Command
}

// newInteractiveClient creates a new interactive client.
func newInteractiveClient(ctx context.Context, shellCommands []*cli.Command, app *cli.App) *interactiveClient {
	return &interactiveClient{
		app:           app,
		shellCommands: shellCommands,
		ctx:           ctx,
	}
}

// completor handles autocompletion for the interactive client.
func (i *interactiveClient) completor(in prompt.Document) []prompt.Suggest {
	// commandPrompts are prompts for commands (no flags)
	commandPrompts := []prompt.Suggest{
		{
			Text:        "help",
			Description: "Shows a list of commands or help for one command",
		},
	}

	// flagPrompts promp flags for each command
	var flagPrompts []prompt.Suggest

	for _, command := range i.shellCommands {
		commandPrompts = append(commandPrompts, prompt.Suggest{
			Text:        command.Name,
			Description: command.Usage,
		})

		for _, flag := range command.Flags {
			docFlag, ok := flag.(cli.DocGenerationFlag)
			if ok {
				flagPrompts = append(flagPrompts, prompt.Suggest{
					Text:        fmt.Sprintf("%s --%s ", command.Name, longestFlag(flag)),
					Description: docFlag.GetUsage(),
				})
			}
		}
	}

	// get the prompts for the current text. if the user has entered part of a command this will be
	// command prompts. if the user has entered a whole command flags will be suggested
	prompts := prompt.FilterHasPrefix(commandPrompts, in.CurrentLineBeforeCursor(), true)
	if len(prompts) == 0 {
		prompts = prompt.FilterHasPrefix(flagPrompts, in.CurrentLineBeforeCursor(), true)

		// right now, autocomplete for the flags will not take into account the first word so
		// autocompleting `start --config` will return `start start --config`. Here, since we've
		// already done the matching we cut off hte command prefix
		for i, currentPrompt := range prompts {
			splitText := strings.Split(currentPrompt.Text, " ")
			currentPrompt.Text = strings.Join(splitText[1:], " ")
			prompts[i] = currentPrompt
		}
	}

	return prompts
}

// longestFlag gets the longest flag from all flag names
// this is used to avoid short (non-descriptive) flags from coming up in the autocomplete.
func longestFlag(flag cli.Flag) (flagName string) {
	maxLength := 0
	for _, name := range flag.Names() {
		flagLength := len(name)
		if flagLength > maxLength {
			maxLength = flagLength
			flagName = name
		}
	}
	return flagName
}

// executor handles executing interactive commands.
func (i *interactiveClient) executor(line string) {
	if line == "" {
		return
	}

	if line == quitCommand || line == quitCommandShort || line == exitCommand {
		os.Exit(0)
	}

	if line == "shell" {
		fmt.Println("cannot start a shell from within a shell!")
		return
	}

	err := i.app.RunContext(i.ctx, strings.Fields("cmd "+line))
	if err != nil {
		fmt.Println("error: ", err)
	}
}

const (
	quitCommand      = "quit"
	quitCommandShort = "q"
	exitCommand      = "exit"
)