forked from GoogleCloudPlatform/professional-services
/
gcpviz.go
166 lines (146 loc) · 5.36 KB
/
gcpviz.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
package main
import (
"bufio"
"bytes"
"context"
"flag"
"io/ioutil"
"log"
"os"
"runtime"
"runtime/pprof"
"strings"
"sync"
"github.com/rosmo/professional-services/tools/gcpviz"
"github.com/dimiro1/banner"
)
var modePtr string
type arrayFlags []string
func (i *arrayFlags) String() string {
return "my string representation"
}
func (i *arrayFlags) Set(value string) error {
*i = append(*i, value)
return nil
}
func main() {
modePtr := flag.String("mode", "", "mode of operation (generate, visualize)")
relationsFilePtr := flag.String("relations-file", "relations.yaml", "location of relations file")
styleFilePtr := flag.String("style-file", "style.yaml", "location of graph style file")
labelsFilePtr := flag.String("labels-file", "labels.yaml", "location of node/edge labels file")
queryFilePtr := flag.String("query-file", "query.js", "location of Gizmo query file")
graphFilePtr := flag.String("graph-file", "graph.db", "location of Graph & Asset database file")
resourceInventoryFilePtr := flag.String("resource-inventory-file", "resource_inventory.json", "location of resource inventory file from Cloud Asset Inventory")
graphTitlePtr := flag.String("graph-title", "", "Title for the graph")
noColorPtr := flag.Bool("no-color", false, "disables color in output")
noBannerPtr := flag.Bool("no-banner", false, "disables banner")
cpuprofile := flag.String("cpuprofile", "", "write cpu profile to `file`")
memprofile := flag.String("memprofile", "", "write memory profile to `file`")
var graphParameters arrayFlags
var queryParameters arrayFlags
flag.Var(&graphParameters, "graph-parameter", "override graph style parameters using SJSON (ie. \"options.overlap=vpsc\")")
flag.Var(&queryParameters, "query-parameter", "additional parameter to pass to Gizmo query (param=value)")
flag.Parse()
banner.Init(os.Stderr, !*noBannerPtr, !*noColorPtr, bytes.NewBufferString(`
{{ .AnsiColor.BrightMagenta }} ██████ ██████ ██████ ██ ██ ██ ███████
{{ .AnsiColor.Magenta }} ██ ██ ██ ██ ██ ██ ██ ███
{{ .AnsiColor.Blue }} ██ ███ ██ ██████ ██ ██ ██ ███
{{ .AnsiColor.Cyan }} ██ ██ ██ ██ ██ ██ ██ ███
{{ .AnsiColor.BrightCyan }} ██████ ██████ ██ ████ ██ ███████
{{ .AnsiColor.Default }}
`))
if *modePtr == "" {
flag.PrintDefaults()
os.Exit(1)
}
if *cpuprofile != "" {
f, err := os.Create(*cpuprofile)
if err != nil {
log.Fatal("Could not create CPU profile: ", err)
}
defer f.Close()
if err := pprof.StartCPUProfile(f); err != nil {
log.Fatal("Could not start CPU profile: ", err)
}
defer pprof.StopCPUProfile()
}
if *memprofile != "" {
f, err := os.Create(*memprofile)
if err != nil {
log.Fatal("Could not create memory profile: ", err)
}
defer f.Close()
runtime.GC() // get up-to-date statistics
if err := pprof.WriteHeapProfile(f); err != nil {
log.Fatal("Could not write memory profile: ", err)
}
}
var overrideParams map[string]string = nil
if len(graphParameters) > 0 {
overrideParams = make(map[string]string, len(graphParameters))
for _, ov := range graphParameters {
param := strings.SplitN(ov, "=", 2)
overrideParams[param[0]] = param[1]
}
}
viz, err := gcpviz.NewGcpViz(*relationsFilePtr, *labelsFilePtr, *styleFilePtr, overrideParams)
if err != nil {
log.Fatalf("Failed to initialize graph engine: %v", err)
}
if *modePtr == "generate" {
err = viz.Create(*graphFilePtr)
if err != nil {
log.Fatalf("Failed to create graph file: %v", err)
}
err = viz.ReadAssetsFromFile(*resourceInventoryFilePtr)
if err != nil {
log.Fatalf("Failed read assets from resource inventory: %v", err)
}
err = viz.EnrichAssets()
if err != nil {
log.Fatalf("Failed create references and enrich assets in resource inventory: %v", err)
}
err = viz.Save()
if err != nil {
log.Fatalf("Failed to save graph file: %v", err)
}
}
if *modePtr == "visualize" {
err = viz.Load(*graphFilePtr)
if err != nil {
log.Fatalf("Failed to load graph file: %v", err)
}
gizmoQuery, err := ioutil.ReadFile(*queryFilePtr)
if err != nil {
log.Fatalf("Failed to load query file: %v", err)
}
var parameters map[string]interface{}
if len(queryParameters) > 0 {
parameters = make(map[string]interface{}, len(queryParameters)+1)
for _, p := range queryParameters {
param := strings.SplitN(p, "=", 2)
if len(param) == 2 {
parameters[param[0]] = param[1]
} else {
log.Fatalf("Invalid query parameter (use param=value format): %s", p)
}
}
} else {
parameters = make(map[string]interface{}, 1)
}
parameters["Title"] = viz.EscapeLabel(*graphTitlePtr)
f := bufio.NewWriter(os.Stdout)
defer f.Flush()
waitGroup := sync.WaitGroup{}
ctx := context.Background()
waitGroup.Add(1)
err = viz.GenerateNodes(&waitGroup, ctx, string(gizmoQuery), parameters, f)
if err != nil {
log.Fatalf("Failed to create graph: %v", err)
}
waitGroup.Wait()
}
if *modePtr != "visualize" && *modePtr != "generate" {
log.Fatal("invalid mode specified, specify either generate or visualize")
}
}