DenyIP-go/ipset.go

package main

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"strconv"
	"strings"
)

func createIPSet(setName string) error {
	cmd := exec.Command("ipset", "create", setName, "hash:ip")
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		// 记录错误信息，但不退出
		log.Printf("failed to execute command: %v, stderr: %s", err, stderr.String())
	}
	return err // 返回错误以便调用者处理
}

func AddIPSet(setName string, ip string) error {
	cmd := exec.Command("ipset", "add", setName, ip)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		// 记录错误信息，但不退出
		log.Printf("failed to add IP to set: %v, stderr: %s", err, stderr.String())
	}
	return err // 返回错误以便调用者处理
}

// NumIPSet returns the number of entries in the specified ipset set.
func NumIPSet(setName string) (int, error) {
	cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list %s | grep \"Number of entries\" | cut -d ':' -f 2 | sed 's/ //g'", setName))

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
		return 0, fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
	}

	output := strings.TrimSpace(stdout.String())
	numEntries, err := strconv.Atoi(output)
	if err != nil {
		log.Printf("failed to parse output as integer: %v, output: %s\n", err, output)
		return 0, fmt.Errorf("failed to parse output as integer: %w, output: %s", err, output)
	}

	return numEntries, nil
}

// IsIpset 检查名为 setName 的 ipset 是否存在，通过返回 0 表示存在，非零表示不存在或其他错误。
func Is_Name_Ipset(setName string) int {
	cmd := exec.Command("ipset", "list", setName)
	err := cmd.Run()

	if err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			// The program has exited with an exit code != 0
			return exitError.ExitCode()
		} else {
			// Another error occurred (e.g., command not found)
			return -1 // 或者你可以选择其他方式来标识这种情况
		}
	}

	// Command executed successfully, the set exists
	return 0
}

func Is_Ip_Ipset(ip string) int {
	cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list | grep \"%s\"", ip))
	err := cmd.Run()

	if err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			// The program has exited with an exit code != 0
			return exitError.ExitCode()
		} else {
			// Another error occurred (e.g., command not found)
			return -1 // 或者你可以选择其他方式来标识这种情况
		}
	}

	// Command executed successfully, the set exists
	return 0
}

// RemoveIPIfInSets 在多个 ipset 中查找并删除某个 IP，成功删除则返回集合名
func RemoveIPIfInSets(prefix string, max int, ip string) (string, error) {
	for i := 0; i < max; i++ {
		setName := fmt.Sprintf("%s%d", prefix, i)

		// 检查 IP 是否在当前集合中
		cmd := exec.Command("ipset", "test", setName, ip)
		output, err := cmd.CombinedOutput()

		if err != nil {
			if strings.Contains(string(output), "is NOT in set") {
				continue // 不在该集合中，尝试下一个
			}
			return "", fmt.Errorf("检测 %s 时出错: %v (%s)", setName, err, output)
		}

		// 存在，执行删除
		cmd = exec.Command("ipset", "del", setName, ip)
		output, err = cmd.CombinedOutput()
		if err != nil {
			return "", fmt.Errorf("从 %s 删除 %s 时出错: %v (%s)", setName, ip, err, output)
		}

		return setName, nil // 删除成功，返回集合名
	}

	return "", nil // 所有集合都没有该 IP，无需删除
}

// 添加 Iptables 规则
func iptables_add(setName string) error {

	cmd := exec.Command("sh", "-c", fmt.Sprintf("iptables -A INPUT -p tcp -m set --match-set %s src -j DROP", setName))

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		//log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
		//err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
	}

	return err
}

// 删除 Iptables 规则
func iptables_del(setName string) error {

	cmd := exec.Command("sh", "-c", fmt.Sprintf("iptables -D INPUT -p tcp -m set --match-set %s src -j DROP", setName))

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		//log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
		//err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
	}

	return err
}

// 打印 Iptables 规则
func iptables_list() error {
	cmd := exec.Command("sh", "-c", "iptables -L -v -n --line-numbers")

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
		err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
	}

	fmt.Print(stdout.String())
	return err
}

/*
func main() {
	// 创建 IPSet，但即使出错也继续执行
	err := createIPSet("root0")
	if err != nil {
		log.Println("创建 IPSet 出错:", err)
	} else {
		fmt.Println("IPSet 创建成功！")
	}

	// 添加 IP 到 IPSet，出错时继续执行
	err = AddIPSet("root0", "1.1.1.1")
	if err != nil {
		log.Println("添加 IP 到 IPSet 出错:", err)
	}

	// 获取 IPSet 条目数，出错时继续执行
	num, err := NumIPSet("root0")
	if err != nil {
		log.Println("获取 IPSet 条目数出错:", err)
	} else {
		fmt.Printf("IPSet 条目数: %d\n", num)
	}

	iptables("root0")
}
*/