优化主进程子进程获取参数

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()
+		IpMutex.Lock()
-		defer ipMutex.Unlock()
+		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 != "" {
+	if *PcapFile != "" {
-		file, err := os.Create(*outputFile)
+		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)

main.go

@@ -20,24 +20,31 @@ var (
 	child  = flag.Bool("child", false, "子进程模式")
 )
 
+// 全局变量
 var (
-	listInterfaces bool
+	interfacesList bool    // 是否列出网络接口
-	interfaceName  *string
+	InterfaceName  *string // 网络接口名称
-	outputFile     *string
+	PcapFile       *string // 输出文件名
-	filter         *string
+	Protocol       *string // BPF 过滤器
 
-	SIG      int
+	SIG      int    // 当前使用的 ipset 集合编号
-	MAX_SIG  = 10
+	MAX_SIG  = 10   // 最大 ipset 集合数量
-	SIG_NAME *string
+	SIG_NAME string // 当前使用的 ipset 集合名称
 
-	// 存储 IPv4 地址的链表
+	IpList  = list.New() // 存储 IPv4 地址的链表
-	ipList  = list.New()
+	IpMutex sync.Mutex   // 保护 ipList 的互斥锁
-	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):
+	case <-time.After(5 * time.Second): // 超时时间调整为5秒
-		fmt.Println("子进程未在规定时间内退出，强制终止...")
+		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 或其他服务管理器。
+	args := []string{}
-	cmd := exec.Command(os.Args[0], "-d=false", "-child=false", "-i", *interfaceName)
+	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 := cmd.Start(); err != nil {
-	if 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:
+		case sig := <-sigChan:
-			fmt.Println("子进程收到终止信号，退出中...")
+			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()        // 锁定互斥锁
+			IpMutex.Lock()        // 锁定互斥锁
-			if ipList.Len() > 0 { // 链表不为空
+			if IpList.Len() > 0 { // 链表不为空
-				e1 := ipList.Front() // 获取链表第一个元素
+				e1 := IpList.Front() // 获取链表第一个元素
-				//fmt.Println(e1.Value) // 输出第一个元素
 
 				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", "", "指定要使用的网络接口")
+	InterfaceName = flag.String("i", "", "指定要使用的网络接口")
-	flag.BoolVar(&listInterfaces, "l", false, "列出可用的网络接口")
+	flag.BoolVar(&interfacesList, "l", false, "列出可用的网络接口")
-	filter = flag.String("f", "tcp", "指定 BPF 过滤器")
+	Protocol = flag.String("f", "tcp", "指定 BPF 过滤器")
-	outputFile = flag.String("o", "", "保存捕获数据的输出文件（可选）")
+	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() // 主进程逻辑
 
 }