From d30161651fa29a73b9045049822885282ed95ad7 Mon Sep 17 00:00:00 2001
From: Akizon77 <i@akz.moe>
Date: Tue, 4 Nov 2025 15:06:53 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20#36=20=E4=BB=8E=E7=BD=91=E5=8D=A1?=
 =?UTF-8?q?=E8=8E=B7=E5=8F=96IP=E5=9C=B0=E5=9D=80;=20=E8=87=AA=E5=AE=9A?=
 =?UTF-8?q?=E4=B9=89IP=E5=9C=B0=E5=9D=80?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/flags/flag.go     |   3 ++
 cmd/root.go           |   3 ++
 monitoring/unit/ip.go | 102 ++++++++++++++++++++++++++++++++++++++----
 3 files changed, 100 insertions(+), 8 deletions(-)

diff --git a/cmd/flags/flag.go b/cmd/flags/flag.go
index 8f9daf7..1b81a91 100644
--- a/cmd/flags/flag.go
+++ b/cmd/flags/flag.go
@@ -22,4 +22,7 @@ var (
 	CustomDNS            string
 	EnableGPU            bool // 启用详细GPU监控
 	ShowWarning          bool // Windows 上显示安全警告，作为子进程运行一次
+	CustomIpv4           string
+	CustomIpv6           string
+	GetIpAddrFromNic     bool // 从网卡获取IP地址
 )
diff --git a/cmd/root.go b/cmd/root.go
index b0f6e31..fc887a8 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -154,5 +154,8 @@ func init() {
 	RootCmd.PersistentFlags().StringVar(&flags.CustomDNS, "custom-dns", "", "Custom DNS server to use (e.g. 8.8.8.8, 114.114.114.114). By default, the program uses the system DNS resolver.")
 	RootCmd.PersistentFlags().BoolVar(&flags.EnableGPU, "gpu", false, "Enable detailed GPU monitoring (usage, memory, multi-GPU support)")
 	RootCmd.PersistentFlags().BoolVar(&flags.ShowWarning, "show-warning", false, "Show security warning on Windows, run once as a subprocess")
+	RootCmd.PersistentFlags().StringVar(&flags.CustomIpv4, "custom-ipv4", "", "Custom IPv4 address to use")
+	RootCmd.PersistentFlags().StringVar(&flags.CustomIpv6, "custom-ipv6", "", "Custom IPv6 address to use")
+	RootCmd.PersistentFlags().BoolVar(&flags.GetIpAddrFromNic, "get-ip-addr-from-nic", false, "Get IP address from network interface")
 	RootCmd.PersistentFlags().ParseErrorsWhitelist.UnknownFlags = true
 }
diff --git a/monitoring/unit/ip.go b/monitoring/unit/ip.go
index 2dfba47..07b5285 100644
--- a/monitoring/unit/ip.go
+++ b/monitoring/unit/ip.go
@@ -9,6 +9,7 @@ import (
 	"regexp"
 	"time"
 
+	"github.com/komari-monitor/komari-agent/cmd/flags"
 	"github.com/komari-monitor/komari-agent/dnsresolver"
 )
 
@@ -119,16 +120,101 @@ func GetIPv6Address() (string, error) {
 }
 
 func GetIPAddress() (ipv4, ipv6 string, err error) {
-	ipv4, err = GetIPv4Address()
-	if err != nil {
-		log.Printf("Get IPV4 Error: %v", err)
-		ipv4 = ""
+
+	if flags.GetIpAddrFromNic {
+		allowNics, err := InterfaceList()
+		if err != nil {
+			log.Printf("Get Interface List Error: %v", err)
+		} else {
+			ipv4, ipv6 = getIPFromInterfaces(allowNics)
+			if ipv4 != "" || ipv6 != "" {
+				log.Printf("Get IP from NIC - IPv4: %s, IPv6: %s", ipv4, ipv6)
+				return ipv4, ipv6, nil
+			}
+		}
 	}
-	ipv6, err = GetIPv6Address()
-	if err != nil {
-		log.Printf("Get IPV6 Error: %v", err)
-		ipv6 = ""
+
+	if flags.CustomIpv4 != "" {
+		ipv4 = flags.CustomIpv4
+	} else {
+		ipv4, err = GetIPv4Address()
+		if err != nil {
+			log.Printf("Get IPV4 Error: %v", err)
+			ipv4 = ""
+		}
+	}
+	if flags.CustomIpv6 != "" {
+		ipv6 = flags.CustomIpv6
+	} else {
+		ipv6, err = GetIPv6Address()
+		if err != nil {
+			log.Printf("Get IPV6 Error: %v", err)
+			ipv6 = ""
+		}
 	}
 
 	return ipv4, ipv6, nil
 }
+
+// getIPFromInterfaces 从指定的网卡接口获取 IPv4 和 IPv6 地址
+func getIPFromInterfaces(nicNames []string) (ipv4, ipv6 string) {
+	interfaces, err := net.Interfaces()
+	if err != nil {
+		log.Printf("Failed to get network interfaces: %v", err)
+		return "", ""
+	}
+	for _, iface := range interfaces {
+		// 检查接口是否在允许列表中
+		if !func(slice []string, item string) bool {
+			for _, s := range slice {
+				if s == item {
+					return true
+				}
+			}
+			return false
+		}(nicNames, iface.Name) {
+			continue
+		}
+
+		// 跳过未启动的接口
+		if iface.Flags&net.FlagUp == 0 {
+			continue
+		}
+
+		addrs, err := iface.Addrs()
+		if err != nil {
+			continue
+		}
+
+		for _, addr := range addrs {
+			var ip net.IP
+			switch v := addr.(type) {
+			case *net.IPNet:
+				ip = v.IP
+			case *net.IPAddr:
+				ip = v.IP
+			}
+
+			if ip == nil || ip.IsLoopback() {
+				continue
+			}
+
+			// 获取 IPv4 地址
+			if ipv4 == "" && ip.To4() != nil {
+				ipv4 = ip.String()
+			}
+
+			// 获取 IPv6 地址（排除链路本地地址）
+			if ipv6 == "" && ip.To4() == nil && !ip.IsLinkLocalUnicast() {
+				ipv6 = ip.String()
+			}
+
+			// 如果已经找到 IPv4 和 IPv6,提前返回
+			if ipv4 != "" && ipv6 != "" {
+				return ipv4, ipv6
+			}
+		}
+	}
+
+	return ipv4, ipv6
+}