127 lines
3.2 KiB
Go
127 lines
3.2 KiB
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")
|
||
}
|
||
*/
|