github.com/Racer159/jackal@v0.32.7-0.20240401174413-0bd2339e4f2e/src/pkg/utils/yaml.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // SPDX-FileCopyrightText: 2021-Present The Jackal Authors
     3  
     4  // Package utils provides generic utility functions.
     5  package utils
     6  
     7  // fork from https://github.com/goccy/go-yaml/blob/master/cmd/ycat/ycat.go
     8  
     9  import (
    10  	"bytes"
    11  	"fmt"
    12  	"io"
    13  	"io/fs"
    14  	"os"
    15  	"regexp"
    16  	"strings"
    17  
    18  	"github.com/Racer159/jackal/src/config"
    19  	"github.com/Racer159/jackal/src/pkg/message"
    20  	"github.com/fatih/color"
    21  	goyaml "github.com/goccy/go-yaml"
    22  	"github.com/goccy/go-yaml/lexer"
    23  	"github.com/goccy/go-yaml/printer"
    24  	"github.com/pterm/pterm"
    25  	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
    26  	"k8s.io/apimachinery/pkg/runtime"
    27  	kubeyaml "k8s.io/apimachinery/pkg/util/yaml"
    28  	k8syaml "sigs.k8s.io/yaml"
    29  )
    30  
    31  const yamlEscape = "\x1b"
    32  
    33  func yamlFormat(attr color.Attribute) string {
    34  	return fmt.Sprintf("%s[%dm", yamlEscape, attr)
    35  }
    36  
    37  // ColorPrintYAML pretty prints a yaml file to the console.
    38  func ColorPrintYAML(data any, hints map[string]string, spaceRootLists bool) {
    39  	text, _ := goyaml.Marshal(data)
    40  	tokens := lexer.Tokenize(string(text))
    41  
    42  	if spaceRootLists {
    43  		for idx := range tokens {
    44  			// Check that this is a dash after a newline (i.e. is a root list)
    45  			if tokens[idx].Origin == "\n-" || (tokens[idx].Origin == "-" && tokens[idx].Prev != nil && strings.HasSuffix(tokens[idx].Prev.Origin, "\n")) {
    46  				tokens[idx].Origin = "\n" + tokens[idx].Origin
    47  			}
    48  		}
    49  	}
    50  
    51  	var p printer.Printer
    52  	p.Bool = func() *printer.Property {
    53  		return &printer.Property{
    54  			Prefix: yamlFormat(color.FgHiWhite),
    55  			Suffix: yamlFormat(color.Reset),
    56  		}
    57  	}
    58  	p.Number = func() *printer.Property {
    59  		return &printer.Property{
    60  			Prefix: yamlFormat(color.FgHiWhite),
    61  			Suffix: yamlFormat(color.Reset),
    62  		}
    63  	}
    64  	p.MapKey = func() *printer.Property {
    65  		return &printer.Property{
    66  			Prefix: yamlFormat(color.FgHiCyan),
    67  			Suffix: yamlFormat(color.Reset),
    68  		}
    69  	}
    70  	p.Anchor = func() *printer.Property {
    71  		return &printer.Property{
    72  			Prefix: yamlFormat(color.FgHiYellow),
    73  			Suffix: yamlFormat(color.Reset),
    74  		}
    75  	}
    76  	p.Alias = func() *printer.Property {
    77  		return &printer.Property{
    78  			Prefix: yamlFormat(color.FgHiYellow),
    79  			Suffix: yamlFormat(color.Reset),
    80  		}
    81  	}
    82  	p.String = func() *printer.Property {
    83  		return &printer.Property{
    84  			Prefix: yamlFormat(color.FgHiMagenta),
    85  			Suffix: yamlFormat(color.Reset),
    86  		}
    87  	}
    88  
    89  	outputYAML := p.PrintTokens(tokens)
    90  
    91  	// Inject the hints into the colorized YAML
    92  	for key, value := range hints {
    93  		outputYAML = strings.Replace(outputYAML, key, value, 1)
    94  	}
    95  
    96  	if config.NoColor {
    97  		// If no color is specified strip any color codes from the output - https://regex101.com/r/YFyIwC/2
    98  		ansiRegex := regexp.MustCompile(`\x1b\[(.*?)m`)
    99  		outputYAML = ansiRegex.ReplaceAllString(outputYAML, "")
   100  	}
   101  
   102  	pterm.Println()
   103  	pterm.Println(outputYAML)
   104  }
   105  
   106  // AddRootListHint adds a hint string for a given root list key and value.
   107  func AddRootListHint(hints map[string]string, listKey string, listValue string, hintText string) map[string]string {
   108  	key := fmt.Sprintf("-%s %s%s:%s %s%s", yamlFormat(color.FgHiCyan), listKey, yamlFormat(color.Reset), yamlFormat(color.FgHiMagenta), listValue, yamlFormat(color.Reset))
   109  	hint := fmt.Sprintf("%s  %s%s", yamlFormat(color.FgHiBlack), hintText, yamlFormat(color.Reset))
   110  	hints[key] = fmt.Sprintf("%s%s", key, hint)
   111  
   112  	return hints
   113  }
   114  
   115  // AddRootHint adds a hint string for a given root key.
   116  func AddRootHint(hints map[string]string, rootKey string, hintText string) map[string]string {
   117  	key := fmt.Sprintf("%s%s%s:", yamlFormat(color.FgHiCyan), rootKey, yamlFormat(color.Reset))
   118  	newKey := fmt.Sprintf("%s%s:%s", yamlFormat(color.FgBlack)+yamlFormat(color.BgWhite), rootKey, yamlFormat(color.Reset))
   119  	hint := fmt.Sprintf("%s  %s%s", yamlFormat(color.FgHiBlack), hintText, yamlFormat(color.Reset))
   120  	hints[key] = fmt.Sprintf("\n%s\n%s%s", message.RuleLine, newKey, hint)
   121  
   122  	return hints
   123  }
   124  
   125  // ReadYaml reads a yaml file and unmarshals it into a given config.
   126  func ReadYaml(path string, destConfig any) error {
   127  	message.Debugf("Reading YAML at %s", path)
   128  
   129  	file, err := os.ReadFile(path)
   130  	if err != nil {
   131  		return err
   132  	}
   133  
   134  	return goyaml.Unmarshal(file, destConfig)
   135  }
   136  
   137  // WriteYaml writes a given config to a yaml file on disk.
   138  func WriteYaml(path string, srcConfig any, perm fs.FileMode) error {
   139  	// Save the parsed output to the config path given
   140  	content, err := goyaml.Marshal(srcConfig)
   141  	if err != nil {
   142  		return err
   143  	}
   144  
   145  	return os.WriteFile(path, content, perm)
   146  }
   147  
   148  // ReloadYamlTemplate marshals a given config, replaces strings and unmarshals it back.
   149  func ReloadYamlTemplate(config any, mappings map[string]string) error {
   150  	text, err := goyaml.Marshal(config)
   151  	if err != nil {
   152  		return err
   153  	}
   154  
   155  	for template, value := range mappings {
   156  		// Prevent user input from escaping the trailing " during yaml marshaling
   157  		lastIdx := len(value) - 1
   158  		if lastIdx > -1 && string(value[lastIdx]) == "\\" {
   159  			value = fmt.Sprintf("%s\\", value)
   160  		}
   161  		// Properly escape " in the yaml text output
   162  		value = strings.ReplaceAll(value, "\"", "\\\"")
   163  		text = []byte(strings.ReplaceAll(string(text), template, value))
   164  	}
   165  
   166  	return goyaml.Unmarshal(text, config)
   167  }
   168  
   169  // FindYamlTemplates finds strings with a given prefix in a config.
   170  func FindYamlTemplates(config any, prefix string, suffix string) (map[string]string, error) {
   171  	mappings := map[string]string{}
   172  
   173  	text, err := goyaml.Marshal(config)
   174  	if err != nil {
   175  		return nil, err
   176  	}
   177  
   178  	// Find all strings that are between the given prefix and suffix
   179  	r := regexp.MustCompile(fmt.Sprintf("%s([A-Z0-9_]+)%s", prefix, suffix))
   180  	matches := r.FindAllStringSubmatch(string(text), -1)
   181  
   182  	for _, match := range matches {
   183  		mappings[match[1]] = ""
   184  	}
   185  
   186  	return mappings, nil
   187  }
   188  
   189  // SplitYAML splits a YAML file into unstructured objects. Returns list of all unstructured objects
   190  // found in the yaml. If an error occurs, returns objects that have been parsed so far too.
   191  // Source: https://github.com/argoproj/gitops-engine/blob/v0.5.2/pkg/utils/kube/kube.go#L286.
   192  func SplitYAML(yamlData []byte) ([]*unstructured.Unstructured, error) {
   193  	var objs []*unstructured.Unstructured
   194  	ymls, err := SplitYAMLToString(yamlData)
   195  	if err != nil {
   196  		return nil, err
   197  	}
   198  	for _, yml := range ymls {
   199  		u := &unstructured.Unstructured{}
   200  		if err := k8syaml.Unmarshal([]byte(yml), u); err != nil {
   201  			return objs, fmt.Errorf("failed to unmarshal manifest: %#v", err)
   202  		}
   203  		objs = append(objs, u)
   204  	}
   205  	return objs, nil
   206  }
   207  
   208  // SplitYAMLToString splits a YAML file into strings. Returns list of yamls
   209  // found in the yaml. If an error occurs, returns objects that have been parsed so far too.
   210  // Source: https://github.com/argoproj/gitops-engine/blob/v0.5.2/pkg/utils/kube/kube.go#L304.
   211  func SplitYAMLToString(yamlData []byte) ([]string, error) {
   212  	// Similar way to what kubectl does
   213  	// https://github.com/kubernetes/cli-runtime/blob/master/pkg/resource/visitor.go#L573-L600
   214  	// Ideally k8s.io/cli-runtime/pkg/resource.Builder should be used instead of this method.
   215  	// E.g. Builder does list unpacking and flattening and this code does not.
   216  	d := kubeyaml.NewYAMLOrJSONDecoder(bytes.NewReader(yamlData), 4096)
   217  	var objs []string
   218  	for {
   219  		ext := runtime.RawExtension{}
   220  		if err := d.Decode(&ext); err != nil {
   221  			if err == io.EOF {
   222  				break
   223  			}
   224  			return objs, fmt.Errorf("failed to unmarshal manifest: %#v", err)
   225  		}
   226  		ext.Raw = bytes.TrimSpace(ext.Raw)
   227  		if len(ext.Raw) == 0 || bytes.Equal(ext.Raw, []byte("null")) {
   228  			continue
   229  		}
   230  		objs = append(objs, string(ext.Raw))
   231  	}
   232  	return objs, nil
   233  }