From 127969d325a6b05f9ab098f46f1cb402ec6998e9 Mon Sep 17 00:00:00 2001
From: Sebastien Blot <sebastien@crowdsec.net>
Date: Mon, 15 Apr 2024 22:06:59 +0200
Subject: [PATCH] add geoip enrich expr helpers

---
 cmd/crowdsec/crowdsec.go       |   8 +++
 pkg/exprhelpers/expr_lib.go    |  23 +++++++
 pkg/exprhelpers/geoip.go       |  63 ++++++++++++++++++
 pkg/exprhelpers/helpers.go     |  30 +++++++++
 pkg/parser/enrich.go           |  21 +-----
 pkg/parser/enrich_date.go      |   6 +-
 pkg/parser/enrich_date_test.go |   2 +-
 pkg/parser/enrich_dns.go       |   6 +-
 pkg/parser/enrich_geoip.go     | 118 ++++++++++++---------------------
 pkg/parser/enrich_unmarshal.go |   6 +-
 pkg/parser/node.go             |   4 +-
 pkg/parser/node_test.go        |   2 +-
 pkg/parser/parsing_test.go     |   6 +-
 pkg/parser/runtime.go          |   2 +-
 pkg/parser/unix_parser.go      |   2 +-
 15 files changed, 184 insertions(+), 115 deletions(-)
 create mode 100644 pkg/exprhelpers/geoip.go

diff --git a/cmd/crowdsec/crowdsec.go b/cmd/crowdsec/crowdsec.go
index f604af1de..ef99f64d3 100644
--- a/cmd/crowdsec/crowdsec.go
+++ b/cmd/crowdsec/crowdsec.go
@@ -19,6 +19,7 @@ import (
 	"github.com/crowdsecurity/crowdsec/pkg/appsec"
 	"github.com/crowdsecurity/crowdsec/pkg/csconfig"
 	"github.com/crowdsecurity/crowdsec/pkg/cwhub"
+	"github.com/crowdsecurity/crowdsec/pkg/exprhelpers"
 	leaky "github.com/crowdsecurity/crowdsec/pkg/leakybucket"
 	"github.com/crowdsecurity/crowdsec/pkg/parser"
 	"github.com/crowdsecurity/crowdsec/pkg/types"
@@ -32,6 +33,13 @@ func initCrowdsec(cConfig *csconfig.Config, hub *cwhub.Hub) (*parser.Parsers, []
 		return nil, nil, fmt.Errorf("while loading context: %w", err)
 	}
 
+	err = exprhelpers.GeoIPInit(hub.GetDataDir())
+
+	if err != nil {
+		//GeoIP databases are not mandatory, do not make crowdsec fail if they are not present
+		log.Warnf("unable to initialize GeoIP: %s", err)
+	}
+
 	// Start loading configs
 	csParsers := parser.NewParsers(hub)
 	if csParsers, err = parser.LoadParsers(cConfig, csParsers); err != nil {
diff --git a/pkg/exprhelpers/expr_lib.go b/pkg/exprhelpers/expr_lib.go
index 520799287..41102bd5b 100644
--- a/pkg/exprhelpers/expr_lib.go
+++ b/pkg/exprhelpers/expr_lib.go
@@ -1,9 +1,11 @@
 package exprhelpers
 
 import (
+	"net"
 	"time"
 
 	"github.com/crowdsecurity/crowdsec/pkg/cticlient"
+	"github.com/oschwald/geoip2-golang"
 )
 
 type exprCustomFunc struct {
@@ -455,6 +457,27 @@ var exprFuncs = []exprCustomFunc{
 			new(func(string) bool),
 		},
 	},
+	{
+		name:     "GeoIPEnrich",
+		function: GeoIPEnrich,
+		signature: []interface{}{
+			new(func(string) *geoip2.City),
+		},
+	},
+	{
+		name:     "GeoIPASNEnrich",
+		function: GeoIPASNEnrich,
+		signature: []interface{}{
+			new(func(string) *geoip2.ASN),
+		},
+	},
+	{
+		name:     "GeoIPRangeEnrich",
+		function: GeoIPRangeEnrich,
+		signature: []interface{}{
+			new(func(string) *net.IPNet),
+		},
+	},
 }
 
 //go 1.20 "CutPrefix":              strings.CutPrefix,
diff --git a/pkg/exprhelpers/geoip.go b/pkg/exprhelpers/geoip.go
new file mode 100644
index 000000000..fb0c344d8
--- /dev/null
+++ b/pkg/exprhelpers/geoip.go
@@ -0,0 +1,63 @@
+package exprhelpers
+
+import (
+	"net"
+)
+
+func GeoIPEnrich(params ...any) (any, error) {
+	if geoIPCityReader == nil {
+		return nil, nil
+	}
+
+	ip := params[0].(string)
+
+	parsedIP := net.ParseIP(ip)
+
+	city, err := geoIPCityReader.City(parsedIP)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return city, nil
+}
+
+func GeoIPASNEnrich(params ...any) (any, error) {
+	if geoIPASNReader == nil {
+		return nil, nil
+	}
+
+	ip := params[0].(string)
+
+	parsedIP := net.ParseIP(ip)
+	asn, err := geoIPASNReader.ASN(parsedIP)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return asn, nil
+}
+
+func GeoIPRangeEnrich(params ...any) (any, error) {
+	if geoIPRangeReader == nil {
+		return nil, nil
+	}
+
+	ip := params[0].(string)
+
+	var dummy interface{}
+
+	parsedIP := net.ParseIP(ip)
+	rangeIP, ok, err := geoIPRangeReader.LookupNetwork(parsedIP, &dummy)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if !ok {
+		return nil, nil
+	}
+
+	return rangeIP, nil
+}
diff --git a/pkg/exprhelpers/helpers.go b/pkg/exprhelpers/helpers.go
index 79a621c7d..e4ecf0a3a 100644
--- a/pkg/exprhelpers/helpers.go
+++ b/pkg/exprhelpers/helpers.go
@@ -20,6 +20,8 @@ import (
 	"github.com/c-robinson/iplib"
 	"github.com/cespare/xxhash/v2"
 	"github.com/davecgh/go-spew/spew"
+	"github.com/oschwald/geoip2-golang"
+	"github.com/oschwald/maxminddb-golang"
 	"github.com/prometheus/client_golang/prometheus"
 	log "github.com/sirupsen/logrus"
 	"github.com/umahmood/haversine"
@@ -55,6 +57,10 @@ var exprFunctionOptions []expr.Option
 
 var keyValuePattern = regexp.MustCompile(`(?P<key>[^=\s]+)=(?:"(?P<quoted_value>[^"\\]*(?:\\.[^"\\]*)*)"|(?P<value>[^=\s]+)|\s*)`)
 
+var geoIPCityReader *geoip2.Reader
+var geoIPASNReader *geoip2.Reader
+var geoIPRangeReader *maxminddb.Reader
+
 func GetExprOptions(ctx map[string]interface{}) []expr.Option {
 	if len(exprFunctionOptions) == 0 {
 		exprFunctionOptions = []expr.Option{}
@@ -72,6 +78,30 @@ func GetExprOptions(ctx map[string]interface{}) []expr.Option {
 	return ret
 }
 
+func GeoIPInit(datadir string) error {
+	var err error
+
+	geoIPCityReader, err = geoip2.Open(filepath.Join(datadir, "GeoLite2-City.mmdb"))
+	if err != nil {
+		log.Errorf("unable to open GeoLite2-City.mmdb : %s", err)
+		return err
+	}
+
+	geoIPASNReader, err = geoip2.Open(filepath.Join(datadir, "GeoLite2-ASN.mmdb"))
+	if err != nil {
+		log.Errorf("unable to open GeoLite2-ASN.mmdb : %s", err)
+		return err
+	}
+
+	geoIPRangeReader, err = maxminddb.Open(filepath.Join(datadir, "GeoLite2-ASN.mmdb"))
+	if err != nil {
+		log.Errorf("unable to open GeoLite2-ASN.mmdb : %s", err)
+		return err
+	}
+
+	return nil
+}
+
 func Init(databaseClient *database.Client) error {
 	dataFile = make(map[string][]string)
 	dataFileRegex = make(map[string][]*regexp.Regexp)
diff --git a/pkg/parser/enrich.go b/pkg/parser/enrich.go
index 5180b9a5f..661410d20 100644
--- a/pkg/parser/enrich.go
+++ b/pkg/parser/enrich.go
@@ -7,7 +7,7 @@ import (
 )
 
 /* should be part of a package shared with enrich/geoip.go */
-type EnrichFunc func(string, *types.Event, interface{}, *log.Entry) (map[string]string, error)
+type EnrichFunc func(string, *types.Event, *log.Entry) (map[string]string, error)
 type InitFunc func(map[string]string) (interface{}, error)
 
 type EnricherCtx struct {
@@ -16,59 +16,42 @@ type EnricherCtx struct {
 
 type Enricher struct {
 	Name       string
-	InitFunc   InitFunc
 	EnrichFunc EnrichFunc
-	Ctx        interface{}
 }
 
 /* mimic plugin loading */
-func Loadplugin(path string) (EnricherCtx, error) {
+func Loadplugin() (EnricherCtx, error) {
 	enricherCtx := EnricherCtx{}
 	enricherCtx.Registered = make(map[string]*Enricher)
 
-	enricherConfig := map[string]string{"datadir": path}
-
 	EnrichersList := []*Enricher{
 		{
 			Name:       "GeoIpCity",
-			InitFunc:   GeoIPCityInit,
 			EnrichFunc: GeoIpCity,
 		},
 		{
 			Name:       "GeoIpASN",
-			InitFunc:   GeoIPASNInit,
 			EnrichFunc: GeoIpASN,
 		},
 		{
 			Name:       "IpToRange",
-			InitFunc:   IpToRangeInit,
 			EnrichFunc: IpToRange,
 		},
 		{
 			Name:       "reverse_dns",
-			InitFunc:   reverseDNSInit,
 			EnrichFunc: reverse_dns,
 		},
 		{
 			Name:       "ParseDate",
-			InitFunc:   parseDateInit,
 			EnrichFunc: ParseDate,
 		},
 		{
 			Name:       "UnmarshalJSON",
-			InitFunc:   unmarshalInit,
 			EnrichFunc: unmarshalJSON,
 		},
 	}
 
 	for _, enricher := range EnrichersList {
-		log.Debugf("Initiating enricher '%s'", enricher.Name)
-		pluginCtx, err := enricher.InitFunc(enricherConfig)
-		if err != nil {
-			log.Errorf("unable to register plugin '%s': %v", enricher.Name, err)
-			continue
-		}
-		enricher.Ctx = pluginCtx
 		log.Infof("Successfully registered enricher '%s'", enricher.Name)
 		enricherCtx.Registered[enricher.Name] = enricher
 	}
diff --git a/pkg/parser/enrich_date.go b/pkg/parser/enrich_date.go
index 20828af90..748a466d7 100644
--- a/pkg/parser/enrich_date.go
+++ b/pkg/parser/enrich_date.go
@@ -56,7 +56,7 @@ func GenDateParse(date string) (string, time.Time) {
 	return "", time.Time{}
 }
 
-func ParseDate(in string, p *types.Event, x interface{}, plog *log.Entry) (map[string]string, error) {
+func ParseDate(in string, p *types.Event, plog *log.Entry) (map[string]string, error) {
 
 	var ret = make(map[string]string)
 	var strDate string
@@ -105,7 +105,3 @@ func ParseDate(in string, p *types.Event, x interface{}, plog *log.Entry) (map[s
 
 	return ret, nil
 }
-
-func parseDateInit(cfg map[string]string) (interface{}, error) {
-	return nil, nil
-}
diff --git a/pkg/parser/enrich_date_test.go b/pkg/parser/enrich_date_test.go
index 084ded525..085ef5ca3 100644
--- a/pkg/parser/enrich_date_test.go
+++ b/pkg/parser/enrich_date_test.go
@@ -48,7 +48,7 @@ func TestDateParse(t *testing.T) {
 	for _, tt := range tests {
 		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
-			strTime, err := ParseDate(tt.evt.StrTime, &tt.evt, nil, logger)
+			strTime, err := ParseDate(tt.evt.StrTime, &tt.evt, logger)
 			cstest.RequireErrorContains(t, err, tt.expectedErr)
 			if tt.expectedErr != "" {
 				return
diff --git a/pkg/parser/enrich_dns.go b/pkg/parser/enrich_dns.go
index f622e6c35..1ff5b0f4f 100644
--- a/pkg/parser/enrich_dns.go
+++ b/pkg/parser/enrich_dns.go
@@ -11,7 +11,7 @@ import (
 /* All plugins must export a list of function pointers for exported symbols */
 //var ExportedFuncs = []string{"reverse_dns"}
 
-func reverse_dns(field string, p *types.Event, ctx interface{}, plog *log.Entry) (map[string]string, error) {
+func reverse_dns(field string, p *types.Event, plog *log.Entry) (map[string]string, error) {
 	ret := make(map[string]string)
 	if field == "" {
 		return nil, nil
@@ -25,7 +25,3 @@ func reverse_dns(field string, p *types.Event, ctx interface{}, plog *log.Entry)
 	ret["reverse_dns"] = rets[0]
 	return ret, nil
 }
-
-func reverseDNSInit(cfg map[string]string) (interface{}, error) {
-	return nil, nil
-}
diff --git a/pkg/parser/enrich_geoip.go b/pkg/parser/enrich_geoip.go
index 0a263c827..f5474ee98 100644
--- a/pkg/parser/enrich_geoip.go
+++ b/pkg/parser/enrich_geoip.go
@@ -6,53 +6,53 @@ import (
 	"strconv"
 
 	"github.com/oschwald/geoip2-golang"
-	"github.com/oschwald/maxminddb-golang"
 	log "github.com/sirupsen/logrus"
 
+	"github.com/crowdsecurity/crowdsec/pkg/exprhelpers"
 	"github.com/crowdsecurity/crowdsec/pkg/types"
 )
 
-func IpToRange(field string, p *types.Event, ctx interface{}, plog *log.Entry) (map[string]string, error) {
-	var dummy interface{}
-	ret := make(map[string]string)
-
-	if field == "" {
-		return nil, nil
-	}
-	ip := net.ParseIP(field)
-	if ip == nil {
-		plog.Infof("Can't parse ip %s, no range enrich", field)
-		return nil, nil
-	}
-	net, ok, err := ctx.(*maxminddb.Reader).LookupNetwork(ip, &dummy)
-	if err != nil {
-		plog.Errorf("Failed to fetch network for %s : %v", ip.String(), err)
-		return nil, nil
-	}
-	if !ok {
-		plog.Debugf("Unable to find range of %s", ip.String())
-		return nil, nil
-	}
-	ret["SourceRange"] = net.String()
-	return ret, nil
-}
-
-func GeoIpASN(field string, p *types.Event, ctx interface{}, plog *log.Entry) (map[string]string, error) {
-	ret := make(map[string]string)
+func IpToRange(field string, p *types.Event, plog *log.Entry) (map[string]string, error) {
 	if field == "" {
 		return nil, nil
 	}
 
-	ip := net.ParseIP(field)
-	if ip == nil {
-		plog.Infof("Can't parse ip %s, no ASN enrich", ip)
-		return nil, nil
-	}
-	record, err := ctx.(*geoip2.Reader).ASN(ip)
+	r, err := exprhelpers.GeoIPRangeEnrich(field)
+
 	if err != nil {
 		plog.Errorf("Unable to enrich ip '%s'", field)
 		return nil, nil //nolint:nilerr
 	}
+
+	if r == nil {
+		plog.Warnf("No range found for ip '%s'", field)
+		return nil, nil //nolint:nilerr
+	}
+
+	record := r.(*net.IPNet)
+
+	ret := make(map[string]string)
+	ret["SourceRange"] = record.String()
+
+	return ret, nil
+}
+
+func GeoIpASN(field string, p *types.Event, plog *log.Entry) (map[string]string, error) {
+	if field == "" {
+		return nil, nil
+	}
+
+	r, err := exprhelpers.GeoIPASNEnrich(field)
+
+	if err != nil {
+		plog.Errorf("Unable to enrich ip '%s'", field)
+		return nil, nil //nolint:nilerr
+	}
+
+	record := r.(*geoip2.ASN)
+
+	ret := make(map[string]string)
+
 	ret["ASNNumber"] = fmt.Sprintf("%d", record.AutonomousSystemNumber)
 	ret["ASNumber"] = fmt.Sprintf("%d", record.AutonomousSystemNumber)
 	ret["ASNOrg"] = record.AutonomousSystemOrganization
@@ -62,21 +62,21 @@ func GeoIpASN(field string, p *types.Event, ctx interface{}, plog *log.Entry) (m
 	return ret, nil
 }
 
-func GeoIpCity(field string, p *types.Event, ctx interface{}, plog *log.Entry) (map[string]string, error) {
-	ret := make(map[string]string)
+func GeoIpCity(field string, p *types.Event, plog *log.Entry) (map[string]string, error) {
 	if field == "" {
 		return nil, nil
 	}
-	ip := net.ParseIP(field)
-	if ip == nil {
-		plog.Infof("Can't parse ip %s, no City enrich", ip)
-		return nil, nil
-	}
-	record, err := ctx.(*geoip2.Reader).City(ip)
+
+	r, err := exprhelpers.GeoIPEnrich(field)
+
 	if err != nil {
-		plog.Debugf("Unable to enrich ip '%s'", ip)
+		plog.Errorf("Unable to enrich ip '%s'", field)
 		return nil, nil //nolint:nilerr
 	}
+
+	record := r.(*geoip2.City)
+	ret := make(map[string]string)
+
 	if record.Country.IsoCode != "" {
 		ret["IsoCode"] = record.Country.IsoCode
 		ret["IsInEU"] = strconv.FormatBool(record.Country.IsInEuropeanUnion)
@@ -88,7 +88,7 @@ func GeoIpCity(field string, p *types.Event, ctx interface{}, plog *log.Entry) (
 		ret["IsInEU"] = strconv.FormatBool(record.RepresentedCountry.IsInEuropeanUnion)
 	} else {
 		ret["IsoCode"] = ""
-		ret["IsInEU"] = strconv.FormatBool(false)
+		ret["IsInEU"] = "false"
 	}
 
 	ret["Latitude"] = fmt.Sprintf("%f", record.Location.Latitude)
@@ -98,33 +98,3 @@ func GeoIpCity(field string, p *types.Event, ctx interface{}, plog *log.Entry) (
 
 	return ret, nil
 }
-
-func GeoIPCityInit(cfg map[string]string) (interface{}, error) {
-	dbCityReader, err := geoip2.Open(cfg["datadir"] + "/GeoLite2-City.mmdb")
-	if err != nil {
-		log.Debugf("couldn't open geoip : %v", err)
-		return nil, err
-	}
-
-	return dbCityReader, nil
-}
-
-func GeoIPASNInit(cfg map[string]string) (interface{}, error) {
-	dbASReader, err := geoip2.Open(cfg["datadir"] + "/GeoLite2-ASN.mmdb")
-	if err != nil {
-		log.Debugf("couldn't open geoip : %v", err)
-		return nil, err
-	}
-
-	return dbASReader, nil
-}
-
-func IpToRangeInit(cfg map[string]string) (interface{}, error) {
-	ipToRangeReader, err := maxminddb.Open(cfg["datadir"] + "/GeoLite2-ASN.mmdb")
-	if err != nil {
-		log.Debugf("couldn't open geoip : %v", err)
-		return nil, err
-	}
-
-	return ipToRangeReader, nil
-}
diff --git a/pkg/parser/enrich_unmarshal.go b/pkg/parser/enrich_unmarshal.go
index dce9c75d4..7ff91b70a 100644
--- a/pkg/parser/enrich_unmarshal.go
+++ b/pkg/parser/enrich_unmarshal.go
@@ -8,7 +8,7 @@ import (
 	"github.com/crowdsecurity/crowdsec/pkg/types"
 )
 
-func unmarshalJSON(field string, p *types.Event, ctx interface{}, plog *log.Entry) (map[string]string, error) {
+func unmarshalJSON(field string, p *types.Event, plog *log.Entry) (map[string]string, error) {
 	err := json.Unmarshal([]byte(p.Line.Raw), &p.Unmarshaled)
 	if err != nil {
 		plog.Errorf("could not unmarshal JSON: %s", err)
@@ -17,7 +17,3 @@ func unmarshalJSON(field string, p *types.Event, ctx interface{}, plog *log.Entr
 	plog.Tracef("unmarshaled JSON: %+v", p.Unmarshaled)
 	return nil, nil
 }
-
-func unmarshalInit(cfg map[string]string) (interface{}, error) {
-	return nil, nil
-}
diff --git a/pkg/parser/node.go b/pkg/parser/node.go
index fe5432ce9..2a2198c3f 100644
--- a/pkg/parser/node.go
+++ b/pkg/parser/node.go
@@ -64,7 +64,7 @@ type Node struct {
 	Data      []*types.DataSource `yaml:"data,omitempty"`
 }
 
-func (n *Node) validate(pctx *UnixParserCtx, ectx EnricherCtx) error {
+func (n *Node) validate(ectx EnricherCtx) error {
 
 	//stage is being set automagically
 	if n.Stage == "" {
@@ -563,7 +563,7 @@ func (n *Node) compile(pctx *UnixParserCtx, ectx EnricherCtx) error {
 		return fmt.Errorf("Node is empty")
 	}
 
-	if err := n.validate(pctx, ectx); err != nil {
+	if err := n.validate(ectx); err != nil {
 		return err
 	}
 
diff --git a/pkg/parser/node_test.go b/pkg/parser/node_test.go
index d85aa82a8..be12176bf 100644
--- a/pkg/parser/node_test.go
+++ b/pkg/parser/node_test.go
@@ -56,7 +56,7 @@ func TestParserConfigs(t *testing.T) {
 			t.Fatalf("Compile: (%d/%d) expected error", idx+1, len(CfgTests))
 		}
 
-		err = CfgTests[idx].NodeCfg.validate(pctx, EnricherCtx{})
+		err = CfgTests[idx].NodeCfg.validate(EnricherCtx{})
 		if CfgTests[idx].Valid == true && err != nil {
 			t.Fatalf("Valid: (%d/%d) expected valid, got : %s", idx+1, len(CfgTests), err)
 		}
diff --git a/pkg/parser/parsing_test.go b/pkg/parser/parsing_test.go
index 04d08cc27..a18aaa6fa 100644
--- a/pkg/parser/parsing_test.go
+++ b/pkg/parser/parsing_test.go
@@ -152,7 +152,11 @@ func prepTests() (*UnixParserCtx, EnricherCtx, error) {
 
 	//Load enrichment
 	datadir := "./test_data/"
-	ectx, err = Loadplugin(datadir)
+	err = exprhelpers.GeoIPInit(datadir)
+	if err != nil {
+		log.Fatalf("unable to initialize GeoIP: %s", err)
+	}
+	ectx, err = Loadplugin()
 	if err != nil {
 		log.Fatalf("failed to load plugin geoip : %v", err)
 	}
diff --git a/pkg/parser/runtime.go b/pkg/parser/runtime.go
index afdf88dc8..1596ef5ff 100644
--- a/pkg/parser/runtime.go
+++ b/pkg/parser/runtime.go
@@ -155,7 +155,7 @@ func (n *Node) ProcessStatics(statics []ExtraField, event *types.Event) error {
 			/*still way too hackish, but : inject all the results in enriched, and */
 			if enricherPlugin, ok := n.EnrichFunctions.Registered[static.Method]; ok {
 				clog.Tracef("Found method '%s'", static.Method)
-				ret, err := enricherPlugin.EnrichFunc(value, event, enricherPlugin.Ctx, n.Logger.WithField("method", static.Method))
+				ret, err := enricherPlugin.EnrichFunc(value, event, n.Logger.WithField("method", static.Method))
 				if err != nil {
 					clog.Errorf("method '%s' returned an error : %v", static.Method, err)
 				}
diff --git a/pkg/parser/unix_parser.go b/pkg/parser/unix_parser.go
index 720bac3d1..280d122ec 100644
--- a/pkg/parser/unix_parser.go
+++ b/pkg/parser/unix_parser.go
@@ -117,7 +117,7 @@ func LoadParsers(cConfig *csconfig.Config, parsers *Parsers) (*Parsers, error) {
 	*/
 	log.Infof("Loading enrich plugins")
 
-	parsers.EnricherCtx, err = Loadplugin(cConfig.ConfigPaths.DataDir)
+	parsers.EnricherCtx, err = Loadplugin()
 	if err != nil {
 		return parsers, fmt.Errorf("failed to load enrich plugin : %v", err)
 	}