iptables-helper/pkg/utils/iptables/handler.go

package iptables

import (
	"iptables-helper/pkg/resp/errorx"
	"iptables-helper/pkg/utils/command"
	"strings"
)

func appendArgs[T string | Chain | PolicyTarget | Action](args []string, argField string, argValue T) []string {
	value := strings.TrimSpace(string(argValue))
	if len(value) <= 0 {
		return args
	}
	args = append(args, argField, value)
	return args
}

func appendArgsWithError[T string | Chain | PolicyTarget | Action](args []string, argField string, argValue T, err error) ([]string, error) {
	value := strings.TrimSpace(string(argValue))
	if len(value) <= 0 && err != nil {
		return args, err
	}
	args = append(args, argField, value)
	return args, nil
}

func AddRule(rule Rule) error {
	var err error
	args := make([]string, 0)

	if args, err = appendArgsWithError(args, "-A", rule.Chain, errorx.NewDefaultError("规则链 Chain 不能为空")); err != nil {
		return err
	}

	args = appendArgs(args, "-j", rule.Jump)
	args = appendArgs(args, "-g", rule.Goto)

	args = appendArgs(args, "-i", rule.InputInterface)
	args = appendArgs(args, "! -i", rule.ExcludeInputInterface)

	args = appendArgs(args, "-o", rule.OutputInterface)
	args = appendArgs(args, "! -o", rule.ExcludeOutputInterface)

	args = appendArgs(args, "-s", rule.Source)
	args = appendArgs(args, "! -s", rule.ExcludeSource)

	args = appendArgs(args, "-d", rule.Destination)
	args = appendArgs(args, "! -d", rule.ExcludeDestination)

	args = appendArgs(args, "-p", rule.Protocol)
	args = appendArgs(args, "! -p", rule.ExcludeProtocol)

	args = appendArgs(args, "-m", rule.Match)

	args = appendArgs(args, "--sport", rule.SrcPort)
	args = appendArgs(args, "--sports", rule.SrcPorts)

	args = appendArgs(args, "--dport", rule.DstPort)
	args = appendArgs(args, "--dports", rule.DstPorts)

	args = appendArgs(args, "--limit", rule.Limit)

	cmd := command.Commander{}
	_, err = cmd.ExecuteWithResult("sudo iptables " + strings.Join(args, " "))
	if err == nil {
		SaveRule()
	}
	return err
}

func DelRuleByCmd(cmd string) error {
	cmd = strings.Split(cmd, ";")[0]
	cmd = strings.Split(cmd, "&")[0]
	commander := command.Commander{}
	_, err := commander.ExecuteWithResult("sudo iptables -D " + cmd)
	if err == nil {
		SaveRule()
	}
	return err
}

func SaveRule() {
	cmd := "sh -c \"iptables-save | tee /etc/iptables.rule\""
	commander := command.Commander{}
	commander.Execute("sudo " + cmd)
}