DenyIP-go/main.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 的互斥锁

	ProcessedIPMap = make(map[string]struct{}) // 使用 map 存储已处理的 IP
	ProcessedMutex sync.Mutex                  // 互斥锁保护 ProcessedIPMap
)

// 启动子进程
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() {
	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:

		}
	}
}

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 := Is_Name_Ipset(IPSET_NAME); return_value == 1 {
		createIPSet(IPSET_NAME)
	}

	// 启动抓包
	go startPacketCapture()

	// 启动IP地域判断，管理
	go func() {
		for {

			IpMutex.Lock()        // 锁定互斥锁
			if IpList.Len() > 0 { // 链表不为空
				e1 := IpList.Front() // 获取链表第一个元素
				ipStr := e1.Value.(net.IP).String()
				region, _ := ip2region(ipStr) // 离线库初步判断地域

				ProcessedMutex.Lock()
				_, processed := ProcessedIPMap[ipStr] // 检查是否已处理
				ProcessedMutex.Unlock()

				if processed { // 如果尚未处理
					log.Printf("\033[33m %s 已经标记为国内，跳过！！！ \033[0m\n", ipStr)
					continue
				}

				if Is_Ip_Ipset(ipStr) != 0 { // IP 不在 ipset 集合中

					if !strings.Contains(region, "中国") && !strings.Contains(region, "内网") { // 离线库判断不在中国内
						log.Printf("\033[33m [%s %s] 离线库为国外, 进一步API判断\033[0m\n", ipStr, region)

						if position, err := curl_(ipStr); err != nil { // 使用 API 判断地域
							log.Printf("获取IP地域出错: %v", err)
						} else {
							log.Printf("\033[31m [%s %s]\033[0m\n\n", ipStr, position) // 打印地域

							if !strings.Contains(position, "中国") && !strings.Contains(position, "内网") { // API 判断为国外
								AddIPSet(IPSET_NAME, ipStr) // 添加 IP 到 ipset 集合
							} else {
								log.Printf("\033[33m %s 离线库为国外, API 判断为国内, 标记为已处理\033[0m\n", ipStr)

								// 将 IP 标记为已处理，国内地址
								ProcessedMutex.Lock()
								ProcessedIPMap[ipStr] = struct{}{}
								ProcessedMutex.Unlock()
							}
						}
					}

				} else { // IP 已在 ipset 集合中
					log.Printf("\033[31m %s 在 ipset 集合中 \033[0m\n", ipStr)
				}

				// 无论是否已处理，都移除该 IP
				IpList.Remove(e1)
			}
			IpMutex.Unlock() // 解锁互斥锁

			log.Printf(" 当前Ip链表长度:%d, Ipset名:%s, 长:%d ProcessedIPMap当前长度:%d\n", IpList.Len(), IPSET_NAME, func() int { // 打印 当前 Ipset 链长度
				_len, _err := NumIPSet(IPSET_NAME)
				if _err == nil {
					return _len
				}
				return 0
			}(), len(ProcessedIPMap))

			// 清理 ProcessedIPMap
			if len(ProcessedIPMap) >= 1024 {
				ProcessedMutex.Lock()                      // 锁定互斥锁
				ProcessedIPMap = make(map[string]struct{}) // 创建一个新的空 map
				ProcessedMutex.Unlock()                    // 解锁
				log.Println("ProcessedIPMap 已清空")

			}

			time.Sleep(1 * time.Second) // 防止高频运行
		}
	}()

	// 启动防火墙管理
	go func() {
		for {
			if ipset_len, _ := NumIPSet(IPSET_NAME); ipset_len >= 65534 {
				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 := Is_Name_Ipset(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 规则\n-s list  打印 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()
		fmt.Printf("\n")
		os.Exit(0)
	}

	if instruction != "" {
		switch instruction {
		case "start":
			for i := 0; i < MAX_IPSET_NAME; i++ {
				_name := fmt.Sprintf("root%d", i)
				iptables_add(_name)
			}

			os.Exit(0)
		case "stop":
			for i := 0; i < MAX_IPSET_NAME; i++ {
				_name := fmt.Sprintf("root%d", i)
				iptables_del(_name)
			}
			os.Exit(0)
		case "l":
			fallthrough
		case "list":
			_ = iptables_list()
			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() // 主进程逻辑

}