From 4704d0daa86affe2e2a17b51df3eef26a59a84fd Mon Sep 17 00:00:00 2001 From: Akizon77 Date: Sun, 20 Jul 2025 07:54:04 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20https://github.com/komari-monitor/komar?= =?UTF-8?q?i/issues/94=20=E6=B5=81=E9=87=8F=E6=8C=89=E6=9C=88=E7=BB=9F?= =?UTF-8?q?=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/flags/flag.go | 1 + cmd/root.go | 1 + monitoring/unit/net.go | 199 +++++++++++++++++++++++++ monitoring/unit/net_test.go | 290 ++++++++++++++++++++++++++++++++++++ 4 files changed, 491 insertions(+) create mode 100644 monitoring/unit/net_test.go diff --git a/cmd/flags/flag.go b/cmd/flags/flag.go index bf2a34d..481d6b8 100644 --- a/cmd/flags/flag.go +++ b/cmd/flags/flag.go @@ -13,4 +13,5 @@ var ( InfoReportInterval int IncludeNics string ExcludeNics string + MonthRotate int ) diff --git a/cmd/root.go b/cmd/root.go index a1e904e..adf7a77 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -70,5 +70,6 @@ func init() { RootCmd.PersistentFlags().IntVar(&flags.InfoReportInterval, "info-report-interval", 5, "Interval in minutes for reporting basic info") RootCmd.PersistentFlags().StringVar(&flags.IncludeNics, "include-nics", "", "Comma-separated list of network interfaces to include") RootCmd.PersistentFlags().StringVar(&flags.ExcludeNics, "exclude-nics", "", "Comma-separated list of network interfaces to exclude") + RootCmd.PersistentFlags().IntVar(&flags.MonthRotate, "month-rotate", 0, "Month reset for network statistics (0 to disable)") RootCmd.PersistentFlags().ParseErrorsWhitelist.UnknownFlags = true } diff --git a/monitoring/unit/net.go b/monitoring/unit/net.go index aa5257c..9f31b32 100644 --- a/monitoring/unit/net.go +++ b/monitoring/unit/net.go @@ -1,7 +1,9 @@ package monitoring import ( + "encoding/json" "fmt" + "os/exec" "strings" "time" @@ -32,10 +34,207 @@ var ( } ) +// VnstatInterface represents a network interface in vnstat output +type VnstatInterface struct { + Name string `json:"name"` + Alias string `json:"alias"` + Created VnstatDate `json:"created"` + Updated VnstatUpdated `json:"updated"` + Traffic VnstatTraffic `json:"traffic"` +} + +// VnstatDate represents date information +type VnstatDate struct { + Date VnstatDateInfo `json:"date"` + Timestamp int64 `json:"timestamp"` +} + +// VnstatUpdated represents updated information +type VnstatUpdated struct { + Date VnstatDateInfo `json:"date"` + Time VnstatTimeInfo `json:"time"` + Timestamp int64 `json:"timestamp"` +} + +// VnstatDateInfo represents date components +type VnstatDateInfo struct { + Year int `json:"year"` + Month int `json:"month"` + Day int `json:"day"` +} + +// VnstatTimeInfo represents time components +type VnstatTimeInfo struct { + Hour int `json:"hour"` + Minute int `json:"minute"` +} + +// VnstatTraffic represents traffic data from vnstat +type VnstatTraffic struct { + Total VnstatTotal `json:"total"` + FiveMinute []VnstatTimeEntry `json:"fiveminute"` + Hour []VnstatTimeEntry `json:"hour"` + Day []VnstatTimeEntry `json:"day"` + Month []VnstatMonthEntry `json:"month"` + Year []VnstatYearEntry `json:"year"` + Top []VnstatTimeEntry `json:"top"` +} + +// VnstatTotal represents total traffic data +type VnstatTotal struct { + Rx uint64 `json:"rx"` + Tx uint64 `json:"tx"` +} + +// VnstatTimeEntry represents a time-based traffic entry +type VnstatTimeEntry struct { + ID int `json:"id"` + Date VnstatDateInfo `json:"date"` + Time VnstatTimeInfo `json:"time,omitempty"` + Timestamp int64 `json:"timestamp"` + Rx uint64 `json:"rx"` + Tx uint64 `json:"tx"` +} + +// VnstatMonthEntry represents a monthly traffic entry +type VnstatMonthEntry struct { + ID int `json:"id"` + Date VnstatDateInfo `json:"date"` + Timestamp int64 `json:"timestamp"` + Rx uint64 `json:"rx"` + Tx uint64 `json:"tx"` +} + +// VnstatYearEntry represents a yearly traffic entry +type VnstatYearEntry struct { + ID int `json:"id"` + Date VnstatDateInfo `json:"date"` + Timestamp int64 `json:"timestamp"` + Rx uint64 `json:"rx"` + Tx uint64 `json:"tx"` +} + +// VnstatOutput represents the complete vnstat JSON output +type VnstatOutput struct { + VnstatVersion string `json:"vnstatversion"` + JsonVersion string `json:"jsonversion"` + Interfaces []VnstatInterface `json:"interfaces"` +} + +func getVnstatData() (map[string]VnstatInterface, error) { + cmd := exec.Command("vnstat", "--json") + output, err := cmd.Output() + if err != nil { + return nil, fmt.Errorf("failed to run vnstat: %w", err) + } + + var vnstatOutput VnstatOutput + if err := json.Unmarshal(output, &vnstatOutput); err != nil { + return nil, fmt.Errorf("failed to parse vnstat output: %w", err) + } + + interfaceMap := make(map[string]VnstatInterface) + for _, iface := range vnstatOutput.Interfaces { + interfaceMap[iface.Name] = iface + } + + return interfaceMap, nil +} + +// calculateMonthlyUsage 计算从月重置日到当前日期的流量使用量 +func calculateMonthlyUsage(iface VnstatInterface, monthRotateDay int) (rx, tx uint64) { + now := time.Now() + currentYear := now.Year() + currentMonth := int(now.Month()) + currentDay := now.Day() + + // 确定统计的起始日期 + var startYear, startMonth, startDay int + if currentDay >= monthRotateDay { + // 当前月的重置日已过，从当前月的重置日开始计算 + startYear = currentYear + startMonth = currentMonth + startDay = monthRotateDay + } else { + // 当前月的重置日未到，从上个月的重置日开始计算 + if currentMonth == 1 { + startYear = currentYear - 1 + startMonth = 12 + } else { + startYear = currentYear + startMonth = currentMonth - 1 + } + startDay = monthRotateDay + } + + startTime := time.Date(startYear, time.Month(startMonth), startDay, 0, 0, 0, 0, time.Local) + + // 统计从起始时间到现在的流量 + for _, entry := range iface.Traffic.Day { + entryTime := time.Date(entry.Date.Year, time.Month(entry.Date.Month), entry.Date.Day, 0, 0, 0, 0, time.Local) + if entryTime.After(startTime) || entryTime.Equal(startTime) { + rx += entry.Rx + tx += entry.Tx + } + } + + return rx, tx +} + +// setVnstatMonthRotate 设置vnstat的月重置日期 +func setVnstatMonthRotate(day int) error { + if day < 1 || day > 31 { + return fmt.Errorf("invalid day: %d, must be between 1 and 31", day) + } + + cmd := exec.Command("vnstat", "--config", fmt.Sprintf("MonthRotate %d", day)) + if err := cmd.Run(); err != nil { + return fmt.Errorf("failed to set vnstat month rotate to day %d: %w", day, err) + } + + return nil +} + func NetworkSpeed() (totalUp, totalDown, upSpeed, downSpeed uint64, err error) { includeNics := parseNics(flags.IncludeNics) excludeNics := parseNics(flags.ExcludeNics) + // 如果设置了月重置（非0），使用vnstat统计totalUp、totalDown + if flags.MonthRotate != 0 { + + vnstatData, err := getVnstatData() + if err != nil { + // 如果vnstat失败，回退到原来的方法，并返回额外的错误信息 + fallbackUp, fallbackDown, fallbackUpSpeed, fallbackDownSpeed, fallbackErr := getNetworkSpeedFallback(includeNics, excludeNics) + if fallbackErr != nil { + return fallbackUp, fallbackDown, fallbackUpSpeed, fallbackDownSpeed, fmt.Errorf("failed to call vnstat: %v; fallback error: %w", err, fallbackErr) + } + return fallbackUp, fallbackDown, fallbackUpSpeed, fallbackDownSpeed, fmt.Errorf("failed to call vnstat: %w", err) + } + + // 使用vnstat数据计算当月（到重置日）的流量使用量 + for interfaceName, interfaceData := range vnstatData { + if shouldInclude(interfaceName, includeNics, excludeNics) { + monthlyRx, monthlyTx := calculateMonthlyUsage(interfaceData, flags.MonthRotate) + totalUp += monthlyTx + totalDown += monthlyRx + } + } + + // 对于实时速度，仍然使用gopsutil方法 + _, _, upSpeed, downSpeed, err = getNetworkSpeedFallback(includeNics, excludeNics) + if err != nil { + return totalUp, totalDown, 0, 0, err + } + + return totalUp, totalDown, upSpeed, downSpeed, nil + } + + // 如果没有设置月重置，使用原来的方法 + return getNetworkSpeedFallback(includeNics, excludeNics) +} + +func getNetworkSpeedFallback(includeNics, excludeNics map[string]struct{}) (totalUp, totalDown, upSpeed, downSpeed uint64, err error) { // 获取第一次网络IO计数器 ioCounters1, err := net.IOCounters(true) if err != nil { diff --git a/monitoring/unit/net_test.go b/monitoring/unit/net_test.go new file mode 100644 index 0000000..9ccf110 --- /dev/null +++ b/monitoring/unit/net_test.go @@ -0,0 +1,290 @@ +package monitoring + +import ( + "strings" + "testing" + + "github.com/komari-monitor/komari-agent/cmd/flags" +) + +func TestConnectionsCount(t *testing.T) { + tcpCount, udpCount, err := ConnectionsCount() + if err != nil { + t.Fatalf("ConnectionsCount failed: %v", err) + } + + if tcpCount < 0 { + t.Errorf("Expected non-negative TCP count, got %d", tcpCount) + } + + if udpCount < 0 { + t.Errorf("Expected non-negative UDP count, got %d", udpCount) + } + + t.Logf("TCP connections: %d, UDP connections: %d", tcpCount, udpCount) +} + +func TestParseNics(t *testing.T) { + tests := []struct { + name string + input string + expected map[string]struct{} + }{ + { + name: "empty string", + input: "", + expected: nil, + }, + { + name: "single nic", + input: "eth0", + expected: map[string]struct{}{ + "eth0": {}, + }, + }, + { + name: "multiple nics", + input: "eth0,wlan0,enp0s3", + expected: map[string]struct{}{ + "eth0": {}, + "wlan0": {}, + "enp0s3": {}, + }, + }, + { + name: "nics with spaces", + input: " eth0 , wlan0 , enp0s3 ", + expected: map[string]struct{}{ + "eth0": {}, + "wlan0": {}, + "enp0s3": {}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := parseNics(tt.input) + + if tt.expected == nil && result != nil { + t.Errorf("Expected nil, got %v", result) + return + } + + if tt.expected != nil && result == nil { + t.Errorf("Expected %v, got nil", tt.expected) + return + } + + if len(result) != len(tt.expected) { + t.Errorf("Expected %d items, got %d", len(tt.expected), len(result)) + return + } + + for key := range tt.expected { + if _, exists := result[key]; !exists { + t.Errorf("Expected key %s not found in result", key) + } + } + }) + } +} + +func TestShouldInclude(t *testing.T) { + tests := []struct { + name string + nicName string + includeNics map[string]struct{} + excludeNics map[string]struct{} + expected bool + }{ + { + name: "loopback interface should be excluded", + nicName: "lo", + includeNics: nil, + excludeNics: nil, + expected: false, + }, + { + name: "docker interface should be excluded", + nicName: "docker0", + includeNics: nil, + excludeNics: nil, + expected: false, + }, + { + name: "normal interface with no filters", + nicName: "eth0", + includeNics: nil, + excludeNics: nil, + expected: true, + }, + { + name: "interface in include list", + nicName: "eth0", + includeNics: map[string]struct{}{ + "eth0": {}, + }, + excludeNics: nil, + expected: true, + }, + { + name: "interface not in include list", + nicName: "wlan0", + includeNics: map[string]struct{}{ + "eth0": {}, + }, + excludeNics: nil, + expected: false, + }, + { + name: "interface in exclude list", + nicName: "eth0", + includeNics: nil, + excludeNics: map[string]struct{}{ + "eth0": {}, + }, + expected: false, + }, + { + name: "interface not in exclude list", + nicName: "wlan0", + includeNics: nil, + excludeNics: map[string]struct{}{ + "eth0": {}, + }, + expected: true, + }, + { + name: "loopback in include list should still be excluded", + nicName: "lo", + includeNics: map[string]struct{}{ + "lo": {}, + }, + excludeNics: nil, + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := shouldInclude(tt.nicName, tt.includeNics, tt.excludeNics) + if result != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, result) + } + }) + } +} + +func TestNetworkSpeedFallback(t *testing.T) { + // 测试回退方法 + includeNics := map[string]struct{}{} + excludeNics := map[string]struct{}{} + + totalUp, totalDown, upSpeed, downSpeed, err := getNetworkSpeedFallback(includeNics, excludeNics) + if err != nil { + t.Fatalf("getNetworkSpeedFallback failed: %v", err) + } + + t.Logf("TotalUp: %d, TotalDown: %d, UpSpeed: %d/s, DownSpeed: %d/s", + totalUp, totalDown, upSpeed, downSpeed) +} + +func TestNetworkSpeedWithoutMonthRotate(t *testing.T) { + + flags.MonthRotate = 1 + + // 设置测试值 + flags.IncludeNics = "" + flags.ExcludeNics = "" + + totalUp, totalDown, upSpeed, downSpeed, err := NetworkSpeed() + if err != nil { + t.Fatalf("NetworkSpeed failed: %v", err) + } + + t.Logf("Without MonthRotate - TotalUp: %d, TotalDown: %d, UpSpeed: %d/s, DownSpeed: %d/s", + totalUp, totalDown, upSpeed, downSpeed) +} + +func TestNetworkSpeedWithMonthRotate(t *testing.T) { + // 保存原始值 + originalMonthRotate := flags.MonthRotate + originalIncludeNics := flags.IncludeNics + originalExcludeNics := flags.ExcludeNics + + // 恢复原始值 + defer func() { + flags.MonthRotate = originalMonthRotate + flags.IncludeNics = originalIncludeNics + flags.ExcludeNics = originalExcludeNics + }() + + // 设置测试值 - 启用月重置 + flags.MonthRotate = 1 + flags.IncludeNics = "" + flags.ExcludeNics = "" + + totalUp, totalDown, upSpeed, downSpeed, err := NetworkSpeed() + + // 如果vnstat不可用，可能会回退到原来的方法，这是正常的 + if err != nil { + if strings.Contains(err.Error(), "failed to call vnstat") { + t.Logf("vnstat not available, this is expected in test environment: %v", err) + return + } + t.Fatalf("NetworkSpeed failed: %v", err) + } + + if totalUp < 0 { + t.Errorf("Expected non-negative totalUp, got %d", totalUp) + } + + if totalDown < 0 { + t.Errorf("Expected non-negative totalDown, got %d", totalDown) + } + + t.Logf("With MonthRotate - TotalUp: %d, TotalDown: %d, UpSpeed: %d/s, DownSpeed: %d/s", + totalUp, totalDown, upSpeed, downSpeed) +} + +func TestGetVnstatData(t *testing.T) { + // 这个测试可能会失败，因为vnstat可能没有安装 + _, err := getVnstatData() + if err != nil { + if strings.Contains(err.Error(), "failed to run vnstat") { + t.Logf("vnstat not available, this is expected: %v", err) + return + } + t.Fatalf("getVnstatData failed unexpectedly: %v", err) + } + + t.Log("vnstat data retrieved successfully") +} + +func TestNetworkSpeedWithNicFilters(t *testing.T) { + // 保存原始值 + originalMonthRotate := flags.MonthRotate + originalIncludeNics := flags.IncludeNics + originalExcludeNics := flags.ExcludeNics + + // 恢复原始值 + defer func() { + flags.MonthRotate = originalMonthRotate + flags.IncludeNics = originalIncludeNics + flags.ExcludeNics = originalExcludeNics + }() + + // 测试排除回环接口 + flags.MonthRotate = 0 + flags.IncludeNics = "" + flags.ExcludeNics = "lo,docker0" + + totalUp, totalDown, upSpeed, downSpeed, err := NetworkSpeed() + if err != nil { + t.Fatalf("NetworkSpeed with excludeNics failed: %v", err) + } + + t.Logf("With excludeNics - TotalUp: %d, TotalDown: %d, UpSpeed: %d/s, DownSpeed: %d/s", + totalUp, totalDown, upSpeed, downSpeed) +}