From a984adc56df8d0829be1c6df53740d272515d2eb Mon Sep 17 00:00:00 2001
From: Akizon77 <hytf233g@gmail.com>
Date: Sun, 20 Jul 2025 07:58:13 +0000
Subject: [PATCH] =?UTF-8?q?feat:=20#5=20=E6=94=AF=E6=8C=81=E6=8C=87?=
 =?UTF-8?q?=E5=AE=9A=E6=8C=82=E8=BD=BD=E8=B7=AF=E5=BE=84?=
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/disk.go | 36 +++++++++++++++++++++++++++---------
 3 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/cmd/flags/flag.go b/cmd/flags/flag.go
index 481d6b8..98c0949 100644
--- a/cmd/flags/flag.go
+++ b/cmd/flags/flag.go
@@ -13,5 +13,6 @@ var (
 	InfoReportInterval  int
 	IncludeNics         string
 	ExcludeNics         string
+	IncludeMountpoints  string
 	MonthRotate         int
 )
diff --git a/cmd/root.go b/cmd/root.go
index adf7a77..4587d87 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -70,6 +70,7 @@ 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().StringVar(&flags.IncludeMountpoints, "include-mountpoint", "", "Semicolon-separated list of mount points to include for disk statistics")
 	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/disk.go b/monitoring/unit/disk.go
index 7e34720..a6e74ab 100644
--- a/monitoring/unit/disk.go
+++ b/monitoring/unit/disk.go
@@ -3,6 +3,7 @@ package monitoring
 import (
 	"strings"
 
+	"github.com/komari-monitor/komari-agent/cmd/flags"
 	"github.com/shirou/gopsutil/v4/disk"
 )
 
@@ -18,15 +19,32 @@ func Disk() DiskInfo {
 		diskinfo.Total = 0
 		diskinfo.Used = 0
 	} else {
-		for _, part := range usage {
-			// 排除临时文件系统和网络驱动器
-			if isPhysicalDisk(part) {
-				u, err := disk.Usage(part.Mountpoint)
-				if err != nil {
-					continue
-				} else {
-					diskinfo.Total += u.Total
-					diskinfo.Used += u.Used
+		// 如果指定了自定义挂载点，只统计指定的挂载点
+		if flags.IncludeMountpoints != "" {
+			includeMounts := strings.Split(flags.IncludeMountpoints, ";")
+			for _, mountpoint := range includeMounts {
+				mountpoint = strings.TrimSpace(mountpoint)
+				if mountpoint != "" {
+					u, err := disk.Usage(mountpoint)
+					if err != nil {
+						continue
+					} else {
+						diskinfo.Total += u.Total
+						diskinfo.Used += u.Used
+					}
+				}
+			}
+		} else {
+			// 使用默认逻辑，排除临时文件系统和网络驱动器
+			for _, part := range usage {
+				if isPhysicalDisk(part) {
+					u, err := disk.Usage(part.Mountpoint)
+					if err != nil {
+						continue
+					} else {
+						diskinfo.Total += u.Total
+						diskinfo.Used += u.Used
+					}
 				}
 			}
 		}