From 3fd43f0d18a52cc0ee6dab0dcdfc91fa1fc87d1d Mon Sep 17 00:00:00 2001
From: Vladimir Avtsenov <vladimir.lsk.cool@gmail.com>
Date: Fri, 30 Aug 2024 04:40:46 +0300
Subject: [PATCH] ipset integration

---
 README.md |  2 +-
 go.mod    |  1 +
 group.go  | 18 +++++++++++++++++-
 kvas2.go  | 10 +++++++++-
 main.go   |  1 +
 5 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 791ee70..c131320 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,7 @@ Realized features:
 - [x] IPTables rules to remap DNS server [1]
 - [ ] Rule composer
 - [ ] List loading/watching (temporary)
-- [ ] IPSet integration
+- [X] IPSet integration
 - [ ] Listing of interfaces
 - [ ] IPTables rules to IPSet [2]
 - [ ] It is not a concept now... REFACTORING TIME!!!
diff --git a/go.mod b/go.mod
index 897c82d..df2b32f 100644
--- a/go.mod
+++ b/go.mod
@@ -5,4 +5,5 @@ go 1.21
 require (
 	github.com/IGLOU-EU/go-wildcard/v2 v2.0.2
 	github.com/coreos/go-iptables v0.7.0
+	github.com/nadoo/ipset v0.5.0
 )
diff --git a/group.go b/group.go
index 9e40748..5e40932 100644
--- a/group.go
+++ b/group.go
@@ -3,6 +3,7 @@ package main
 import (
 	"errors"
 	"fmt"
+	"github.com/nadoo/ipset"
 	"strconv"
 
 	"kvas2-go/models"
@@ -17,7 +18,8 @@ type GroupOptions struct {
 
 type Group struct {
 	*models.Group
-	options GroupOptions
+	ipsetName string
+	options   GroupOptions
 }
 
 func (g *Group) Enable() error {
@@ -45,6 +47,15 @@ func (g *Group) Enable() error {
 		return errors.New(string(out))
 	}
 
+	err = ipset.Destroy(g.ipsetName)
+	if err != nil {
+		return fmt.Errorf("failed to destroy ipset: %w", err)
+	}
+	err = ipset.Create(g.ipsetName)
+	if err != nil {
+		return fmt.Errorf("failed to create ipset: %w", err)
+	}
+
 	g.options.Enabled = true
 	g.options.FWMark = fwmark
 	g.options.Table = table
@@ -67,6 +78,11 @@ func (g *Group) Disable() error {
 		return errors.New(string(out))
 	}
 
+	err = ipset.Destroy(g.ipsetName)
+	if err != nil {
+		return fmt.Errorf("failed to destroy ipset: %w", err)
+	}
+
 	g.options.Enabled = false
 	g.options.FWMark = 0
 	g.options.Table = 0
diff --git a/kvas2.go b/kvas2.go
index c43645b..768e140 100644
--- a/kvas2.go
+++ b/kvas2.go
@@ -10,6 +10,8 @@ import (
 	"kvas2-go/models"
 	"kvas2-go/pkg/dns-proxy"
 	"kvas2-go/pkg/iptables-helper"
+
+	"github.com/nadoo/ipset"
 )
 
 var (
@@ -20,6 +22,7 @@ var (
 type Config struct {
 	MinimalTTL             time.Duration
 	ChainPostfix           string
+	IpSetPostfix           string
 	TargetDNSServerAddress string
 	ListenPort             uint16
 }
@@ -113,7 +116,8 @@ func (a *App) AppendGroup(group *models.Group) error {
 	}
 
 	a.Groups[group.ID] = &Group{
-		Group: group,
+		Group:     group,
+		ipsetName: fmt.Sprintf("%s%d", a.Config.IpSetPostfix, group.ID),
 	}
 
 	if a.isRunning {
@@ -183,6 +187,10 @@ func (a *App) handleMessage(msg *dnsProxy.Message) {
 func New(config Config) (*App, error) {
 	var err error
 
+	if err = ipset.Init(); err != nil {
+		return nil, fmt.Errorf("failed to initialize ipset: %w", err)
+	}
+
 	app := &App{}
 
 	app.Config = config
diff --git a/main.go b/main.go
index 58f8f98..d08140c 100644
--- a/main.go
+++ b/main.go
@@ -14,6 +14,7 @@ func main() {
 	app, err := New(Config{
 		MinimalTTL:             time.Hour,
 		ChainPostfix:           "KVAS2_",
+		IpSetPostfix:           "kvas2_",
 		TargetDNSServerAddress: "127.0.0.1:53",
 		ListenPort:             7548,
 	})