From 572bea6a718be862bfcd6729e510d95e0ad37a6f Mon Sep 17 00:00:00 2001
From: Viktor Liu <viktor@netbird.io>
Date: Tue, 24 Dec 2024 11:18:58 +0100
Subject: [PATCH] Simplify debug code

---
 client/server/debug.go | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/client/server/debug.go b/client/server/debug.go
index a33c65279..50c62ea9d 100644
--- a/client/server/debug.go
+++ b/client/server/debug.go
@@ -24,6 +24,7 @@ import (
 	"google.golang.org/protobuf/encoding/protojson"
 
 	"github.com/netbirdio/netbird/client/anonymize"
+	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -368,18 +369,7 @@ func (s *Server) addFirewallStats(req *proto.DebugBundleRequest, anonymizer *ano
 	}
 
 	if req.GetAnonymize() {
-		for _, stat := range stats {
-			if stat.SourceIP != nil {
-				if ip, ok := netip.AddrFromSlice(stat.SourceIP); ok {
-					stat.SourceIP = anonymizer.AnonymizeIP(ip).AsSlice()
-				}
-			}
-			if stat.DestIP != nil {
-				if ip, ok := netip.AddrFromSlice(stat.DestIP); ok {
-					stat.DestIP = anonymizer.AnonymizeIP(ip).AsSlice()
-				}
-			}
-		}
+		anonymizeStatIPs(stats, anonymizer)
 	}
 
 	jsonBytes, err := json.MarshalIndent(stats, "", "  ")
@@ -978,3 +968,18 @@ func anonymizeSlice(v []any, anonymizer *anonymize.Anonymizer) []any {
 	}
 	return v
 }
+
+func anonymizeStatIPs(stats []*firewall.FlowStats, anonymizer *anonymize.Anonymizer) {
+	for _, stat := range stats {
+		if stat.SourceIP != nil {
+			if ip, ok := netip.AddrFromSlice(stat.SourceIP); ok {
+				stat.SourceIP = anonymizer.AnonymizeIP(ip).AsSlice()
+			}
+		}
+		if stat.DestIP != nil {
+			if ip, ok := netip.AddrFromSlice(stat.DestIP); ok {
+				stat.DestIP = anonymizer.AnonymizeIP(ip).AsSlice()
+			}
+		}
+	}
+}