/
hitmap.go
107 lines (96 loc) · 3.13 KB
/
hitmap.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
package output
import (
"sort"
"strconv"
"strings"
"github.com/olekukonko/tablewriter"
)
// HitMap is a map of strings to ints. It's used to track hits for a given rule
// and/or file. The index is the ruleID or filename. The value is the number of
// times the rule was hit or the number of findings in the file.
type HitMap map[string]int
// ruleHitMap creates a HitMap where the key is ruleID and the value is the
// number of hits for that rule.
func ruleHitMap(o Output) HitMap {
// Create the hitmap.
hm := make(HitMap)
// Go through the results and create the hitmap.
for _, result := range o.Results {
hm[result.RuleID()]++
}
return hm
}
// fileHitMap creates a HitMap where the key is file path and the value is
// the number of hits in that file.
func fileHitMap(o Output) HitMap {
// Create the hitmap.
hm := make(HitMap)
// Go through the results and create the hitmap.
for _, result := range o.Results {
hm[result.FilePath()]++
}
return hm
}
// -----
// HitMapRow is a single row in the HitMap.
type HitMapRow struct {
Key string
Value string
}
// SortedData returns the HitMap data as a series of rows. If sortByCount is true,
// the data will be sorted by count, otherwise they will be sorted by key.
//
// Note: Sorting by count is in descending order because most of the time we
// assume we want rules/files with higher hits first. Sorting by key is
// ascending order alphabetically.
func (m HitMap) SortedData(sortByCount bool) []HitMapRow {
// First get a slice of keys.
mapKeys := make([]string, len(m))
// We're using an index instead of append because we know the size of the
// array and it's probably a bit faster.
index := 0
for k := range m {
mapKeys[index] = k
index++
}
// Sort by value.
if sortByCount {
// We can sort using sort.Slice and providing our own less function that
// does the comparison by value of keys. Note: We're sorting in
// descending order here because most of the time we assume we want
// findings/files with higher hits first.
sort.Slice(mapKeys, func(i, j int) bool {
return m[mapKeys[i]] > m[mapKeys[j]]
})
} else {
// Sort by key. This is in ascending order.
sort.Strings(mapKeys)
}
// mapKeys is sorted. We can create the data.
data := make([]HitMapRow, len(m))
for index, key := range mapKeys {
data[index] = HitMapRow{Key: key, Value: strconv.Itoa(m[key])}
}
return data
}
// Returns the HitMap data as a text table. This is useful for text printing.
// It's the caller's responsibility to pass the correct number of headers.
func (m HitMap) ToStringTable(headers []string, sortByCount bool) string {
// Convert the HitMap data into a [][]string to pass to tablewriter.
data := make([][]string, len(m))
for index, row := range m.SortedData(sortByCount) {
data[index] = []string{row.Key, row.Value}
}
// String builder to hold the result.
var final strings.Builder
// Create the table writer and set the destination to the string builder.
table := tablewriter.NewWriter(&final)
// Set the headers.
table.SetHeader(headers)
// Append the data.
table.AppendBulk(data)
// Render the table.
table.Render()
// Return the data.
return final.String()
}