DenyIP-go/main.go

package main

import (
	"container/list"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

var (
	Daemon = flag.Bool("d", false, "守护进程模式")
	child  = flag.Bool("child", false, "子进程模式")
)

var (
	listInterfaces bool
	interfaceName  *string
	outputFile     *string
	filter         *string

	SIG      int
	MAX_SIG  = 10
	SIG_NAME *string

	// 存储 IPv4 地址的链表
	ipList  = list.New()
	ipMutex sync.Mutex // 用于保护 ipList 的互斥锁
)

// 启动子进程
func startChildProcess() (*exec.Cmd, error) {
	cmd := exec.Command(os.Args[0], "-child=true", "-i", *interfaceName)
	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("停止子进程失败: %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):
		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, os.Interrupt, syscall.SIGTERM)

	<-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)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Start()
	if err != nil {
		log.Fatalf("无法启动新实例: %v", err)
	}
	fmt.Println("新实例已启动，当前进程将退出")
	os.Exit(0)
}

func runChildProcess() { // 子进程逻辑
	//fmt.Println("子进程启动中...")
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	ticker := time.NewTicker(3 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-sigChan:
			fmt.Println("子进程收到终止信号，退出中...")
			return
		case <-ticker.C:
			//fmt.Println("子进程!!!")
		}
	}
}

func runMainProcess() { // 主进程逻辑

	fmt.Println("主进程启动...")

	cmd, err := startChildProcess()
	if err != nil {
		log.Fatalf("子进程启动失败: %v", err)
	}

	SIG = 0
	SIG_NAME := fmt.Sprintf("root%d", SIG)
	if return_value := IsIpset(SIG_NAME); return_value == 1 {
		createIPSet(SIG_NAME)
	}

	// 启动抓包
	go startPacketCapture()

	// 启动IP地域判断，管理
	go func() {
		for {
			ipMutex.Lock()        // 锁定互斥锁
			if ipList.Len() > 0 { // 链表不为空
				e1 := ipList.Front() // 获取链表第一个元素
				//fmt.Println(e1.Value) // 输出第一个元素

				region, _ := ip2region(e1.Value.(net.IP).String())

				log.Printf("当前 Ipset 链 %s %d\n", SIG_NAME, func() int {
					_len, _err := NumIPSet(SIG_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(SIG_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(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)

				// 创建新的 ipset 集合
				SIG++

				if SIG >= MAX_SIG {
					//log.Printf("已创建 %d 个集合!!!", MAX_SIG)
					log.Printf("\033[31m 已创建 %d 个集合!!! \033[0m\n", MAX_SIG)
				}

				SIG_NAME = fmt.Sprintf("root%d", SIG)
				if return_value := IsIpset(SIG_NAME); return_value == 1 {
					createIPSet(SIG_NAME)
				}

			}

			time.Sleep(7 * time.Second)
		}
	}()

	// 等待信号并清理
	waitForSignalAndCleanUp(cmd)
}

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", "", "保存捕获数据的输出文件（可选）")
	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_SIG; i++ {
				_name := fmt.Sprintf("root%d", i)
				iptables_add(_name)
			}

			os.Exit(0)
		case "stop":
			fmt.Println("停止 Iptables 规则")
			for i := 0; i < MAX_SIG; i++ {
				_name := fmt.Sprintf("root%d", i)
				iptables_del(_name)
			}
			os.Exit(0)
		default:
			log.Fatalf("未知的操作: %s. 请使用 'start' 或 'stop'.", instruction)
		}
	}

	if listInterfaces {
		printAvailableInterfaces()
		os.Exit(0)
	}

	if *interfaceName == "" {
		log.Fatal("请使用 -i 标志指定网络接口，或者使用 -l 列出接口。")
	}
}

func main() {
	handleCmd()

	// 守护进程模式
	if *Daemon {
		startDaemon()
	}

	// 子进程逻辑
	if *child {
		runChildProcess()
		return
	}

	runMainProcess()

}