allow configuration via TOML

build-rpm.sh

@@ -5,6 +5,7 @@ VERSION=`git describe --tag | sed 's/^v//'`
 ITERATION=1
 BINARY=munchclient
 DEFAULTS_FILE="$BINARY"
+CONFIG_FILE="$BINARY.toml"
 SERVICE_FILE="$BINARY.service"
 PKG_TYPE=rpm
 DESCR="ScraperWall traffic collector"
@@ -17,19 +18,17 @@ make
 
 install -v -m 755 $BINARY $DESTDIR/usr/bin/
 install -v -m 644 defaults/$DEFAULTS_FILE $DESTDIR/etc/default/
-#install -v -m 755 init.d/$BINARY $DESTDIR/etc/init.d/
-
-#install -m 644 $SERVICE_FILE $DESTDIR/etc/systemd/system/
-#install -m 644 $MAXMIND_CITY_DB $DESTDIR/usr/share/munch/
-#install -m 644 $MAXMIND_ASN_DB $DESTDIR/usr/share/munch/
+install -v -m 644 munchclient.toml /etc/
 
 
 fpm -s dir -t $PKG_TYPE -C $DESTDIR --name $BINARY \
   --version $VERSION \
   --iteration $ITERATION \
   --description "$DESCR" \
-  --config-files etc/default/$BINARY \
-  --rpm-init init.d/$BINARY \
+  --config-files "etc/default/$BINARY" \
+  --config-files "etc/$CONFIG_FILE" \
+  --rpm-init "init.d/$BINARY" \
+  -p rmps \
 	--rpm-sign
   # --deb-systemd $SERVICE_FILE \
 

init.d/munchclient

@@ -25,8 +25,8 @@ exec=/usr/bin/munchclient
 lockfile=/var/lock/subsys/$prog
 
 # Source config
-if [ -f /etc/sysconfig/$prog ] ; then
-    . /etc/sysconfig/$prog
+if [ -f /etc/default/$prog ] ; then
+    . /etc/default/$prog
 fi
 
 start() {
@@ -35,7 +35,7 @@ start() {
 	umask 077
 
         echo -n $"Starting ScraperWall request collector: "
-        daemon --pidfile="$PIDFILE" $exec -i "$PIDFILE" $SYSLOGD_OPTIONS
+        daemon --pidfile="$PIDFILE" $exec -config /etc/munchclient.toml -pidfile "$PIDFILE"
         RETVAL=$?
         echo
         [ $RETVAL -eq 0 ] && touch $lockfile

main.go

@@ -13,9 +13,11 @@ import (
 	"strings"
 	"time"
 
+	"github.com/BurntSushi/toml"
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/pcap"
+	"github.com/kr/pretty"
 	"github.com/nats-io/nats"
 	"github.com/nats-io/nats/encoders/protobuf"
 
@@ -36,12 +38,14 @@ var (
 	requestsFile          = flag.String("requests", "", "CSV file containing requests (IP and URL)")
 	useXForwardedAsSource = flag.Bool("use-x-forwarded", false, "Use the IP address in X-Forwarded-For as source")
 	pidFile               = flag.String("pidfile", "/var/run/munchclient.pid", "The location of the PID file")
+	configFile            = flag.String("config", "", "The location of the TOML config file")
 	doVersion             = flag.Bool("version", false, "Show version information")
 
 	natsEC  *nats.EncodedConn
 	count   uint64
 	timeout = -1 * time.Second
 	ipPriv  *ip.IP
+	config  Config
 
 	// Version contains the program Version, e.g. 1.0.1
 	Version string
@@ -50,23 +54,72 @@ var (
 	BuildDate string
 )
 
+// Config contains the program configuration
+type Config struct {
+	Live                  bool
+	Interface             string
+	SnapshotLen           int
+	Filter                string
+	Promiscuous           bool
+	NatsURL               string
+	NatsQueue             string
+	SleepFor              time.Duration
+	RequestsFile          string
+	UseXForwardedAsSource bool
+	PidFile               string
+}
+
 func init() {
 	flag.Parse()
 
 	nats.RegisterEncoder(protobuf.PROTOBUF_ENCODER, &protobuf.ProtobufEncoder{})
 }
 
+func loadConfig() {
+
+	// initialize with values from the command line / environment
+	config.Live = *doLiveCapture
+	config.Interface = *iface
+	config.SnapshotLen = *snapshotLen
+	config.Filter = *filter
+	config.Promiscuous = *promiscuous
+	config.NatsURL = *natsURL
+	config.NatsQueue = *natsQueue
+	config.SleepFor = *sleepFor
+	config.RequestsFile = *requestsFile
+	config.UseXForwardedAsSource = *useXForwardedAsSource
+	config.PidFile = *pidFile
+
+	if *configFile == "" {
+		return
+	}
+
+	_, err := os.Stat(*configFile)
+	if err != nil {
+		log.Printf("%s: %s\n", *configFile, err)
+		return
+	}
+
+	if _, err = toml.DecodeFile(*configFile, &config); err != nil {
+		log.Printf("%s: %s\n", *configFile, err)
+	}
+
+	pretty.Println(config)
+}
+
 func main() {
 	if *doVersion {
 		version()
 		os.Exit(0)
 	}
 
-	if err := pidfile.Write(*pidFile); err != nil {
-		log.Fatal("munchclient is already running, exiting!")
+	loadConfig()
+
+	if err := pidfile.Write(config.PidFile); err != nil {
+		log.Fatalf("munchclient: %s\n", err)
 		os.Exit(1)
 	}
-	defer pidfile.Remove(*pidFile)
+	defer pidfile.Remove(config.PidFile)
 
 	go func(c *uint64) {
 		for {
@@ -78,11 +131,11 @@ func main() {
 
 	// NATS
 	//
-	if *natsURL == "" {
+	if config.NatsURL == "" {
 		log.Fatal("No NATS URL specified (-nats-url)!")
 	}
 
-	natsConn, err := nats.Connect(*natsURL)
+	natsConn, err := nats.Connect(config.NatsURL)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -93,10 +146,10 @@ func main() {
 	}
 
 	// What should I do?
-	if *requestsFile != "" {
+	if config.RequestsFile != "" {
 		replayFile()
-	} else if *doLiveCapture {
-		fmt.Printf("live capture (%s, %s) to %s\n", *iface, *filter, *natsURL)
+	} else if config.Live {
+		fmt.Printf("live capture (%s, %s) to %s\n", config.Interface, config.Filter, config.NatsURL)
 		liveCapture()
 	}
 }
@@ -106,13 +159,13 @@ func liveCapture() {
 
 	// PCAP setup
 	//
-	handle, err := pcap.OpenLive(*iface, int32(*snapshotLen), *promiscuous, timeout)
+	handle, err := pcap.OpenLive(config.Interface, int32(config.SnapshotLen), config.Promiscuous, timeout)
 	if err != nil {
 		log.Fatal(err)
 	}
 	defer handle.Close()
 
-	err = handle.SetBPFFilter(*filter)
+	err = handle.SetBPFFilter(config.Filter)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -210,7 +263,7 @@ func processPacket(packet gopacket.Packet) {
 	}
 	data.Source = data.IpSrc
 
-	if *useXForwardedAsSource && data.XForwardedFor != "" {
+	if config.UseXForwardedAsSource && data.XForwardedFor != "" {
 		if strings.Contains(data.XForwardedFor, ",") {
 			ips := strings.Split(data.XForwardedFor, ",")
 			for i := len(ips) - 1; i >= 0; i-- {
@@ -234,7 +287,7 @@ func processPacket(packet gopacket.Packet) {
 		data.Source = data.XRealIP
 	}
 
-	natsEC.Publish(*natsQueue, &data)
+	natsEC.Publish(config.NatsQueue, &data)
 }
 
 // replayFile takes a file containing a list of requests (SourceIP Url) and queues the requests
@@ -247,16 +300,16 @@ func replayFile() {
 	var endTs time.Time
 
 	for {
-		fh, err := os.Open(*requestsFile)
+		fh, err := os.Open(config.RequestsFile)
 		if err != nil {
-			log.Fatalf("Failed to open request file '%s': %s", *requestsFile, err)
+			log.Fatalf("Failed to open request file '%s': %s", config.RequestsFile, err)
 		}
 
 		c := csv.NewReader(fh)
 		c.Comma = ' '
 
 		for {
-			if *sleepFor > time.Nanosecond {
+			if config.SleepFor > time.Nanosecond {
 				startTs = time.Now()
 			}
 
@@ -278,13 +331,13 @@ func replayFile() {
 			req.Host = "www.scraperwall.com"
 			req.CreatedAt = time.Now().UnixNano()
 
-			natsEC.Publish(*natsQueue, &req)
+			natsEC.Publish(config.NatsQueue, &req)
 
 			count++
-			if *sleepFor >= time.Nanosecond {
+			if config.SleepFor >= time.Nanosecond {
 				endTs = time.Now()
-				if endTs.Before(startTs.Add(*sleepFor)) {
-					time.Sleep(*sleepFor - endTs.Sub(startTs))
+				if endTs.Before(startTs.Add(config.SleepFor)) {
+					time.Sleep(config.SleepFor - endTs.Sub(startTs))
 				}
 			}
 		}

munchclient.toml

@@ -0,0 +1,11 @@
+Live = true
+Interface = "eth0"
+SnapshotLen = 8192
+Filter = "tcp port 80 and dst 10.1.1.1"
+Promiscuous = false
+NatsURL = "nats://127.0.0.1:4222"
+NatsQueue = "requests"
+UseXForwardedAsSource = true
+PidFile = "/tmp/munchclient.pid"
+# SleepFor = 100µs
+# RequestsFile = requests.csv