Merge pull request #93 from DCSO/unicorn

add testdata submission for flow reports

Author: satta

README.md

@@ -213,7 +213,7 @@ FEVER can optionally inject in-band test data into downstream submissions, such
    }
    ```
    will be created and forwarded.
-*  For passive DNS observation submissions, use the `pdns.test-domain` config item to insert a dummy entry for that domain, e.g. for `pdns.tests-domain` set to `heartbeat.fever-heartbeat`:
+* For passive DNS observation submissions, use the `pdns.test-domain` config item to insert a dummy entry for that domain, e.g. for `pdns.test-domain` set to `heartbeat.fever-heartbeat`:
    ```json
    {
      "timestamp_start": "2021-12-07T18:18:00.029197078Z",
@@ -230,11 +230,34 @@ FEVER can optionally inject in-band test data into downstream submissions, such
            }
          ]
        },
-   ...
+      ...
      }
    }
    ```
-
+*  For flow report submission, use the `flowreport.testdata*` config items to insert a dummy flow for that specific IPs and ports, e.g. for :
+   ```yaml
+   flowreport:
+     # ...
+     testdata-srcip: 0.0.0.1
+     testdata-destip: 0.0.0.2
+     testdata-destport: 99999
+   ```
+   we would get
+   ```json
+   {
+      "sensor-id": "XXX",
+      "time-start": "2021-12-08T13:53:36.442182896+01:00",
+      "time-end": "2021-12-08T13:53:46.490743527+01:00",
+      "tuples": {
+          "0.0.0.1_0.0.0.2_99999": {
+              "count": 1,
+              "total_bytes_toclient": 23,
+              "total_bytes_toserver": 42
+          }
+      },
+      ...
+   }
+   ```
 
 ## Author/Contact
 

cmd/fever/cmds/run.go

@@ -441,6 +441,12 @@ func mainfunc(cmd *cobra.Command, args []string) {
 			}).Info("compression of flow stats")
 		}
 		ua := processing.MakeUnicornAggregator(submitter, unicornSleep, dummyMode)
+		testSrcIP := viper.GetString("flowreport.testdata-srcip")
+		testDestIP := viper.GetString("flowreport.testdata-destip")
+		testDestPort := viper.GetInt64("flowreport.testdata-destport")
+		if testSrcIP != "" && testDestIP != "" {
+			ua.EnableTestFlow(testSrcIP, testDestIP, testDestPort)
+		}
 		dispatcher.RegisterHandler(ua)
 		ua.Run()
 		defer func() {

fever.yaml

@@ -80,6 +80,11 @@ flowreport:
   submission-exchange: aggregations
   # Set to true to disable gzip compression for uploads.
   nocompress: false
+  # If both srcip and destip are non-empty, inject an extra flow record for
+  # these towards the given destination port.
+  #testdata-srcip: 0.0.0.1
+  #testdata-destip: 0.0.0.2
+  #testdata-destport: 99999
 
 # Configuration for metrics (i.e. InfluxDB) submission.
 metrics:

processing/unicorn_aggregator.go

@@ -6,6 +6,7 @@ package processing
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"os"
 	"strconv"
 	"sync"
@@ -41,6 +42,9 @@ type UnicornAggregator struct {
 	StringBuf            bytes.Buffer
 	UnicornTuplesMutex   sync.RWMutex `json:"-"`
 	UnicornProxyMapMutex sync.RWMutex `json:"-"`
+	TestFlowSrcIP        string
+	TestFlowDestIP       string
+	TestFlowDestPort     int64
 }
 
 // MakeUnicornAggregate creates a new empty UnicornAggregate object.
@@ -59,12 +63,13 @@ func MakeUnicornAggregator(statsSubmitter util.StatsSubmitter,
 		Logger: log.WithFields(log.Fields{
 			"domain": "aggregate",
 		}),
-		Submitter:    statsSubmitter,
-		DummyMode:    dummyMode,
-		SubmitPeriod: submitPeriod,
-		CloseChan:    make(chan bool),
-		ClosedChan:   make(chan bool),
-		Aggregate:    *MakeUnicornAggregate(),
+		Submitter:        statsSubmitter,
+		DummyMode:        dummyMode,
+		SubmitPeriod:     submitPeriod,
+		CloseChan:        make(chan bool),
+		ClosedChan:       make(chan bool),
+		Aggregate:        *MakeUnicornAggregate(),
+		TestFlowDestPort: 99999,
 	}
 	return a
 }
@@ -86,6 +91,16 @@ func (a *UnicornAggregator) stop() {
 }
 
 func (a *UnicornAggregator) submit(submitter util.StatsSubmitter, dummyMode bool) {
+	if a.TestFlowSrcIP != "" && a.TestFlowDestIP != "" {
+		// Inject test flow into aggregation
+		a.CountFlowTuple(
+			fmt.Sprintf("%s_%s_%d", a.TestFlowSrcIP,
+				a.TestFlowDestIP, a.TestFlowDestPort),
+			23,
+			42,
+			20, // count 20 to ensure some limits are met downstream
+		)
+	}
 	// Lock the current measurements for submission. Since this is a blocking
 	// operation, we don't want this to depend on how long submitter.Submit()
 	// takes but keep it independent of that. Hence we take the time to create
@@ -122,14 +137,16 @@ func (a *UnicornAggregator) submit(submitter util.StatsSubmitter, dummyMode bool
 
 }
 
-// CountFlowTuple increments the flow tuple counter for the given key.
+// CountFlowTuple increments the flow tuple counter for the given key. If addCnt
+// is >1, then the caller is responsible for providing the correct (sub-total)
+// counts for bytestoclient and bytestoserver.
 func (a *UnicornAggregator) CountFlowTuple(key string, bytestoclient int64,
-	bytestoserver int64) {
+	bytestoserver int64, addCnt int64) {
 	a.UnicornTuplesMutex.Lock()
 	if _, ok := a.Aggregate.FlowTuples[key]; !ok {
 		a.Aggregate.FlowTuples[key] = make(map[string]int64)
 	}
-	a.Aggregate.FlowTuples[key]["count"]++
+	a.Aggregate.FlowTuples[key]["count"] += addCnt
 	a.Aggregate.FlowTuples[key]["total_bytes_toclient"] += bytestoclient
 	a.Aggregate.FlowTuples[key]["total_bytes_toserver"] += bytestoserver
 	a.UnicornTuplesMutex.Unlock()
@@ -187,7 +204,7 @@ func (a *UnicornAggregator) Consume(e *types.Entry) error {
 		a.StringBuf.Write([]byte("_"))
 		a.StringBuf.Write([]byte(strconv.Itoa(int(e.DestPort))))
 		a.CountFlowTuple(a.StringBuf.String(), e.BytesToClient,
-			e.BytesToServer)
+			e.BytesToServer, 1)
 		a.StringBuf.Reset()
 	}
 
@@ -210,3 +227,10 @@ func (a *UnicornAggregator) GetName() string {
 func (a *UnicornAggregator) GetEventTypes() []string {
 	return []string{"http", "flow"}
 }
+
+// EnableTestFlow adds a dummy flow with the given specs to each aggregation
+func (a *UnicornAggregator) EnableTestFlow(srcip, dstip string, dstport int64) {
+	a.TestFlowSrcIP = srcip
+	a.TestFlowDestIP = dstip
+	a.TestFlowDestPort = dstport
+}

processing/unicorn_aggregator_test.go

@@ -184,6 +184,56 @@ func TestUnicornAggregator(t *testing.T) {
 	}
 }
 
+func TestUnicornAggregatorWithTestdata(t *testing.T) {
+	rand.Seed(time.Now().UTC().UnixNano())
+	dsub := &testSubmitter{
+		Data: make([]string, 0),
+	}
+	f := MakeUnicornAggregator(dsub, 500*time.Millisecond, false)
+	f.EnableTestFlow("1.2.3.4", "5.6.7.8", 33333)
+	f.Run()
+
+	for {
+		if dsub.GetTotalAggs() < 1 {
+			log.Debug(dsub.GetTotalAggs())
+			time.Sleep(100 * time.Millisecond)
+		} else {
+			break
+		}
+	}
+
+	consumeWaitChan := make(chan bool)
+	f.Stop(consumeWaitChan)
+	<-consumeWaitChan
+
+	var d UnicornAggregate
+
+	err := json.Unmarshal([]byte(dsub.Data[0]), &d)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v, ok := d.FlowTuples["1.2.3.4_5.6.7.8_33333"]; ok {
+		if val, ok := v["count"]; ok {
+			if val != 20 {
+				t.Fatalf("wrong value: %v", val)
+			}
+		}
+		if val, ok := v["total_bytes_toclient"]; ok {
+			if val != 23 {
+				t.Fatalf("wrong value: %v", val)
+			}
+		}
+		if val, ok := v["total_bytes_toserver"]; ok {
+			if val != 42 {
+				t.Fatalf("wrong value: %v", val)
+			}
+		}
+	} else {
+		t.Fatalf("missing key in map: %v", d.FlowTuples)
+	}
+
+}
+
 func TestUnicornAggregatorWithDispatch(t *testing.T) {
 	rand.Seed(time.Now().UTC().UnixNano())
 	dsub := &testSubmitter{