golang统计虚拟网卡指定端口进出流量

前置条件

1. libpcap(apt install libpcap-dev)||(yum install libpcap-devel)
2. gopacket(go get github.com/google/gopacket)

上代码

package main

import (
	"fmt"
	"os"
	"strconv"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/pcap"
	"github.com/oswaldoooo/cmicro/sys"
)

type packet_info struct {
	src_host, dst_host [4]byte
	src_port, dst_port uint32
}
//手动解析一下原始链路层数据包，这里不清楚可以去wireshark抓包比对一下位置
func parsepacket(src []byte) packet_info {
	var ans packet_info
	copy(ans.src_host[:], src[26:30])
	copy(ans.dst_host[:], src[30:34])
	ans.src_port = uint32(src[34])*256 + uint32(src[35])
	ans.dst_port = uint32(src[36])*256 + uint32(src[37])
	return ans
}
//ip4地址为4个byte，byte实际为uint8，所以这里需要转换一下
func format_packet(src [4]byte) string {
	return strconv.FormatUint(uint64(src[0]), 10) + "." + strconv.FormatUint(uint64(src[1]), 10) + "." + strconv.FormatUint(uint64(src[2]), 10) + "." + strconv.FormatUint(uint64(src[3]), 10)
}

var in_band_used = make(map[string]**uint64)
var out_band_used = make(map[string]**uint64)

func init() {
	for _, ele := range os.Args[1:] {
		if _, ok := in_band_used[ele]; !ok {
			in_band_used[ele] = new(*uint64)
			ptr := sys.Shm_Open(ele+"in.shm", in_band_used[ele])
			if ptr == nil {
				fmt.Println(ele + " create failed")
				delete(in_band_used, ele)
			} else {
				fmt.Println(ele+" in band registered,used", **in_band_used[ele])
			}

		}
		if _, ok := out_band_used[ele]; !ok {
			out_band_used[ele] = new(*uint64)
			ptr := sys.Shm_Open(ele+"out.shm", out_band_used[ele])
			if ptr == nil {
				fmt.Println(ele + " create failed")
				delete(out_band_used, ele)
			} else {
				fmt.Println(ele+" out band registered,used", **out_band_used[ele])
			}
		}
	}

}
func main() {
	hand, err := pcap.OpenLive("ens160", 8192, true, 100*time.Millisecond)
	var (
		data []byte
		cap  gopacket.CaptureInfo
		pack packet_info
	)
	if err != nil {
		fmt.Println("[error]", err.Error())
	}
	var aimport uint32 = 9000
	if aimport == 0 {
		fmt.Println("aim port is not correct")
		os.Exit(1)
	}
	var okone, oktwo bool
	var inptr, outptr **uint64
	for {
		data, cap, err = hand.ZeroCopyReadPacketData()
		if err == nil {
			pack = parsepacket(data)
			inptr, okone = in_band_used[format_packet(pack.src_host)+":"+strconv.FormatUint(uint64(pack.src_port), 10)]
			outptr, oktwo = out_band_used[format_packet(pack.dst_host)+":"+strconv.FormatUint(uint64(pack.dst_port), 10)]
			if okone || oktwo {
				fmt.Printf("get target pakcet: src add %s:%d,dst add %s:%d,data length %d\n", format_packet(pack.src_host), pack.src_port, format_packet(pack.dst_host), pack.dst_port, cap.Length)
			}
			if okone {
				**inptr += uint64(cap.Length)
				fmt.Printf("get in band %s\n", format_packet(pack.src_host)+":"+strconv.FormatUint(uint64(pack.src_port), 10))
			} else if oktwo {
				**outptr += uint64(cap.Length)
				fmt.Printf("get out band %s\n", format_packet(pack.dst_host)+":"+strconv.FormatUint(uint64(pack.dst_port), 10))
			}

		}
	}
}

测试结果

统计结果