mirror of
https://github.com/fankes/komari-agent.git
synced 2025-10-18 18:49:23 +08:00
feat: https://github.com/komari-monitor/komari/issues/94 流量按月统计
This commit is contained in:
Akizon77
@@ -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 {
|
||||
|
||||
290
monitoring/unit/net_test.go
Normal file
290
monitoring/unit/net_test.go
Normal file
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user