Skip to content

Commit

Permalink
Add shell completion
Browse files Browse the repository at this point in the history
Signed-off-by: Vitor Savian <[email protected]>
  • Loading branch information
vitorsavian committed Jul 12, 2024
1 parent 3134a11 commit 71e0867
Show file tree
Hide file tree
Showing 4 changed files with 156 additions and 0 deletions.
1 change: 1 addition & 0 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ func main() {
cmds.NewCertCommand(),
cmds.NewSecretsEncryptCommand(),
cmds.NewTokenCommand(),
cmds.NewCompletionCommand(),
}

if err := app.Run(configfilearg.MustParse(os.Args)); err != nil {
Expand Down
1 change: 1 addition & 0 deletions pkg/cli/cmds/cmds_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ func Test_NewCommands(t *testing.T) {
NewCertCommand(),
NewSecretsEncryptCommand(),
NewTokenCommand(),
NewCompletionCommand(),
}

for _, command := range app.Commands {
Expand Down
153 changes: 153 additions & 0 deletions pkg/cli/cmds/completion.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
package cmds

import (
"fmt"
"os"

"github.com/k3s-io/k3s/pkg/cli/cmds"
"github.com/k3s-io/k3s/pkg/version"
"github.com/urfave/cli"
)

var (
completionFlags = []cli.Flag{
&cli.BoolFlag{
Name: "kubectl",
Usage: "(kubectl) export kubeconfig",
},
&cli.BoolFlag{
Name: "crictl",
Usage: "(crictl) export crictl config file",
},
}

k3sCompletionBase = mustCmdFromK3S(cmds.NewCompletionCommand(Run), K3SFlagSet{
"i": copyFlag,
})
)

func NewCompletionCommand() cli.Command {
cmd := k3sCompletionBase
cmd.Flags = append(cmd.Flags, completionFlags...)
return cmd
}

func isKubectlSet(kubectl bool) string {
if kubectl {
return " --kubectl"
}
return ""
}

func isCrictlSet(crictl bool) string {
if crictl {
return " --crictl"
}
return ""
}

func Run(ctx *cli.Context) error {
if ctx.NArg() < 1 {
return fmt.Errorf("must provide a valid SHELL argument")
}
shell := ctx.Args()[0]
completetionScript, err := genCompletionScript(shell, ctx.Bool("kubectl"), ctx.Bool("crictl"))
if err != nil {
return err
}
if ctx.Bool("i") {
return writeToRC(shell, ctx.Bool("kubectl"), ctx.Bool("crictl"))
}
fmt.Println(completetionScript)
return nil
}

func genCompletionScript(shell string, kubectl, crictl bool) (string, error) {
var completionScript string
if shell == "bash" {
completionScript = fmt.Sprintf(`#! /bin/bash
_cli_bash_autocomplete() {
if [[ "${COMP_WORDS[0]}" != "source" ]]; then
local cur opts base
COMPREPLY=()
cur="${COMP_WORDS[COMP_CWORD]}"
if [[ "$cur" == "-"* ]]; then
opts=$( ${COMP_WORDS[@]:0:$COMP_CWORD} ${cur} --generate-bash-completion )
else
opts=$( ${COMP_WORDS[@]:0:$COMP_CWORD} --generate-bash-completion )
fi
COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
return 0
fi
}
complete -o bashdefault -o default -o nospace -F _cli_bash_autocomplete %s
`, version.Program)
} else if shell == "zsh" {
completionScript = fmt.Sprintf(`#compdef %[1]s
_cli_zsh_autocomplete() {
local -a opts
local cur
cur=${words[-1]}
if [[ "$cur" == "-"* ]]; then
opts=("${(@f)$(_CLI_ZSH_AUTOCOMPLETE_HACK=1 ${words[@]:0:#words[@]-1} ${cur} --generate-bash-completion)}")
else
opts=("${(@f)$(_CLI_ZSH_AUTOCOMPLETE_HACK=1 ${words[@]:0:#words[@]-1} --generate-bash-completion)}")
fi
if [[ "${opts[1]}" != "" ]]; then
_describe 'values' opts
else
_files
fi
return
}
compdef _cli_zsh_autocomplete %[1]s
`, version.Program)
} else {
return "", fmt.Errorf("unknown shell: %s", shell)
}

if kubectl {
completionScript = fmt.Sprintf(`%s
export KUBECONFIG=/etc/rancher/rke2/rke2.yaml
`, completionScript)
}

if crictl {
completionScript = fmt.Sprintf(`%s
export CRI_CONFIG_FILE=/var/lib/rancher/rke2/agent/etc/crictl.yaml
`, completionScript)
}

return completionScript, nil
}

func writeToRC(shell string, kubectl, crictl bool) error {
rcFileName := ""
if shell == "bash" {
rcFileName = "/.bashrc"
} else if shell == "zsh" {
rcFileName = "/.zshrc"
}

home, err := os.UserHomeDir()
if err != nil {
return nil
}
rcFileName = home + rcFileName
f, err := os.OpenFile(rcFileName, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
return err
}
defer f.Close()
bashEntry := fmt.Sprintf("# >> %[1]s command completion (start)\n. <(%[1]s completion %[2]s%[3]s%[4]s)\n# >> %[1]s command completion (end)", version.Program, shell, isKubectlSet(kubectl), isCrictlSet(crictl))
if _, err := f.WriteString(bashEntry); err != nil {
return err
}
fmt.Printf("Autocomplete for %s added to: %s\n", shell, rcFileName)
return nil
}
1 change: 1 addition & 0 deletions pkg/cli/cmds/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,7 @@ func validateCloudProviderName(clx *cli.Context, role CLIRole) {
func NewApp() *cli.App {
app := cli.NewApp()
app.Name = appName
app.EnableBashCompletion = true
app.Usage = "Rancher Kubernetes Engine 2"
app.Version = fmt.Sprintf("%s (%s)", version.Version, version.GitCommit)
cli.VersionPrinter = func(c *cli.Context) {
Expand Down

0 comments on commit 71e0867

Please sign in to comment.