feat: https://github.com/komari-monitor/komari/issues/94 流量按月统计

cmd/flags/flag.go

@@ -13,4 +13,5 @@ var (
 	InfoReportInterval  int
 	IncludeNics         string
 	ExcludeNics         string
+	MonthRotate         int
 )

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
 }

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 {

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)
+}