eruhs/core/capture.go

package core

import (
    "context"
    "eruhs/core/parser"
    "github.com/dustin/go-humanize"
    "github.com/google/gopacket"
    "github.com/google/gopacket/pcap"
    "time"
)

type DataPack struct {
    TotalLength int
    LayerStack  []parser.Layer
}

func (d *DataPack) AddLayer(layer parser.Layer) {
    d.LayerStack = append(d.LayerStack, layer)
}

type Capture struct {
    handle      *pcap.Handle
    started     bool
    stopFunc    func()
    callback    func(pack *DataPack)
    errCallback func(err error)
}

func NewCapture(callback func(pack *DataPack), errCallback func(err error)) *Capture {
    return &Capture{
        callback:    callback,
        errCallback: errCallback,
        started:     false,
    }
}

// StartCaptureOnDevice 启动抓包
func (c *Capture) StartCaptureOnDevice(device string) {
    if c.started {
        return
    } else {
        c.started = true
    }

    handle, err := pcap.OpenLive(device, 4*humanize.KByte, false, 1*time.Second)
    if err != nil {
        c.errCallback(err)
        return
    }

    c.handle = handle
    ctx, cancel := context.WithCancel(context.Background())
    c.stopFunc = cancel

    go func(ctx context.Context) {
        defer c.handle.Close()
        packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
        for packet := range packetSource.Packets() {
            select {
            case <-ctx.Done():
                return
            default:
                parsedPack, err := c.handlePacket(packet)
                if err != nil {
                    if parsedPack != nil {
                        c.callback(parsedPack)
                    } else {
                        c.errCallback(err)
                    }
                }

                c.callback(parsedPack)
            }
        }
    }(ctx)
}

func (c *Capture) Pause() {

}

func (c *Capture) Stop() {
    c.stopFunc()
    c.started = false
}