فهرست منبع

Download datafile (#895)

* add the ability to download datafile on cscli hub upgrade on files are missing
* fix stuff + lint
* fix error management

Co-authored-by: sabban <15465465+sabban@users.noreply.github.com>
Manuel Sabban 3 سال پیش
والد
کامیت
4dbbd4b3c4
3فایلهای تغییر یافته به همراه52 افزوده شده و 8 حذف شده
  1. 4 0
      cmd/crowdsec-cli/utils.go
  2. 47 8
      pkg/cwhub/download.go
  3. 1 0
      pkg/parser/enrich_geoip.go

+ 4 - 0
cmd/crowdsec-cli/utils.go

@@ -223,7 +223,11 @@ func UpgradeConfig(itemType string, name string, force bool) {
 		found = true
 		if v.UpToDate {
 			log.Infof("%s : up-to-date", v.Name)
+
 			if !force {
+				if err = cwhub.DownloadDataIfNeeded(csConfig.Hub, v, false); err != nil {
+					log.Fatalf("%s : download failed : %v", v.Name, err)
+				}
 				continue
 			}
 		}

+ 47 - 8
pkg/cwhub/download.go

@@ -3,6 +3,7 @@ package cwhub
 import (
 	"bytes"
 	"crypto/sha256"
+	"path"
 	"path/filepath"
 
 	//"errors"
@@ -134,7 +135,7 @@ func DownloadItem(hub *csconfig.Hub, target Item, overwrite bool) (Item, error)
 		}
 		if target.UpToDate {
 			log.Debugf("%s : up-to-date, not updated", target.Name)
-			return target, nil
+			//  We still have to check if data files are present
 		}
 	}
 	req, err := http.NewRequest("GET", fmt.Sprintf(RawFileURLTemplate, HubBranch, target.RemotePath), nil)
@@ -204,7 +205,35 @@ func DownloadItem(hub *csconfig.Hub, target Item, overwrite bool) (Item, error)
 	target.Tainted = false
 	target.UpToDate = true
 
-	dec := yaml.NewDecoder(bytes.NewReader(body))
+	if err = downloadData(dataFolder, overwrite, bytes.NewReader(body)); err != nil {
+		return target, errors.Wrapf(err, "while downloading data for %s", target.FileName)
+	}
+
+	hubIdx[target.Type][target.Name] = target
+	return target, nil
+}
+
+func DownloadDataIfNeeded(hub *csconfig.Hub, target Item, force bool) error {
+	var (
+		dataFolder = hub.DataDir
+		itemFile   *os.File
+		err        error
+	)
+	itemFilePath := fmt.Sprintf("%s/%s", hub.HubDir, target.RemotePath)
+
+	if itemFile, err = os.Open(itemFilePath); err != nil {
+		return errors.Wrapf(err, "while opening %s", itemFilePath)
+	}
+	if err = downloadData(dataFolder, force, itemFile); err != nil {
+		return errors.Wrapf(err, "while downloading data for %s", itemFilePath)
+	}
+	return nil
+}
+
+func downloadData(dataFolder string, force bool, reader io.Reader) error {
+	var err error
+	dec := yaml.NewDecoder(reader)
+
 	for {
 		data := &types.DataSet{}
 		err = dec.Decode(data)
@@ -212,14 +241,24 @@ func DownloadItem(hub *csconfig.Hub, target Item, overwrite bool) (Item, error)
 			if err == io.EOF {
 				break
 			} else {
-				return target, errors.Wrap(err, "while reading file")
+				return errors.Wrap(err, "while reading file")
 			}
 		}
-		err = types.GetData(data.Data, dataFolder)
-		if err != nil {
-			return target, errors.Wrap(err, "while getting data")
+
+		download := false
+		if !force {
+			for _, dataS := range data.Data {
+				if _, err := os.Stat(path.Join(dataFolder, dataS.DestPath)); os.IsNotExist(err) {
+					download = true
+				}
+			}
+		}
+		if download || force {
+			err = types.GetData(data.Data, dataFolder)
+			if err != nil {
+				return errors.Wrap(err, "while getting data")
+			}
 		}
 	}
-	hubIdx[target.Type][target.Name] = target
-	return target, nil
+	return nil
 }

+ 1 - 0
pkg/parser/enrich_geoip.go

@@ -94,6 +94,7 @@ func GeoIpCity(field string, p *types.Event, ctx interface{}) (map[string]string
 func GeoIpInit(cfg map[string]string) (interface{}, error) {
 	var ctx GeoIpEnricherCtx
 	var err error
+
 	ctx.dbc, err = geoip2.Open(cfg["datadir"] + "/GeoLite2-City.mmdb")
 	if err != nil {
 		log.Debugf("couldn't open geoip : %v", err)