DenyIP-go/ipset.go

package main

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

// 创建 ipset 集合
func createIPSet(setName string) error {
	// 建议增加 maxelem，防止频繁扩容。默认 65536 有点小。
	// 加上 -exist 防止报错
	cmd := exec.Command("ipset", "create", setName, "hash:ip", "maxelem", "200000", "-exist")
	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("ipset create failed: %v, stderr: %s", err, stderr.String())
	}
	return nil
}

// 添加 IP 到集合
func AddIPSet(setName string, ip string) error {
	cmd := exec.Command("ipset", "add", setName, ip)
	// 不关注输出，只关注 exit code
	return cmd.Run()
}

// NumIPSet 获取集合中 IP 的数量
// 注意：频繁调用此函数开销较大（因为它通过 shell 解析文本），建议仅在监控循环中低频调用
func NumIPSet(setName string) (int, error) {
	// 这里必须用 sh -c 因为涉及管道符 |
	cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list %s | grep 'Number of entries' | cut -d ':' -f 2", setName))

	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	err := cmd.Run()
	if err != nil {
		return 0, err
	}

	output := strings.TrimSpace(stdout.String())
	numEntries, err := strconv.Atoi(output)
	if err != nil {
		return 0, fmt.Errorf("解析数量失败: %w, 输出: %s", err, output)
	}

	return numEntries, nil
}

// Is_Name_Ipset 检查集合是否存在
// 返回 0 表示存在 (ExitCode 0)，其他表示不存在
func Is_Name_Ipset(setName string) int {
	cmd := exec.Command("ipset", "list", setName, "-n") // -n 只输出名字，减少开销
	err := cmd.Run()
	if err != nil {
		return 1 // 不存在或出错
	}
	return 0 // 存在
}

// Is_Ip_Ipset 检查 IP 是否在 ANY root 集合中
// 优化：不再使用 grep 全局搜索，而是遍历 root0...rootN 使用 ipset test
func Is_Ip_Ipset(ip string) int {
	// 遍历所有可能的集合名
	for i := 0; i < MAX_IPSET_NAME; i++ {
		setName := fmt.Sprintf("root%d", i)

		// 快速检查集合是否存在（可选，为了严谨）
		// 如果 ipset test 在集合不存在时会报错，所以我们直接运行 test
		// 如果 ip 在集合中，exit code 为 0
		cmd := exec.Command("ipset", "test", setName, ip)
		err := cmd.Run()

		if err == nil {
			return 1 // 找到了，IP 在该集合中
		}
	}

	return 0 // 所有集合都找过了，不在其中
}

// RemoveIPIfInSets 在所有集合中查找并删除某个 IP
func RemoveIPIfInSets(prefix string, max int, ip string) (string, error) {
	for i := 0; i < max; i++ {
		setName := fmt.Sprintf("%s%d", prefix, i)

		// 1. 先测试是否存在 (避免无意义的 del 调用报错)
		testCmd := exec.Command("ipset", "test", setName, ip)
		if err := testCmd.Run(); err != nil {
			// exit code != 0 说明不在这个集合里，继续找下一个
			continue
		}

		// 2. 存在则删除
		delCmd := exec.Command("ipset", "del", setName, ip)
		if err := delCmd.Run(); err != nil {
			return "", fmt.Errorf("从 %s 删除 %s 失败: %v", setName, ip, err)
		}

		return setName, nil // 成功删除
	}

	return "", nil // 未找到
}

// iptables_add 添加规则
// 修正：使用 -I (Insert) 而不是 -A (Append)，确保规则在最前面生效
func iptables_add(setName string) error {
	// 检查规则是否已经存在，避免重复添加
	checkCmd := exec.Command("iptables", "-C", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
	if err := checkCmd.Run(); err == nil {
		// 规则已存在，直接返回
		return nil
	}

	// 添加规则
	cmd := exec.Command("iptables", "-I", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		log.Printf("添加 iptables 规则失败 [%s]: %v, stderr: %s", setName, err, stderr.String())
	}
	return err
}

// iptables_del 删除规则
func iptables_del(setName string) error {
	// 循环删除，防止有重复规则
	for {
		cmd := exec.Command("iptables", "-D", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
		if err := cmd.Run(); err != nil {
			// 删除失败通常意味着规则不存在了，跳出循环
			break
		}
	}
	return nil
}

// iptables_list 打印规则
func iptables_list() error {
	cmd := exec.Command("iptables", "-L", "INPUT", "-v", "-n", "--line-numbers")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

/*
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")
}
*/