Merge pull request #90 from DCSO/alertify-add-fields

Ensure that alertified events also contain added fields

Author: satta

cmd/fever/cmds/alertify.go

@@ -140,7 +140,11 @@ func alertify(cmd *cobra.Command, args []string) {
 	limit := viper.GetUint("alert-limit")
 	extrakey := viper.GetString("extra-key")
 
+	addFields := viper.GetStringMapString("add-fields")
 	a := makeAlertifyAlertifier(prefix, extrakey)
+	if err := a.SetAddedFields(addFields); err != nil {
+		log.Fatal(err)
+	}
 	for e := range eventChan {
 		err := emitAlertsForEvent(a, e, ioc, os.Stdout, uint64(limit))
 		if err != nil {

db/slurper_ejdb.go

@@ -1,3 +1,4 @@
+//go:build ignore
 // +build ignore
 
 package db

processing/forward_handler.go

@@ -5,7 +5,6 @@ package processing
 
 import (
 	"crypto/tls"
-	"fmt"
 	"sync"
 	"time"
 
@@ -103,31 +102,11 @@ func (fh *ForwardHandler) EnableRDNS(expiryPeriod time.Duration) {
 // AddFields enables the addition of a custom set of top-level fields to the
 // forwarded JSON.
 func (fh *ForwardHandler) AddFields(fields map[string]string) error {
-	j := ""
-	// We preprocess the JSON to be able to only use fast string operations
-	// later. This code progressively builds a JSON snippet by adding JSON
-	// key-value pairs for each added field, e.g. `, "foo":"bar"`.
-	for k, v := range fields {
-		// Escape the fields to make sure we do not mess up the JSON when
-		// encountering weird symbols in field names or values.
-		kval, err := util.EscapeJSON(k)
-		if err != nil {
-			fh.Logger.Warningf("cannot escape value: %s", v)
-			return err
-		}
-		vval, err := util.EscapeJSON(v)
-		if err != nil {
-			fh.Logger.Warningf("cannot escape value: %s", v)
-			return err
-		}
-		j += fmt.Sprintf(",%s:%s", kval, vval)
+	addedFields, err := util.PreprocessAddedFields(fields)
+	if err != nil {
+		return err
 	}
-	// We finish the list of key-value pairs with a final brace:
-	// `, "foo":"bar"}`. This string can now just replace the final brace in a
-	// given JSON string. If there were no added fields, we just leave the
-	// output at the final brace.
-	j += "}"
-	fh.AddedFields = j
+	fh.AddedFields = addedFields
 	return nil
 }
 

util/add_fields_preprocess.go

@@ -0,0 +1,36 @@
+package util
+
+import (
+	"fmt"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// PreprocessAddedFields preprocesses the added fields to be able to only use
+// fast string operations to add them to JSON text later. This code
+// progressively builds a JSON snippet by adding JSON key-value pairs for each
+// added field, e.g. `, "foo":"bar"`.
+func PreprocessAddedFields(fields map[string]string) (string, error) {
+	j := ""
+	for k, v := range fields {
+		// Escape the fields to make sure we do not mess up the JSON when
+		// encountering weird symbols in field names or values.
+		kval, err := EscapeJSON(k)
+		if err != nil {
+			log.Warningf("cannot escape value: %s", v)
+			return "", err
+		}
+		vval, err := EscapeJSON(v)
+		if err != nil {
+			log.Warningf("cannot escape value: %s", v)
+			return "", err
+		}
+		j += fmt.Sprintf(",%s:%s", kval, vval)
+	}
+	// We finish the list of key-value pairs with a final brace:
+	// `, "foo":"bar"}`. This string can now just replace the final brace in a
+	// given JSON string. If there were no added fields, we just leave the
+	// output at the final brace.
+	j += "}"
+	return j, nil
+}

util/add_fields_preprocess_test.go

@@ -0,0 +1,45 @@
+package util
+
+import "testing"
+
+func TestPreprocessAddedFields(t *testing.T) {
+	type args struct {
+		fields map[string]string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    string
+		wantErr bool
+	}{
+		{
+			name: "empty fieldset",
+			args: args{
+				fields: map[string]string{},
+			},
+			want: "}",
+		},
+		{
+			name: "fieldset present",
+			args: args{
+				fields: map[string]string{
+					"foo": "bar",
+					"baz": "quux",
+				},
+			},
+			want: `,"foo":"bar","baz":"quux"}`,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := PreprocessAddedFields(tt.args.fields)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("PreprocessAddedFields() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("PreprocessAddedFields() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

util/alertifier.go

@@ -32,6 +32,7 @@ type AlertJSONProvider interface {
 type Alertifier struct {
 	alertPrefix   string
 	extraModifier ExtraModifier
+	addedFields   string
 	matchTypes    map[string]AlertJSONProvider
 }
 
@@ -66,6 +67,17 @@ func (a *Alertifier) SetExtraModifier(em ExtraModifier) {
 	a.extraModifier = em
 }
 
+// SetAddedFields adds string key-value pairs to be added as extra JSON
+// values.
+func (a *Alertifier) SetAddedFields(fields map[string]string) error {
+	af, err := PreprocessAddedFields(fields)
+	if err != nil {
+		return err
+	}
+	a.addedFields = af
+	return nil
+}
+
 // MakeAlert generates a new Entry representing an `alert` event based on the
 // given input metadata event. It uses the information from the Alertifier as
 // well as the given IoC to craft an `alert` sub-object in the resulting
@@ -141,6 +153,14 @@ func (a *Alertifier) MakeAlert(inputEvent types.Entry, ioc string,
 	if err != nil {
 		return nil, err
 	}
+	// Append added fields string, if present
+	if len(a.addedFields) > 1 {
+		j := l
+		jlen := len(j)
+		j = j[:jlen-1]
+		j = append(j, a.addedFields...)
+		l = j
+	}
 	// update returned entry
 	newEntry.Timestamp = eventTimestampFormatted
 	newEntry.JSONLine = string(l)