jarl/cmd/jarl/complete.go

package main

// Shell completion. Cobra completes commands and flags by itself
// ('jarl completion bash'); this adds the values: server and group names
// from the configuration file, and zone, view, ACL, TSIG key, user and
// record type names from the server. Lookups on the server use the
// configured token and give up quickly, completing nothing, when the
// server can't be reached.

import (
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"jarl/internal/api"
	"jarl/internal/model"
	"jarl/internal/rr"
)

const completeTimeout = 3 * time.Second

// setupCompletion adds value completion to the command tree.
func setupCompletion(root *cobra.Command) {
	root.RegisterFlagCompletionFunc("server", completeServers(true))
	root.RegisterFlagCompletionFunc("output", cobra.FixedCompletions([]string{"table", "json", "yaml", "id", "field="}, cobra.ShellCompDirectiveNoFileComp|cobra.ShellCompDirectiveNoSpace))
	var walk func(c *cobra.Command)
	walk = func(c *cobra.Command) {
		for _, name := range []string{"view", "user", "zone"} {
			if c.LocalFlags().Lookup(name) != nil {
				c.RegisterFlagCompletionFunc(name, completeFromServer(name))
			}
		}
		if c.ValidArgsFunction == nil && c.Runnable() {
			if fn := argCompletion(c); fn != nil {
				c.ValidArgsFunction = fn
			}
		}
		for _, sub := range c.Commands() {
			walk(sub)
		}
	}
	walk(root)
}

// argCompletion completes a command's arguments from the placeholders in
// its usage line ("add ZONE NAME TYPE DATA...").
func argCompletion(c *cobra.Command) cobra.CompletionFunc {
	fields := strings.Fields(c.Use)
	if len(fields) < 2 {
		return nil
	}
	params := fields[1:]
	parent := ""
	if c.HasParent() {
		parent = c.Parent().Name()
	}
	return func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
		if len(args) >= len(params) && !strings.HasSuffix(params[len(params)-1], "...") {
			return nil, cobra.ShellCompDirectiveNoFileComp
		}
		p := params[min(len(args), len(params)-1)]
		p = strings.Trim(strings.TrimSuffix(p, "..."), "[]")
		if i := strings.Index(p, "|"); i >= 0 {
			p = p[i+1:] // "[ID|NAME]": complete names
		}
		switch {
		case p == "FILE":
			return nil, cobra.ShellCompDirectiveDefault
		case p == "ZONE":
			return lookup(cmd, "zone")
		case p == "VIEW":
			return lookup(cmd, "view")
		case p == "COLOR":
			return append([]string{"default"}, model.ViewColors...), cobra.ShellCompDirectiveNoFileComp
		case p == "ACL":
			return lookup(cmd, "acl")
		case p == "LIST":
			return lookup(cmd, "blocklist")
		case p == "USER":
			return lookup(cmd, "user")
		case p == "KEY" && parent == "tsig":
			return lookup(cmd, "tsig")
		case p == "TYPE":
			var out []cobra.Completion
			for _, t := range rr.Types() {
				out = append(out, cobra.CompletionWithDesc(t.Type, t.Description))
			}
			return out, cobra.ShellCompDirectiveNoFileComp
		case p == "NAME" && (parent == "server" || parent == "config" && cmd.Name() == "use"):
			return completeServers(false)(cmd, args, toComplete)
		case p == "NAME" && parent == "group" && len(args) == 0:
			return completeGroups(cmd)
		case p == "MEMBER":
			all, d := completeServers(true)(cmd, args, toComplete)
			var out []cobra.Completion
			for _, c := range all {
				name, _, _ := strings.Cut(c, "\t")
				if name != "all" && name != args[0] {
					out = append(out, c)
				}
			}
			return out, d
		}
		return nil, cobra.ShellCompDirectiveNoFileComp
	}
}

// completeServers completes server names (and with groups, group names
// and "all") from the configuration file.
func completeServers(groups bool) cobra.CompletionFunc {
	return func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
		cfg, err := loadCLIConfig()
		if err != nil {
			return nil, cobra.ShellCompDirectiveNoFileComp
		}
		// After a comma, complete the next item of a list.
		prefix := ""
		if i := strings.LastIndex(toComplete, ","); i >= 0 {
			prefix = toComplete[:i+1]
		}
		var out []cobra.Completion
		for _, name := range sortedKeys(cfg.Servers) {
			out = append(out, cobra.CompletionWithDesc(prefix+name, cfg.Servers[name].URL))
		}
		if groups {
			for _, name := range sortedKeys(cfg.Groups) {
				out = append(out, cobra.CompletionWithDesc(prefix+name, "group: "+strings.Join(cfg.Groups[name], ", ")))
			}
			if len(cfg.Servers) > 0 {
				out = append(out, cobra.CompletionWithDesc(prefix+"all", "every server"))
			}
			if _, err := os.Stat(socketPath()); err == nil {
				out = append(out, cobra.CompletionWithDesc(prefix+localName, "jarld on this machine (local socket)"))
			}
		}
		return out, cobra.ShellCompDirectiveNoFileComp
	}
}

func completeGroups(cmd *cobra.Command) ([]cobra.Completion, cobra.ShellCompDirective) {
	cfg, err := loadCLIConfig()
	if err != nil {
		return nil, cobra.ShellCompDirectiveNoFileComp
	}
	return sortedKeys(cfg.Groups), cobra.ShellCompDirectiveNoFileComp
}

func completeFromServer(kind string) cobra.CompletionFunc {
	return func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
		return lookup(cmd, kind)
	}
}

// lookup asks the server for names of a kind of object.
func lookup(cmd *cobra.Command, kind string) ([]cobra.Completion, cobra.ShellCompDirective) {
	none := cobra.ShellCompDirectiveNoFileComp
	c, err := newClient()
	if err != nil {
		return nil, none
	}
	c.HTTP.Timeout = completeTimeout
	seen := map[string]bool{}
	var out []cobra.Completion
	addName := func(name, desc string) {
		if name != "" && !seen[name] {
			seen[name] = true
			out = append(out, cobra.CompletionWithDesc(name, desc))
		}
	}
	switch kind {
	case "zone":
		var zones []api.Zone
		if c.Get(ctx(), "/zones", nil, &zones) != nil {
			return nil, none
		}
		// A zone in several views is one name; its description lists them.
		views := map[string][]string{}
		for _, z := range zones {
			name := dot(z.Name)
			if z.View != "" {
				views[name] = append(views[name], z.View)
			}
			addName(name, z.Type)
		}
		for i, c := range out {
			if vs := views[c[:strings.IndexByte(c, '\t')]]; len(vs) > 0 {
				label := "view "
				if len(vs) > 1 {
					label = "views "
				}
				out[i] = c + ", " + label + strings.Join(vs, ", ")
			}
		}
	case "view":
		var views []api.View
		if c.Get(ctx(), "/views", nil, &views) != nil {
			return nil, none
		}
		for _, v := range views {
			addName(v.Name, v.Comment)
		}
	case "blocklist":
		var lists []api.Blocklist
		if c.Get(ctx(), "/blocking/lists", nil, &lists) != nil {
			return nil, none
		}
		for _, b := range lists {
			addName(b.Name, b.URL)
		}
	case "acl":
		var acls []api.ACL
		if c.Get(ctx(), "/acls", nil, &acls) != nil {
			return nil, none
		}
		for _, a := range acls {
			addName(a.Name, a.Comment)
		}
	case "tsig":
		var keys []api.TSIGKey
		if c.Get(ctx(), "/tsig-keys", nil, &keys) != nil {
			return nil, none
		}
		for _, k := range keys {
			addName(k.Name, k.Algorithm)
		}
	case "user":
		var users []api.User
		if c.Get(ctx(), "/users", nil, &users) != nil {
			return nil, none
		}
		for _, u := range users {
			addName(u.Username, u.Comment)
		}
	}
	return out, none
}