ProjectOctopus/agent-go/status/Network.go

package status

import (
	"fmt"
	"github.com/go-openapi/errors"
	"github.com/shirou/gopsutil/v3/net"
	net2 "net"
	"regexp"
	"strings"
	"time"
)

type NetworkMetric struct {
	net.IOCountersStat

	SendSpeed float64
	RecvSpeed float64
}

type NetworkStatus struct {
	Name         string   `json:"name,omitempty"`
	InternalIPv4 []string `json:"internal_ip_v4,omitempty"`
	InternalIPv6 []string `json:"internal_ip_v6,omitempty"`
	ExternalIPv4 []string `json:"external_ip_v4,omitempty"`
	ExternalIPv6 []string `json:"external_ip_v6,omitempty"`
	Mac          string   `json:"mac,omitempty"`
	Sent         uint64   `json:"sent,omitempty"`
	Recv         uint64   `json:"recv,omitempty"`
	SentRate     string   `json:"sent_rate,omitempty"`
	RecvRate     string   `json:"recv_rate,omitempty"`
}

func GetNetworkMetric() ([]NetworkMetric, error) {
	interfaceStatList, err := net.Interfaces()
	if err != nil {
		log.ErrorF("[GetNetworkMetric] - get interface list error %v\n", err)
		return nil, err
	}
	contains := false
	for _, interfaceStat := range interfaceStatList {
		if MatchNetInterfaceRight(interfaceStat.Name) {
			contains = true
			break
		}
	}
	if !contains {
		msg := "[GetNetworkMetric]- can not find desired net interface !"
		log.Warn(msg)
		return nil, errors.New(233, msg)
	}

	ioCountersStats, err := net.IOCounters(false)
	if err != nil {
		log.ErrorF("[GetNetworkMetric] - get io counter list error %v\n", err)
		return nil, err
	}

	duration := time.Second * 2
	time.Sleep(duration)

	afterIoCounterList, err := net.IOCounters(false)
	if err != nil {
		log.ErrorF("[GetNetworkMetric] - get io counter list second error %v\n", err)
		return nil, err
	}

	var result []NetworkMetric

	for _, interfaceStat := range interfaceStatList {
		if MatchNetInterfaceRight(interfaceStat.Name) {

			for _, after := range afterIoCounterList {
				if strings.Contains(after.Name, interfaceStat.Name) {
					for _, before := range ioCountersStats {
						if strings.Contains(before.Name, interfaceStat.Name) {
							// match after and before interface

							// assign
							networkMetric := NetworkMetric{
								IOCountersStat: after,
								SendSpeed:      0,
								RecvSpeed:      0,
							}

							networkMetric.SendSpeed = (float64(after.BytesSent - before.BytesSent)) / duration.Seconds()
							networkMetric.RecvSpeed = (float64(after.BytesRecv - before.BytesRecv)) / duration.Seconds()

							result = append(result, networkMetric)
						}
					}
				}
			}

		}
	}

	return result, nil
}

func GetNetworkStatus() (*NetworkStatus, error) {
	interfaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}
	var mainInterface net.InterfaceStat
	for _, intf := range interfaces {
		if strings.HasPrefix(intf.Name, "ens") || strings.HasPrefix(intf.Name, "eth") || strings.HasPrefix(intf.Name, "eno") {
			mainInterface = intf
			break
		}
	}
	var allAddrs []string
	log.DebugF("all main interface address are %v", mainInterface.Addrs)
	for _, addr := range mainInterface.Addrs {
		allAddrs = append(allAddrs, addr.Addr)
	}
	ipv4List, ipv6List := GetInternelIpAddrs(allAddrs)
	log.DebugF("ipv4 list are => %v, ipv6 list are => %v", ipv4List, ipv6List)

	counters, err := net.IOCounters(true)
	if err != nil {
		return nil, err
	}
	var sent uint64
	var recv uint64
	for _, counter := range counters {
		if counter.Name == mainInterface.Name {
			sent = counter.BytesSent
			recv = counter.BytesRecv
			break
		}
	}

	// 休眠3秒
	time.Sleep(3 * time.Second)

	var sentAfter uint64
	var recvAfter uint64
	countersAfter, err := net.IOCounters(true)
	if err != nil {
		return nil, err
	}
	for _, counter := range countersAfter {
		if counter.Name == mainInterface.Name {
			sentAfter = counter.BytesSent
			recvAfter = counter.BytesRecv
			break
		}
	}
	sendRate := fmt.Sprintf(formatBytes((sentAfter-sent)/3) + "/s")
	recvRate := fmt.Sprintf(formatBytes((recvAfter-recv)/3) + "/s")

	info := &NetworkStatus{
		Name:         mainInterface.Name,
		InternalIPv4: ipv4List,
		InternalIPv6: ipv6List,
		ExternalIPv4: nil,
		ExternalIPv6: nil,
		Mac:          mainInterface.HardwareAddr,
		Sent:         sent,
		Recv:         recv,
		SentRate:     sendRate,
		RecvRate:     recvRate,
	}
	return info, nil
}

func formatBytes(bytes uint64) string {
	const unit = 1024
	if bytes < unit {
		return fmt.Sprintf("%d B", bytes)
	}
	div, exp := int64(unit), 0
	for n := bytes / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
}

func GetInternelIpAddrs(addresses []string) ([]string, []string) {
	var ipv4 []string
	var ipv6 []string
	for _, addr := range addresses {
		// it parse (0.0.0.0) not cidr
		ip, _, err := net2.ParseCIDR(addr)
		if err != nil {
			continue
		}
		if ip.To4() != nil {
			ipv4 = append(ipv4, addr)
		} else if ip.To16() != nil {
			ipv6 = append(ipv6, addr)
		}
	}
	return ipv4, ipv6
}

func MatchNetInterfaceRight(interfaceName string) bool {

	match, _ := regexp.MatchString(`^(eth|enp|wlan|ens|eno)\d+`, interfaceName)
	if match {
		return true
	}

	match, _ = regexp.MatchString(`^(lo)`, interfaceName)
	if match {
		return true
	}

	return match
}