From f30c25b7e8123027e07dc05c62e8a55ae6b2287f Mon Sep 17 00:00:00 2001
From: Alexandr Dubovikov <alexandr.dubovikov@gmail.com>
Date: Wed, 30 Oct 2024 23:10:49 +0100
Subject: [PATCH] fix: added replace token for incoming collector

---
 config/config.go     |  1 +
 main.go              |  1 +
 publish/publisher.go | 15 +++++++++++++++
 3 files changed, 17 insertions(+)

diff --git a/config/config.go b/config/config.go
index 4419e1b..19c5e4f 100644
--- a/config/config.go
+++ b/config/config.go
@@ -24,6 +24,7 @@ type Config struct {
 	Zip                bool
 	HepCollector       string
 	CollectOnlySip     bool
+	ReplaceToken       bool
 	HepServer          string
 	HepNodePW          string
 	HepNodeID          uint
diff --git a/main.go b/main.go
index 4a950df..fe26b8e 100644
--- a/main.go
+++ b/main.go
@@ -44,6 +44,7 @@ func createFlags() {
 	flag.StringVar(&config.Cfg.HEPBufferFile, "hep-buffer-file", "HEP-Buffer.dump", "filename and location for hep-buffer file")
 	flag.StringVar(&config.Cfg.PrometheusIPPort, "prometheus", ":8090", "prometheus metrics - ip:port. By default all IPs")
 	flag.BoolVar(&config.Cfg.CollectOnlySip, "collectonlysip", false, "collect only sip")
+	flag.BoolVar(&config.Cfg.ReplaceToken, "replacetoken", false, "replace token for collector")
 	flag.BoolVar(&config.Cfg.Reassembly, "tcpassembly", false, "If true, tcpassembly will be enabled")
 	flag.BoolVar(&config.Cfg.SipAssembly, "sipassembly", false, "If true, sipassembly will be enabled")
 	flag.UintVar(&config.Cfg.SendRetries, "tcpsendretries", 0, "Number of retries for sending before giving up and reconnecting")
diff --git a/publish/publisher.go b/publish/publisher.go
index 64f92f4..45a7c07 100644
--- a/publish/publisher.go
+++ b/publish/publisher.go
@@ -70,6 +70,21 @@ func (pub *Publisher) Start(pq chan *decoder.Packet) {
 
 		//Version == 100 just for forwarding...
 		if pkt.Version == 100 {
+
+			if config.Cfg.ReplaceToken {
+				msg, err := DecodeHEP(pkt.Payload)
+				if err == nil {
+					msg.NodePW = config.Cfg.HepNodePW
+					pkt.Payload, err = msg.Marshal()
+					if err != nil {
+						pub.output(pkt.Payload)
+					}
+				} else {
+					logp.Warn("Bad HEP: %v", err)
+				}
+				continue
+			}
+
 			pub.output(pkt.Payload)
 			logp.Debug("publisher", "sent hep message from collector")
 		} else if pkt.Version == 0 {