mirror of
https://github.com/fankes/komari-agent.git
synced 2025-10-18 10:39:24 +08:00
409 lines
12 KiB
Go
409 lines
12 KiB
Go
//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
|
||
}
|