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"
)

// 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"`
}

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

	if node != "" {
		filters = append(filters, filter{"=", "certname", node})
	}
	if factName != "" {
		filters = append(filters, filter{"=", "name", factName})
	} else if showRole {
		filters = append(filters, filter{"=", "name", roleFact})
	}

	if match != "" {
		filters = append(filters, filter{"=", "value", match})
	} else if partialMatch != "" {
		filters = append(filters, filter{"~", "value", partialMatch})
	}

	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 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
}

func isTerminal(f *os.File) bool {
	fi, err := f.Stat()
	if err != nil {
		return false
	}
	return (fi.Mode() & os.ModeCharDevice) != 0
}

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, partialMatch string, showRole, 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")
	}

	var collected []fact
	var stdinLines []string

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

	if nodeName == "" && !isTerminal(os.Stdin) {
		scanner := bufio.NewScanner(os.Stdin)
		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:
		hostFactMap := map[string]map[string]interface{}{}
		for _, f := range collected {
			if _, ok := hostFactMap[f.Certname]; !ok {
				hostFactMap[f.Certname] = map[string]interface{}{}
			}
			key := factName
			if key == "" {
				if showRole {
					key = cfg.RoleFact
				} else {
					key = "value"
				}
			}
			hostFactMap[f.Certname][key] = valueAny(f.Value)
		}
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		enc.SetEscapeHTML(false)
		enc.Encode(hostFactMap)

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

	case ansible:
		hosts := map[string]interface{}{}
		for _, line := range returnData {
			host := strings.Fields(line)[0]
			hosts[host] = map[string]interface{}{}
		}
		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
		partialMatch string
		nodeOnly     bool
		valueOnly    bool
		count        bool
		ansible      bool
		jsonMode     bool
		allFacts     bool
		puppetDBURL  string
	)

	rootCmd := &cobra.Command{
		Use:   appName,
		Short: "Query PuppetDB for nodes.",
		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, match, partialMatch, showRole, 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")
	f.BoolVarP(&showRole, "role", "R", false, "Show role fact ("+defaultRoleFact+" by default)")
	f.StringVarP(&match, "match", "m", "", "Exact value match")
	f.StringVar(&partialMatch, "pm", "", "Partial match on value")
	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)

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