diff --git a/cap.go b/cap.go index 36ecdea..f8f735b 100644 --- a/cap.go +++ b/cap.go @@ -29,7 +29,7 @@ func printAvailableInterfaces() { // 判断 IP 是否在链表中 func isIPInList(ip net.IP) bool { - for e := ipList.Front(); e != nil; e = e.Next() { + for e := IpList.Front(); e != nil; e = e.Next() { if e.Value.(net.IP).Equal(ip) { return true } @@ -43,10 +43,10 @@ func printPacketInfo(packet gopacket.Packet) { if ipLayer != nil { ip, _ := ipLayer.(*layers.IPv4) - ipMutex.Lock() - defer ipMutex.Unlock() + IpMutex.Lock() + defer IpMutex.Unlock() if !isIPInList(ip.SrcIP) { - ipList.PushBack(ip.SrcIP) + IpList.PushBack(ip.SrcIP) log.Printf("\033[31m 已添加源 IP: %s 到链表 \033[0m\n", ip.SrcIP) } @@ -54,12 +54,10 @@ func printPacketInfo(packet gopacket.Packet) { } func startPacketCapture() { - handle, err := pcap.OpenLive(*interfaceName, 65535, true, pcap.BlockForever) + handle, err := pcap.OpenLive(*InterfaceName, 65535, true, pcap.BlockForever) if err != nil { - log.Fatalf("打开网络接口 %s 出错: %v", *interfaceName, err) + log.Fatalf("打开网络接口 %s 出错: %v", *InterfaceName, err) } - defer handle.Close() - defer func() { fmt.Println("清理资源...") if handle != nil { @@ -67,14 +65,14 @@ func startPacketCapture() { } }() - err = handle.SetBPFFilter(*filter) + err = handle.SetBPFFilter(*Protocol) if err != nil { log.Fatalf("设置 BPF 过滤器出错: %v", err) } var pcapWriter *pcapgo.Writer - if *outputFile != "" { - file, err := os.Create(*outputFile) + if *PcapFile != "" { + file, err := os.Create(*PcapFile) if err != nil { log.Fatalf("创建输出文件出错: %v", err) } @@ -88,7 +86,7 @@ func startPacketCapture() { } packetSource := gopacket.NewPacketSource(handle, handle.LinkType()) - fmt.Printf("正在监听网络接口 %s,使用过滤器 '%s'...\n", *interfaceName, *filter) + fmt.Printf("正在监听网络接口 %s,使用过滤器 '%s'...\n", *InterfaceName, *Protocol) sigChan := make(chan os.Signal, 1) signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) diff --git a/denyip b/denyip index 79d5060..7c2fe03 100644 Binary files a/denyip and b/denyip differ diff --git a/main.go b/main.go index 6e4f14d..1108f40 100644 --- a/main.go +++ b/main.go @@ -20,24 +20,31 @@ var ( child = flag.Bool("child", false, "子进程模式") ) +// 全局变量 var ( - listInterfaces bool - interfaceName *string - outputFile *string - filter *string + interfacesList bool // 是否列出网络接口 + InterfaceName *string // 网络接口名称 + PcapFile *string // 输出文件名 + Protocol *string // BPF 过滤器 - SIG int - MAX_SIG = 10 - SIG_NAME *string + SIG int // 当前使用的 ipset 集合编号 + MAX_SIG = 10 // 最大 ipset 集合数量 + SIG_NAME string // 当前使用的 ipset 集合名称 - // 存储 IPv4 地址的链表 - ipList = list.New() - ipMutex sync.Mutex // 用于保护 ipList 的互斥锁 + IpList = list.New() // 存储 IPv4 地址的链表 + IpMutex sync.Mutex // 保护 ipList 的互斥锁 ) // 启动子进程 func startChildProcess() (*exec.Cmd, error) { - cmd := exec.Command(os.Args[0], "-child=true", "-i", *interfaceName) + 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 @@ -56,10 +63,9 @@ func stopChildProcess(cmd *exec.Cmd) error { // 尝试优雅地停止子进程 if err := cmd.Process.Signal(syscall.SIGTERM); err != nil { - return fmt.Errorf("停止子进程失败: %w", err) + return fmt.Errorf("发送 SIGTERM 信号失败: %w", err) } - // 等待一段时间让子进程退出 done := make(chan error, 1) go func() { done <- cmd.Wait() @@ -70,8 +76,8 @@ func stopChildProcess(cmd *exec.Cmd) error { if err != nil { return fmt.Errorf("等待子进程退出时出错: %w", err) } - case <-time.After(5 * time.Second): - fmt.Println("子进程未在规定时间内退出,强制终止...") + case <-time.After(5 * time.Second): // 超时时间调整为5秒 + fmt.Println("子进程未在规定时间内退出,尝试强制终止...") if err := cmd.Process.Kill(); err != nil { return fmt.Errorf("强制终止子进程失败: %w", err) } @@ -84,46 +90,63 @@ func stopChildProcess(cmd *exec.Cmd) error { // 等待信号并优雅退出 func waitForSignalAndCleanUp(cmd *exec.Cmd) { sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + 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 + } + } - <-sigChan - fmt.Println("主进程收到终止信号,准备清理子进程...") if err := stopChildProcess(cmd); err != nil { log.Fatalf("清理子进程时遇到错误: %v", err) } fmt.Println("主进程退出") } -// 守护进程模式(简化版) +// 守护进程模式 func startDaemon() { - // 注意:此方法不会创建真正的守护进程,它只是以分离的方式启动一个新的进程。 - // 对于创建真正的守护进程,你应该使用 systemd 或其他服务管理器。 - cmd := exec.Command(os.Args[0], "-d=false", "-child=false", "-i", *interfaceName) + // 创建一个新的实例并让当前进程退出。注意这并不是守护进程的标准实现。 + 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 - err := cmd.Start() - if err != nil { + if err := cmd.Start(); err != nil { log.Fatalf("无法启动新实例: %v", err) } fmt.Println("新实例已启动,当前进程将退出") os.Exit(0) } -func runChildProcess() { // 子进程逻辑 - //fmt.Println("子进程启动中...") +// 子进程逻辑 +func runChildProcess() { + fmt.Println("子进程启动中...") + sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT) ticker := time.NewTicker(3 * time.Second) defer ticker.Stop() for { select { - case <-sigChan: - fmt.Println("子进程收到终止信号,退出中...") + case sig := <-sigChan: + fmt.Printf("子进程收到信号: %v,准备退出...\n", sig) return case <-ticker.C: - //fmt.Println("子进程!!!") + fmt.Println("子进程正在运行...") + fmt.Println("子进程接收到的参数:", os.Args) } } } @@ -138,7 +161,7 @@ func runMainProcess() { // 主进程逻辑 } SIG = 0 - SIG_NAME := fmt.Sprintf("root%d", SIG) + SIG_NAME = fmt.Sprintf("root%d", SIG) if return_value := IsIpset(SIG_NAME); return_value == 1 { createIPSet(SIG_NAME) } @@ -149,10 +172,9 @@ func runMainProcess() { // 主进程逻辑 // 启动IP地域判断,管理 go func() { for { - ipMutex.Lock() // 锁定互斥锁 - if ipList.Len() > 0 { // 链表不为空 - e1 := ipList.Front() // 获取链表第一个元素 - //fmt.Println(e1.Value) // 输出第一个元素 + IpMutex.Lock() // 锁定互斥锁 + if IpList.Len() > 0 { // 链表不为空 + e1 := IpList.Front() // 获取链表第一个元素 region, _ := ip2region(e1.Value.(net.IP).String()) @@ -171,16 +193,16 @@ func runMainProcess() { // 主进程逻辑 log.Printf("curl_ %s 位置 %s\n", e1.Value.(net.IP).String(), position) // 打印地域 AddIPSet(SIG_NAME, e1.Value.(net.IP).String()) // 添加 Ip 到 ipset 集合 - ipList.Remove(e1) // 移除第一个元素 + IpList.Remove(e1) // 移除第一个元素 } } else { // 这时是国内地址 - ipList.Remove(e1) + IpList.Remove(e1) } - log.Printf("Ip 链表长度:%d\n", ipList.Len()) // 链表长度 + log.Printf("Ip 链表长度:%d\n", IpList.Len()) // 链表长度 time.Sleep(2 * time.Second) } - ipMutex.Unlock() // 链表解锁互斥锁 + IpMutex.Unlock() // 链表解锁互斥锁 } @@ -188,9 +210,7 @@ func runMainProcess() { // 主进程逻辑 // 启动防火墙管理 go func() { - for { - if ipset_len, _ := NumIPSet(SIG_NAME); ipset_len >= 65534 { //log.Printf("ipset %s 列表已满 %dd\n", SIG_NAME, ipset_len) log.Printf("\033[31m ipset %s 列表已满 %d \033[0m\n", SIG_NAME, ipset_len) @@ -222,10 +242,10 @@ func handleCmd() { // 定义命令行标志 var instruction string var help bool - interfaceName = flag.String("i", "", "指定要使用的网络接口") - flag.BoolVar(&listInterfaces, "l", false, "列出可用的网络接口") - filter = flag.String("f", "tcp", "指定 BPF 过滤器") - outputFile = flag.String("o", "", "保存捕获数据的输出文件(可选)") + 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, "帮助信息") @@ -264,12 +284,12 @@ func handleCmd() { } } - if listInterfaces { + if interfacesList { printAvailableInterfaces() os.Exit(0) } - if *interfaceName == "" { + if *InterfaceName == "" { log.Fatal("请使用 -i 标志指定网络接口,或者使用 -l 列出接口。") } } @@ -284,10 +304,10 @@ func main() { // 子进程逻辑 if *child { - runChildProcess() + runChildProcess() // 子进程逻辑 return } - runMainProcess() + runMainProcess() // 主进程逻辑 }