spry/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/peterbourgon/mergemap"

	sprig "github.com/Masterminds/sprig/v3"
	yaml "gopkg.in/yaml.v3"
)

type StrSet map[string]struct{}

func (s *StrSet) Add(str string) {
	(*s)[str] = struct{}{}
}

func (s *StrSet) Has(str string) bool {
	if _, ok := (*s)[str]; ok {
		return true
	}
	return false
}

func (s *StrSet) AddHas(str string) bool {
	if ok := s.Has(str); ok {
		return ok
	}
	s.Add(str)
	return false
}

type valFileSlice []string

func (v *valFileSlice) String() string {
	return fmt.Sprintf("%v", *v)
}

func (v *valFileSlice) Set(value string) error {
	*v = append(*v, value)
	return nil
}

type valEnt struct {
	address []string
	v       string
}

type valSlice []valEnt

func (v *valSlice) String() string {
	return fmt.Sprintf("%v", *v)
}

func (v *valSlice) Set(value string) error {
	pieces := strings.Split(value, "=")
	if len(pieces) < 2 {
		return fmt.Errorf("Invalid value setting.")
	}
	*v = append(*v, valEnt{
		strings.Split(pieces[0], "."),
		pieces[1],
	})
	return nil
}

var flags struct {
	valuesFiles     valFileSlice
	values          valSlice
	outputDelimiter string
	helmCompat      bool
}

func init() {
	flag.Var(&flags.values, "set", "Set a specific value (foo.bar=spam)")
	flag.Var(&flags.valuesFiles, "f", "Specify a values file (JSON or YAML)")
	flag.StringVar(
		&flags.outputDelimiter,
		"d", "---",
		"Output delimiter between template files rendered",
	)
	flag.BoolVar(
		&flags.helmCompat,
		"helm", false,
		"Helm compatibility: provided values will be nested under .Values",
	)
	flag.Parse()

	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s [Options] [Template Files]:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Fprintf(flag.CommandLine.Output(), "\n")
		fmt.Fprintf(
			flag.CommandLine.Output(),
			"Template Files can also be directories to recurse for templates.\n",
		)
		fmt.Fprintf(flag.CommandLine.Output(), "\n")
		fmt.Fprintf(
			flag.CommandLine.Output(),
			"If no template files are provided it will attempt read from STDIN.\n",
		)
		fmt.Fprintf(
			flag.CommandLine.Output(),
			"If no input is available from STDIN it will print this usage message instead.\n",
		)
		fmt.Fprintf(
			flag.CommandLine.Output(),
			"Pass only '-' to Template Files to force waiting for input.\n",
		)
		fmt.Fprintf(flag.CommandLine.Output(), "\n")
	}
}

func checkErr(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %+v", err)
		os.Exit(1)
	}
}

func expandFiles(files []string) (expFiles []string) {
	expFiles = []string{}
	fileSet := StrSet{}
	directories := []string{}

	// Ensure we don't process any provided values files as templates.
	for _, path := range flags.valuesFiles {
		abs, err := filepath.Abs(path)
		if err != nil {
			fileSet.Add(path)
			break
		}
		fileSet.Add(abs)
	}

	for _, path := range files {
		fileInfo, err := os.Stat(path)
		checkErr(err)

		if fileInfo.IsDir() {
			// When provided path is a directory, defer loading files from it.
			directories = append(directories, path)
		} else {
			abs, err := filepath.Abs(path)
			if err != nil {
				// On failure to find the absolute path, just use the provided path.
				abs = path
			}
			// Deduplicate files.
			if ok := fileSet.AddHas(abs); !ok {
				expFiles = append(expFiles, path)
			}
		}
		// Walk each provided directory to find templates.
		for _, path := range directories {
			filepath.WalkDir(path, func(path string, d fs.DirEntry, err error) error {
				if err != nil {
					fmt.Fprintf(os.Stderr, "WARN: %+v", err)
					return nil
				}
				if d.IsDir() {
					// Don't recurse hidden directories.
					if d.Name() != "." && strings.HasPrefix(d.Name(), ".") {
						return filepath.SkipDir
					}
					return nil
				}
				// Skip hidden files.
				if strings.HasPrefix(d.Name(), ".") {
					return nil
				}
				abs, err := filepath.Abs(path)
				if err != nil {
					// On failure to find the absolute path, just use the provided path.
					abs = path
				}
				// Deduplicate files.
				if ok := fileSet.AddHas(abs); !ok {
					expFiles = append(expFiles, path)
				}
				return nil
			})
		}
	}
	return
}

func mergeValues() (res map[string]interface{}) {
	res = map[string]interface{}{}
	result := res
	if flags.helmCompat {
		res["Values"] = map[string]interface{}{}
		result = res["Values"].(map[string]interface{})
	}
	for _, vf := range flags.valuesFiles {
		vfContent, err := os.ReadFile(vf)
		checkErr(err)
		if strings.HasSuffix(vf, ".json") {
			err = json.Unmarshal(vfContent, &result)
			checkErr(err)
		} else {
			err = yaml.Unmarshal(vfContent, result)
			checkErr(err)
		}
	}
	for _, v := range flags.values {
		valueMap := map[string]interface{}{}
		subMap := valueMap
		for i, node := range v.address {
			if i == len(v.address)-1 {
				subMap[node] = v.v
				break
			}
			subMap[node] = map[string]interface{}{}
			subMap = subMap[node].(map[string]interface{})
		}
		mergemap.Merge(result, valueMap)
	}
	return
}

func main() {
	var err error
	tpl := template.New("base").Funcs(sprig.FuncMap())

	values := mergeValues()
	// fmt.Fprintf(os.Stderr, "%+v\n", values)

	if len(flag.Args()) == 0 {
		// With no provided arguments read from STDIN if anything is immediately available through it.
		// otherwise print the usage and exit.
		ch := make(chan struct{})
		go func() {
			select {
			case <-ch:
				return
			case <-time.After(100 * time.Millisecond):
				flag.Usage()
				os.Exit(2)
			}
		}()
		var data []byte
		data, err = io.ReadAll(os.Stdin)
		// Tell the waiting routine that we got input on STDIN.
		ch <- struct{}{}
		checkErr(err)
		tpl, err = tpl.Parse(string(data))
		checkErr(err)
	} else if len(flag.Args()) == 1 && flag.Args()[0] == "-" {
		// If only a single non-flag argument is provided and it's a -, wait for STDIN indefinitely.
		var data []byte
		data, err = io.ReadAll(os.Stdin)
		checkErr(err)
		tpl, err = tpl.Parse(string(data))
		checkErr(err)
	} else {
		tplFiles := expandFiles(flag.Args())
		// Otherwise use the provided tplFiles/directories for templates to render.
		// fmt.Fprintf(os.Stderr, "%+v\n", tplFiles)
		tpl, err = tpl.ParseFiles(tplFiles...)
		checkErr(err)
		for i, file := range tplFiles {
			if i > 0 {
				fmt.Println(flags.outputDelimiter)
			}
			checkErr(tpl.Funcs(sprig.FuncMap()).ExecuteTemplate(os.Stdout, filepath.Base(file), values))
		}
		return
	}
	checkErr(tpl.Execute(os.Stdout, values))
}