支持自定义规则链管理

cmd/service/main.go

@@ -1,6 +1,7 @@
 package service
 
 import (
+	"errors"
 	"github.com/fsnotify/fsnotify"
 	"github.com/spf13/viper"
 	"iptables-helper/internel/app"
@@ -19,11 +20,10 @@ func Main() {
 	}
 
 	if err := viper.ReadInConfig(); err != nil {
-		if _, ok := err.(viper.ConfigFileNotFoundError); ok {
+		var configFileNotFoundError viper.ConfigFileNotFoundError
+		if errors.As(err, &configFileNotFoundError) {
 			_ = toml.GenerateConfig()
 			logger.Log().Fatalf("未找到配置文件, 已生成示例配置文件于运行路径下")
-		} else {
-			logger.Log().Fatalf("配置解析失败 %s", err)
 		}
 	}
 

internel/app/app.go

@@ -2,13 +2,17 @@ package app
 
 import (
 	"encoding/json"
+	"fmt"
 	"github.com/gofiber/fiber/v2"
 	"github.com/spf13/viper"
+	globalConf "iptables-helper/internel/conf"
 	"iptables-helper/internel/middleware"
 	"iptables-helper/internel/route"
 	"iptables-helper/pkg/config"
 	"iptables-helper/pkg/logger"
+	"iptables-helper/pkg/utils/command"
 	fib "iptables-helper/pkg/utils/fiber"
+	"strings"
 	"time"
 )
 
@@ -29,14 +33,35 @@ func CreateApp(c *config.Conf) *fiber.App {
 
 func Run() {
 	err := viper.Unmarshal(conf)
+	globalConf.Conf = conf
 	if err != nil {
 		logger.Log().Fatalf("配置文件解析失败: %s, 请检查配置是否有误", err)
 	}
 	Shutdown()
-
 	// 初始化数据源
 	fib.Exec(func() {
 		InitApp(conf)
+
+		if !conf.Custom.UseCustomChain {
+			logger.Log().Infof("使用全局管理")
+		} else {
+			customChain := strings.TrimSpace(conf.Custom.CustomChain)
+			logger.Log().Infof("使用自定义链管理 (管理自定义链: %s)", customChain)
+			if len(customChain) > 0 {
+				commander := command.Commander{}
+				_, err := commander.ExecuteWithResult(fmt.Sprintf("sudo iptables -N %s", customChain))
+				if err == nil {
+					// 入口
+					commander.Execute(fmt.Sprintf("sudo iptables -A INPUT -j %s", customChain))
+					// 出口
+					commander.Execute(fmt.Sprintf("sudo iptables -A OUTPUT -j %s", customChain))
+					// 转发
+					commander.Execute(fmt.Sprintf("sudo iptables -A FORWARD -j %s", customChain))
+				}
+			} else {
+				logger.Log().Fatal("防火墙配置失败: 自定义链名称不能为空")
+			}
+		}
 	})
 
 	// 创建 fiber 服务器

internel/conf/conf.go

@@ -0,0 +1,5 @@
+package conf
+
+import "iptables-helper/pkg/config"
+
+var Conf = new(config.Conf)

internel/controller/controller.go

@@ -3,6 +3,7 @@ package controller
 import (
 	"github.com/gofiber/fiber/v2"
 	"github.com/shirou/gopsutil/net"
+	"iptables-helper/internel/conf"
 	response "iptables-helper/pkg/resp"
 	"iptables-helper/pkg/resp/errorx"
 	"iptables-helper/pkg/utils/command"
@@ -30,7 +31,26 @@ func getRuleInfo(api fiber.Router) {
 	api.Get("/info", func(ctx *fiber.Ctx) error {
 		cmder := command.Commander{}
 		result, _ := cmder.ExecuteWithResult("sudo iptables -S")
-		return ctx.JSON(response.NewResponse(iptables.Parse(result)))
+		data := iptables.Parse(result)
+		if conf.Conf.Custom.UseCustomChain {
+			chains := make([]iptables.Chain, 0)
+			for _, chain := range data.Chains {
+				if string(chain) == conf.Conf.Custom.CustomChain {
+					chains = append(chains, chain)
+					break
+				}
+			}
+			data.Chains = chains
+
+			rules := make([]iptables.Rule, 0)
+			for _, rule := range data.Rules {
+				if string(rule.Chain) == conf.Conf.Custom.CustomChain {
+					rules = append(rules, rule)
+				}
+			}
+			data.Rules = rules
+		}
+		return ctx.JSON(response.NewResponse(data))
 	})
 }
 

pkg/config/config.go

@@ -2,6 +2,7 @@ package config
 
 type Conf struct {
 	Server *ServerConfig         `comment:"服务器配置"`
+	Custom *CustomFireWallConfig `comment:"自定义配置"`
 }
 
 type ServerConfig struct {
@@ -17,6 +18,11 @@ type ServerConfig struct {
 	EnableSwag      bool `yaml:"enableSwag" comment:"是否启用 swag 访问路径: /swagger"`
 }
 
+type CustomFireWallConfig struct {
+	UseCustomChain bool   `yaml:"useCustomChain" comment:"是否只使用某个自定义规则链进行管理, 否则管理所有规则链， 是则只管理指定的链"`
+	CustomChain    string `yaml:"customChain" comment:"自定义规则链名称"`
+}
+
 func DefaultConfig() *Conf {
 	return &Conf{
 		Server: &ServerConfig{
@@ -27,5 +33,9 @@ func DefaultConfig() *Conf {
 			ServerHeader:    "SkServer",
 			EnableRoutesMsg: false,
 		},
+		Custom: &CustomFireWallConfig{
+			UseCustomChain: false,
+			CustomChain:    "",
+		},
 	}
 }