mirror of
https://github.com/fankes/komari-agent.git
synced 2025-10-18 18:49:23 +08:00
feat: 添加 Windows 安全警告功能
This commit is contained in:
Akizon77
408
cmd/warn_windows.go
Normal file
408
cmd/warn_windows.go
Normal file
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user