add downlaod data support
This commit is contained in:
parent
2611f7fa23
commit
3425c60c53
3 changed files with 86 additions and 11 deletions
|
@ -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)
|
||||
|
|
|
@ -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
69
pkg/leakybucket/utils.go
Normal 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
|
||||
}
|
Loading…
Add table
Reference in a new issue