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 IsIpset(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 iptables(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
}

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