// Copyright © 2022 Roberto Hidalgo // SPDX-License-Identifier: Apache-2.0 package command import ( "bytes" "context" "fmt" "os" "strings" "text/template" "time" "git.rob.mx/nidito/chinampa/pkg/exec" "git.rob.mx/nidito/chinampa/pkg/render" "github.com/sirupsen/logrus" "github.com/spf13/cobra" "gopkg.in/yaml.v3" ) // ValueType represent the kinds of or option. type ValueType string const ( // ValueTypeDefault is the empty string, maps to ValueTypeString. ValueTypeDefault ValueType = "" // ValueTypeString a value treated like a string. ValueTypeString ValueType = "string" // ValueTypeBoolean is a value treated like a boolean. ValueTypeBoolean ValueType = "bool" // ValueTypeBoolean is a value treated like an integer. ValueTypeInt ValueType = "int" ) type SourceCommand struct { Path []string Args string } type CompletionFunc func(cmd *Command, currentValue string, config string) (values []string, flag cobra.ShellCompDirective, err error) // ValueSource represents the source for an auto-completed and/or validated option/argument. type ValueSource struct { // Directories prompts for directories with the given prefix. Directories *string `json:"dirs,omitempty" yaml:"dirs,omitempty" validate:"omitempty,excluded_with=Command Files Func Script Static"` // Files prompts for files with the given extensions Files *[]string `json:"files,omitempty" yaml:"files,omitempty" validate:"omitempty,excluded_with=Command Func Directories Script Static"` // Script runs the provided command with `bash -c "$script"` and returns an option for every line of stdout. Script string `json:"script,omitempty" yaml:"script,omitempty" validate:"omitempty,excluded_with=Command Directories Files Func Static"` // Static returns the given list. Static *[]string `json:"static,omitempty" yaml:"static,omitempty" validate:"omitempty,excluded_with=Command Directories Files Func Script"` // Command runs a subcommand and returns an option for every line of stdout. Command *SourceCommand `json:"command,omitempty" yaml:"command,omitempty" validate:"omitempty,excluded_with=Directories Files Func Script Static"` // Func runs a function Func CompletionFunc `json:"-" yaml:"-" validate:"omitempty,excluded_with=Command Directories Files Script Static"` // Timeout is the maximum amount of time we will wait for a Script, Command, or Func before giving up on completions/validations. Timeout int `json:"timeout,omitempty" yaml:"timeout,omitempty" validate:"omitempty,excluded_with=Directories Files Static"` // Suggestion if provided will only suggest autocomplete values but will not perform validation of a given value Suggestion bool `json:"suggest-only" yaml:"suggest-only" validate:"omitempty"` // nolint:tagliatelle // SuggestRaw if provided the shell will not add a space after autocompleting SuggestRaw bool `json:"suggest-raw" yaml:"suggest-raw" validate:"omitempty"` // nolint:tagliatelle command *Command `json:"-" yaml:"-" validate:"-"` computed *[]string flag cobra.ShellCompDirective custom string // The app-defined key's value } // Validates tells if a value needs to be validated. func (vs *ValueSource) Validates() bool { if vs.Directories != nil || vs.Files != nil { return false } return !vs.Suggestion } // Resolve returns the values for autocomplete and validation. func (vs *ValueSource) Resolve(currentValue string) (values []string, flag cobra.ShellCompDirective, err error) { if vs.computed != nil { return *vs.computed, vs.flag, nil } if vs.Timeout == 0 { vs.Timeout = 5 } flag = cobra.ShellCompDirectiveDefault timeout := time.Duration(vs.Timeout) logrus.Errorf("resolving: %s", vs) switch { case vs.Static != nil: values = *vs.Static case vs.Files != nil: flag = cobra.ShellCompDirectiveFilterFileExt values = *vs.Files case vs.Directories != nil: flag = cobra.ShellCompDirectiveFilterDirs values = []string{*vs.Directories} case vs.Func != nil: ctx, cancel := context.WithTimeout(context.Background(), timeout*time.Second) defer cancel() done := make(chan error, 1) panicChan := make(chan any, 1) go func() { defer func() { if p := recover(); p != nil { panicChan <- p } }() values, flag, err = vs.Func(vs.command, currentValue, vs.custom) done <- err }() select { case err = <-done: return case p := <-panicChan: panic(p) case <-ctx.Done(): flag = cobra.ShellCompDirectiveError err = ctx.Err() return } case vs.Command != nil: if vs.command == nil { return nil, cobra.ShellCompDirectiveError, fmt.Errorf("bug: command is nil") } argString, err := vs.command.ResolveTemplate(vs.Command.Args, currentValue) if err != nil { return nil, cobra.ShellCompDirectiveError, err } args := strings.Split(argString, " ") sub, _, err := vs.command.Cobra.Root().Find(vs.Command.Path) if err != nil { return nil, cobra.ShellCompDirectiveError, fmt.Errorf("could not find a command named %s", vs.Command.Path) } ctx, cancel := context.WithTimeout(context.Background(), timeout*time.Second) defer cancel() // The cancel should be deferred so resources are cleaned up sub.SetArgs(args) var stdout bytes.Buffer sub.SetOut(&stdout) var stderr bytes.Buffer sub.SetErr(&stderr) err = sub.ExecuteContext(ctx) if err != nil { return nil, cobra.ShellCompDirectiveError, err } values = strings.Split(stdout.String(), "\n") flag = cobra.ShellCompDirectiveDefault case vs.Script != "": if vs.command == nil { return nil, cobra.ShellCompDirectiveError, fmt.Errorf("bug: command is nil") } cmd, err := vs.command.ResolveTemplate(vs.Script, currentValue) if err != nil { return nil, cobra.ShellCompDirectiveError, err } args := append([]string{"/bin/bash", "-c"}, cmd) values, flag, err = exec.Exec(vs.command.FullName(), args, os.Environ(), timeout*time.Second) if err != nil { return nil, flag, err } default: return nil, flag, fmt.Errorf("Empty value source") } vs.computed = &values if vs.SuggestRaw { flag |= cobra.ShellCompDirectiveNoSpace } vs.flag = flag return values, flag, err } type AutocompleteTemplate struct { Args map[string]string Opts map[string]string } func (tpl *AutocompleteTemplate) Opt(name string) string { if val, ok := tpl.Opts[name]; ok { return fmt.Sprintf("--%s %s", name, val) } return "" } func (tpl *AutocompleteTemplate) Arg(name string) string { return tpl.Args[name] } func (cmd *Command) ResolveTemplate(templateString string, currentValue string) (string, error) { var buf bytes.Buffer tplData := &AutocompleteTemplate{ Args: cmd.Arguments.AllKnownStr(), Opts: cmd.Options.AllKnownStr(), } fnMap := template.FuncMap{ "Opt": tplData.Opt, "Arg": tplData.Arg, "Current": func() string { return currentValue }, } for k, v := range render.TemplateFuncs { fnMap[k] = v } tpl, err := template.New("subcommand").Funcs(fnMap).Parse(templateString) if err != nil { return "", err } err = tpl.Execute(&buf, tplData) if err != nil { return "", err } return buf.String(), nil } func (vs *ValueSource) UnmarshalYAML(node *yaml.Node) error { vs.Timeout = 0 vs.Suggestion = false vs.SuggestRaw = false vs.computed = nil intermediate := map[string]yaml.Node{} if err := node.Decode(&intermediate); err != nil { logrus.Errorf("could not decode valuesource: %s", err) return err } if t, ok := intermediate["timeout"]; ok { if err := t.Decode(&vs.Timeout); err != nil { logrus.Errorf("could not decode timeout: %s", err) return err } delete(intermediate, "timeout") } if t, ok := intermediate["suggest-only"]; ok { if err := t.Decode(&vs.Suggestion); err != nil { logrus.Errorf("could not decode suggest-only: %s", err) return err } delete(intermediate, "suggest-only") } if t, ok := intermediate["suggest-raw"]; ok { if err := t.Decode(&vs.SuggestRaw); err != nil { logrus.Errorf("could not decode suggest-raw: %s", err) return err } delete(intermediate, "suggest-raw") } for key, node := range intermediate { if cfn, ok := customCompleters[key]; ok { if err := node.Decode(&vs.custom); err != nil { logrus.Errorf("could not decode custom: %s", err) return err } vs.Func = cfn break } switch key { case "dirs": var res string if err := node.Decode(&res); err != nil { return err } vs.Directories = &res case "files": res := []string{} if err := node.Decode(&res); err != nil { return err } vs.Files = &res case "script": if err := node.Decode(&vs.Script); err != nil { return err } case "static": static := []string{} if err := node.Decode(&static); err != nil { logrus.Errorf("could not decode static: %s", err) return err } vs.Static = &static case "command": if err := node.Decode(&vs.Command); err != nil { return err } default: return fmt.Errorf("unknown value source key: %s", key) } } return nil } var customCompleters = map[string]CompletionFunc{} func RegisterValueSource(key string, completion CompletionFunc) { customCompleters[key] = completion }