317 lines
7.7 KiB
Go
317 lines
7.7 KiB
Go
package main
|
||
|
||
import (
|
||
"container/list"
|
||
"flag"
|
||
"fmt"
|
||
"log"
|
||
"net"
|
||
"os"
|
||
"os/exec"
|
||
"os/signal"
|
||
"runtime"
|
||
"strings"
|
||
"sync"
|
||
"syscall"
|
||
"time"
|
||
)
|
||
|
||
var (
|
||
daemon = flag.Bool("d", false, "守护进程模式")
|
||
child = flag.Bool("child", false, "子进程模式")
|
||
)
|
||
|
||
// 全局变量
|
||
var (
|
||
InterfacesList bool // 是否列出网络接口
|
||
InterfaceName *string // 网络接口名称
|
||
PcapFile *string // 输出文件名
|
||
Protocol *string // BPF 过滤器
|
||
|
||
IPSET_NUMBER int // 当前使用的 ipset 集合编号
|
||
MAX_IPSET_NAME = 10 // 最大 ipset 集合数量
|
||
IPSET_NAME string // 当前使用的 ipset 集合名称
|
||
|
||
IpList = list.New() // 存储 IPv4 地址的链表
|
||
IpMutex sync.Mutex // 保护 ipList 的互斥锁
|
||
)
|
||
|
||
// 启动子进程
|
||
func startChildProcess() (*exec.Cmd, error) {
|
||
args := []string{}
|
||
for _, arg := range os.Args[1:] {
|
||
if !strings.HasPrefix(arg, "-child") { // 只过滤 -child 标志
|
||
args = append(args, arg)
|
||
}
|
||
}
|
||
args = append(args, "-child=true")
|
||
cmd := exec.Command(os.Args[0], args...)
|
||
cmd.Stdout = os.Stdout
|
||
cmd.Stderr = os.Stderr
|
||
|
||
if err := cmd.Start(); err != nil {
|
||
return nil, fmt.Errorf("启动子进程失败: %w", err)
|
||
}
|
||
fmt.Printf("子进程已启动, PID: %d\n", cmd.Process.Pid)
|
||
return cmd, nil
|
||
}
|
||
|
||
// 停止子进程
|
||
func stopChildProcess(cmd *exec.Cmd) error {
|
||
if cmd == nil || cmd.Process == nil {
|
||
return fmt.Errorf("子进程无效")
|
||
}
|
||
|
||
// 尝试优雅地停止子进程
|
||
if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
|
||
return fmt.Errorf("发送 SIGTERM 信号失败: %w", err)
|
||
}
|
||
|
||
done := make(chan error, 1)
|
||
go func() {
|
||
done <- cmd.Wait()
|
||
}()
|
||
|
||
select {
|
||
case err := <-done:
|
||
if err != nil {
|
||
return fmt.Errorf("等待子进程退出时出错: %w", err)
|
||
}
|
||
case <-time.After(5 * time.Second): // 超时时间调整为5秒
|
||
fmt.Println("子进程未在规定时间内退出，尝试强制终止...")
|
||
if err := cmd.Process.Kill(); err != nil {
|
||
return fmt.Errorf("强制终止子进程失败: %w", err)
|
||
}
|
||
}
|
||
|
||
fmt.Printf("子进程已停止, PID: %d\n", cmd.Process.Pid)
|
||
return nil
|
||
}
|
||
|
||
// 等待信号并优雅退出
|
||
func waitForSignalAndCleanUp(cmd *exec.Cmd) {
|
||
sigChan := make(chan os.Signal, 1)
|
||
signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
|
||
|
||
for sig := range sigChan {
|
||
fmt.Printf("主进程收到信号: %v\n", sig)
|
||
if cmd != nil && cmd.Process != nil {
|
||
_ = cmd.Process.Signal(sig) // 转发信号到子进程
|
||
}
|
||
if sig == syscall.SIGINT || sig == syscall.SIGTERM {
|
||
break
|
||
}
|
||
}
|
||
|
||
if err := stopChildProcess(cmd); err != nil {
|
||
log.Fatalf("清理子进程时遇到错误: %v", err)
|
||
}
|
||
fmt.Println("主进程退出")
|
||
}
|
||
|
||
// 守护进程模式
|
||
func startDaemon() {
|
||
// 创建一个新的实例并让当前进程退出。注意这并不是守护进程的标准实现。
|
||
args := []string{}
|
||
for _, arg := range os.Args[1:] {
|
||
if !strings.HasPrefix(arg, "-d") && !strings.HasPrefix(arg, "-child") { // 过滤掉 -d 和 -child 标志
|
||
args = append(args, arg)
|
||
}
|
||
}
|
||
args = append(args, "-d=false", "-child=false")
|
||
|
||
cmd := exec.Command(os.Args[0], args...)
|
||
cmd.Stdout = os.Stdout
|
||
cmd.Stderr = os.Stderr
|
||
if err := cmd.Start(); err != nil {
|
||
log.Fatalf("无法启动新实例: %v", err)
|
||
}
|
||
fmt.Println("新实例已启动，当前进程将退出")
|
||
os.Exit(0)
|
||
}
|
||
|
||
// 子进程逻辑
|
||
func runChildProcess() {
|
||
fmt.Println("子进程启动中...")
|
||
|
||
sigChan := make(chan os.Signal, 1)
|
||
signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
|
||
|
||
ticker := time.NewTicker(3 * time.Second)
|
||
defer ticker.Stop()
|
||
|
||
for {
|
||
select {
|
||
case sig := <-sigChan:
|
||
fmt.Printf("子进程收到信号: %v，准备退出...\n", sig)
|
||
return
|
||
case <-ticker.C:
|
||
fmt.Println("子进程正在运行...")
|
||
fmt.Println("子进程接收到的参数:", os.Args)
|
||
}
|
||
}
|
||
}
|
||
|
||
func runMainProcess() { // 主进程逻辑
|
||
|
||
fmt.Println("主进程启动...")
|
||
|
||
cmd, err := startChildProcess()
|
||
if err != nil {
|
||
log.Fatalf("子进程启动失败: %v", err)
|
||
}
|
||
|
||
IPSET_NUMBER = 0
|
||
IPSET_NAME = fmt.Sprintf("root%d", IPSET_NUMBER)
|
||
if return_value := IsIpset(IPSET_NAME); return_value == 1 {
|
||
createIPSet(IPSET_NAME)
|
||
}
|
||
|
||
// 启动抓包
|
||
go startPacketCapture()
|
||
|
||
// 启动IP地域判断，管理
|
||
go func() {
|
||
for {
|
||
IpMutex.Lock() // 锁定互斥锁
|
||
if IpList.Len() > 0 { // 链表不为空
|
||
e1 := IpList.Front() // 获取链表第一个元素
|
||
|
||
region, _ := ip2region(e1.Value.(net.IP).String())
|
||
|
||
log.Printf("当前 Ipset 链 %s %d\n", IPSET_NAME, func() int {
|
||
_len, _err := NumIPSet(IPSET_NAME)
|
||
if _err == nil {
|
||
return _len
|
||
}
|
||
return 0
|
||
}())
|
||
|
||
if !strings.Contains(region, "中国") && !strings.Contains(region, "内网") {
|
||
if position, err := curl_(e1.Value.(net.IP).String()); err != nil { //判断地域
|
||
log.Printf("获取Ip地域出错: %v", err)
|
||
} else {
|
||
log.Printf("curl_ %s 位置 %s\n", e1.Value.(net.IP).String(), position) // 打印地域
|
||
AddIPSet(IPSET_NAME, e1.Value.(net.IP).String()) // 添加 Ip 到 ipset 集合
|
||
|
||
IpList.Remove(e1) // 移除第一个元素
|
||
}
|
||
} else { // 这时是国内地址
|
||
IpList.Remove(e1)
|
||
}
|
||
|
||
log.Printf("Ip 链表长度:%d\n", IpList.Len()) // 链表长度
|
||
time.Sleep(2 * time.Second)
|
||
}
|
||
IpMutex.Unlock() // 链表解锁互斥锁
|
||
|
||
}
|
||
|
||
}()
|
||
|
||
// 启动防火墙管理
|
||
go func() {
|
||
for {
|
||
if ipset_len, _ := NumIPSet(IPSET_NAME); ipset_len >= 65534 {
|
||
//log.Printf("ipset %s 列表已满 %dd\n", IPSET_NAME, ipset_len)
|
||
log.Printf("\033[31m ipset %s 列表已满 %d \033[0m\n", IPSET_NAME, ipset_len)
|
||
|
||
// 创建新的 ipset 集合
|
||
IPSET_NUMBER++
|
||
|
||
if IPSET_NUMBER >= MAX_IPSET_NAME {
|
||
//log.Printf("已创建 %d 个集合!!!", MAX_IPSET_NAME)
|
||
log.Printf("\033[31m 已创建 %d 个集合!!! \033[0m\n", MAX_IPSET_NAME)
|
||
}
|
||
|
||
IPSET_NAME = fmt.Sprintf("root%d", IPSET_NUMBER)
|
||
if return_value := IsIpset(IPSET_NAME); return_value == 1 {
|
||
createIPSet(IPSET_NAME)
|
||
}
|
||
|
||
}
|
||
|
||
time.Sleep(7 * time.Second)
|
||
}
|
||
}()
|
||
|
||
// 等待信号并清理
|
||
waitForSignalAndCleanUp(cmd)
|
||
}
|
||
|
||
func handleCmd() {
|
||
// 定义命令行标志
|
||
var instruction string
|
||
var help bool
|
||
InterfaceName = flag.String("i", "", "指定要使用的网络接口")
|
||
flag.BoolVar(&InterfacesList, "l", false, "列出可用的网络接口")
|
||
Protocol = flag.String("f", "tcp", "指定 BPF 过滤器")
|
||
PcapFile = flag.String("o", "", "保存捕获数据的输出文件（可选）")
|
||
flag.StringVar(&instruction, "s", "", "-s start 启动 Iptables 规则\n-s stop 停止 Iptables 规则")
|
||
flag.BoolVar(&help, "h", false, "")
|
||
flag.BoolVar(&help, "help", false, "帮助信息")
|
||
|
||
flag.Parse()
|
||
|
||
if help {
|
||
fmt.Printf(
|
||
"\t\tDenyip firewall\n" +
|
||
"\tVersion 0.1\n" +
|
||
"\tE-mail: aixiao@aixiao.me\n" +
|
||
"\tDate: 20250102\n")
|
||
|
||
flag.Usage()
|
||
os.Exit(0)
|
||
}
|
||
|
||
if instruction != "" {
|
||
switch instruction {
|
||
case "start":
|
||
fmt.Println("启动 Iptables 规则")
|
||
for i := 0; i < MAX_IPSET_NAME; i++ {
|
||
_name := fmt.Sprintf("root%d", i)
|
||
iptables_add(_name)
|
||
}
|
||
|
||
os.Exit(0)
|
||
case "stop":
|
||
fmt.Println("停止 Iptables 规则")
|
||
for i := 0; i < MAX_IPSET_NAME; i++ {
|
||
_name := fmt.Sprintf("root%d", i)
|
||
iptables_del(_name)
|
||
}
|
||
os.Exit(0)
|
||
default:
|
||
log.Fatalf("未知的操作: %s. 请使用 'start' 或 'stop'.", instruction)
|
||
}
|
||
}
|
||
|
||
if InterfacesList {
|
||
printAvailableInterfaces()
|
||
os.Exit(0)
|
||
}
|
||
|
||
if *InterfaceName == "" {
|
||
log.Fatal("请使用 -i 标志指定网络接口，或者使用 -l 列出接口。")
|
||
}
|
||
}
|
||
|
||
func main() {
|
||
runtime.GOMAXPROCS(runtime.NumCPU()) // 设置最大CPU核数
|
||
handleCmd()
|
||
|
||
// 守护进程模式
|
||
if *daemon {
|
||
startDaemon()
|
||
}
|
||
|
||
// 子进程逻辑
|
||
if *child {
|
||
runChildProcess() // 子进程逻辑
|
||
return
|
||
}
|
||
|
||
runMainProcess() // 主进程逻辑
|
||
|
||
}
|