From dbef6669b1685d5feecf2416c0ff01bbf0b836b1 Mon Sep 17 00:00:00 2001 From: zxb <919411476@qq.com> Date: Fri, 3 Nov 2023 14:07:50 +0800 Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E8=87=AA=E5=AE=9A=E4=B9=89?= =?UTF-8?q?=E8=A7=84=E5=88=99=E9=93=BE=E7=AE=A1=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/service/main.go | 6 +++--- internel/app/app.go | 27 ++++++++++++++++++++++++++- internel/conf/conf.go | 5 +++++ internel/controller/controller.go | 22 +++++++++++++++++++++- pkg/config/config.go | 12 +++++++++++- 5 files changed, 66 insertions(+), 6 deletions(-) create mode 100644 internel/conf/conf.go diff --git a/cmd/service/main.go b/cmd/service/main.go index dfad90e..b6f1ef5 100644 --- a/cmd/service/main.go +++ b/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) } } diff --git a/internel/app/app.go b/internel/app/app.go index 0600990..f2b7073 100644 --- a/internel/app/app.go +++ b/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 服务器 diff --git a/internel/conf/conf.go b/internel/conf/conf.go new file mode 100644 index 0000000..d3ff427 --- /dev/null +++ b/internel/conf/conf.go @@ -0,0 +1,5 @@ +package conf + +import "iptables-helper/pkg/config" + +var Conf = new(config.Conf) diff --git a/internel/controller/controller.go b/internel/controller/controller.go index 79461d5..18177b5 100644 --- a/internel/controller/controller.go +++ b/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)) }) } diff --git a/pkg/config/config.go b/pkg/config/config.go index b90563d..0953efb 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -1,7 +1,8 @@ package config type Conf struct { - Server *ServerConfig `comment:"服务器配置"` + 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: "", + }, } }