From 0b43814c0754bd77366608abfbb72607b545c0cb Mon Sep 17 00:00:00 2001
From: Vladimir Avtsenov <vladimir.lsk.cool@gmail.com>
Date: Thu, 5 Sep 2024 05:42:10 +0300
Subject: [PATCH] iptables rules to ipset

---
 README.md |  2 +-
 group.go  | 63 +++++++++++++++++++++++++++++++++++++++++++++++++++++--
 kvas2.go  |  6 +++---
 3 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 26cc8ab..59afc7f 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,7 @@ Realized features:
 - [x] IPTables rules for rebind DNS server port [1]
 - [X] IPSet integration
 - [X] IP integration
-- [ ] IPTables rules to IPSet [2]
+- [X] IPTables rules to IPSet [2]
 - [ ] Rule composer (CRUD)
 - [ ] GORM integration
 - [X] Listing of interfaces
diff --git a/group.go b/group.go
index 71a05fd..098cd77 100644
--- a/group.go
+++ b/group.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"strconv"
 	"time"
 
 	"kvas2-go/models"
@@ -55,7 +56,7 @@ DomainSearch:
 	return nil
 }
 
-func (g *Group) Enable() error {
+func (g *Group) Enable(a *App) error {
 	if g.options.Enabled {
 		return nil
 	}
@@ -150,18 +151,76 @@ func (g *Group) Enable() error {
 		return fmt.Errorf("failed to create ipset: %w", err)
 	}
 
+	preroutingChainName := fmt.Sprintf("%sROUTING_%d_PREROUTING", a.Config.ChainPostfix, g.ID)
+
+	err = a.IPTables.ClearChain("mangle", preroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to clear chain: %w", err)
+	}
+
+	err = a.IPTables.AppendUnique("mangle", preroutingChainName, "-m", "set", "--match-set", g.ipsetName, "dst", "-j", "MARK", "--set-mark", strconv.Itoa(int(mark)))
+	if err != nil {
+		return fmt.Errorf("failed to create rule: %w", err)
+	}
+
+	err = a.IPTables.AppendUnique("mangle", "PREROUTING", "-j", preroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to linking chain: %w", err)
+	}
+
+	postroutingChainName := fmt.Sprintf("%sROUTING_%d_POSTROUTING", a.Config.ChainPostfix, g.ID)
+
+	err = a.IPTables.ClearChain("nat", postroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to clear chain: %w", err)
+	}
+
+	err = a.IPTables.AppendUnique("nat", postroutingChainName, "-o", g.Interface, "-j", "MASQUERADE")
+	if err != nil {
+		return fmt.Errorf("failed to create rule: %w", err)
+	}
+
+	err = a.IPTables.AppendUnique("nat", "POSTROUTING", "-j", postroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to linking chain: %w", err)
+	}
+
 	g.options.Enabled = true
 
 	return nil
 }
 
-func (g *Group) Disable() error {
+func (g *Group) Disable(a *App) error {
 	if !g.options.Enabled {
 		return nil
 	}
 
 	var err error
 
+	preroutingChainName := fmt.Sprintf("%sROUTING_%d_PREROUTING", a.Config.ChainPostfix, g.ID)
+
+	err = a.IPTables.DeleteIfExists("mangle", "PREROUTING", "-j", preroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to unlinking chain: %w", err)
+	}
+
+	err = a.IPTables.ClearAndDeleteChain("mangle", preroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to delete chain: %w", err)
+	}
+
+	postroutingChainName := fmt.Sprintf("%sROUTING_%d_POSTROUTING", a.Config.ChainPostfix, g.ID)
+
+	err = a.IPTables.DeleteIfExists("nat", "POSTROUTING", "-j", postroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to unlinking chain: %w", err)
+	}
+
+	err = a.IPTables.ClearAndDeleteChain("nat", postroutingChainName)
+	if err != nil {
+		return fmt.Errorf("failed to delete chain: %w", err)
+	}
+
 	if g.options.ipRule != nil {
 		err = netlink.RuleDel(g.options.ipRule)
 		if err != nil {
diff --git a/kvas2.go b/kvas2.go
index da2c250..f9108db 100644
--- a/kvas2.go
+++ b/kvas2.go
@@ -94,7 +94,7 @@ func (a *App) Listen(ctx context.Context) []error {
 	}
 
 	for idx, _ := range a.Groups {
-		err = a.Groups[idx].Enable()
+		err = a.Groups[idx].Enable(a)
 		if err != nil {
 			handleError(fmt.Errorf("failed to enable group: %w", err))
 			return errs
@@ -113,7 +113,7 @@ func (a *App) Listen(ctx context.Context) []error {
 	}
 
 	for idx, _ := range a.Groups {
-		err = a.Groups[idx].Disable()
+		err = a.Groups[idx].Disable(a)
 		if err != nil {
 			handleError(fmt.Errorf("failed to disable group: %w", err))
 			return errs
@@ -146,7 +146,7 @@ func (a *App) AppendGroup(group *models.Group) error {
 	}
 
 	if a.isRunning {
-		err := a.Groups[group.ID].Enable()
+		err := a.Groups[group.ID].Enable(a)
 		if err != nil {
 			return fmt.Errorf("failed to enable appended group: %w", err)
 		}