DenyIP-go/cap.go

package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/pcapgo"
)

// 打印可用的网络接口信息
func printAvailableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("可用的网络接口:")
	for _, device := range devices {
		fmt.Printf("名称: %s\n描述: %s\n地址: %+v\n\n", device.Name, device.Description, device.Addresses)
	}
}

// 判断 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
}

// 打印捕获到的每个数据包的信息
func printPacketInfo(packet gopacket.Packet) {
	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)

		}
	}
}

func startPacketCapture() {
	handle, err := pcap.OpenLive(*InterfaceName, 65535, true, pcap.BlockForever)
	if err != nil {
		log.Fatalf("打开网络接口 %s 出错: %v", *InterfaceName, err)
	}
	defer func() {
		fmt.Println("清理资源...")
		if handle != nil {
			handle.Close()
		}
	}()

	err = handle.SetBPFFilter(*Protocol)
	if err != nil {
		log.Fatalf("设置 BPF 过滤器出错: %v", err)
	}

	var pcapWriter *pcapgo.Writer
	if *PcapFile != "" {
		file, err := os.Create(*PcapFile)
		if err != nil {
			log.Fatalf("创建输出文件出错: %v", err)
		}
		defer file.Close()

		pcapWriter = pcapgo.NewWriter(file)
		err = pcapWriter.WriteFileHeader(65535, layers.LinkTypeEthernet)
		if err != nil {
			log.Fatalf("写入全局头部出错: %v", err)
		}
	}

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	fmt.Printf("正在监听网络接口 %s，使用过滤器 '%s'...\n", *InterfaceName, *Protocol)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	go func() {
		for packet := range packetSource.Packets() {
			printPacketInfo(packet)
			if pcapWriter != nil {
				err := pcapWriter.WritePacket(packet.Metadata().CaptureInfo, packet.Data())
				if err != nil {
					log.Printf("写入数据包出错: %v", err)
				}
			}
		}
	}()

	<-sigChan
	fmt.Println("\n停止抓包...")
}