From a78756f101b1c17f627feea0c5227f4ca1196a38 Mon Sep 17 00:00:00 2001 From: Akizon77 Date: Fri, 10 Oct 2025 00:43:06 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20Windows=20?= =?UTF-8?q?=E5=AE=89=E5=85=A8=E8=AD=A6=E5=91=8A=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/flags/flag.go | 3 +- cmd/root.go | 11 + cmd/warn.go | 13 + cmd/warn_windows.go | 408 +++++++++++++++++++++++++++ go.mod | 4 +- go.sum | 4 + monitoring/unit/gpu_detailed_test.go | 8 +- 7 files changed, 445 insertions(+), 6 deletions(-) create mode 100644 cmd/warn.go create mode 100644 cmd/warn_windows.go diff --git a/cmd/flags/flag.go b/cmd/flags/flag.go index 7c9deac..8f9daf7 100644 --- a/cmd/flags/flag.go +++ b/cmd/flags/flag.go @@ -20,5 +20,6 @@ var ( CFAccessClientSecret string MemoryIncludeCache bool CustomDNS string - EnableGPU bool // 启用详细GPU监控 + EnableGPU bool // 启用详细GPU监控 + ShowWarning bool // Windows 上显示安全警告，作为子进程运行一次 ) diff --git a/cmd/root.go b/cmd/root.go index d15d645..1885cbf 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -19,6 +19,16 @@ var RootCmd = &cobra.Command{ Short: "komari agent", Long: `komari agent`, Run: func(cmd *cobra.Command, args []string) { + + if flags.ShowWarning { + ShowToast() + os.Exit(0) + } + + if !flags.DisableWebSsh { + go WarnKomariRunning() + } + log.Println("Komari Agent", update.CurrentVersion) log.Println("Github Repo:", update.Repo) @@ -113,5 +123,6 @@ func init() { RootCmd.PersistentFlags().BoolVar(&flags.MemoryIncludeCache, "memory-include-cache", false, "Include cache/buffer in memory usage") 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().ParseErrorsWhitelist.UnknownFlags = true } diff --git a/cmd/warn.go b/cmd/warn.go new file mode 100644 index 0000000..a1fffb0 --- /dev/null +++ b/cmd/warn.go @@ -0,0 +1,13 @@ +//go:build !windows + +package cmd + +func WarnKomariRunning() { + // No-op on non-Windows platforms + return +} + +func ShowToast() { + // No-op on non-Windows platforms + return +} diff --git a/cmd/warn_windows.go b/cmd/warn_windows.go new file mode 100644 index 0000000..3f719f6 --- /dev/null +++ b/cmd/warn_windows.go @@ -0,0 +1,408 @@ +//go:build windows + +package cmd + +import ( + "errors" + "fmt" + "log" + "os" + "path/filepath" + "strings" + "sync" + "time" + "unsafe" + + toast "gopkg.in/toast.v1" + + "github.com/go-ole/go-ole" + "github.com/go-ole/go-ole/oleutil" + "golang.org/x/sys/windows" +) + +// WarnKomariRunning +// 作为 SYSTEM（Session 0）运行时： +// 1) 轮询已登录的交互会话（WTSActive） +// 2) 对新检测到的会话，以该用户身份在其会话内启动当前进程（追加 --show-warning 参数） +// 3) 用户态子进程会进入 ShowToast() 分支并发送 Toast +func WarnKomariRunning() { + + // 启用权限 + if err := enablePrivileges([]string{"SeAssignPrimaryTokenPrivilege", "SeIncreaseQuotaPrivilege"}); err != nil { + log.Printf("[warn] enabling privileges failed: %v", err) + } + + seen := map[uint32]struct{}{} + var mu sync.Mutex + + sessions := []uint32{} + for _, sid := range sessions { + seen[sid] = struct{}{} + } + + // 轮询新登录 + ticker := time.NewTicker(3 * time.Second) + defer ticker.Stop() + + for range ticker.C { + current, err := enumerateActiveSessions() + if err != nil { + log.Printf("[warn] enumerateActiveSessions error: %v", err) + continue + } + + // 将 current 列表转换为集合，便于清理旧会话 + currentSet := make(map[uint32]struct{}, len(current)) + for _, sid := range current { + currentSet[sid] = struct{}{} + } + + // 找到新出现的会话 -> 在该会话启动进程 + for _, sid := range current { + mu.Lock() + _, known := seen[sid] + if !known { + seen[sid] = struct{}{} + mu.Unlock() + if err := launchSelfInSession(sid, []string{"--show-warning"}); err != nil { + log.Printf("[warn] launch in session %d failed: %v", sid, err) + } else { + log.Printf("[info] launched toast helper in session %d", sid) + } + } else { + mu.Unlock() + } + } + + // 清理不再存在的会话，避免 map 膨胀 + mu.Lock() + for sid := range seen { + if _, ok := currentSet[sid]; !ok { + delete(seen, sid) + } + } + mu.Unlock() + } +} + +// ShowToast 在用户态中执行 +func ShowToast() { + title := "Komari is Running" + message := "The remote control software \"Komari\" is running, which allows others to control your computer. If this was not initiated by you, please terminate the program immediately." + + const aumid = "Komari.Monitor.Agent" + const linkName = "Komari Warning (Auto Delete Later)" + + if err := ensureStartMenuShortcut(aumid, linkName); err != nil { + log.Printf("[warn] ensureStartMenuShortcut failed: %v", err) + } + + n := toast.Notification{ + AppID: aumid, + Title: title, + Message: message, + Actions: []toast.Action{ + {Type: "protocol", Label: "Help", Arguments: "https://komari-document.pages.dev/faq/uninstall.html"}, + }, + } + if err := n.Push(); err != nil { + log.Printf("[warn] toast push failed: %v", err) + } + + // 等待 15 秒后删除快捷方式 + shortcutPath := getStartMenuShortcutPath(linkName) + time.Sleep(15 * time.Second) + if err := os.Remove(shortcutPath); err != nil { + if !os.IsNotExist(err) { + log.Printf("[warn] remove shortcut failed: %v", err) + } + } +} + +// ensureStartMenuShortcut 使用 WScript.Shell 创建 .lnk 并设置 AppUserModelID +func ensureStartMenuShortcut(aumid, linkName string) error { + programs := getStartMenuProgramsDir() + if err := os.MkdirAll(programs, 0o755); err != nil { + return err + } + shortcutPath := filepath.Join(programs, sanitizeFileName(linkName)+".lnk") + if _, err := os.Stat(shortcutPath); err == nil { + return nil + } + + if hr := ole.CoInitializeEx(0, ole.COINIT_APARTMENTTHREADED); hr != nil { + // S_OK (0) 或 S_FALSE (1) 都视为成功；go-ole 将非零 HRESULT 作为 error 返回 + // 当返回错误时，我们再进行 Uninitialize 保护即可 + // 这里直接继续执行，由后续操作决定是否可用 + } + defer ole.CoUninitialize() + + unknown, err := oleutil.CreateObject("WScript.Shell") + if err != nil { + return fmt.Errorf("CreateObject WScript.Shell: %w", err) + } + defer unknown.Release() + + shell, err := unknown.QueryInterface(ole.IID_IDispatch) + if err != nil { + return fmt.Errorf("QueryInterface IDispatch: %w", err) + } + defer shell.Release() + + cs, err := oleutil.CallMethod(shell, "CreateShortcut", shortcutPath) + if err != nil { + return fmt.Errorf("CreateShortcut: %w", err) + } + shortcut := cs.ToIDispatch() + defer shortcut.Release() + + exePath, _ := os.Executable() + exeDir := filepath.Dir(exePath) + + if _, err = oleutil.PutProperty(shortcut, "TargetPath", exePath); err != nil { + return fmt.Errorf("set TargetPath: %w", err) + } + if _, err = oleutil.PutProperty(shortcut, "WorkingDirectory", exeDir); err != nil { + return fmt.Errorf("set WorkingDirectory: %w", err) + } + _, _ = oleutil.PutProperty(shortcut, "Description", "Komari Agent") + // 设置 AUMID + if _, err = oleutil.PutProperty(shortcut, "AppUserModelID", aumid); err != nil { + // 某些系统该属性不存在时，依然尝试保存；Toast 可能仍然显示 + log.Printf("[warn] set AppUserModelID failed: %v", err) + } + + if _, err = oleutil.CallMethod(shortcut, "Save"); err != nil { + return fmt.Errorf("shortcut.Save: %w", err) + } + return nil +} + +// 返回当前用户开始菜单 Programs 目录 +func getStartMenuProgramsDir() string { + return filepath.Join(os.Getenv("APPDATA"), "Microsoft", "Windows", "Start Menu", "Programs") +} + +// 获取快捷方式完整路径 +func getStartMenuShortcutPath(linkName string) string { + return filepath.Join(getStartMenuProgramsDir(), sanitizeFileName(linkName)+".lnk") +} + +func sanitizeFileName(name string) string { + replacer := strings.NewReplacer("\\", "_", "/", "_", ":", "_", "*", "_", "?", "_", "\"", "_", "<", "_", ">", "_", "|", "_") + return replacer.Replace(name) +} + +// enumerateActiveSessions 列出当前处于交互活动状态的会话 ID +func enumerateActiveSessions() ([]uint32, error) { + type wtsSessionInfo struct { + SessionID uint32 + WinStation *uint16 + State uint32 + } + + wtsapi := windows.NewLazySystemDLL("wtsapi32.dll") + procEnum := wtsapi.NewProc("WTSEnumerateSessionsW") + procFree := wtsapi.NewProc("WTSFreeMemory") + + var ( + server windows.Handle // WTS_CURRENT_SERVER_HANDLE == 0 + pinfo *wtsSessionInfo + count uint32 + version uint32 = 1 + ) + r1, _, err := procEnum.Call( + uintptr(server), + 0, + uintptr(version), + uintptr(unsafe.Pointer(&pinfo)), + uintptr(unsafe.Pointer(&count)), + ) + if r1 == 0 { + return nil, fmt.Errorf("WTSEnumerateSessionsW: %w", err) + } + defer procFree.Call(uintptr(unsafe.Pointer(pinfo))) + + // WTS_CONNECTSTATE_CLASS + const WTSActive = 0 + + // 遍历结构数组 + res := make([]uint32, 0, count) + infos := unsafe.Slice(pinfo, int(count)) + for i := 0; i < len(infos); i++ { + info := &infos[i] + if info.State == WTSActive { + if hasUserName(info.SessionID) { + res = append(res, info.SessionID) + } + } + } + return res, nil +} + +// hasUserName 检查会话是否有用户名（避免将空会话当作登录） +func hasUserName(sessionID uint32) bool { + const WTSUserName = 5 + wtsapi := windows.NewLazySystemDLL("wtsapi32.dll") + procQuery := wtsapi.NewProc("WTSQuerySessionInformationW") + procFree := wtsapi.NewProc("WTSFreeMemory") + + var buf *uint16 + var blen uint32 + r1, _, _ := procQuery.Call(0, uintptr(sessionID), uintptr(WTSUserName), uintptr(unsafe.Pointer(&buf)), uintptr(unsafe.Pointer(&blen))) + if r1 == 0 || buf == nil { + return false + } + defer procFree.Call(uintptr(unsafe.Pointer(buf))) + name := windows.UTF16PtrToString(buf) + return strings.TrimSpace(name) != "" +} + +// launchSelfInSession 在指定会话中以该用户身份启动当前进程并追加 args +func launchSelfInSession(sessionID uint32, extraArgs []string) error { + // 获取用户令牌 + userToken, err := queryUserToken(sessionID) + if err != nil { + return fmt.Errorf("queryUserToken: %w", err) + } + defer userToken.Close() + + primary, err := duplicateTokenPrimary(userToken) + if err != nil { + return fmt.Errorf("duplicateToken: %w", err) + } + defer primary.Close() + + exePath, _ := os.Executable() + // 仅保留进程名，去掉已有的 --show-warning，避免递归 + baseArgs := filterArgs(os.Args[1:], "--show-warning") + fullArgs := append([]string{quoteIfNeeded(exePath)}, baseArgs...) + fullArgs = append(fullArgs, extraArgs...) + cmdlineStr := strings.Join(fullArgs, " ") + cmdline, err := windows.UTF16PtrFromString(cmdlineStr) + if err != nil { + return fmt.Errorf("UTF16PtrFromString: %w", err) + } + + env, err := createEnvironmentBlock(primary) + if err != nil { + return fmt.Errorf("createEnvironmentBlock: %w", err) + } + defer destroyEnvironmentBlock(env) + + var si windows.StartupInfo + si.Cb = uint32(unsafe.Sizeof(si)) + si.Flags = 0 + si.ShowWindow = 0 + // 指定桌面，确保窗口可见 + desktop, _ := windows.UTF16PtrFromString("winsta0\\default") + si.Desktop = desktop + + var pi windows.ProcessInformation + // CREATE_UNICODE_ENVIRONMENT | DETACHED_PROCESS + const CREATE_UNICODE_ENVIRONMENT = 0x00000400 + const DETACHED_PROCESS = 0x00000008 + + err = windows.CreateProcessAsUser(primary, nil, cmdline, nil, nil, false, CREATE_UNICODE_ENVIRONMENT|DETACHED_PROCESS, env, nil, &si, &pi) + if err != nil { + return fmt.Errorf("CreateProcessAsUser: %w", err) + } + windows.CloseHandle(pi.Thread) + windows.CloseHandle(pi.Process) + return nil +} + +// enablePrivileges 尝试启用一组权限 +func enablePrivileges(names []string) error { + var errs []string + var token windows.Token + if err := windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &token); err != nil { + return err + } + defer token.Close() + for _, name := range names { + if err := setPrivilege(token, name, true); err != nil { + errs = append(errs, fmt.Sprintf("%s: %v", name, err)) + } + } + if len(errs) > 0 { + return errors.New(strings.Join(errs, "; ")) + } + return nil +} + +func setPrivilege(token windows.Token, privName string, enable bool) error { + var luid windows.LUID + nameUTF16, _ := windows.UTF16PtrFromString(privName) + if err := windows.LookupPrivilegeValue(nil, nameUTF16, &luid); err != nil { + return err + } + tp := windows.Tokenprivileges{ + PrivilegeCount: 1, + Privileges: [1]windows.LUIDAndAttributes{ + {Luid: luid, Attributes: 0}, + }, + } + if enable { + tp.Privileges[0].Attributes = windows.SE_PRIVILEGE_ENABLED + } + return windows.AdjustTokenPrivileges(token, false, &tp, 0, nil, nil) +} + +// queryUserToken 调用 WTSQueryUserToken 获取指定会话的用户令牌（模拟令牌） +func queryUserToken(sessionID uint32) (windows.Token, error) { + wtsapi := windows.NewLazySystemDLL("wtsapi32.dll") + proc := wtsapi.NewProc("WTSQueryUserToken") + var h windows.Handle + r1, _, err := proc.Call(uintptr(sessionID), uintptr(unsafe.Pointer(&h))) + if r1 == 0 { + return 0, fmt.Errorf("WTSQueryUserToken: %w", err) + } + return windows.Token(h), nil +} + +// duplicateTokenPrimary 将模拟令牌复制为主令牌，以供 CreateProcessAsUser 使用 +func duplicateTokenPrimary(token windows.Token) (windows.Token, error) { + var primary windows.Token + err := windows.DuplicateTokenEx(token, windows.TOKEN_ALL_ACCESS, nil, windows.SecurityIdentification, windows.TokenPrimary, &primary) + return primary, err +} + +// createEnvironmentBlock 为用户令牌创建环境块 +func createEnvironmentBlock(token windows.Token) (*uint16, error) { + userenv := windows.NewLazySystemDLL("userenv.dll") + proc := userenv.NewProc("CreateEnvironmentBlock") + var env *uint16 + r1, _, err := proc.Call(uintptr(unsafe.Pointer(&env)), uintptr(token), 0) + if r1 == 0 { + return nil, fmt.Errorf("CreateEnvironmentBlock: %w", err) + } + return env, nil +} + +func destroyEnvironmentBlock(env *uint16) { + if env == nil { + return + } + userenv := windows.NewLazySystemDLL("userenv.dll") + proc := userenv.NewProc("DestroyEnvironmentBlock") + _, _, _ = proc.Call(uintptr(unsafe.Pointer(env))) +} + +func quoteIfNeeded(s string) string { + if strings.ContainsAny(s, " \t\"") { + return "\"" + strings.ReplaceAll(s, "\"", "\\\"") + "\"" + } + return s +} + +func filterArgs(args []string, drop string) []string { + out := make([]string, 0, len(args)) + for i := 0; i < len(args); i++ { + if args[i] == drop { + continue + } + out = append(out, args[i]) + } + return out +} diff --git a/go.mod b/go.mod index 96ae14f..ba63b28 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/UserExistsError/conpty v0.1.4 github.com/blang/semver v3.5.1+incompatible github.com/creack/pty v1.1.24 + github.com/go-ole/go-ole v1.2.6 github.com/gorilla/websocket v1.5.3 github.com/klauspost/cpuid/v2 v2.3.0 github.com/prometheus-community/pro-bing v0.7.0 @@ -13,17 +14,18 @@ require ( github.com/shirou/gopsutil/v4 v4.25.6 github.com/spf13/cobra v1.9.1 golang.org/x/sys v0.33.0 + gopkg.in/toast.v1 v1.0.0-20180812000517-0a84660828b2 ) require ( github.com/ebitengine/purego v0.8.4 // indirect - github.com/go-ole/go-ole v1.2.6 // indirect github.com/google/go-github/v30 v30.1.0 // indirect github.com/google/go-querystring v1.0.0 // indirect github.com/google/uuid v1.6.0 // indirect github.com/inconshreveable/go-update v0.0.0-20160112193335-8152e7eb6ccf // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect + github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d // indirect github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect github.com/spf13/pflag v1.0.6 // indirect github.com/tcnksm/go-gitconfig v0.1.2 // indirect diff --git a/go.sum b/go.sum index 4966574..eb643db 100644 --- a/go.sum +++ b/go.sum @@ -37,6 +37,8 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= +github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d h1:VhgPp6v9qf9Agr/56bj7Y/xa04UccTW04VP0Qed4vnQ= +github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d/go.mod h1:YUTz3bUH2ZwIWBy3CJBeOBEugqcmXREj14T+iG/4k4U= github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/gomega v1.4.2 h1:3mYCb7aPxS/RU7TI1y4rkEn1oKmPRjNJLNEXgw7MH2I= github.com/onsi/gomega v1.4.2/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= @@ -104,6 +106,8 @@ google.golang.org/appengine v1.3.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/toast.v1 v1.0.0-20180812000517-0a84660828b2 h1:MZF6J7CV6s/h0HBkfqebrYfKCVEo5iN+wzE4QhV3Evo= +gopkg.in/toast.v1 v1.0.0-20180812000517-0a84660828b2/go.mod h1:s1Sn2yZos05Qfs7NKt867Xe18emOmtsO3eAKbDaon0o= gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE= gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= diff --git a/monitoring/unit/gpu_detailed_test.go b/monitoring/unit/gpu_detailed_test.go index d625840..7ed817b 100644 --- a/monitoring/unit/gpu_detailed_test.go +++ b/monitoring/unit/gpu_detailed_test.go @@ -51,14 +51,14 @@ func TestDetailedGPUInfo(t *testing.T) { t.Logf(" Name: %s", info.Name) t.Logf(" Memory Total: %d MB", info.MemoryTotal) t.Logf(" Memory Used: %d MB", info.MemoryUsed) - t.Logf(" Memory Free: %d MB", info.MemoryFree) + //t.Logf(" Memory Free: %d MB", info.MemoryFree) t.Logf(" Utilization: %.1f%%", info.Utilization) t.Logf(" Temperature: %d°C", info.Temperature) // 验证数据的合理性 - if info.MemoryTotal > 0 && info.MemoryUsed+info.MemoryFree != info.MemoryTotal { - t.Logf("Warning: Memory usage calculation may be inconsistent for %s", info.Name) - } + //if info.MemoryTotal > 0 && info.MemoryUsed+info.MemoryFree != info.MemoryTotal { + // t.Logf("Warning: Memory usage calculation may be inconsistent for %s", info.Name) + //} if info.Utilization < 0 || info.Utilization > 100 { t.Errorf("Invalid utilization value for %s: %.1f%%", info.Name, info.Utilization)