From e0baad92febaadd9e60283b16e4c712dee28c1c4 Mon Sep 17 00:00:00 2001
From: Paulo Nascimento <paulo.nascimento@medwing.com>
Date: Sat, 20 Mar 2021 18:57:27 -0300
Subject: [PATCH] Usignregex to determine if it os ipv6 or not

---
 pkg/iptables/iptables.go | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/pkg/iptables/iptables.go b/pkg/iptables/iptables.go
index 37a935e..c3308d0 100644
--- a/pkg/iptables/iptables.go
+++ b/pkg/iptables/iptables.go
@@ -17,6 +17,7 @@ package iptables
 import (
 	"fmt"
 	"net"
+	"regexp"
 	"sync"
 	"time"
 
@@ -353,6 +354,9 @@ func (c *Controller) Set(rules []Rule) error {
 	c.Lock()
 	defer c.Unlock()
 	var i int
+
+	ipv6Regex,_ := regexp.Compile("[-]d\\s(.*:.*\\s[-]m\\scomment)")
+
 	for ; i < len(rules); i++ {
 		if i < len(c.rules) {
 			if rules[i].String() != c.rules[i].String() {
@@ -362,15 +366,19 @@ func (c *Controller) Set(rules []Rule) error {
 			}
 		}
 		if i >= len(c.rules) {
-			var proto = ProtocolIPv4 //rules[i].Proto()
+			proto := ProtocolIPv4
 
-			var protocolName = "ipv4"
+			ruleString := rules[i].String()
+			if ipv6Regex.MatchString(ruleString) {
+				proto = ProtocolIPv6
+			}
+
+			protocolName := "ipv4"
 
 			if proto == ProtocolIPv6 {
 				protocolName = "ipv6"
 			}
 
-			var ruleString = rules[i].String()
 			level.Debug(c.logger).Log("msg", "Applying Firewall Rule...", "Rule", ruleString, "Protocol", protocolName)
 			if err := rules[i].Add(c.client(proto)); err != nil {
 				return fmt.Errorf("failed to add rule: %v", err)