Merge pull request #5 from myENA/feature/AddedCustomRateLimits

Feature/added custom rate limits

common/common.go

@@ -6,14 +6,27 @@ import (
 	"encoding/binary"
 	"fmt"
 	"github.com/mostlygeek/arp"
+	"gopkg.in/yaml.v2"
 	"log"
 	"math/big"
 	"net"
+	"os"
 	"os/exec"
 	"runtime"
 	"strings"
 )
 
+type CustomDeviceRateLimits struct {
+	RateLimits []CustomDeviceRateLimit `yaml:"rateLimits"`
+}
+
+type CustomDeviceRateLimit struct {
+	DeviceIp  string `yaml:"deviceIp"`
+	RateLimit int    `yaml:"rateLimit"`
+}
+
+var customDeviceRateLimits CustomDeviceRateLimits
+
 func CryptoRandomNumber(max int64) int64 {
 	n, err := crand.Int(crand.Reader, big.NewInt(max))
 	if err != nil {
@@ -109,3 +122,41 @@ func DarwinMACFormat(macString string) string {
 	builder.WriteString(macString[len(macString)-group:])
 	return builder.String()
 }
+
+func InitCustomRateLimits(rateLimitPath string) error {
+	var customRateLimits CustomDeviceRateLimits
+	configFileName := rateLimitPath
+	source, err := os.ReadFile(configFileName)
+	if err != nil {
+		fmt.Println("failed reading custom device rate limits")
+		return err
+	}
+	err = yaml.Unmarshal(source, &customRateLimits)
+	if err != nil {
+		log.Fatalf("error: %v", err)
+		fmt.Println("failed unmarshal the custom rate limits")
+		return err
+	}
+	customDeviceRateLimits = customRateLimits
+	return nil
+}
+
+func HasCustomRateLimit(ip string) (bool, int) {
+	for _, customRate := range customDeviceRateLimits.RateLimits {
+		if strings.HasSuffix(customRate.DeviceIp, "/24") {
+			_, ipv4Net, err := net.ParseCIDR(customRate.DeviceIp)
+			if err != nil {
+				return false, 0
+			}
+			if ipv4Net.Contains(net.ParseIP(ip)) {
+				return true, customRate.RateLimit
+			} else {
+				return false, 0
+			}
+		}
+		if customRate.DeviceIp == ip {
+			return true, customRate.RateLimit
+		}
+	}
+	return false, 0
+}

flowproxy.go

@@ -3,6 +3,7 @@ package main
 import (
 	"flag"
 	"fmt"
+	"github.com/myENA/flowproxy/common"
 	"github.com/myENA/flowproxy/proxy"
 	"github.com/myENA/flowproxy/tproxy"
 	"os"
@@ -75,6 +76,12 @@ func main() {
 	tproxyVerbose := tproxyCmd.Bool("verbose", false, "Whether to log every flow received. "+
 		"Warning can be a lot")
 
+	initCustomRatesErr := common.InitCustomRateLimits("rateLimits.yaml")
+	if initCustomRatesErr != nil {
+		fmt.Println("failed to load the custom rate limits")
+		os.Exit(1)
+	}
+
 	// Start parsing command line args
 	if len(os.Args) < 2 {
 		printHelpHeader()

go.mod

@@ -5,8 +5,11 @@ go 1.21.3
 require (
 	github.com/google/gopacket v1.1.19
 	golang.org/x/time v0.5.0
+	gopkg.in/yaml.v2 v2.4.0
 )
 
+require golang.org/x/net v0.0.0-20190620200207-3b0461eec859 // indirect
+
 require (
 	github.com/mostlygeek/arp v0.0.0-20170424181311-541a2129847a
 	golang.org/x/sys v0.1.0 // indirect

rateLimits.yaml

@@ -0,0 +1,2 @@
+rateLimits:
+  - {deviceIp: "192.168.56.106", rateLimit: 20}

tproxy/netflow-dummy-packets/rateLimitsTestData.yaml

@@ -0,0 +1,2 @@
+rateLimits:
+  - {deviceIp: "192.168.56.106", rateLimit: 20}

tproxy/tproxy.go

@@ -264,6 +264,9 @@ func parseNetflow(ctx context.Context, wg *sync.WaitGroup, proxyChan <-chan gopa
 					srcIP = ip.SrcIP
 				}
 			}
+
+			hasCustomRateLimit, customRateLimit := common.HasCustomRateLimit(srcIP.String())
+
 			udpLayer := packet.Layer(layers.LayerTypeUDP)
 			payload := udpLayer.LayerPayload()
 			ok9, err9 := netflow.IsValidNetFlow(payload, 9)
@@ -275,7 +278,9 @@ func parseNetflow(ctx context.Context, wg *sync.WaitGroup, proxyChan <-chan gopa
 				deviceManager.SeenDevice(srcIP.String())
 			} else {
 				deviceManager.AddDevice(srcIP.String())
-				if rateLimit {
+				if hasCustomRateLimit {
+					deviceManager.SetSampleRate(srcIP.String(), customRateLimit)
+				} else if rateLimit {
 					deviceManager.SetSampleRate(srcIP.String(), rate)
 				}
 			}
@@ -300,7 +305,12 @@ func parseNetflow(ctx context.Context, wg *sync.WaitGroup, proxyChan <-chan gopa
 					dataChan <- packet
 					continue
 				}
-				if rateLimit {
+				if hasCustomRateLimit {
+					if deviceManager.CheckSampleRate(srcIP.String(), int(dataCount)) {
+						rStats.Netflow9.DataSent = rStats.Netflow9.DataSent + dataCount
+						dataChan <- packet
+					}
+				} else if rateLimit {
 					if deviceManager.CheckSampleRate(srcIP.String(), int(dataCount)) {
 						rStats.Netflow9.DataSent = rStats.Netflow9.DataSent + dataCount
 						dataChan <- packet