node-lookup/main.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"sort"
	"strings"
	"syscall"

	"github.com/spf13/cobra"
	"gopkg.in/yaml.v3"
)

const (
	defaultPuppetDBURL = "http://puppetdbapi.service.consul:8080/pdb/query/v4/facts"
	defaultRoleFact    = "enc_role"
	configFileName     = "config.yaml"
	appName            = "node-lookup"
)

var version = "dev"

// config holds all configurable values. Fields map 1:1 to config file keys,
// env vars (NODE_LOOKUP_*), and (where applicable) CLI flags.
type config struct {
	PuppetDBURL string `yaml:"puppetdb_url"`
	RoleFact    string `yaml:"role_fact"`
}

func defaultConfig() config {
	return config{
		PuppetDBURL: defaultPuppetDBURL,
		RoleFact:    defaultRoleFact,
	}
}

// configDir returns the XDG_CONFIG_HOME/node-lookup directory.
func configDir() string {
	base := os.Getenv("XDG_CONFIG_HOME")
	if base == "" {
		home, _ := os.UserHomeDir()
		base = filepath.Join(home, ".config")
	}
	return filepath.Join(base, appName)
}

// configPath returns the full path to the config file.
func configPath() string {
	return filepath.Join(configDir(), configFileName)
}

// loadConfig reads the config file (if present), then applies env var overrides.
// Precedence (lowest → highest): defaults < config file < env vars < CLI flags.
func loadConfig() (config, error) {
	cfg := defaultConfig()

	path := configPath()
	data, err := os.ReadFile(path)
	if err != nil && !os.IsNotExist(err) {
		return cfg, fmt.Errorf("reading config %s: %w", path, err)
	}
	if err == nil {
		if err := yaml.Unmarshal(data, &cfg); err != nil {
			return cfg, fmt.Errorf("parsing config %s: %w", path, err)
		}
	}

	if v := os.Getenv("NODE_LOOKUP_URL"); v != "" {
		cfg.PuppetDBURL = v
	}
	if v := os.Getenv("NODE_LOOKUP_ROLE_FACT"); v != "" {
		cfg.RoleFact = v
	}

	return cfg, nil
}

// writeDefaultConfig creates the config directory and writes a default config file.
func writeDefaultConfig() error {
	dir := configDir()
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return fmt.Errorf("creating config dir: %w", err)
	}

	path := configPath()
	if _, err := os.Stat(path); err == nil {
		return fmt.Errorf("config already exists at %s", path)
	}

	cfg := defaultConfig()
	data, _ := yaml.Marshal(cfg)
	header := []byte("# node-lookup configuration\n# Fields can be overridden with env vars: NODE_LOOKUP_URL, NODE_LOOKUP_ROLE_FACT\n\n")
	if err := os.WriteFile(path, append(header, data...), 0o644); err != nil {
		return fmt.Errorf("writing config: %w", err)
	}
	fmt.Println("Config written to", path)
	return nil
}

type fact struct {
	Certname string          `json:"certname"`
	Name     string          `json:"name"`
	Value    json.RawMessage `json:"value"`
}

// splitFactNames splits a comma-separated -F value into trimmed, non-empty
// names, so `-F ipaddress,enc_role` queries both facts.
func splitFactNames(factName string) []string {
	var names []string
	for _, n := range strings.Split(factName, ",") {
		if n = strings.TrimSpace(n); n != "" {
			names = append(names, n)
		}
	}
	return names
}

// nameFilter returns a PQL filter matching any of the given fact names: a plain
// equality for one name, an "or" over per-name equalities for several.
func nameFilter(names []string) []interface{} {
	if len(names) == 1 {
		return []interface{}{"=", "name", names[0]}
	}
	or := []interface{}{"or"}
	for _, n := range names {
		or = append(or, []interface{}{"=", "name", n})
	}
	return or
}

func buildQuery(node, factName, match, roleFact string, showRole, partial, inverse bool) string {
	type filter = []interface{}
	var filters []filter

	if node != "" {
		filters = append(filters, filter{"=", "certname", node})
	}
	if names := splitFactNames(factName); len(names) > 0 {
		filters = append(filters, nameFilter(names))
	} else if showRole {
		filters = append(filters, filter{"=", "name", roleFact})
	}

	if match != "" {
		op := "="
		if partial {
			op = "~"
		}
		inner := filter{op, "value", match}
		if inverse {
			filters = append(filters, filter{"not", inner})
		} else {
			filters = append(filters, inner)
		}
	}

	if len(filters) == 0 {
		b, _ := json.Marshal([]interface{}{"=", "name", roleFact})
		return string(b)
	}

	combined := make([]interface{}, 0, len(filters)+1)
	combined = append(combined, "and")
	for _, f := range filters {
		combined = append(combined, f)
	}
	b, _ := json.Marshal(combined)
	return string(b)
}

func queryPuppetDB(puppetDBURL, query string) ([]fact, error) {
	params := url.Values{}
	params.Set("query", query)
	resp, err := http.Get(puppetDBURL + "?" + params.Encode())
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	var facts []fact
	if err := json.NewDecoder(resp.Body).Decode(&facts); err != nil {
		return nil, fmt.Errorf("decode error: %w", err)
	}
	return facts, nil
}

// valueString returns the fact value as a plain string (unquotes JSON strings,
// leaves numbers/bools/objects/arrays as compact JSON).
func valueString(raw json.RawMessage) string {
	var s string
	if err := json.Unmarshal(raw, &s); err == nil {
		return s
	}
	return string(raw)
}

func valueAny(raw json.RawMessage) interface{} {
	var v interface{}
	_ = json.Unmarshal(raw, &v)
	return v
}

func processResults(facts []fact) []string {
	out := make([]string, 0, len(facts))
	for _, f := range facts {
		out = append(out, f.Certname+" "+valueString(f.Value))
	}
	sort.Strings(out)
	return out
}

func countResults(lines []string) []string {
	counts := map[string]int{}
	for _, line := range lines {
		parts := strings.SplitN(line, " ", 2)
		val := line
		if len(parts) == 2 {
			val = parts[1]
		}
		counts[val]++
	}
	out := make([]string, 0, len(counts))
	for val, n := range counts {
		out = append(out, fmt.Sprintf("%d: %s", n, val))
	}
	sort.Strings(out)
	return out
}

// stdinReader returns a buffered reader over f and true only when f actually
// carries piped/redirected data to consume as node names. Terminals and
// character devices such as /dev/null return false, and an empty pipe or empty
// file (immediate EOF on peek) also returns false. This means running without a
// TTY — e.g. invoked by an agent or CI where stdin is /dev/null or a closed
// pipe — falls through to a normal query instead of silently consuming empty
// input and printing nothing.
func stdinReader(f *os.File) (*bufio.Reader, bool) {
	fi, err := f.Stat()
	if err != nil {
		return nil, false
	}
	if (fi.Mode() & os.ModeCharDevice) != 0 {
		return nil, false // terminal or /dev/null
	}
	r := bufio.NewReader(f)
	if _, err := r.Peek(1); err != nil {
		return nil, false // empty pipe / empty file (EOF)
	}
	return r, true
}

// matchValue resolves the value to match against. The -m/--match flag wins; if
// it is empty, the (optional) positional argument is used instead. The
// positional fallback exists so combined shorthands like `-pm k8s` work — pflag
// leaves the space-separated `k8s` as a positional rather than attaching it to
// the grouped -m flag.
func matchValue(flagMatch string, args []string) string {
	if flagMatch != "" {
		return flagMatch
	}
	if len(args) > 0 {
		return args[0]
	}
	return ""
}

// factsByHost groups collected facts into {certname: {factname: value}}. Each
// value is keyed by the fact's real name, so multiple -F facts each appear
// under the host; it falls back to the -F string, the role fact (with -R), or
// "value" only when a result carries no name. Shared by the -j and -A outputs.
func factsByHost(collected []fact, factName, roleFact string, showRole bool) map[string]map[string]interface{} {
	out := map[string]map[string]interface{}{}
	for _, f := range collected {
		if _, ok := out[f.Certname]; !ok {
			out[f.Certname] = map[string]interface{}{}
		}
		key := f.Name
		if key == "" {
			if key = factName; key == "" {
				if showRole {
					key = roleFact
				} else {
					key = "value"
				}
			}
		}
		out[f.Certname][key] = valueAny(f.Value)
	}
	return out
}

func allFactsForNode(puppetDBURL, node string) ([]fact, error) {
	query, _ := json.Marshal([]interface{}{"=", "certname", node})
	return queryPuppetDB(puppetDBURL, string(query))
}

func run(cfg config, nodeName, factName, match string, showRole, partial, inverse, nodeOnly, valueOnly, count, ansible, jsonMode, allFacts bool) error {
	signal.Ignore(syscall.SIGPIPE)

	if allFacts {
		if nodeName == "" {
			return fmt.Errorf("-a requires -n")
		}
		facts, err := allFactsForNode(cfg.PuppetDBURL, nodeName)
		if err != nil {
			return err
		}
		sort.Slice(facts, func(i, j int) bool { return facts[i].Name < facts[j].Name })
		for _, f := range facts {
			fmt.Printf("%-40s %s\n", f.Name, valueString(f.Value))
		}
		return nil
	}

	if (nodeOnly || valueOnly || count || ansible) && !showRole && factName == "" {
		return fmt.Errorf("-R or -F must be used with -1, -2, -C, or -A")
	}
	if (match != "" || partial || inverse) && !showRole && factName == "" {
		return fmt.Errorf("-R or -F must be used with -m, -p, or -i")
	}

	var collected []fact
	var stdinLines []string

	doQuery := func(node string) error {
		query := buildQuery(node, factName, match, cfg.RoleFact, showRole, partial, inverse)
		facts, err := queryPuppetDB(cfg.PuppetDBURL, query)
		if err != nil {
			return err
		}
		collected = append(collected, facts...)
		return nil
	}

	if reader, ok := stdinReader(os.Stdin); ok && nodeName == "" {
		scanner := bufio.NewScanner(reader)
		if count {
			for scanner.Scan() {
				stdinLines = append(stdinLines, scanner.Text())
			}
		} else {
			for scanner.Scan() {
				fields := strings.Fields(scanner.Text())
				if len(fields) == 0 {
					continue
				}
				if err := doQuery(fields[0]); err != nil {
					fmt.Fprintln(os.Stderr, "error:", err)
				}
			}
		}
	} else {
		if err := doQuery(nodeName); err != nil {
			return err
		}
	}

	returnData := processResults(collected)

	switch {
	case jsonMode:
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		enc.SetEscapeHTML(false)
		_ = enc.Encode(factsByHost(collected, factName, cfg.RoleFact, showRole))

	case count:
		values := stdinLines
		if len(values) == 0 {
			values = returnData
		}
		fmt.Println(strings.Join(countResults(values), "\n"))

	case ansible:
		// Attach each host's queried fact(s) as inventory host vars, e.g.
		// `-F ipaddress,enc_role -A` yields hosts with ipaddress + enc_role set.
		hosts := map[string]interface{}{}
		for host, vars := range factsByHost(collected, factName, cfg.RoleFact, showRole) {
			hosts[host] = vars
		}
		inventory := map[string]interface{}{
			"all": map[string]interface{}{"hosts": hosts},
		}
		b, _ := yaml.Marshal(inventory)
		_, _ = os.Stdout.Write(b)

	case nodeOnly:
		for _, line := range returnData {
			fmt.Println(strings.Fields(line)[0])
		}

	case valueOnly:
		for _, line := range returnData {
			parts := strings.SplitN(line, " ", 2)
			if len(parts) == 2 {
				fmt.Println(parts[1])
			}
		}

	default:
		fmt.Println(strings.Join(returnData, "\n"))
	}

	return nil
}

func main() {
	cfg, err := loadConfig()
	if err != nil {
		fmt.Fprintln(os.Stderr, "config error:", err)
		os.Exit(1)
	}

	var (
		nodeName    string
		factName    string
		showRole    bool
		match       string
		partial     bool
		inverse     bool
		nodeOnly    bool
		valueOnly   bool
		count       bool
		ansible     bool
		jsonMode    bool
		allFacts    bool
		puppetDBURL string
	)

	rootCmd := &cobra.Command{
		Use:   appName + " [value]",
		Short: "Query PuppetDB for nodes.",
		// Accept an optional positional match value in addition to -m. This makes
		// combined shorthands like `-pm k8s` work: pflag does not attach a
		// space-separated value to a string flag grouped with a bool flag (only
		// `-pm=k8s` or `-p -m k8s` do), so `k8s` arrives here as a positional
		// argument instead. Falling back to it keeps the ergonomic form working.
		Args: cobra.MaximumNArgs(1),
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			if cmd.Flags().Changed("url") {
				cfg.PuppetDBURL = puppetDBURL
			}
			return nil
		},
		RunE: func(cmd *cobra.Command, args []string) error {
			return run(cfg, nodeName, factName, matchValue(match, args), showRole, partial, inverse, nodeOnly, valueOnly, count, ansible, jsonMode, allFacts)
		},
		SilenceUsage: true,
	}

	f := rootCmd.Flags()
	f.StringVarP(&nodeName, "node", "n", "", "Node name")
	f.StringVarP(&factName, "fact", "F", "", "Fact name (comma-separated for several, e.g. -F ipaddress,enc_role)")
	f.BoolVarP(&showRole, "role", "R", false, "Show role fact ("+defaultRoleFact+" by default)")
	f.StringVarP(&match, "match", "m", "", "Value to match (use with -p and/or -i)")
	f.BoolVarP(&partial, "partial", "p", false, "Partial/regex match modifier (combine with -m)")
	f.BoolVarP(&inverse, "inverse", "i", false, "Inverse match modifier (combine with -m)")
	f.BoolVarP(&nodeOnly, "nodeonly", "1", false, "Show only the node name")
	f.BoolVarP(&valueOnly, "valueonly", "2", false, "Show only the value")
	f.BoolVarP(&count, "count", "C", false, "Count fact occurrences")
	f.BoolVarP(&ansible, "ansible", "A", false, "Output as Ansible inventory")
	f.BoolVarP(&jsonMode, "json", "j", false, "Emit valid JSON for all output")
	f.BoolVarP(&allFacts, "all", "a", false, "Show all facts for a node (requires -n)")
	rootCmd.PersistentFlags().StringVar(&puppetDBURL, "url", cfg.PuppetDBURL, "PuppetDB facts URL (overrides config and NODE_LOOKUP_URL)")

	configCmd := &cobra.Command{
		Use:   "config",
		Short: "Manage configuration",
	}

	configInitCmd := &cobra.Command{
		Use:   "init",
		Short: "Write a default config file to " + configPath(),
		RunE: func(cmd *cobra.Command, args []string) error {
			return writeDefaultConfig()
		},
		SilenceUsage: true,
	}

	configShowCmd := &cobra.Command{
		Use:   "show",
		Short: "Print the active configuration",
		RunE: func(cmd *cobra.Command, args []string) error {
			fmt.Printf("config file : %s\n", configPath())
			fmt.Printf("puppetdb_url: %s\n", cfg.PuppetDBURL)
			fmt.Printf("role_fact   : %s\n", cfg.RoleFact)
			return nil
		},
		SilenceUsage: true,
	}

	configCmd.AddCommand(configInitCmd, configShowCmd)
	rootCmd.AddCommand(configCmd)

	versionCmd := &cobra.Command{
		Use:          "version",
		Short:        "Print the version",
		Run:          func(cmd *cobra.Command, args []string) { fmt.Println(version) },
		SilenceUsage: true,
	}
	rootCmd.AddCommand(versionCmd)

	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}