feat: 信令跟踪功能接口

src/modules/trace/service/packet.go

@@ -0,0 +1,191 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strings"
+	"sync"
+	"time"
+
+	"be.ems/src/framework/logger"
+	"be.ems/src/framework/vo"
+	"github.com/gopacket/gopacket"
+	"github.com/gopacket/gopacket/pcap"
+	"github.com/gopacket/gopacket/pcapgo"
+)
+
+// 实例化服务层 Packet 结构体
+var NewPacket = &Packet{
+	taskMap: sync.Map{},
+}
+
+// 信令跟踪 服务层处理
+type Packet struct {
+	taskMap sync.Map // 捕获任务
+}
+
+// task  任务信息
+type task struct {
+	TaskNo  string          // 任务编号
+	Handle  *pcap.Handle    // 捕获句柄
+	File    *os.File        // 捕获信息输出文件句柄
+	Writer  *pcapgo.Writer  // 捕获信息输出句柄
+	Expire  time.Time       // 过期时间
+	context context.Context // 上下文 控制完成结束
+}
+
+// NetworkDevices 获取网卡设备信息
+func (s *Packet) NetworkDevices() []vo.TreeSelect {
+	arr := make([]vo.TreeSelect, 0)
+	devices, err := pcap.FindAllDevs()
+	if err != nil {
+		logger.Errorf("interfaces find all devices err: %s", err.Error())
+		return arr
+	}
+
+	for _, device := range devices {
+		if len(device.Addresses) == 0 {
+			continue
+		}
+
+		lable := device.Description
+		if lable == "" {
+			lable = device.Name
+		}
+		item := vo.TreeSelect{
+			ID:       device.Name,
+			Label:    lable,
+			Children: []vo.TreeSelect{},
+		}
+
+		for _, address := range device.Addresses {
+			if address.IP != nil {
+				ip := address.IP.String()
+				item.Children = append(item.Children, vo.TreeSelect{ID: ip, Label: ip})
+			}
+		}
+		arr = append(arr, item)
+	}
+
+	return arr
+}
+
+// verifyDevice 检查网卡设备是否存在
+func (s *Packet) verifyDevice(str string) (string, bool) {
+	devices, err := pcap.FindAllDevs()
+	if err != nil {
+		logger.Errorf("interfaces find all devices err: %s", err.Error())
+		return "", false
+	}
+	for _, device := range devices {
+		if len(device.Addresses) == 0 {
+			continue
+		}
+		if device.Name == str {
+			return device.Name, true
+		}
+		for _, address := range device.Addresses {
+			if address.IP.String() == str {
+				return device.Name, true
+			}
+		}
+	}
+	return "", false
+}
+
+// LiveStart 开始捕获数据
+func (s *Packet) LiveStart(taskNo, deviceName, outputFile string) (string, error) {
+	if _, ok := s.taskMap.Load(taskNo); ok {
+		return "", fmt.Errorf("task no. %s already exist", taskNo)
+	}
+
+	// Verify the specified network interface exists
+	device, deviceOk := s.verifyDevice(deviceName)
+	if !deviceOk {
+		return "", fmt.Errorf("network device not exist: %s", deviceName)
+	}
+
+	snapshotLength := 262144
+
+	// open device
+	handle, err := pcap.OpenLive(device, int32(snapshotLength), true, pcap.BlockForever)
+	if err != nil {
+		logger.Errorf("open live err: %s", err.Error())
+		if strings.Contains(err.Error(), "operation not permitted") {
+			return "", fmt.Errorf("you don't have permission to capture on that/these device(s)")
+		}
+		return "", err
+	}
+
+	// write a new file
+	var w *pcapgo.Writer
+	var f *os.File
+	if outputFile != "" {
+		f, err = os.Create(outputFile)
+		if err != nil {
+			return "", err
+		}
+		w = pcapgo.NewWriter(f)
+		w.WriteFileHeader(uint32(snapshotLength), handle.LinkType())
+	}
+
+	// capture packets
+	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
+	packetSource.Lazy = false
+	packetSource.NoCopy = true
+	packetSource.DecodeStreamsAsDatagrams = true
+
+	// save tasks
+	ctx := context.Background()
+	task := &task{
+		TaskNo:  taskNo,
+		Handle:  handle,
+		File:    f,
+		Writer:  w,
+		Expire:  time.Now().Add(time.Second * 120),
+		context: ctx,
+	}
+	s.taskMap.Store(taskNo, task)
+
+	// start capture
+	go func() {
+		for packet := range packetSource.PacketsCtx(ctx) {
+			if packet == nil {
+				continue
+			}
+			if packet.Metadata().Timestamp.Before(time.Now()) {
+				continue
+			}
+			if w != nil {
+				w.WritePacket(packet.Metadata().CaptureInfo, packet.Data())
+			}
+			fmt.Println(packet.Dump())
+		}
+	}()
+	return "task initiated", nil
+}
+
+// LiveFilter 捕获过滤
+func (s *Packet) LiveFilter(taskNo, expr string) error {
+	info, ok := s.taskMap.Load(taskNo)
+	if !ok {
+		return fmt.Errorf("task no. %s not exist", taskNo)
+	}
+	return info.(*task).Handle.SetBPFFilter(expr)
+}
+
+// LiveStop 停止捕获数据
+func (s *Packet) LiveStop(taskNo string) error {
+	info, ok := s.taskMap.Load(taskNo)
+	if !ok {
+		return fmt.Errorf("task no. %s not exist", taskNo)
+	}
+	info.(*task).context.Done()
+	info.(*task).Handle.Close()
+	if info.(task).File != nil {
+		info.(task).File.Close()
+	}
+	s.taskMap.Delete(taskNo)
+	return nil
+}