diff --git a/README.md b/README.md index b69ae1b..84775f0 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ ```bash apt-get install ipset -apt-get install libcap-dev libpcap-dev libdbus-1-dev +apt-get install libcap-dev libpcap-dev libdbus-1-dev libsystemd-dev git clone https://git.aixiao.me/aixiao/DenyIP-go.git cd DenyIP-go diff --git a/cap.go b/cap.go index 26d9352..6860fdb 100644 --- a/cap.go +++ b/cap.go @@ -3,7 +3,6 @@ package main import ( "fmt" "log" - "net" "os" "os/signal" "syscall" @@ -14,7 +13,7 @@ import ( "github.com/google/gopacket/pcapgo" ) -// 打印可用的网络接口信息 +// 打印可用的网络接口信息 (保持不变) func printAvailableInterfaces() { devices, err := pcap.FindAllDevs() if err != nil { @@ -27,54 +26,61 @@ func printAvailableInterfaces() { } } -// 判断 IP 是否在链表中 -func isIPInList(ip net.IP) bool { - for e := IpList.Front(); e != nil; e = e.Next() { - if e.Value.(net.IP).Equal(ip) { - return true - } - } - return false -} +// ---------------------------------------------------------------------- +// 移除 isIPInList 函数 (不再需要) +// ---------------------------------------------------------------------- -// 打印捕获到的每个数据包的信息 -func printPacketInfo(packet gopacket.Packet) { +// 处理捕获到的数据包 (原 printPacketInfo 的优化版) +func handlePacket(packet gopacket.Packet) { + // 1. 快速提取网络层 (IPv4) + // 使用 Layer() 比 Layers() 稍微快一点点，但如果追求极致性能，建议使用 ZeroCopy 模式 ipLayer := packet.Layer(layers.LayerTypeIPv4) - - if ipLayer != nil { - ip, _ := ipLayer.(*layers.IPv4) - IpMutex.Lock() - defer IpMutex.Unlock() - if !isIPInList(ip.SrcIP) { - IpList.PushBack(ip.SrcIP) - log.Printf("\033[31m 已添加源 IP: %s 到链表 \033[0m\n", ip.SrcIP) - - } + if ipLayer == nil { + return } + ip, ok := ipLayer.(*layers.IPv4) + if !ok { + return + } + + // 2. 提取源 IP 字符串 + srcIP := ip.SrcIP.String() + + // 3. 将 IP 推入全局处理队列 + // 这里不再加锁，也不判断是否存在链表中，直接交给 PushIPToQueue 的非阻塞逻辑处理 + // PushIPToQueue 会负责去重 (PendingIPs) 和限流 + PushIPToQueue(srcIP) + + // 注意：为了性能，去掉了这里的日志打印。 + // 在高流量攻击下，每秒打印几千行日志会导致 IO 阻塞，拖慢抓包速度。 + // 具体的封禁日志会在 processIP (消费者) 中打印。 } // startPacketCapture 启动数据包捕获 func startPacketCapture() { + if *InterfaceName == "" { + log.Fatal("未指定网络接口，请使用 -i 参数") + } + // 打开指定的网络接口进行实时数据包捕获 - handle, err := pcap.OpenLive(*InterfaceName, 65535, true, pcap.BlockForever) + // snaplen 设置为 1600 足够捕获头部信息，没必要设为 65535，可以稍微节省内存 + handle, err := pcap.OpenLive(*InterfaceName, 1600, true, pcap.BlockForever) if err != nil { log.Fatalf("打开网络接口 %s 出错: %v", *InterfaceName, err) } - // 确保在函数退出时关闭句柄，释放资源 defer func() { - fmt.Println("清理资源...") - if handle != nil { - handle.Close() - } + fmt.Println("清理抓包资源...") + handle.Close() }() - // 设置 BPF（Berkeley Packet Filter）过滤器，以便只捕获指定协议的数据包 - err = handle.SetBPFFilter(*Protocol) - if err != nil { - log.Fatalf("设置 BPF 过滤器出错: %v", err) + // 设置 BPF 过滤器 + if *Protocol != "" { + if err = handle.SetBPFFilter(*Protocol); err != nil { + log.Fatalf("设置 BPF 过滤器出错: %v", err) + } } - // 如果指定了输出文件，则创建文件并初始化 pcapgo.Writer + // --- PCAP 文件保存逻辑 (如果不需要存包，这部分其实是最大的性能瓶颈) --- var pcapWriter *pcapgo.Writer if *PcapFile != "" { file, err := os.Create(*PcapFile) @@ -84,28 +90,35 @@ func startPacketCapture() { defer file.Close() pcapWriter = pcapgo.NewWriter(file) - // 写入 pcap 文件头部，指定最大捕获长度和链路层类型 - err = pcapWriter.WriteFileHeader(65535, layers.LinkTypeEthernet) - if err != nil { + // 写入文件头 + if err = pcapWriter.WriteFileHeader(1600, layers.LinkTypeEthernet); err != nil { log.Fatalf("写入全局头部出错: %v", err) } } - // 创建数据包源，用于从网络接口读取数据包 + // 使用 ZeroCopyPacketDataSource 可以减少内存拷贝，提升性能 (Linux Only) + // 如果在非 Linux 环境编译报错，请改回 gopacket.NewPacketSource packetSource := gopacket.NewPacketSource(handle, handle.LinkType()) + // packetSource.NoCopy = true // 开启 NoCopy 模式 (可选，稍微危险但更快) + log.Printf(" 正在监听网络接口 %s, 使用过滤器 '%s'...\n", *InterfaceName, *Protocol) - // 创建信号通道，用于捕获中断信号 + // 创建信号通道 sigChan := make(chan os.Signal, 1) signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) - // 启动一个 goroutine 来处理捕获到的数据包 + // 启动 Goroutine 处理数据包 go func() { + // 直接遍历 channel for packet := range packetSource.Packets() { - // 打印数据包信息 - printPacketInfo(packet) - // 如果指定了输出文件，则将数据包写入文件 + + // 1. 核心业务：提取 IP 进队列 + handlePacket(packet) + + // 2. 可选业务：存盘 + // 警告：在高并发攻击下，写文件 IO 会成为瓶颈。建议生产环境非必要不开启 -o 参数 if pcapWriter != nil { + // 注意：CaptureInfo 和 Data 必须在 packet 被复用前读取 err := pcapWriter.WritePacket(packet.Metadata().CaptureInfo, packet.Data()) if err != nil { log.Printf("写入数据包出错: %v", err) @@ -114,7 +127,7 @@ func startPacketCapture() { } }() - // 等待中断信号，收到信号后停止捕获 + // 阻塞等待信号 <-sigChan fmt.Println("\n停止抓包...") } diff --git a/denyip b/denyip index 660a286..095c143 100644 Binary files a/denyip and b/denyip differ diff --git a/denyip.upx b/denyip.upx new file mode 100644 index 0000000..177d398 Binary files /dev/null and b/denyip.upx differ diff --git a/embed.go b/embed.go index 7779e4f..976982b 100644 --- a/embed.go +++ b/embed.go @@ -59,9 +59,9 @@ func embed_ip2region() { // 你可以在这里使用 ip2region.xdb 文件了，例如： if _, err := os.Stat("ip2region/ip2region.xdb"); err == nil { - log.Println(" ✅ 确认 ip2region.xdb 已成功写出") + log.Println(" 确认 ip2region.xdb 已成功写出") // 这里可以调用 ip2region 逻辑加载它 } else { - log.Println(" ❌ 找不到 ip2region.xdb") + log.Println(" 找不到 ip2region.xdb") } } diff --git a/go.mod b/go.mod index 832650a..a3d653e 100644 --- a/go.mod +++ b/go.mod @@ -1,10 +1,10 @@ module denyip -go 1.25.4 +go 1.25.5 require ( github.com/google/gopacket v1.1.19 - github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251207115101-d4b8f9f841b9 + github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251215094000-f12b9765b373 ) require ( diff --git a/go.sum b/go.sum index 5bb33f6..faf2a7d 100644 --- a/go.sum +++ b/go.sum @@ -12,6 +12,8 @@ github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251124080701-096d68ea7 github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251124080701-096d68ea7706/go.mod h1:+mNMTBuDMdEGhWzoQgc6kBdqeaQpWh5ba8zqmp2MxCU= github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251207115101-d4b8f9f841b9 h1:0IngVEHYqJUpjrnY9T1dZ2AMIbsI/sCUxxg77eGXXes= github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251207115101-d4b8f9f841b9/go.mod h1:+mNMTBuDMdEGhWzoQgc6kBdqeaQpWh5ba8zqmp2MxCU= +github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251215094000-f12b9765b373 h1:JrfpSwKuIr0RoPoJlYTlbZknElX41AuAwiWhhNF7Ff0= +github.com/lionsoul2014/ip2region/binding/golang v0.0.0-20251215094000-f12b9765b373/go.mod h1:+mNMTBuDMdEGhWzoQgc6kBdqeaQpWh5ba8zqmp2MxCU= github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= diff --git a/ipset.go b/ipset.go index edbbd56..44bfb0b 100644 --- a/ipset.go +++ b/ipset.go @@ -4,181 +4,154 @@ import ( "bytes" "fmt" "log" + "os" "os/exec" "strconv" "strings" ) +// 创建 ipset 集合 func createIPSet(setName string) error { - cmd := exec.Command("ipset", "create", setName, "hash:ip") - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout + // 建议增加 maxelem，防止频繁扩容。默认 65536 有点小。 + // 加上 -exist 防止报错 + cmd := exec.Command("ipset", "create", setName, "hash:ip", "maxelem", "200000", "-exist") + var stderr bytes.Buffer cmd.Stderr = &stderr err := cmd.Run() if err != nil { - // 记录错误信息，但不退出 - log.Printf("failed to execute command: %v, stderr: %s", err, stderr.String()) + return fmt.Errorf("ipset create failed: %v, stderr: %s", err, stderr.String()) } - return err // 返回错误以便调用者处理 + return nil } +// 添加 IP 到集合 func AddIPSet(setName string, ip string) error { cmd := exec.Command("ipset", "add", setName, ip) - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - err := cmd.Run() - if err != nil { - // 记录错误信息，但不退出 - log.Printf("failed to add IP to set: %v, stderr: %s", err, stderr.String()) - } - return err // 返回错误以便调用者处理 + // 不关注输出，只关注 exit code + return cmd.Run() } -// NumIPSet returns the number of entries in the specified ipset set. +// NumIPSet 获取集合中 IP 的数量 +// 注意：频繁调用此函数开销较大（因为它通过 shell 解析文本），建议仅在监控循环中低频调用 func NumIPSet(setName string) (int, error) { - cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list %s | grep \"Number of entries\" | cut -d ':' -f 2 | sed 's/ //g'", setName)) + // 这里必须用 sh -c 因为涉及管道符 | + cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list %s | grep 'Number of entries' | cut -d ':' -f 2", setName)) - var stdout, stderr bytes.Buffer + var stdout bytes.Buffer cmd.Stdout = &stdout - cmd.Stderr = &stderr err := cmd.Run() if err != nil { - log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String()) - return 0, fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String()) + return 0, err } output := strings.TrimSpace(stdout.String()) numEntries, err := strconv.Atoi(output) if err != nil { - log.Printf("failed to parse output as integer: %v, output: %s\n", err, output) - return 0, fmt.Errorf("failed to parse output as integer: %w, output: %s", err, output) + return 0, fmt.Errorf("解析数量失败: %w, 输出: %s", err, output) } return numEntries, nil } -// IsIpset 检查名为 setName 的 ipset 是否存在，通过返回 0 表示存在，非零表示不存在或其他错误。 +// Is_Name_Ipset 检查集合是否存在 +// 返回 0 表示存在 (ExitCode 0)，其他表示不存在 func Is_Name_Ipset(setName string) int { - cmd := exec.Command("ipset", "list", setName) + cmd := exec.Command("ipset", "list", setName, "-n") // -n 只输出名字，减少开销 err := cmd.Run() - if err != nil { - if exitError, ok := err.(*exec.ExitError); ok { - // The program has exited with an exit code != 0 - return exitError.ExitCode() - } else { - // Another error occurred (e.g., command not found) - return -1 // 或者你可以选择其他方式来标识这种情况 - } + return 1 // 不存在或出错 } - - // Command executed successfully, the set exists - return 0 + return 0 // 存在 } +// Is_Ip_Ipset 检查 IP 是否在 ANY root 集合中 +// 优化：不再使用 grep 全局搜索，而是遍历 root0...rootN 使用 ipset test func Is_Ip_Ipset(ip string) int { - cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list | grep \"%s\"", ip)) - err := cmd.Run() + // 遍历所有可能的集合名 + for i := 0; i < MAX_IPSET_NAME; i++ { + setName := fmt.Sprintf("root%d", i) - if err != nil { - if exitError, ok := err.(*exec.ExitError); ok { - // The program has exited with an exit code != 0 - return exitError.ExitCode() - } else { - // Another error occurred (e.g., command not found) - return -1 // 或者你可以选择其他方式来标识这种情况 + // 快速检查集合是否存在（可选，为了严谨） + // 如果 ipset test 在集合不存在时会报错，所以我们直接运行 test + // 如果 ip 在集合中，exit code 为 0 + cmd := exec.Command("ipset", "test", setName, ip) + err := cmd.Run() + + if err == nil { + return 1 // 找到了，IP 在该集合中 } } - // Command executed successfully, the set exists - return 0 + return 0 // 所有集合都找过了，不在其中 } -// RemoveIPIfInSets 在多个 ipset 中查找并删除某个 IP，成功删除则返回集合名 +// RemoveIPIfInSets 在所有集合中查找并删除某个 IP func RemoveIPIfInSets(prefix string, max int, ip string) (string, error) { for i := 0; i < max; i++ { setName := fmt.Sprintf("%s%d", prefix, i) - // 检查 IP 是否在当前集合中 - cmd := exec.Command("ipset", "test", setName, ip) - output, err := cmd.CombinedOutput() - - if err != nil { - if strings.Contains(string(output), "is NOT in set") { - continue // 不在该集合中，尝试下一个 - } - return "", fmt.Errorf("检测 %s 时出错: %v (%s)", setName, err, output) + // 1. 先测试是否存在 (避免无意义的 del 调用报错) + testCmd := exec.Command("ipset", "test", setName, ip) + if err := testCmd.Run(); err != nil { + // exit code != 0 说明不在这个集合里，继续找下一个 + continue } - // 存在，执行删除 - cmd = exec.Command("ipset", "del", setName, ip) - output, err = cmd.CombinedOutput() - if err != nil { - return "", fmt.Errorf("从 %s 删除 %s 时出错: %v (%s)", setName, ip, err, output) + // 2. 存在则删除 + delCmd := exec.Command("ipset", "del", setName, ip) + if err := delCmd.Run(); err != nil { + return "", fmt.Errorf("从 %s 删除 %s 失败: %v", setName, ip, err) } - return setName, nil // 删除成功，返回集合名 + return setName, nil // 成功删除 } - return "", nil // 所有集合都没有该 IP，无需删除 + return "", nil // 未找到 } -// 添加 Iptables 规则 +// iptables_add 添加规则 +// 修正：使用 -I (Insert) 而不是 -A (Append)，确保规则在最前面生效 func iptables_add(setName string) error { + // 检查规则是否已经存在，避免重复添加 + checkCmd := exec.Command("iptables", "-C", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP") + if err := checkCmd.Run(); err == nil { + // 规则已存在，直接返回 + return nil + } - cmd := exec.Command("sh", "-c", fmt.Sprintf("iptables -A INPUT -m set --match-set %s src -j DROP", setName)) - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout + // 添加规则 + cmd := exec.Command("iptables", "-I", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP") + var stderr bytes.Buffer cmd.Stderr = &stderr err := cmd.Run() if err != nil { - //log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String()) - //err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String()) + log.Printf("添加 iptables 规则失败 [%s]: %v, stderr: %s", setName, err, stderr.String()) } - return err } -// 删除 Iptables 规则 +// iptables_del 删除规则 func iptables_del(setName string) error { - - cmd := exec.Command("sh", "-c", fmt.Sprintf("iptables -D INPUT -m set --match-set %s src -j DROP", setName)) - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - err := cmd.Run() - if err != nil { - //log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String()) - //err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String()) + // 循环删除，防止有重复规则 + for { + cmd := exec.Command("iptables", "-D", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP") + if err := cmd.Run(); err != nil { + // 删除失败通常意味着规则不存在了，跳出循环 + break + } } - - return err + return nil } -// 打印 Iptables 规则 +// iptables_list 打印规则 func iptables_list() error { - cmd := exec.Command("sh", "-c", "iptables -L -v -n --line-numbers") - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - err := cmd.Run() - if err != nil { - log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String()) - err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String()) - } - - fmt.Print(stdout.String()) - return err + cmd := exec.Command("iptables", "-L", "INPUT", "-v", "-n", "--line-numbers") + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + return cmd.Run() } /* diff --git a/main.go b/main.go index 02e944d..07d596f 100644 --- a/main.go +++ b/main.go @@ -1,7 +1,6 @@ package main import ( - "container/list" "encoding/json" "flag" "fmt" @@ -13,11 +12,12 @@ import ( "runtime" "strings" "sync" + "sync/atomic" "syscall" "time" ) -var BuildDate = "unknown" // 由编译时注入 +var BuildDate = "unknown" func init() { // 强制使用 Go 的纯用户态 DNS 解析器 @@ -28,34 +28,311 @@ func init() { var ( daemon = flag.Bool("d", false, "守护进程模式") - child = flag.Bool("child", false, "子进程模式, (不要使用！！！)") + child = flag.Bool("child", false, "子进程模式") ) -// 全局变量 +// 全局配置变量 var ( - InterfacesList bool // 是否列出网络接口 - InterfaceName *string // 网络接口名称 - PcapFile *string // 输出文件名 - Protocol *string // BPF 过滤器 + InterfacesList bool + InterfaceName *string + PcapFile *string + Protocol *string - IPSET_NUMBER int // 当前使用的 ipset 集合编号 - MAX_IPSET_NAME = 100 // 最大 ipset 集合数量 - IPSET_NAME string // 当前使用的 ipset 集合名称 + IPSET_NUMBER int + MAX_IPSET_NAME = 100 + IPSET_NAME string - IpList = list.New() // 存储 IPv4 地址的链表 - IpMutex sync.Mutex // 保护 ipList 的互斥锁 + // --- 优化部分开始 --- - ProcessedIPMap = map[string]time.Time{} // 使用 map 存储已处理的 IP - ProcessedMutex sync.Mutex // 互斥锁保护 ProcessedIPMap + // 替换 List 为带缓冲的 Channel + // 容量设为 5000，足以应对大多数突发流量 + IpChannel = make(chan string, 5000) + + // 用于入队去重的 Map，防止同一个 IP 在处理中时重复入队 + // key: IP string, value: struct{} + PendingIPs sync.Map + PendingCount int64 + + // --- 优化部分结束 --- + + ProcessedIPMap = map[string]time.Time{} + ProcessedMutex sync.Mutex local_ipv4_addr string ) -// 启动子进程 +// --- 核心逻辑优化：生产者 (入队) --- +// PushIPToQueue 安全地将 IP 放入队列 +// 该函数应由 startPacketCapture 调用 +func PushIPToQueue(ipStr string) { + // 0. 基础校验 + if ipStr == "" { + return + } + + // 1. 快速检查：如果已经在 ProcessedIPMap (已知的国内IP或白名单)，直接忽略 + ProcessedMutex.Lock() + _, processed := ProcessedIPMap[ipStr] + ProcessedMutex.Unlock() + if processed { + return + } + + // 2. 去重检查：如果已经在 PendingIPs (队列中或正在处理)，跳过 + // LoadOrStore: 如果 key 存在，返回 true；否则写入并返回 false + if _, loaded := PendingIPs.LoadOrStore(ipStr, struct{}{}); loaded { + return + } + + // 【计数增加】只有确定要入队时才增加 + atomic.AddInt64(&PendingCount, 1) + + // 3. 非阻塞入队 + select { + case IpChannel <- ipStr: + // 成功入队 + default: + // 队列已满，丢弃该包 + // 【回滚状态】从 Map 删除，并减少计数 + PendingIPs.Delete(ipStr) + atomic.AddInt64(&PendingCount, -1) + // 可选：log.Println("警告：处理队列已满，丢弃 IP:", ipStr) + } +} + +// --- 核心逻辑优化：消费者 (处理 IP) --- +func processIP(ipStr string) { + // PushIPToQueue 已经拦截了空字符串，这里其实不需要再判断 + // 但为了代码健壮性，如果必须判断，defer 必须在 return 之前定义 + + // 确保函数结束时从 Pending 状态移除并减少计数 + defer func() { + PendingIPs.Delete(ipStr) + atomic.AddInt64(&PendingCount, -1) + }() + + if ipStr == "" { + return + } + + // 再次检查 ProcessedIPMap (防止排队期间被其他协程处理了) + ProcessedMutex.Lock() + _, processed := ProcessedIPMap[ipStr] + ProcessedMutex.Unlock() + if processed { + return + } + + // 检查白名单 + // --- 修改开始：使用读写锁检查白名单 --- + whiteListLock.RLock() // 加读锁 + _, isWhitelisted := whiteList[ipStr] // 检查是否存在 + whiteListLock.RUnlock() // 解读锁 + + if isWhitelisted { + log.Printf("\033[33m %s 在白名单中, 跳过 \033[0m\n", ipStr) + // 尝试从 ipset 移除 + RemoveIPIfInSets("root", MAX_IPSET_NAME, ipStr) + return + } + // --- 修改结束 --- + + REGION := "中国 内网" + + // 如果 IP 已经在 ipset 中，通常无需处理 + if Is_Ip_Ipset(ipStr) == 1 { + //log.Printf("\033[31m %s 已在 ipset 集合中 \033[0m\n", ipStr) + return + } + + // 1. 离线库判断 (快速) + region, _ := ip2region(ipStr) + + // 如果不包含 "中国" 也不包含 "内网"，则判定为疑似国外 + if !ContainsPart(region, REGION) { + log.Printf("\033[33m [%s %s] 离线库为国外, 进一步API判断\033[0m\n", ipStr, region) + + // 2. 在线 API 判断 (慢速) + position, err := curl_(ipStr) + if err != nil { + log.Printf("获取IP地域出错: %v", err) + return // API 失败暂时跳过，等待下次重试 + } + + log.Printf("\033[31m [%s %s]\033[0m\n", ipStr, position) + + if !ContainsPart(position, REGION) { + // --- 确认为国外 --- + AddIPSet(IPSET_NAME, ipStr) + log.Printf("\033[31m [封禁] 已添加国外 IP: %s \033[0m\n", ipStr) + } else { + // --- 确认为国内 --- + log.Printf("\033[32m %s API 修正为国内, 标记放行\033[0m\n", ipStr) + ProcessedMutex.Lock() + ProcessedIPMap[ipStr] = time.Now() + ProcessedMutex.Unlock() + } + } else { + // 离线库确认为国内，标记放行 + ProcessedMutex.Lock() + ProcessedIPMap[ipStr] = time.Now() + ProcessedMutex.Unlock() + } +} + +func RunMainProcess() { + log.Println(" 主进程启动...") + + WriteLocalAddr() + + cmd, err := StartChildProcess() + if err != nil { + log.Fatalf("子进程启动失败: %v", err) + } + + IPSET_NUMBER = 0 + IPSET_NAME = fmt.Sprintf("root%d", IPSET_NUMBER) + if Is_Name_Ipset(IPSET_NAME) == 0 { // 假设 0 表示不存在 + createIPSet(IPSET_NAME) + } + + // 1. 启动抓包 (生产者) + go startPacketCapture() + + // 2. 启动 Worker Pool (消费者) + // 根据机器性能调整 worker 数量，建议 20-50 + numWorkers := 40 + log.Printf(" 启动 %d 个并发 Worker 处理 IP...", numWorkers) + for i := 0; i < numWorkers; i++ { + go func(id int) { + for ip := range IpChannel { + processIP(ip) + } + }(i) + } + + // 定期保存 Map 数据 (替代原来在循环里保存) + go func() { + ticker := time.NewTicker(1 * time.Minute) + for range ticker.C { + if err := saveMapToFile("cn.json"); err != nil { + log.Printf(" 自动保存 Map 失败: %v", err) + } + } + }() + + // 过期清理 ProcessedIPMap + go func() { + ticker := time.NewTicker(1 * time.Minute) + for range ticker.C { + now := time.Now() + ProcessedMutex.Lock() + count := 0 + for ip, t := range ProcessedIPMap { + if t.Year() == 1971 { + continue + } + if now.Sub(t) > 30*time.Minute { + delete(ProcessedIPMap, ip) + count++ + } + } + ProcessedMutex.Unlock() + if count > 0 { + log.Printf(" 已清理 %d 个过期 ProcessedIPMap 项", count) + } + } + }() + + // 白名单刷新 + go func() { + for { + time.Sleep(10 * time.Minute) + if err := LoadWhiteList("whitelist.txt"); err != nil { + log.Printf(" 刷新白名单失败: %v", err) + } + } + }() + + // 防火墙扩容管理 + // 防火墙扩容管理 + go func() { + for { + time.Sleep(1 * time.Second) + + // 1. 获取当前集合长度，必须处理错误 + ipset_len, err := NumIPSet(IPSET_NAME) + if err != nil { + // 如果是因为集合不存在导致的错误，尝试创建它 + if Is_Name_Ipset(IPSET_NAME) != 0 { + log.Printf("检测到集合 %s 不存在，正在初始化...", IPSET_NAME) + createIPSet(IPSET_NAME) + iptables_add(IPSET_NAME) // 重点：创建后必须同步添加 iptables 规则 + } + continue + } + + // 2. 检查是否需要扩容 + // 注意：ipset 默认 maxelem 是 65536，达到这个值 add 操作就会失败 + if ipset_len >= 65530 { + log.Printf("\033[31m ipset %s 列表已满 %d，准备扩容... \033[0m\n", IPSET_NAME, ipset_len) + + IPSET_NUMBER++ + if IPSET_NUMBER >= MAX_IPSET_NAME { + log.Printf("\033[31m 警告：已达到最大集合数量限制!!! \033[0m\n") + // 这里可以根据需求决定是否 return 或采取其他措施 + } + + newSetName := fmt.Sprintf("root%d", IPSET_NUMBER) + + // 3. 只有不存在时才创建 (注意判断逻辑: != 0 表示不存在) + if Is_Name_Ipset(newSetName) != 0 { + log.Printf("\033[32m 正在创建并应用新集合: %s \033[0m\n", newSetName) + if err := createIPSet(newSetName); err == nil { + // 4. 关键：创建新集合后，必须立刻将其加入 iptables 拦截规则 + iptables_add(newSetName) + // 5. 切换全局变量 + IPSET_NAME = newSetName + } else { + log.Printf("创建集合失败: %v", err) + } + } else { + // 如果集合已经存在（可能是上次运行留下的），直接切换过去 + IPSET_NAME = newSetName + iptables_add(newSetName) // 确保规则存在 + } + } + } + }() + + // 打印日志 + go func() { + for { + time.Sleep(7 * time.Second) + + // 获取 ipset 数量 + ipset_len, _ := NumIPSet(IPSET_NAME) + + // 安全地获取 ProcessedIPMap 的长度（不要打印内容，不要在无锁状态下读取！） + ProcessedMutex.Lock() + processedLen := len(ProcessedIPMap) + ProcessedMutex.Unlock() + + // 获取 Pending 数量 + pendingCount := atomic.LoadInt64(&PendingCount) + + log.Printf("\033[32m [状态监控] IPSet(%s): %d | 已处理缓存: %d | 待处理积压: %d \033[0m\n", + IPSET_NAME, ipset_len, processedLen, pendingCount) + } + }() + + waitForSignalAndCleanUp(cmd) +} + func StartChildProcess() (*exec.Cmd, error) { args := []string{} for _, arg := range os.Args[1:] { - if !strings.HasPrefix(arg, "-child") { // 只过滤 -child 标志 + if !strings.HasPrefix(arg, "-child") { args = append(args, arg) } } @@ -71,374 +348,88 @@ func StartChildProcess() (*exec.Cmd, error) { 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) + return fmt.Errorf("SIGTERM 失败: %w", err) } - done := make(chan error, 1) - go func() { - done <- cmd.Wait() - }() - + 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) - } + return err + case <-time.After(2 * time.Second): + _ = cmd.Process.Kill() } - - 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 - } + sig := <-sigChan + fmt.Printf("主进程收到信号: %v\n", sig) + if cmd != nil && cmd.Process != nil { + _ = cmd.Process.Signal(sig) } - - if err := StopChildProcess(cmd); err != nil { - log.Fatalf("清理子进程时遇到错误: %v", err) - } - fmt.Println("主进程退出") + StopChildProcess(cmd) + saveMapToFile("cn.json") } -// 守护进程模式 func StartDaemon() { - // 创建一个新的实例并让当前进程退出。注意这并不是守护进程的标准实现。 args := []string{} for _, arg := range os.Args[1:] { - if !strings.HasPrefix(arg, "-d") && !strings.HasPrefix(arg, "-child") { // 过滤掉 -d 和 -child 标志 + if !strings.HasPrefix(arg, "-d") && !strings.HasPrefix(arg, "-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("新实例已启动，当前进程将退出") + cmd.Start() 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: - - } - } + signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM) + <-sigChan } -// 保存ProcessedIPMap到文件 func saveMapToFile(filePath string) error { ProcessedMutex.Lock() defer ProcessedMutex.Unlock() - file, err := os.Create(filePath) if err != nil { - return fmt.Errorf("创建文件失败: %w", err) + return err } defer file.Close() - - encoder := json.NewEncoder(file) - if err := encoder.Encode(ProcessedIPMap); err != nil { - return fmt.Errorf("编码 ProcessedIPMap 失败: %w", err) - } - - //log.Println("ProcessedIPMap 已成功保存到文件") - return nil + return json.NewEncoder(file).Encode(ProcessedIPMap) } -func loadFromFile(filePath string, logMessage string) error { +func loadFromFile(filePath string) error { ProcessedMutex.Lock() defer ProcessedMutex.Unlock() - file, err := os.Open(filePath) if err != nil { - if os.IsNotExist(err) { - log.Println("文件不存在，跳过加载 Map") - return nil - } - return fmt.Errorf("打开文件失败: %w", err) + return nil } defer file.Close() - - // 尝试用新格式解码 - decoder := json.NewDecoder(file) - temp := make(map[string]time.Time) - if err := decoder.Decode(&temp); err == nil { - ProcessedIPMap = temp - log.Println(logMessage) - return nil - } - - // 如果失败，尝试旧格式 - file.Seek(0, 0) // 重置读取位置 - decoder = json.NewDecoder(file) - oldTemp := make(map[string]string) - if err := decoder.Decode(&oldTemp); err == nil { - for ip := range oldTemp { - //ProcessedIPMap[ip] = time.Now() // 给旧 IP 打个当前时间戳 - ProcessedIPMap[ip] = time.Date(1971, 1, 1, 0, 0, 0, 0, time.UTC) // 标记为永不过期 - } - log.Println(logMessage + "（从旧格式转换）") - return nil - } - - return fmt.Errorf("解码 Map 失败: %w", err) + return json.NewDecoder(file).Decode(&ProcessedIPMap) } func InitMap() { - if err := loadFromFile("cn.json", " Map 已成功从文件加载"); err != nil { - log.Fatalf(" 加载 Map 失败: %v", err) - } - - defer func() { - // 程序退出时保存数据 - if err := saveMapToFile("cn.json"); err != nil { - log.Printf(" 保存 Map 失败: %v", err) - } - }() - + loadFromFile("cn.json") } func WriteLocalAddr() { - local_ipv4_addr = GetLocalIpv4Addr() // 本机地址 - // 将本机外网地址加入到已处理集合中 + local_ipv4_addr = GetLocalIpv4Addr() if local_ipv4_addr != "NULL" { - //log.Printf("\033[33m %s 本机地址 \033[0m\n", ipStr) ProcessedMutex.Lock() - //ProcessedIPMap[local_ipv4_addr] = time.Now() - ProcessedIPMap[local_ipv4_addr] = time.Date(1971, 1, 1, 0, 0, 0, 0, time.UTC) // 标记为永不过期 + ProcessedIPMap[local_ipv4_addr] = time.Date(1971, 1, 1, 0, 0, 0, 0, time.UTC) ProcessedMutex.Unlock() } - - // 写入json文件 - if err := saveMapToFile("cn.json"); err != nil { - log.Printf("实时保存 Map 失败: %v", err) - } - -} - -func RunMainProcess() { // 主进程逻辑 - - log.Println(" 主进程启动...") - - WriteLocalAddr() // 将本机外网地址加入到已处理集合中 - - 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) - IpList.Remove(e1) - goto next - } - - // 检查是否在白名单中 - if _, ip_ := whiteList[ipStr]; ip_ { - log.Printf("\033[33m %s 跳过白名单, 跳过！！！ \033[0m\n", ipStr) - IpList.Remove(e1) - - setName, err := RemoveIPIfInSets("root", MAX_IPSET_NAME, ipStr) - if err != nil { - log.Printf(" %s 删除 IP 出错: %v\n", ipStr, err) - } else if setName != "" { - log.Printf(" %s 已从 %s 中移除 \n", ipStr, setName) - } else { - log.Printf(" %s 不在任何 IPSet 中，无需移除\n", ipStr) - } - - goto next - } - - REGION := "中国 内网" // 默认地域 - - if Is_Ip_Ipset(ipStr) != 0 { // IP 不在 ipset 集合中 - - //if !strings.Contains(region, "中国") && !strings.Contains(region, "内网") { // 离线库判断不在中国内 - if !ContainsPart(region, 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 判断为国外 - if !ContainsPart(region, REGION) { - AddIPSet(IPSET_NAME, ipStr) // 添加 IP 到 ipset 集合 - - // 钉钉告警，废弃！钉钉可能限制文本长度，和发送次数！ - // warning_ding(ipStr, position) // 警告 IP 地域 - } else { - log.Printf("\033[33m %s 离线库为国外, API 判断为国内, 标记为已处理\033[0m\n", ipStr) - - // 将 IP 标记为已处理，国内地址 - ProcessedMutex.Lock() - ProcessedIPMap[ipStr] = time.Now() - ProcessedMutex.Unlock() - - // 写入json文件 - if err := saveMapToFile("cn.json"); err != nil { - log.Printf(" 实时保存 Map 失败: %v", err) - } - } - } - } - - } else { // IP 已在 ipset 集合中 - log.Printf("\033[31m %s 在 ipset 集合中 \033[0m\n", ipStr) - } - - // 无论是否已处理，都移除该 IP - IpList.Remove(e1) - } - next: - IpMutex.Unlock() // 解锁互斥锁 - - // 打印当前info - log.Printf(" 当前Ip链表长度:%d, Ipset名:%s, 长:%d, ProcessedIPMap:[%s]当前长度:%d\n", IpList.Len(), IPSET_NAME, func() int { // 打印 当前 Ipset 链长度 - _len, _err := NumIPSet(IPSET_NAME) - if _err == nil { - return _len - } - return 0 - }(), ProcessedIPMap, len(ProcessedIPMap)) - - time.Sleep(1 * time.Second) // 防止高频运行 - } - }() - - // 启动一个 goroutine（后台任务），用于定期清理 ProcessedIPMap 中过期的 IP 条目 - go func() { - for { - time.Sleep(1 * time.Minute) - - // 获取当前时间，用于与 map 中的时间戳进行比较 - now := time.Now() - - // 加锁，确保并发安全地访问全局变量 ProcessedIPMap - ProcessedMutex.Lock() - - // 遍历 ProcessedIPMap，检查每个 IP 的记录是否已超过 24 小时 - for ip, t := range ProcessedIPMap { - - if t.Year() == 1971 { - continue // 不清理标记为“永不过期”的 IP - } - - // 如果当前时间减去记录时间大于 1 小时，则删除该条目 - if now.Sub(t) > 30*time.Minute { - delete(ProcessedIPMap, ip) - } - } - - // 解锁，允许其他 goroutine 访问 ProcessedIPMap - ProcessedMutex.Unlock() - - // 打印日志，表示本次清理已完成 - log.Println(" 已清理过期 ProcessedIPMap 项") - } - }() - - // 定时重新加载白名单 - go func() { - for { - time.Sleep(1 * time.Minute) // 每 10 分钟自动刷新 - if err := LoadWhiteList("whitelist.txt"); err != nil { - log.Printf(" 刷新白名单失败: %v", err) - } - } - }() - - // 启动防火墙管理 - go func() { - for { - if ipset_len, _ := NumIPSet(IPSET_NAME); ipset_len >= 65535 { - 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(3 * time.Second) - } - }() - - // 等待信号并清理 - waitForSignalAndCleanUp(cmd) } func HandleCmd() { @@ -449,7 +440,11 @@ func HandleCmd() { flag.BoolVar(&InterfacesList, "l", false, "列出可用的网络接口") Protocol = flag.String("f", "'tcp' or 'udp' or 'tcp or udp'", "指定 BPF 过滤器") PcapFile = flag.String("o", "", "保存捕获数据的输出文件（可选）") - flag.StringVar(&instruction, "s", "", "-s start 启动 Iptables 规则\n-s stop 停止 Iptables 规则\n-s list 打印 Iptables 规则\n-s reload 重启 Iptables 规则") + flag.StringVar(&instruction, "s", "", + "-s start 启动 Iptables 规则\n"+ + "-s stop 停止 Iptables 规则\n"+ + "-s list 打印 Iptables 规则\n"+ + "-s reload 重启 Iptables 规则") flag.BoolVar(&help, "h", false, "") flag.BoolVar(&help, "help", false, "帮助信息") flag.Parse() @@ -519,34 +514,21 @@ func HandleCmd() { } func main() { - runtime.GOMAXPROCS(runtime.NumCPU()) // 设置最大CPU核数 + runtime.GOMAXPROCS(runtime.NumCPU()) HandleCmd() CheckCommandExists("iptables") CheckCommandExists("ipset") - embed_ip2region() - // 加载白名单 - err := LoadWhiteList("whitelist.txt") - if err != nil { - log.Fatalf(" whiteList Map 加载白名单失败: %v", err) - } else { - log.Println(" whiteList Map 白名单加载成功") - } - - // 守护进程模式 if *daemon { StartDaemon() } - - // 子进程逻辑 if *child { - RunChildProcess() // 子进程逻辑 + RunChildProcess() return } InitMap() - RunMainProcess() // 主进程逻辑 - + RunMainProcess() }