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 }