From 25b3c678f3c75162b0dcbb3798fbd02fec7e94f3 Mon Sep 17 00:00:00 2001
From: Vladimir Avtsenov <vladimir.lsk.cool@gmail.com>
Date: Sat, 24 Aug 2024 21:26:49 +0300
Subject: [PATCH] iptables rules to remap DNS server

---
 README.md              |  2 +-
 dns-proxy/dns-proxy.go | 39 ++++++++++++++++++++++++++++++++++++++-
 go.mod                 |  4 +++-
 main.go                | 24 ++++++++++++++++++++----
 4 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 11dc51d..cf304a7 100644
--- a/README.md
+++ b/README.md
@@ -5,7 +5,7 @@ Better implementation of [KVAS](https://github.com/qzeleza/kvas)
 Roadmap:
 - [x] DNS Proxy
 - [x] DNS Records table
-- [ ] IPTables rules to remap DNS server [1]
+- [x] IPTables rules to remap DNS server [1]
 - [ ] Rule composer
 - [ ] List loading/watching (temporary)
 - [ ] IPSet integration
diff --git a/dns-proxy/dns-proxy.go b/dns-proxy/dns-proxy.go
index 3bc109e..c408a33 100644
--- a/dns-proxy/dns-proxy.go
+++ b/dns-proxy/dns-proxy.go
@@ -3,6 +3,7 @@ package dnsProxy
 import (
 	"encoding/hex"
 	"fmt"
+	"github.com/coreos/go-iptables/iptables"
 	"log"
 	"net"
 	"time"
@@ -23,7 +24,23 @@ type DNSProxy struct {
 }
 
 func (p DNSProxy) Close() error {
-	return p.udpConn.Close()
+	ipt, err := iptables.New()
+	if err != nil {
+		log.Fatalf("iptables init fail: %v", err)
+	}
+
+	err = ipt.DeleteIfExists("nat", "PREROUTING", "-j", "KVAS2_DNSOVERRIDE")
+	if err != nil {
+		log.Fatalf("failed to attaching chain: %v", err)
+	}
+
+	err = ipt.ClearAndDeleteChain("nat", "KVAS2_DNSOVERRIDE")
+	if err != nil {
+		log.Fatalf("failed to delete chain: %v", err)
+	}
+
+	return nil
+	//return p.udpConn.Close()
 }
 
 func (p DNSProxy) sendToUpstream(isTCP bool, request []byte) ([]byte, error) {
@@ -90,6 +107,26 @@ func (p DNSProxy) handleDNSRequest(clientAddr *net.UDPAddr, buffer []byte) {
 func (p DNSProxy) Listen() error {
 	var err error
 
+	ipt, err := iptables.New()
+	if err != nil {
+		log.Fatalf("iptables init fail: %v", err)
+	}
+
+	err = ipt.ClearChain("nat", "KVAS2_DNSOVERRIDE")
+	if err != nil {
+		log.Fatalf("failed to clean chain: %v", err)
+	}
+
+	err = ipt.AppendUnique("nat", "KVAS2_DNSOVERRIDE", "-p", "udp", "--dport", "53", "-j", "REDIRECT", "--to-port", "7548")
+	if err != nil {
+		log.Fatalf("failed to create rule: %v", err)
+	}
+
+	err = ipt.InsertUnique("nat", "PREROUTING", 1, "-j", "KVAS2_DNSOVERRIDE")
+	if err != nil {
+		log.Fatalf("failed to attaching chain: %v", err)
+	}
+
 	udpAddr, err := net.ResolveUDPAddr("udp", p.listenAddr)
 	if err != nil {
 		return fmt.Errorf("failed to resolve UDP address: %v", err)
diff --git a/go.mod b/go.mod
index 514b814..00593c0 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module kvas2-go
 
-go 1.21
\ No newline at end of file
+go 1.21
+
+require github.com/coreos/go-iptables v0.7.0
diff --git a/main.go b/main.go
index 1ce4300..65b02b8 100644
--- a/main.go
+++ b/main.go
@@ -5,6 +5,9 @@ import (
 	dnsProxy "kvas2-go/dns-proxy"
 	ruleComposer "kvas2-go/rule-composer"
 	"log"
+	"os"
+	"os/signal"
+	"syscall"
 )
 
 var (
@@ -46,9 +49,22 @@ func main() {
 			}
 		}
 	}
-	err := proxy.Listen()
-	if err != nil {
-		log.Fatal(err)
+
+	go func() {
+		err := proxy.Listen()
+		if err != nil {
+			log.Fatal(err)
+		}
+	}()
+
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+
+	for {
+		select {
+		case <-c:
+			proxy.Close()
+			return
+		}
 	}
-	defer proxy.Close()
 }