Skip to content

Commit 2f03b1a

Browse files
author
Daniel Jimenez
committed
Refactor notify pkg
1 parent c80c752 commit 2f03b1a

File tree

9 files changed

+220
-83
lines changed

9 files changed

+220
-83
lines changed

‎pkg/notify/notifier.go‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,25 @@ Copyright 2020 Adevinta
44

55
package notify
66

7+
import (
8+
"github.com/adevinta/vulnerability-db/pkg/store"
9+
)
10+
11+
// FindingNotification represents a notification associated with a finding state change.
12+
type FindingNotification struct {
13+
store.FindingExpanded
14+
// TODO: Tag field is defined here in order for the FindingNotification struct to be
15+
// backward compatible with the "old" vulnerability notifications format so integrations
16+
// with external components that still use the old representation (e.g.: Hermes) are
17+
// maintained and can be isolated in the notify pkg. Once these integrations are modified
18+
// to use the new notification format we'll have to decide if we make this tag field
19+
// serializable into JSON and keep using the tag field as identifier, or we migrate to use
20+
// the current team identifier.
21+
Tag string `json:"-"`
22+
}
23+
724
// Notifier is a connector that allows
825
// to push messages to a notification service.
926
type Notifier interface {
10-
Push(mssg interface{}) error
27+
PushFinding(f FindingNotification) error
1128
}

‎pkg/notify/sns.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,14 +55,14 @@ func NewSNSNotifier(conf SNSConfig, logger *log.Logger) (*SNSNotifier, error) {
5555
return notifier, nil
5656
}
5757

58-
// Push pushes a notification to the configured sns topic.
59-
func (n *SNSNotifier) Push(message interface{}) error {
58+
// PushFinding pushes a finding notification to the configured sns topic.
59+
func (n *SNSNotifier) PushFinding(f FindingNotification) error {
6060
if !n.conf.Enabled {
6161
return nil
6262
}
6363

6464
n.logger.Info("Pushing notification to SNS")
65-
content, err := json.Marshal(&message)
65+
content, err := json.Marshal(f)
6666
if err != nil {
6767
return err
6868
}

‎pkg/notify/sns_test.go‎

Lines changed: 149 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,18 @@ package notify
66

77
import (
88
"testing"
9+
"unicode"
910

1011
"github.com/aws/aws-sdk-go/aws"
1112
"github.com/google/go-cmp/cmp"
1213
"github.com/google/go-cmp/cmp/cmpopts"
14+
"github.com/lib/pq"
1315

1416
"github.com/aws/aws-sdk-go/service/sns"
1517
"github.com/aws/aws-sdk-go/service/sns/snsiface"
1618
log "github.com/sirupsen/logrus"
19+
20+
"github.com/adevinta/vulnerability-db/pkg/store"
1721
)
1822

1923
type snsMock struct {
@@ -26,7 +30,7 @@ func (m *snsMock) Publish(s *sns.PublishInput) (*sns.PublishOutput, error) {
2630
return nil, nil
2731
}
2832

29-
func TestSNSNotifier_Push(t *testing.T) {
33+
func TestSNSNotifier_PushFinding(t *testing.T) {
3034
type fields struct {
3135
conf SNSConfig
3236
sns *snsMock
@@ -36,7 +40,7 @@ func TestSNSNotifier_Push(t *testing.T) {
3640
tests := []struct {
3741
name string
3842
fields fields
39-
message map[string]interface{}
43+
f FindingNotification
4044
want *sns.PublishInput
4145
wantErr bool
4246
}{
@@ -50,27 +54,165 @@ func TestSNSNotifier_Push(t *testing.T) {
5054
Enabled: true,
5155
},
5256
},
53-
message: map[string]interface{}{"a": "b"},
57+
f: FindingNotification{
58+
FindingExpanded: store.FindingExpanded{
59+
Finding: store.Finding{
60+
ID: "FindingID-1",
61+
AffectedResource: "AffectedResource-1",
62+
Score: 9,
63+
Status: "OPEN",
64+
Details: "Details-1",
65+
ImpactDetails: "ImpactDetails-1",
66+
},
67+
Issue: store.IssueLabels{
68+
Issue: store.Issue{
69+
ID: "IssueID-1",
70+
Summary: "Summary-1",
71+
CWEID: 1,
72+
Description: "Description-1",
73+
Recommendations: pq.StringArray{
74+
"Recommendation-1",
75+
"Recommendation-2",
76+
},
77+
ReferenceLinks: pq.StringArray{
78+
"ReferenceLink-1",
79+
"ReferenceLink-2",
80+
},
81+
},
82+
Labels: []string{
83+
"Label-1",
84+
"Label-2",
85+
},
86+
},
87+
Target: store.TargetTeams{
88+
Target: store.Target{
89+
ID: "TargetID-1",
90+
Identifier: "Identifier-1",
91+
},
92+
Teams: []string{
93+
"Team-1",
94+
"Team-2",
95+
},
96+
},
97+
Source: store.Source{
98+
ID: "SourceID-1",
99+
Instance: "SourceInstance-1",
100+
Options: "SourceOptions-1",
101+
SourceFamily: store.SourceFamily{
102+
Name: "SourceName-1",
103+
Component: "SourceComponent-1",
104+
},
105+
},
106+
Resources: store.Resources{
107+
{
108+
Name: "ResourceName-1",
109+
Attributes: []string{
110+
"Attr-1",
111+
"Attr-2",
112+
},
113+
Resources: []map[string]string{
114+
{
115+
"Attr-1": "1",
116+
"Attr-2": "2",
117+
},
118+
},
119+
},
120+
},
121+
TotalExposure: 10,
122+
CurrentExposure: 5,
123+
},
124+
Tag: "tag-1",
125+
},
54126
want: &sns.PublishInput{
55-
Message: aws.String(`{"a":"b"}`),
127+
Message: aws.String(removeSpaces(
128+
`{
129+
"id": "FindingID-1",
130+
"affected_resource": "AffectedResource-1",
131+
"score": 9,
132+
"status": "OPEN",
133+
"details": "Details-1",
134+
"impact_details": "ImpactDetails-1",
135+
"issue": {
136+
"id": "IssueID-1",
137+
"summary": "Summary-1",
138+
"cwe_id": 1,
139+
"description": "Description-1",
140+
"recommendations": [
141+
"Recommendation-1",
142+
"Recommendation-2"
143+
],
144+
"reference_links": [
145+
"ReferenceLink-1",
146+
"ReferenceLink-2"
147+
],
148+
"labels": [
149+
"Label-1",
150+
"Label-2"
151+
]
152+
},
153+
"target": {
154+
"id": "TargetID-1",
155+
"identifier": "Identifier-1",
156+
"teams": [
157+
"Team-1",
158+
"Team-2"
159+
]
160+
},
161+
"source": {
162+
"id": "SourceID-1",
163+
"instance": "SourceInstance-1",
164+
"options": "SourceOptions-1",
165+
"time": "0001-01-01T00:00:00Z",
166+
"name": "SourceName-1",
167+
"component": "SourceComponent-1"
168+
},
169+
"resources": [
170+
{
171+
"name": "ResourceName-1",
172+
"attributes": [
173+
"Attr-1",
174+
"Attr-2"
175+
],
176+
"resources": [
177+
{
178+
"Attr-1": "1",
179+
"Attr-2": "2"
180+
}
181+
]
182+
}
183+
],
184+
"total_exposure": 10,
185+
"current_exposure": 5
186+
}`)),
56187
TopicArn: aws.String("arn:aTopic"),
57188
},
58189
},
59190
}
191+
60192
for _, tt := range tests {
61193
t.Run(tt.name, func(t *testing.T) {
62194
s := &SNSNotifier{
63195
conf: tt.fields.conf,
64196
sns: tt.fields.sns,
65197
logger: tt.fields.logger,
66198
}
67-
if err := s.Push(tt.message); (err != nil) != tt.wantErr {
68-
t.Errorf("SNSNotifier.Push() error = %v, wantErr %v", err, tt.wantErr)
199+
if err := s.PushFinding(tt.f); (err != nil) != tt.wantErr {
200+
t.Errorf("got error: %v but wanted: %v", err, tt.wantErr)
69201
}
70202
diff := cmp.Diff(tt.want, tt.fields.sns.notification, cmpopts.IgnoreUnexported(sns.PublishInput{}))
71203
if diff != "" {
72-
t.Errorf("want!= got. Diffs:%s", diff)
204+
t.Errorf("SNS payload does not match with expected one. Diff:\n%s", diff)
73205
}
74206
})
75207
}
76208
}
209+
210+
func removeSpaces(s string) string {
211+
rr := make([]rune, 0, len(s))
212+
for _, r := range s {
213+
if !unicode.IsSpace(r) {
214+
rr = append(rr, r)
215+
}
216+
}
217+
return string(rr)
218+
}

‎pkg/processor/processor.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,7 @@ func (p *CheckProcessor) notifyFindings(findingsState []store.FindingState, tag
296296

297297
// Notify finding state update
298298
for _, f := range findings {
299-
err = p.notifier.Push(FindingNotification{
299+
err = p.notifier.PushFinding(notify.FindingNotification{
300300
FindingExpanded: f,
301301
Tag: tag,
302302
})

‎pkg/processor/processor_test.go‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,18 +28,17 @@ type mockNotifier struct {
2828
calls uint8
2929
}
3030

31-
func (n *mockNotifier) Push(mssg interface{}) error {
31+
func (n *mockNotifier) PushFinding(f notify.FindingNotification) error {
3232
n.calls++
3333
return nil
3434
}
3535

3636
type inMemMockNotifier struct {
37-
sent []FindingNotification
37+
sent []notify.FindingNotification
3838
}
3939

40-
func (n *inMemMockNotifier) Push(mssg interface{}) error {
41-
notif, _ := mssg.(FindingNotification)
42-
n.sent = append(n.sent, notif)
40+
func (n *inMemMockNotifier) PushFinding(f notify.FindingNotification) error {
41+
n.sent = append(n.sent, f)
4342
return nil
4443
}
4544

@@ -221,7 +220,7 @@ func TestNotifyOpenFindings(t *testing.T) {
221220
name string
222221
fields fields
223222
input input
224-
expectedNotifs []FindingNotification
223+
expectedNotifs []notify.FindingNotification
225224
}{
226225
{
227226
name: "Should send notificatons adding check tag",
@@ -242,7 +241,7 @@ func TestNotifyOpenFindings(t *testing.T) {
242241
},
243242
tag: "test-tag",
244243
},
245-
expectedNotifs: []FindingNotification{
244+
expectedNotifs: []notify.FindingNotification{
246245
{
247246
FindingExpanded: store.FindingExpanded{
248247
Finding: store.Finding{ID: "1"},

‎pkg/processor/types.go‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@ Copyright 2020 Adevinta
44

55
package processor
66

7-
import "github.com/adevinta/vulnerability-db/pkg/store"
8-
97
// Notification is the message received in the queue from the Vulcan Core SNS topic.
108
type Notification struct {
119
Type string `json:"Type"`
@@ -42,19 +40,6 @@ type CheckMessage struct {
4240
Tag string `json:"tag"`
4341
}
4442

45-
// FindingNotification represents a notification associated with a finding state change.
46-
type FindingNotification struct {
47-
store.FindingExpanded
48-
// TODO: Tag field is defined here in order for the FindingNotification struct to be
49-
// backward compatible with the "old" vulnerability notifications format so integrations
50-
// with external components that still use the old representation (e.g.: Hermes) are
51-
// maintained and can be isolated in the notify pkg. Once these integrations are modified
52-
// to use the new notification format we'll have to decide if we make this tag field
53-
// serializable into JSON and keep using the tag field as identifier, or we migrate to use
54-
// the current team identifier.
55-
Tag string `json:"-"`
56-
}
57-
5843
type vulnerability struct {
5944
ID string `json:"id"`
6045
Summary string `json:"summary"`

‎pkg/store/findings.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -84,12 +84,12 @@ type FindingExposure struct {
8484
// and source data.
8585
type FindingExpanded struct {
8686
Finding
87-
TotalExposure int64 `json:"total_exposure"`
88-
CurrentExposure int64 `json:"current_exposure,omitempty"`
89-
Resources Resources `json:"resources"`
9087
Issue IssueLabels `json:"issue"`
9188
Target TargetTeams `json:"target"`
9289
Source Source `json:"source"`
90+
Resources Resources `json:"resources"`
91+
TotalExposure int64 `json:"total_exposure"`
92+
CurrentExposure int64 `json:"current_exposure,omitempty"`
9393
}
9494

9595
// Resources defines the structure of a the resources of a finding.

‎pkg/store/targets.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,14 @@ import "context"
99
// Target represents the target
1010
// scope for a check execution.
1111
type Target struct {
12-
ID string
13-
Identifier string
12+
ID string `json:"id"`
13+
Identifier string `json:"identifier"`
1414
}
1515

1616
// TargetTeams represents a target along with its associated teams.
1717
type TargetTeams struct {
1818
Target
19-
Teams []string
19+
Teams []string `json:"teams"`
2020
}
2121

2222
func (db *psqlxStore) CreateTarget(t Target) (*Target, error) {

0 commit comments

Comments
 (0)