rootsystem/test.go

package main

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"html"
	"io"
	"io/ioutil"
	"os/exec"
	"regexp"
	"strings"

	errors "git.sequentialread.com/forest/pkg-errors"
)

type XMLAttr struct {
	Name      string
	Value     string
	Namespace string
	Local     string
}

type XMLNode struct {
	XMLName  xml.Name
	Attrs    []*XMLAttr `xml:"-"`
	Content  []byte     `xml:",innerxml"`
	Children []*XMLNode `xml:",any"`
	Parent   *XMLNode   `xml:"-"`
}

type XMLQuery struct {
	DirectChild bool
	NodeType    string
	Class       string
	Attr        string
}

func (node *XMLNode) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	attrPointers := []*XMLAttr{}
	for _, attr := range start.Attr {
		// I dont really give a shit about XML namespaces but I do have to search for xlink:title which is separate from normal title
		// hardcoded that this is happening inside the svg namespace so all non svg namespaces attrs will be specified like ns:attr
		name := attr.Name.Local
		if attr.Name.Space != "" {
			// attr.Name.Space is a url like http://www.w3.org/1999/xlink or ttp://www.w3.org/2000/svg so we are just grabbing the last part
			split := strings.Split(attr.Name.Space, "/")
			ns := split[len(split)-1]
			if ns != "svg" {
				name = fmt.Sprintf("%s:%s", ns, name)
			}
		}
		myAttr := XMLAttr{
			Name:      name,
			Value:     attr.Value,
			Namespace: attr.Name.Space,
			Local:     attr.Name.Local,
		}
		attrPointers = append(attrPointers, &myAttr)
	}
	node.Attrs = attrPointers

	// this little derived type wont directly have this UnmarshalXML method,
	// thus prevent the decoder from thinking it has to invoke our custom UnmarshalXML again...
	type dontCauseInfiniteLoop XMLNode

	return d.DecodeElement((*dontCauseInfiniteLoop)(node), &start)
}

//https://stackoverflow.com/questions/39780433/golang-xml-customize-output

// // MarshalXML generate XML output for PrecsontructedInfo
// func (preconstructed PreconstructedInfo) MarshalXML(e *xml.Encoder, start xml.StartElement) (err error) {
// 	if (PreconstructedInfo{} == preconstructed) {
// 			return nil
// 	}
// 	if preconstructed.Decks > 0 {
// 			start.Attr = []xml.Attr{xml.Attr{Name: xml.Name{Local: "decks"}, Value: strconv.Itoa(preconstructed.Size)}}
// 	}
// 	if strings.Compare(preconstructed.Type, "") != 0 {
// 			start.Attr = append(start.Attr, xml.Attr{Name: xml.Name{Local: "type"}, Value: preconstructed.Type})
// 	}

// 	err = e.EncodeToken(start)
// 	e.EncodeElement(preconstructed.Size, xml.StartElement{Name: xml.Name{Local: "size"}})
// 	return e.EncodeToken(xml.EndElement{Name: start.Name})
// }

func (node XMLNode) MarshalXML(e *xml.Encoder, start xml.StartElement) error {

	attrsToEncode := []xml.Attr{}
	for _, attr := range node.Attrs {
		// this encoder will print some stupid shit if you give it an xmlns attr:
		// for example: <svg  ...  xmlns:_ _xmlns:xlink="http://www.w3.org/1999/xlink">
		if attr.Namespace != "xmlns" {
			attrsToEncode = append(attrsToEncode, xml.Attr{
				Name: xml.Name{
					Local: attr.Local,
					Space: attr.Namespace,
				},
				Value: attr.Value,
			})
		}

	}
	start.Attr = attrsToEncode
	start.Name = node.XMLName

	// this little derived type wont directly have this UnmarshalXML method,
	// thus prevent the decoder from thinking it has to invoke our custom UnmarshalXML again...
	type dontCauseInfiniteLoop XMLNode

	return e.EncodeElement((dontCauseInfiniteLoop)(node), start)
}

func main() {
	bytes, err := ioutil.ReadFile("test-dot.txt")
	if err != nil {
		panic(err)
	}

	dot := string(bytes)

	exitCode, dotStdout, dotStderr, err := shellExecInputPipe(".", &dot, "dot", "-Tsvg")
	err = errorFromShellExecResult("dot -Tsvg", exitCode, dotStdout, dotStderr, err)
	if err != nil {
		panic(err)
	}

	svgString := string(dotStdout)

	// the svg output from dot has html comments that contain metadata preceeding each entity.
	// for example:
	// <!-- param_ssh_private_key_filepath -->
	// <g id="node16" class="node">
	// We will attach that metadata to the svg elements as a data property
	// Note that this will not work perfectly for the "cluster"s (aka modules)
	// they are handled separately after the SVG has been inserted.
	// Yes parsing XML with regex is "evil"  but keep in mind, this XML was generated by a program, not written by hand.
	// So it is a lot more predictable. Also, I couldn't figure out how to parse the comments with the go std lib XML parser.
	commentRegex := regexp.MustCompile(`<!--\s*(.+)\s*-->\s*\n\s*(<[a-z]+\s)`)
	svgString = commentRegex.ReplaceAllStringFunc(svgString, func(matched string) string {
		matches := commentRegex.FindStringSubmatch(matched)
		return fmt.Sprintf("%s dot=\"%s\" ", matches[2], html.UnescapeString(matches[1]))
	})

	ioutil.WriteFile("./test-gen.svg", []byte(svgString), 0777)

	var svgDoc XMLNode

	err = xml.Unmarshal(dotStdout, &svgDoc)
	if err != nil {
		panic(err)
	}

	var setParentRecurse func(*XMLNode)
	setParentRecurse = func(node *XMLNode) {
		for _, child := range node.Children {
			child.Parent = node
			setParentRecurse(child)
		}
	}
	setParentRecurse(&svgDoc)

	// remove the white background from the SVG
	svgDoc.WithQuerySelector(
		[]XMLQuery{XMLQuery{NodeType: "g", Class: "graph"}, XMLQuery{NodeType: "polygon", DirectChild: true}},
		func(node *XMLNode) {
			node.SetAttr("fill", "none")
		},
	)

	// remove the white background from the edges
	svgDoc.WithQuerySelector(
		[]XMLQuery{XMLQuery{NodeType: "g", Class: "edge"}, XMLQuery{NodeType: "path"}},
		func(node *XMLNode) {
			node.SetAttr("fill", "none")
		},
	)

	// correctly set the dot property for the "cluster"s (aka modules)
	svgDoc.WithQuerySelector(
		[]XMLQuery{XMLQuery{NodeType: "a", Attr: "xlink:title"}},
		func(node *XMLNode) {
			if node.Parent != nil && node.Parent.Parent != nil {
				title := html.UnescapeString(node.GetAttr("xlink:title"))
				title = regexp.MustCompile(`[.-]`).ReplaceAllString(title, "_")
				node.Parent.Parent.SetAttr("dot", title)
			}
		},
	)

	// correctly set the dot property for resources or "node"s
	svgDoc.WithQuerySelector(
		[]XMLQuery{XMLQuery{NodeType: "g", Class: "node"}, XMLQuery{NodeType: "title"}},
		func(node *XMLNode) {
			if node.Parent != nil {
				title := html.UnescapeString(string(node.Content))
				node.Parent.SetAttr("dot", title)
			}
		},
	)

	// correctly set the dot property for edges
	svgDoc.WithQuerySelector(
		[]XMLQuery{XMLQuery{NodeType: "g", Class: "edge"}, XMLQuery{NodeType: "title"}},
		func(node *XMLNode) {
			if node.Parent != nil {
				title := html.UnescapeString(string(node.Content))
				title = strings.Replace(title, "->", "_to_", 1)
				node.Parent.SetAttr("dot", title)
			}
		},
	)

	// // in order to get dot to render it correctly, the connections are actually between resources,
	// // not between modules. So we have to trim the resource name off of the dot attribute
	// // on the connections to make them match what we have in the status JSON object.
	// const moduleNames = Array.from(svg.querySelectorAll("g.cluster")).map(x => x.dataset['dot']);
	// Array.from(svg.querySelectorAll("g.edge")).forEach(edge => {
	// 	const dot = edge.dataset['dot'];
	// 	const fromTo = dot.split('->');
	// 	if(fromTo.length != 2) {
	// 		console.error("malformed dot on edge: fromTo.length != 2", + dot);
	// 	}
	// 	moduleNames.forEach(moduleName => {
	// 		if(fromTo[0].indexOf(moduleName) == 0) {
	// 			edge.dataset['dot'] = `${moduleName}->${fromTo[1]}`;
	// 		}
	// 		if(fromTo[1].indexOf(moduleName) == 0) {
	// 			edge.dataset['dot'] = `${fromTo[0]}->${moduleName}`;
	// 		}
	// 	});
	// });

	// svgString := strings.ReplaceAll(string(dotStdout), `fill="none"`, `fill="#ffffff"`)
	// matchTitleTag := regexp.MustCompile(`<title>[^<]*</title>`)
	// svgString = matchTitleTag.ReplaceAllString(svgString, "")
	// svgString = strings.ReplaceAll(svgString, `Times,serif`, `-apple-system,system-ui,BlinkMacSystemFont,Ubuntu,Roboto,Segoe UI,sans-serif`)

	// testJson, _ := json.MarshalIndent(simpleStatus, "", " ")
	// ioutil.WriteFile("./test-status.json", testJson, 0777)

	// this encoder doesn't want to include my custom attributes...
	// so we will hijack the id attribute! Record the value of the custom attribute so we can
	// set the value of Id later.

	var cleanUpSvgRecurse func(*XMLNode)
	cleanUpSvgRecurse = func(node *XMLNode) {
		// We cant serialize cyclical relationships like parent -> child -> parent
		node.Parent = nil

		// the XML serializer does not like my custom attributes like `data-dot` or `dot`.
		// so we will just hijack the id attribute instead, poggers
		dot := node.GetAttr("dot")
		if dot != "" {
			node.SetAttr("id", dot)
		}

		// if we don't zero out the content for non-leaf nodes, we will get duplicates of everything
		if len(node.Children) != 0 {
			node.Content = []byte{}
		}

		for _, child := range node.Children {
			cleanUpSvgRecurse(child)
		}
	}
	cleanUpSvgRecurse(&svgDoc)

	testJson2, err := json.MarshalIndent(svgDoc, "", " ")
	if err != nil {
		panic(err)
	}
	ioutil.WriteFile("./test-processed.json", testJson2, 0777)

	svgBytes, err := xml.MarshalIndent(svgDoc, "", "  ")

	// it goes crazy spamming xmlns everywhere, lets just remove all of them.
	svgBytes = regexp.MustCompile(`xmlns[^=]*="[^"]+"`).ReplaceAll(svgBytes, []byte{})

	// now we add the xmlns back to just the root element
	startOfSvgTagWithXMLNS := []byte(`<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" `)
	svgBytes = regexp.MustCompile(`^\s*<svg\s*`).ReplaceAll(svgBytes, startOfSvgTagWithXMLNS)

	svgHeader := []byte(`<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
 "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
<!-- Generated by graphviz & server.garden rootsystem -->
`)

	svgBytes = append(svgHeader, svgBytes...)

	// if err != nil {
	// 	panic(err)
	// }
	ioutil.WriteFile("./test-processed.svg", svgBytes, 0777)
}

func padStringForDot(str string) string {
	pad := strings.Repeat(" ", int(float32(len(str))/4))
	return fmt.Sprintf("%s%s%s", pad, str, pad)
}

// TODO which user to run this command as ??
// Dont run as root. (this process probably runs as root)
func shellExec(workingDirectory string, executable string, arguments ...string) (int, []byte, []byte, error) {
	return shellExecInputPipe(workingDirectory, nil, executable, arguments...)
}

func shellExecInputPipe(workingDirectory string, input *string, executable string, arguments ...string) (int, []byte, []byte, error) {
	process := exec.Command(executable, arguments...)
	process.Dir = workingDirectory

	if input != nil {
		stdin, err := process.StdinPipe()
		if err != nil {
			return -1, []byte{}, []byte{}, errors.Wrap(err, "process.StdinPipe() returned")
		}

		go func() {
			defer stdin.Close()
			io.WriteString(stdin, *input)
		}()
	}

	var processStdoutBuffer, processStderrBuffer bytes.Buffer
	process.Stdout = &processStdoutBuffer
	process.Stderr = &processStderrBuffer
	err := process.Start()
	if err != nil {
		err = errors.Wrapf(err, "can't ShellExec(%s %s), process.Start() returned", executable, strings.Join(arguments, " "))
		return process.ProcessState.ExitCode(), []byte(""), []byte(""), err
	}

	err = process.Wait()
	if err != nil {
		err = errors.Wrapf(err, "can't ShellExec(%s %s), process.Wait() returned", executable, strings.Join(arguments, " "))
	}

	return process.ProcessState.ExitCode(), processStdoutBuffer.Bytes(), processStderrBuffer.Bytes(), err
}

func errorFromShellExecResult(command string, exitCode int, stdout []byte, stderr []byte, err error) error {
	if exitCode != 0 || err != nil {
		errorString := "nil"
		if err != nil {
			errorString = err.Error()
		}
		return fmt.Errorf(
			"%s failed with exit code %d, stdout: \n----\n%s\n----\nstderr: \n----\n%s\n----\nerror: %s",
			command, exitCode, stdout, stderr, errorString,
		)
	}
	return nil
}

func (doc *XMLNode) QuerySelectorAll(querySteps []XMLQuery) []*XMLNode {

	iterate := func(search []*XMLNode, q XMLQuery) []*XMLNode {
		toReturn := []*XMLNode{}
		var recurse func(*XMLNode, bool)
		recurse = func(node *XMLNode, isDirectChild bool) {
			match := true
			if q.NodeType != "" && node.XMLName.Local != q.NodeType {
				match = false
			}
			if q.Class != "" {
				found := false
				for _, attr := range node.Attrs {
					classes := strings.Split(attr.Value, " ")
					if attr.Name == "class" {
						for _, class := range classes {
							if class == q.Class {
								found = true
							}
						}
					}
				}
				if !found {
					match = false
				}
			}
			if q.Attr != "" {
				found := false
				for _, attr := range node.Attrs {
					if attr.Name == q.Attr {
						found = true
					}
				}
				if !found {
					match = false
				}
			}
			if match && (isDirectChild || !q.DirectChild) {
				toReturn = append(toReturn, node)
			}
			if isDirectChild || !q.DirectChild {
				for _, child := range node.Children {
					recurse(child, false)
				}
			}
		}
		for _, node := range search {
			for _, child := range node.Children {
				recurse(child, true)
			}
		}

		return toReturn
	}

	matched := []*XMLNode{doc}
	for _, step := range querySteps {
		matched = iterate(matched, step)
	}

	return matched
}

func (doc *XMLNode) WithQuerySelector(querySteps []XMLQuery, withResult func(*XMLNode)) {
	allResults := doc.QuerySelectorAll(querySteps)
	for _, res := range allResults {
		withResult(res)
	}
}

func (node *XMLNode) SetAttr(attr string, value string) {
	has := false
	for _, a := range node.Attrs {
		if a.Name == attr {
			a.Value = value
			has = true
		}
	}
	if !has {
		newAttr := XMLAttr{
			Name:  attr,
			Value: value,
		}
		node.Attrs = append(node.Attrs, &newAttr)
	}
}

func (node *XMLNode) GetAttr(attr string) string {
	for _, a := range node.Attrs {
		if a.Name == attr {
			return a.Value
		}
	}
	return ""
}