add downlaod data support

This commit is contained in:
AlteredCoder 2020-05-22 17:11:18 +02:00
parent 2611f7fa23
commit 3425c60c53
3 changed files with 86 additions and 11 deletions

View file

@ -105,7 +105,7 @@ func main() {
log.Infof("Loading grok library")
/* load base regexps for two grok parsers */
parserCTX, err = p.Init(map[string]interface{}{"patterns": cConfig.ConfigFolder + string("/patterns/")})
parserCTX, err = p.Init(map[string]interface{}{"patterns": cConfig.ConfigFolder + string("/patterns/"), "data": cConfig.DataFolder})
if err != nil {
log.Errorf("failed to initialize parser : %v", err)
return
@ -204,11 +204,11 @@ func main() {
for _, scenarios := range CustomScenarios {
bucketFiles = append(bucketFiles, scenarios.Filename)
}
holders, outputEventChan, err = leaky.LoadBuckets(bucketFiles)
holders, outputEventChan, err = leaky.LoadBuckets(bucketFiles, cConfig.DataFolder)
} else {
log.Infof("Loading scenarios")
holders, outputEventChan, err = leaky.Init(map[string]string{"patterns": cConfig.ConfigFolder + "/scenarios/"})
holders, outputEventChan, err = leaky.Init(map[string]string{"patterns": cConfig.ConfigFolder + "/scenarios/", "data": cConfig.DataFolder})
}
if err != nil {
log.Fatalf("Scenario loading failed : %v", err)

View file

@ -51,6 +51,7 @@ type BucketFactory struct {
CacheSize int `yaml:"cache_size"` //CacheSize, if > 0, limits the size of in-memory cache of the bucket
Profiling bool `yaml:"profiling"` //Profiling, if true, will make the bucket record pours/overflows/etc.
OverflowFilter string `yaml:"overflow_filter"` //OverflowFilter if present, is a filter that must return true for the overflow to go through
Data []*dataSource `yaml:"data,omitempty"` //If data are needed for the filter to work (ie downloading file)
BucketName string `yaml:"-"`
Filename string `yaml:"-"`
RunTimeFilter *vm.Program `json:"-"`
@ -101,15 +102,14 @@ func ValidateFactory(b *BucketFactory) error {
/* Init recursively process yaml files from a directory and loads them as BucketFactory */
func Init(cfg map[string]string) ([]BucketFactory, chan types.Event, error) {
return LoadBucketDir(cfg["patterns"])
return LoadBucketDir(cfg["patterns"], cfg["data"])
}
func LoadBuckets(files []string) ([]BucketFactory, chan types.Event, error) {
func LoadBuckets(files []string, dataDir string) ([]BucketFactory, chan types.Event, error) {
var (
ret []BucketFactory = []BucketFactory{}
response chan types.Event
)
var seed namegenerator.Generator = namegenerator.NewNameGenerator(time.Now().UTC().UnixNano())
response = make(chan types.Event, 1)
for _, f := range files {
@ -160,7 +160,7 @@ func LoadBuckets(files []string) ([]BucketFactory, chan types.Event, error) {
g.Filename = filepath.Clean(f)
g.BucketName = seed.Generate()
g.ret = response
err = LoadBucket(&g)
err = LoadBucket(&g, dataDir)
if err != nil {
log.Errorf("Failed to load bucket : %v", err)
return nil, nil, fmt.Errorf("loadBucket failed : %v", err)
@ -172,7 +172,7 @@ func LoadBuckets(files []string) ([]BucketFactory, chan types.Event, error) {
return ret, response, nil
}
func LoadBucketDir(dir string) ([]BucketFactory, chan types.Event, error) {
func LoadBucketDir(dir string, dataDir string) ([]BucketFactory, chan types.Event, error) {
var (
filenames []string
)
@ -183,11 +183,11 @@ func LoadBucketDir(dir string) ([]BucketFactory, chan types.Event, error) {
for _, f := range files {
filenames = append(filenames, dir+f.Name())
}
return LoadBuckets(filenames)
return LoadBuckets(filenames, dataDir)
}
/* Init recursively process yaml files from a directory and loads them as BucketFactory */
func LoadBucket(g *BucketFactory) error {
func LoadBucket(g *BucketFactory, dataDir string) error {
var err error
if g.Debug {
var clog = logrus.New()
@ -222,7 +222,7 @@ func LoadBucket(g *BucketFactory) error {
if g.Filter == "" {
g.logger.Warningf("Bucket without filter, abort.")
return fmt.Errorf("bucket without filter directive.")
return fmt.Errorf("bucket doesn't have filter")
}
g.RunTimeFilter, err = expr.Compile(g.Filter, expr.Env(exprhelpers.GetExprEnv(map[string]interface{}{"evt": &types.Event{}})))
if err != nil {
@ -236,6 +236,12 @@ func LoadBucket(g *BucketFactory) error {
}
}
if len(g.Data) > 0 {
if err := getData(g.Data, dataDir); err != nil {
return fmt.Errorf("unable to download data: %s", err)
}
}
g.logger.Infof("Adding %s bucket", g.Type)
//return the Holder correponding to the type of bucket
g.processors = []Processor{}

69
pkg/leakybucket/utils.go Normal file
View file

@ -0,0 +1,69 @@
package leakybucket
import (
"fmt"
"io/ioutil"
"net/http"
"os"
"path"
log "github.com/sirupsen/logrus"
)
type dataSource struct {
SourceURL string `yaml:"source_url"`
DestPath string `yaml:"dest_file"`
}
func downloadFile(url string, destPath string) error {
log.Debugf("downloading %s in %s", url, destPath)
req, err := http.NewRequest("GET", url, nil)
if err != nil {
return err
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
return err
}
if resp.StatusCode != 200 {
return fmt.Errorf("download response 'HTTP %d' : %s", resp.StatusCode, string(body))
}
file, err := os.OpenFile(destPath, os.O_RDWR|os.O_CREATE, 0644)
if err != nil {
return err
}
_, err = file.WriteString(string(body))
if err != nil {
return err
}
err = file.Sync()
if err != nil {
return err
}
return nil
}
func getData(data []*dataSource, dataDir string) error {
for _, dataS := range data {
destPath := path.Join(dataDir, dataS.DestPath)
log.Infof("downloading data '%s' in '%s'", dataS.SourceURL, destPath)
err := downloadFile(dataS.SourceURL, destPath)
if err != nil {
return err
}
}
return nil
}