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"

	htl "gitlab.com/CRThaze/helm-tmpl-lang"
)

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
	helm                  bool
	helmStrict            bool
	helmDNSFunc           bool
	helmLint              bool
	helmNestValues        bool
	outputFilenamePrefix  string
	outputFilenameSTDERR  bool
	outputDelimiterSTDERR 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.helm,
		"helm", false,
		"Use Helm templating language (superset of Sprig).",
	)
	flag.BoolVar(
		&flags.helmStrict,
		"helmStrict", false,
		"When using Helm Tpl Lang, use strict rendering.",
	)
	flag.BoolVar(
		&flags.helmDNSFunc,
		"helmDNS", true,
		"When using Helm Tpl Lang, support DNS resolution.",
	)
	flag.BoolVar(
		&flags.helmLint,
		"helmLint", false,
		"When using Helm Tpl Lang, enable the Linting Mode.",
	)
	flag.BoolVar(
		&flags.helmNestValues,
		"helmNestValues", true,
		"When using Helm Tpl Lang, Nest provided values under '.Values'",
	)
	flag.StringVar(
		&flags.outputFilenamePrefix,
		"fp", "# ",
		"Prefix for the filename in the output.",
	)
	flag.BoolVar(
		&flags.outputDelimiterSTDERR,
		"delimErr", false,
		"Whether to print output delimiter to STDERR.",
	)
	flag.BoolVar(
		&flags.outputFilenameSTDERR,
		"filenameErr", true,
		"Whether to print filename to STDERR.",
	)
	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")
		fmt.Fprintf(
			flag.CommandLine.Output(),
			"Helm Mode: use the Helm Templating Language/Engine\n"+
				"with the following caveats:\n"+
				"\t- 'lookup' function unavailable.\n",
		)
		fmt.Fprintf(flag.CommandLine.Output(), "\n")
	}
}

func checkErr(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %+v\n", 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.helm && flags.helmNestValues {
		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
	var ht *htl.Templater
	lo := &htl.LangOpts{
		Strict:    flags.helmStrict,
		LintMode:  flags.helmLint,
		EnableDNS: flags.helmDNSFunc,
	}
	tpl := template.New("base").Funcs(sprig.FuncMap())

	values := mergeValues()

	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)

		if flags.helm {
			ht = htl.NewTemplaterWithBytes(data, nil)
			fmt.Fprintf(os.Stderr, "WARN: Helm Compatibility only partially implemented.\n")
			tpl, err = tpl.Parse(string(data))
		} else {
			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)

		if flags.helm {
			ht = htl.NewTemplaterWithBytes(data, lo)
		} else {
			tpl, err = tpl.Parse(string(data))
		}
		checkErr(err)
	} else {
		// Otherwise use the provided tplFiles/directories for templates to render.
		tplFiles := expandFiles(flag.Args())
		if flags.helm {
			ht, err = htl.NewTemplaterWithFiles(tplFiles, lo)
			checkErr(err)
			rendered, err := ht.Render(values)
			checkErr(err)
			first := true
			for filename, content := range rendered {
				if first {
					first = false
				} else {
					out := os.Stdout
					if flags.outputDelimiterSTDERR {
						out = os.Stderr
					}
					fmt.Fprintf(out, "%s\n", flags.outputDelimiter)

					out = os.Stdout
					if flags.outputFilenameSTDERR {
						out = os.Stderr
					}
					fmt.Fprintf(out, "%s%s\n", flags.outputFilenamePrefix, filename)
				}
				fmt.Fprint(os.Stdout, content)
			}
		} else {
			tpl, err = tpl.ParseFiles(tplFiles...)
			checkErr(err)
			for i, file := range tplFiles {
				if i > 0 {
					out := os.Stdout
					if flags.outputDelimiterSTDERR {
						out = os.Stderr
					}
					fmt.Fprintf(out, "%s\n", flags.outputDelimiter)

					out = os.Stdout
					if flags.outputFilenameSTDERR {
						out = os.Stderr
					}
					fmt.Fprintf(out, "%s%s\n", flags.outputFilenamePrefix, file)
				}
				checkErr(tpl.ExecuteTemplate(os.Stdout, filepath.Base(file), values))
			}
		}
		return
	}
	if ht != nil {
		rendered, err := ht.Render(values)
		checkErr(err)
		fmt.Fprintln(os.Stdout, rendered)
	} else {
		checkErr(tpl.Execute(os.Stdout, values))
	}
}