/
load.go
117 lines (110 loc) · 3.31 KB
/
load.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
package putil
import (
"fmt"
"log"
"os"
"path/filepath"
"strings"
v1 "github.com/ocurity/dracon/api/proto/v1"
"google.golang.org/protobuf/proto"
)
// LoadToolResponse loads raw results.
func LoadToolResponse(inPath string) ([]*v1.LaunchToolResponse, error) {
responses := []*v1.LaunchToolResponse{}
if err := filepath.Walk(inPath, func(path string, f os.FileInfo, err error) error {
if f == nil {
return fmt.Errorf("path %s doesn't exist", path)
}
if !f.IsDir() && (strings.HasSuffix(f.Name(), ".pb")) {
pbBytes, err := os.ReadFile(path)
if err != nil {
return err
}
res := v1.LaunchToolResponse{}
if err := proto.Unmarshal(pbBytes, &res); err != nil {
log.Printf("skipping %s as unable to unmarshal", path)
} else {
responses = append(responses, &res)
}
}
return nil
}); err != nil {
return responses, err
}
return responses, nil
}
// LoadTaggedToolResponse loads raw results that have been tagged, it's used by the enrichers.
func LoadTaggedToolResponse(inPath string) ([]*v1.LaunchToolResponse, error) {
responses := []*v1.LaunchToolResponse{}
if err := filepath.Walk(inPath, func(path string, f os.FileInfo, err error) error {
if f == nil {
return fmt.Errorf("path %s doesn't exist", path)
}
if !f.IsDir() && (strings.HasSuffix(f.Name(), ".tagged.pb")) {
pbBytes, err := os.ReadFile(path)
if err != nil {
return err
}
res := v1.LaunchToolResponse{}
if err := proto.Unmarshal(pbBytes, &res); err != nil {
log.Printf("skipping %s as unable to unmarshal", path)
} else {
responses = append(responses, &res)
}
}
return nil
}); err != nil {
return responses, err
}
return responses, nil
}
// LoadEnrichedToolResponse loads enriched results from the enricher.
func LoadEnrichedToolResponse(inPath string) ([]*v1.EnrichedLaunchToolResponse, error) {
responses := []*v1.EnrichedLaunchToolResponse{}
if err := filepath.Walk(inPath, func(path string, f os.FileInfo, err error) error {
if f == nil {
return fmt.Errorf("path %s doesn't exist", path)
}
if !f.IsDir() && (strings.HasSuffix(f.Name(), ".enriched.aggregated.pb")) {
pbBytes, err := os.ReadFile(path)
if err != nil {
return err
}
res := v1.EnrichedLaunchToolResponse{}
if err := proto.Unmarshal(pbBytes, &res); err != nil {
log.Printf("skipping %s as unable to unmarshal", path)
} else {
responses = append(responses, &res)
}
}
return nil
}); err != nil {
return responses, err
}
return responses, nil
}
// LoadEnrichedNonAggregatedToolResponse loads enriched but not aggregated results from the enricher.
func LoadEnrichedNonAggregatedToolResponse(inPath string) ([]*v1.EnrichedLaunchToolResponse, error) {
responses := []*v1.EnrichedLaunchToolResponse{}
if err := filepath.Walk(inPath, func(path string, f os.FileInfo, err error) error {
if f == nil {
return fmt.Errorf("path %s doesn't exist", path)
}
if !f.IsDir() && (strings.HasSuffix(f.Name(), ".enriched.pb")) {
pbBytes, err := os.ReadFile(path)
if err != nil {
return err
}
res := v1.EnrichedLaunchToolResponse{}
if err := proto.Unmarshal(pbBytes, &res); err != nil {
log.Printf("skipping %s as unable to unmarshal", path)
} else {
responses = append(responses, &res)
}
}
return nil
}); err != nil {
return responses, err
}
return responses, nil
}