-
Notifications
You must be signed in to change notification settings - Fork 1
/
elastic.go
84 lines (71 loc) · 1.56 KB
/
elastic.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
package vior
import (
"context"
"log"
"strings"
"github.com/olivere/elastic"
)
// ElasticStorage is an example of the storage for CSP reports,
// that implements DataStorage interface
type ElasticStorage struct {
Pipe chan *ExtReport
URL string
Client *elastic.Client
IdxName string
DocType string
Ctx context.Context
}
// NewElasticStorage bootstraps and initializes ElasticStorage
func NewElasticStorage(url string, idxname string, doctype string) (*ElasticStorage, error) {
e := &ElasticStorage{
Pipe: make(chan *ExtReport),
URL: url,
IdxName: idxname,
DocType: doctype,
Ctx: context.Background(),
}
if err := e.Init(); err != nil {
return nil, err
}
return e, nil
}
// Init initializes Elastic client, creates index and starts
// goroutine that pops reports from incomming channel
func (e *ElasticStorage) Init() error {
ec, err := elastic.NewClient(
elastic.SetURL(e.URL),
)
if err != nil {
return err
}
e.Client = ec
_, err = e.Client.CreateIndex(e.IdxName).Do(e.Ctx)
if err != nil {
if !strings.Contains(err.Error(), "already_exists_exception") {
return err
}
}
go func() {
for rep := range e.Pipe {
go e.Save(rep)
}
}()
return nil
}
// Save saves the report in Elastic
func (e *ElasticStorage) Save(r *ExtReport) error {
_, err := e.Client.Index().
Index(e.IdxName).
Type(e.DocType).
BodyJson(*r).
Do(e.Ctx)
if err != nil {
log.Println(err)
return err
}
return nil
}
// GetPipe returns a pipe to write reports to
func (e *ElasticStorage) GetPipe() chan *ExtReport {
return e.Pipe
}