init

internal/capture/capture.go

@@ -0,0 +1,162 @@
+package capture
+
+import (
+	"fmt"
+	"log"
+	"sync"
+	"time"
+
+	"equipment-analyzer/internal/model"
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	"github.com/google/gopacket/pcap"
+)
+
+type PacketCapture struct {
+	handle       *pcap.Handle
+	isCapturing  bool
+	stopChan     chan bool
+	mutex        sync.RWMutex
+	tcpProcessor *TCPProcessor
+	dataChan     chan *model.TCPData
+	errorChan    chan error
+}
+
+type Config struct {
+	InterfaceName string
+	Filter        string
+	Timeout       time.Duration
+	BufferSize    int
+}
+
+func NewPacketCapture() *PacketCapture {
+	return &PacketCapture{
+		stopChan:     make(chan bool),
+		tcpProcessor: NewTCPProcessor(),
+		dataChan:     make(chan *model.TCPData, 1000),
+		errorChan:    make(chan error, 100),
+	}
+}
+
+func (pc *PacketCapture) Start(config Config) error {
+	pc.mutex.Lock()
+	defer pc.mutex.Unlock()
+
+	if pc.isCapturing {
+		return fmt.Errorf("capture already running")
+	}
+
+	// 打开网络接口
+	handle, err := pcap.OpenLive(
+		config.InterfaceName,
+		int32(config.BufferSize),
+		true, // promiscuous
+		config.Timeout,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to open interface: %v", err)
+	}
+
+	// 设置过滤器
+	if err := handle.SetBPFFilter(config.Filter); err != nil {
+		handle.Close()
+		return fmt.Errorf("failed to set filter: %v", err)
+	}
+
+	pc.handle = handle
+	pc.isCapturing = true
+
+	// 启动抓包协程
+	go pc.captureLoop()
+
+	return nil
+}
+
+func (pc *PacketCapture) Stop() {
+	pc.mutex.Lock()
+	defer pc.mutex.Unlock()
+
+	if !pc.isCapturing {
+		return
+	}
+
+	pc.isCapturing = false
+	close(pc.stopChan)
+
+	if pc.handle != nil {
+		pc.handle.Close()
+	}
+}
+
+func (pc *PacketCapture) IsCapturing() bool {
+	pc.mutex.RLock()
+	defer pc.mutex.RUnlock()
+	return pc.isCapturing
+}
+
+func (pc *PacketCapture) captureLoop() {
+	packetSource := gopacket.NewPacketSource(pc.handle, pc.handle.LinkType())
+
+	log.Println("[抓包] 开始监听数据包...")
+
+	for {
+		select {
+		case <-pc.stopChan:
+			return
+		default:
+			packet, err := packetSource.NextPacket()
+			if err != nil {
+				if err.Error() == "Timeout Expired" {
+					// 静默跳过超时
+					continue
+				}
+				log.Printf("Error reading packet: %v", err)
+				continue
+			}
+
+			// 处理TCP包
+			pc.processTCPPacket(packet)
+		}
+	}
+}
+
+func (pc *PacketCapture) processTCPPacket(packet gopacket.Packet) {
+	tcpLayer := packet.Layer(layers.LayerTypeTCP)
+	if tcpLayer == nil {
+		return
+	}
+
+	tcp, ok := tcpLayer.(*layers.TCP)
+	if !ok {
+		return
+	}
+
+	// 提取TCP负载
+	if len(tcp.Payload) == 0 {
+		return
+	}
+
+	// 创建TCP数据包
+	tcpData := &model.TCPData{
+		Payload: tcp.Payload,
+		Seq:     uint32(tcp.Seq),
+		Ack:     uint32(tcp.Ack),
+		SrcPort: uint16(tcp.SrcPort),
+		DstPort: uint16(tcp.DstPort),
+	}
+
+	// 发送给TCP处理器
+	pc.tcpProcessor.ProcessPacket(tcpData)
+}
+
+func (pc *PacketCapture) GetCapturedData() []string {
+	return pc.tcpProcessor.GetFinalBuffer()
+}
+
+func (pc *PacketCapture) ProcessAllData() {
+	pc.tcpProcessor.ProcessAllData()
+}
+
+func (pc *PacketCapture) Clear() {
+	pc.tcpProcessor.Clear()
+}