eruhs/app.go

package main

import (
    "context"
    "encoding/hex"
    "eruhs/core"
    "eruhs/core/parser"
    "eruhs/describer"
    "fmt"
    "github.com/wailsapp/wails/v2/pkg/runtime"
    "time"
)

// App struct
type App struct {
    ctx     context.Context
    capture *core.Capture
}

// NewApp creates a new App application struct
func NewApp() *App {
    return &App{}
}

// startup is called when the app starts. The context is saved
// so we can call the runtime methods
func (a *App) startup(ctx context.Context) {
    a.ctx = ctx
}

func (a *App) GetNetworkDeviceList() ([]core.Device, error) {
    devices, err := core.GetAllDevice()
    if err != nil {
        return nil, err
    }

    return devices, nil
}

var globalPackCnt = 0
var startTime = time.Now().UnixMilli()
var prevSendTime = time.Now().UnixMilli()
var batchSendBuffer []UIDataPackView
var packetMap = map[int]*core.DataPack{}

func (a *App) StartCapture(device string) {
    if a.capture == nil {
        a.capture = core.NewCapture(packCallback(a.ctx), errCallback(a.ctx))
    }

    globalPackCnt = 0
    startTime = time.Now().UnixMilli()
    prevSendTime = time.Now().UnixMilli()
    batchSendBuffer = nil
    batchSendBuffer = make([]UIDataPackView, 0, 1000)
    runtime.LogPrint(a.ctx, "StartCapture")

    a.capture.StartCaptureOnDevice(device)
}

func (a *App) PauseCapture() {

}

func (a *App) StopCapture() {
    a.capture.Stop()
}

func (a *App) GetDetail(id int) map[int]LayerDescribe {
    pack := packetMap[id]
    if pack == nil {
        return map[int]LayerDescribe{}
    }

    result := map[int]LayerDescribe{}
    for _, layer := range pack.LayerStack {
        switch layer.LayerType() {
        case parser.LayerTypeEthernet:
            eth := layer.(*parser.Ethernet)
            result[0] = LayerDescribe{
                HeaderRaw: hex.Dump(eth.Header),
                LayerName: "ETH",
                Describe:  (&describer.EthernetDescriber{Eth: eth}).GetDetailDescribe(),
            }
        case parser.LayerTypeARP:
            arp := layer.(*parser.ARP)
            result[1] = LayerDescribe{
                HeaderRaw: hex.Dump(arp.Header),
                LayerName: "ARP",
                Describe:  (&describer.ARPDescriber{ARP: arp}).GetDetailDescribe(),
            }
        case parser.LayerTypeIPv4:
            ipv4 := layer.(*parser.IPv4)
            result[1] = LayerDescribe{
                HeaderRaw: hex.Dump(ipv4.Header),
                LayerName: "IPv4",
                Describe:  (&describer.IPv4Describer{IPv4: ipv4}).GetDetailDescribe(),
            }
        case parser.LayerTypeTCP:
            tcp := layer.(*parser.TCP)
            result[2] = LayerDescribe{
                HeaderRaw:             hex.Dump(tcp.Header),
                ApplicationPayloadRaw: hex.Dump(tcp.Payload),
                LayerName:             "TCP",
                Describe:              (&describer.TCPDescriber{TCP: tcp}).GetDetailDescribe(),
            }
        case parser.LayerTypeUDP:
            udp := layer.(*parser.UDP)
            result[2] = LayerDescribe{
                HeaderRaw:             hex.Dump(udp.Header),
                ApplicationPayloadRaw: hex.Dump(udp.Payload),
                LayerName:             "UDP",
                Describe:              (&describer.UDPDescriber{UDP: udp}).GetDetailDescribe(),
            }
        default:
        }
    }

    return result
}

func errCallback(ctx context.Context) func(err error) {
    return func(err error) {
        fmt.Println(err)
        runtime.EventsEmit(ctx, "onErrOccurrence", err.Error())
    }
}

func packCallback(ctx context.Context) func(pack *core.DataPack) {
    return func(pack *core.DataPack) {
        globalPackCnt++
        packetMap[globalPackCnt] = pack
        dataPackView := UIDataPackView{
            Id:     globalPackCnt,
            Time:   time.Duration(time.Now().UnixMilli() - startTime),
            Source: "",
            Dest:   "",
            Proto:  "",
            Length: pack.TotalLength,
        }

        for _, layer := range pack.LayerStack {
            switch layer.LayerType() {
            case parser.LayerTypeEthernet:
                eth := layer.(*parser.Ethernet)
                dataPackView.Source = eth.Source.String()
                dataPackView.Dest = eth.Destination.String()
                dataPackView.Proto = "ETH"
            case parser.LayerTypeARP:
                dataPackView.Proto = "ARP"
            case parser.LayerTypeIPv4:
                ipv4 := layer.(*parser.IPv4)
                dataPackView.Source = ipv4.SrcIP.String()
                dataPackView.Dest = ipv4.DstIP.String()
                dataPackView.Proto = "IPv4"
            case parser.LayerTypeTCP:
                tcp := layer.(*parser.TCP)
                dataPackView.Source = fmt.Sprintf("%s:%d", dataPackView.Source, tcp.SrcPort)
                dataPackView.Dest = fmt.Sprintf("%s:%d", dataPackView.Dest, tcp.DstPort)
                dataPackView.Proto = "TCP"
            case parser.LayerTypeUDP:
                udp := layer.(*parser.UDP)
                dataPackView.Source = fmt.Sprintf("%s:%d", dataPackView.Source, udp.SrcPort)
                dataPackView.Dest = fmt.Sprintf("%s:%d", dataPackView.Dest, udp.DstPort)
                dataPackView.Proto = "UDP"
            default:
            }
        }

        if (time.Now().UnixMilli() - prevSendTime) > int64(0.2*1000) {
            runtime.EventsEmit(ctx, "onNewDataPacketBatch", batchSendBuffer)
            batchSendBuffer = nil
            batchSendBuffer = make([]UIDataPackView, 0, 1000)
            prevSendTime = time.Now().UnixMilli()
        } else {
            batchSendBuffer = append(batchSendBuffer, dataPackView)
        }

    }
}