-
Notifications
You must be signed in to change notification settings - Fork 21
Expand file tree
/
Copy pathspur.go
More file actions
157 lines (135 loc) · 4.06 KB
/
spur.go
File metadata and controls
157 lines (135 loc) · 4.06 KB
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
/*
* ZAnnotate Copyright 2026 Regents of the University of Michigan
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
* use this file except in compliance with the License. You may obtain a copy
* of the License at http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
* implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package zannotate
import (
"bytes"
"encoding/json"
"errors"
"flag"
"fmt"
"io"
"net"
"net/http"
"os"
"strings"
"time"
log "github.com/sirupsen/logrus"
)
const SpurApiUrl = "https://api.spur.us/v2/context/"
type SpurAnnotatorFactory struct {
BasePluginConf
apiKey string // Spur API Key, pulled from env var
timeoutSecs int
}
type SpurAnnotator struct {
Factory *SpurAnnotatorFactory
Id int
client *http.Client
}
// Spur Annotator Factory (Global)
func (a *SpurAnnotatorFactory) MakeAnnotator(i int) Annotator {
var v SpurAnnotator
v.Factory = a
v.Id = i
return &v
}
func (a *SpurAnnotatorFactory) Initialize(_ *GlobalConf) error {
// Check for API key
a.apiKey = os.Getenv("SPUR_API_KEY")
if len(a.apiKey) == 0 {
return errors.New("SPUR_API_KEY environment variable not set. Please use 'export SPUR_API_KEY=your_api_key' to set it")
}
return nil
}
func (a *SpurAnnotatorFactory) GetWorkers() int {
return a.Threads
}
func (a *SpurAnnotatorFactory) Close() error {
return nil
}
func (a *SpurAnnotatorFactory) IsEnabled() bool {
return a.Enabled
}
func (a *SpurAnnotatorFactory) AddFlags(flags *flag.FlagSet) {
flags.BoolVar(&a.Enabled, "spur", false, "enrich with Spur's threat intelligence data")
flags.IntVar(&a.Threads, "spur-threads", 100, "how many threads to use for Spur lookups")
flags.IntVar(&a.timeoutSecs, "spur-timeout", 2, "timeout for each Spur query, in seconds")
}
// Spur Annotator (Per-Worker)
func (a *SpurAnnotator) Initialize() error {
a.client = &http.Client{
Timeout: time.Duration(a.Factory.timeoutSecs) * time.Second,
}
return nil
}
func (a *SpurAnnotator) GetFieldName() string {
return "spur"
}
// Annotate performs a Spur data lookup for the given IP address and returns the results.
// If an error occurs or a lookup fails, it returns nil
func (a *SpurAnnotator) Annotate(ip net.IP) interface{} {
req, err := http.NewRequest(
http.MethodGet,
fmt.Sprintf("%s%s", SpurApiUrl, ip.String()),
nil,
)
if err != nil {
log.Errorf("failed to create Spur HTTP request for IP %s: %v", ip.String(), err)
return nil
}
req.Header.Set("Token", a.Factory.apiKey) // Set the API key in the request header
resp, err := a.client.Do(req)
if err != nil {
log.Errorf("http request to Spur API failed for IP %s: %v", ip.String(), err)
return nil
}
defer func(Body io.ReadCloser) {
err := Body.Close()
if err != nil {
log.Errorf("failed to close Spur API response body: %v", err)
}
}(resp.Body)
body, err := io.ReadAll(resp.Body)
if err != nil {
log.Errorf("failed to read response body for IP %s: %v", ip.String(), err)
return nil
}
switch resp.StatusCode {
case http.StatusOK:
trimmed, _ := strings.CutSuffix(string(body), "\n") // Remove trailing newline if present, cleans up output
return json.RawMessage(trimmed)
case http.StatusUnauthorized:
// retrieve error message from body if possible
var msg string
if json.Valid(body) {
var compact bytes.Buffer
if err := json.Compact(&compact, body); err == nil {
msg = compact.String()
} else {
// fallback to raw trimmed text if compaction fails
msg = strings.TrimSpace(string(body))
}
}
// wouldn't be able to recover from an auth error, so log fatal with details and exit
log.Fatalf("error from Spur API for IP %s. Spur responded with %s: double-check your API key", ip.String(), msg)
}
return nil
}
func (a *SpurAnnotator) Close() error {
return nil
}
func init() {
s := new(SpurAnnotatorFactory)
RegisterAnnotator(s)
}