ips option

memodi · Dec 3, 2021 · d719205 · d719205
1 parent 3b5422b
commit d719205
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 17 deletions.
diff --git a/nflow-generator.go b/nflow-generator.go
@@ -10,6 +10,7 @@ import (
 	"math/rand"
 	"net"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/jessevdk/go-flags"
@@ -37,11 +38,11 @@ var opts struct {
 	CollectorPort string `short:"p" long:"port" description:"port number of the target netflow collector"`
 	SpikeProto    string `short:"s" long:"spike" description:"run a second thread generating a spike for the specified protocol"`
 	FalseIndex    bool   `short:"f" long:"false-index" description:"generate false SNMP interface indexes, otherwise set to 0"`
+	IPs           string `short:"i" long:"ips" description:"uses specific list of ips, comma separated"`
 	Help          bool   `short:"h" long:"help" description:"show nflow-generator help"`
 }
 
 func main() {
-
 	_, err := flags.Parse(&opts)
 	if err != nil {
 		showUsage()
@@ -66,7 +67,14 @@ func main() {
 	}
 	log.Infof("sending netflow data to a collector ip: %s and port: %s",
 		opts.CollectorIP, opts.CollectorPort)
-
+	var ips []string
+	if len(opts.IPs) > 0 {
+		ips = strings.Split(opts.IPs, ",")
+		log.Info("specified ips:")
+		for _, ip := range ips {
+			log.Infof("%s", ip)
+		}
+	}
 	for {
 		rand.Seed(time.Now().Unix())
 		n := randomNum(50, 1000)
@@ -75,14 +83,14 @@ func main() {
 			GenerateSpike()
 		}
 		if n > 900 {
-			data := GenerateNetflow(8)
+			data := GenerateNetflow(8, ips)
 			buffer := BuildNFlowPayload(data)
 			_, err := conn.Write(buffer.Bytes())
 			if err != nil {
 				log.Fatal("Error connecting to the target collector: ", err)
 			}
 		} else {
-			data := GenerateNetflow(16)
+			data := GenerateNetflow(16, ips)
 			buffer := BuildNFlowPayload(data)
 			_, err := conn.Write(buffer.Bytes())
 			if err != nil {
@@ -130,6 +138,7 @@ Application Options:
         p2p - generates udp/6681
         bittorrent - generates udp/6682
   -f, --false-index generate a false snmp index values of 1 or 2. The default is 0. (Optional)
+  -i, --ips uses specific list of ips, comma separated (Optional)
 
 Example Usage:
 
@@ -139,6 +148,9 @@ Example Usage:
     -generate default flows to device 172.16.86.138, port 9995
     ./nflow-generator -t 172.16.86.138 -p 9995 
 
+    -generate default flows between ips 172.16.86.1, 172.16.86.2, 172.16.86.3 to device 172.16.86.138, port 9995
+    ./nflow-generator -t 172.16.86.138 -p 9995 -i 172.16.86.1,172.16.86.2,172.16.86.3
+
     -generate default flows along with a spike in the specified protocol:
     ./nflow-generator -t 172.16.86.138 -p 9995 -s ssh
 

diff --git a/nflow_payload.go b/nflow_payload.go
@@ -94,17 +94,27 @@ func BuildNFlowPayload(data Netflow) bytes.Buffer {
 }
 
 //Generate a netflow packet w/ user-defined record count
-func GenerateNetflow(recordCount int) Netflow {
+func GenerateNetflow(recordCount int, ips []string) Netflow {
 	data := new(Netflow)
 	header := CreateNFlowHeader(recordCount)
-	records := []NetflowPayload{}
+	var records []NetflowPayload
 	if recordCount == 8 {
 		// overwrite payload to add some variations for traffic spikes.
 		records = CreateVariablePayload(recordCount)
 	} else {
 		records = CreateNFlowPayload(recordCount)
 	}
 
+	//override ips from list if specified
+	if len(ips) > 0 {
+		rand.Seed(time.Now().Unix())
+		for i := 0; i < len(records); i++ {
+			records[i].SrcIP = IPtoUint32(ips[rand.Int()%len(ips)])
+			records[i].DstIP = IPtoUint32(ips[rand.Int()%len(ips)])
+			records[i].NextHopIP = IPtoUint32(ips[rand.Int()%len(ips)])
+		}
+	}
+
 	data.Header = header
 	data.Records = records
 	return *data
@@ -116,7 +126,7 @@ func CreateNFlowHeader(recordCount int) NetflowHeader {
 	t := time.Now().UnixNano()
 	sec := t / int64(time.Second)
 	nsec := t - sec*int64(time.Second)
-	sysUptime = uint32((t-StartTime) / int64(time.Millisecond))+1000
+	sysUptime = uint32((t-StartTime)/int64(time.Millisecond)) + 1000
 	flowSequence++
 
 	// log.Infof("Time: %d; Seconds: %d; Nanoseconds: %d\n", t, sec, nsec)
@@ -528,13 +538,13 @@ func CreateRandomFlow() NetflowPayload {
 }
 
 // patch up the common fields of the packets
-func FillCommonFields (
-		payload *NetflowPayload, 
-		numPktOct int, 
-		ipProtocol int, 
-		srcPrefixMask int) NetflowPayload {
+func FillCommonFields(
+	payload *NetflowPayload,
+	numPktOct int,
+	ipProtocol int,
+	srcPrefixMask int) NetflowPayload {
 
-// Fill template with values not filled by caller
+	// Fill template with values not filled by caller
 	// payload.SrcIP = IPtoUint32("10.154.20.12")
 	// payload.DstIP = IPtoUint32("77.12.190.94")
 	// payload.NextHopIP = IPtoUint32("150.20.145.1")
@@ -557,10 +567,10 @@ func FillCommonFields (
 	payload.Padding2 = 0
 
 	// now handle computed values
-	if !opts.FalseIndex {                       // default interfaces are zero
+	if !opts.FalseIndex { // default interfaces are zero
 		payload.SnmpInIndex = 0
 		payload.SnmpOutIndex = 0
-	} else if payload.SrcIP > payload.DstIP {   // false-index
+	} else if payload.SrcIP > payload.DstIP { // false-index
 		payload.SnmpInIndex = 1
 		payload.SnmpOutIndex = 2
 	} else {
@@ -569,8 +579,8 @@ func FillCommonFields (
 	}
 
 	uptime := int(sysUptime)
-	payload.SysUptimeEnd = uint32(uptime - randomNum(10,500))
-	payload.SysUptimeStart = payload.SysUptimeEnd - uint32(randomNum(10,500))
+	payload.SysUptimeEnd = uint32(uptime - randomNum(10, 500))
+	payload.SysUptimeStart = payload.SysUptimeEnd - uint32(randomNum(10, 500))
 
 	// log.Infof("S&D : %x %x %d, %d", payload.SrcIP, payload.DstIP, payload.DstPort, payload.SnmpInIndex)
 	// log.Infof("Time: %d %d %d", sysUptime, payload.SysUptimeStart, payload.SysUptimeEnd)