3319 lines
140 KiB
Python
3319 lines
140 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
NetDesk Host - unattended remote desktop agent for Windows (full access).
|
|
|
|
Modes:
|
|
python netdesk_host.py --server wss://VPS/ws [--listen 8901] # simple in-process mode
|
|
python netdesk_host.py --install # install+start boot-time service broker (run as admin once)
|
|
python netdesk_host.py --uninstall
|
|
python netdesk_host.py --broker # service broker (normally started by the scheduled task)
|
|
python netdesk_host.py --session-agent --broker ws://127.0.0.1:PORT/agent --token TOK
|
|
# per-desktop worker, spawned by the broker
|
|
|
|
Service architecture (like TeamViewer/RustDesk):
|
|
broker (SYSTEM, session 0, starts at boot)
|
|
- owns the outbound connection to the NetDesk server
|
|
- spawns a desktop agent into the active console session as SYSTEM
|
|
-> agent can capture and inject on the user desktop, the logon screen
|
|
and UAC prompts (full access, before/without any user being logged on)
|
|
- relays traffic server <-> agent (agent holds all E2E keys, broker is a
|
|
dumb pipe)
|
|
|
|
Features: H.264 capture/encode, SendInput keyboard+mouse (desktop-following),
|
|
clipboard sync, file manager with folder upload, chat, remote power, process
|
|
manager, privacy mode (monitor off), lock-on-disconnect, optional system audio
|
|
(PyAudioWPatch), optional TOTP 2FA, wake display on connect.
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import logging
|
|
import mss
|
|
import numpy as np
|
|
import os
|
|
import queue
|
|
import secrets
|
|
import shutil
|
|
import string
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from fractions import Fraction
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import aiohttp
|
|
from aiohttp import web, WSMsgType
|
|
from cryptography.hazmat.primitives.asymmetric import ec
|
|
from cryptography.hazmat.primitives import serialization
|
|
|
|
# WebRTC (aiortc): optional P2P/UDP fast path. When it's missing or the NAT
|
|
# won't allow it, everything keeps working over the WSS relay exactly as before.
|
|
try:
|
|
from aiortc import (RTCPeerConnection, RTCSessionDescription,
|
|
RTCConfiguration, RTCIceServer)
|
|
RTC_AVAILABLE = True
|
|
except Exception:
|
|
RTC_AVAILABLE = False
|
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
from cryptography.hazmat.primitives import hashes
|
|
|
|
# DXGI Desktop Duplication (dxcam): optional, Windows-only fast-capture path.
|
|
# ~10x faster than mss GDI readback (4ms vs ~45ms/frame at 1080p) with built-in
|
|
# damage detection (grab() returns None when the screen is unchanged). When it's
|
|
# missing or fails (secure desktop, RDP, odd GPU), capture falls back to mss.
|
|
try:
|
|
import dxcam as _dxcam
|
|
_DX_AVAILABLE = (sys.platform == "win32")
|
|
except Exception:
|
|
_dxcam = None
|
|
_DX_AVAILABLE = False
|
|
|
|
log = logging.getLogger("netdesk.host")
|
|
VERSION = "1.5.0-netdesk"
|
|
BUILTIN_SERVER = "wss://netdesk.space/ws"
|
|
IS_WIN = sys.platform == "win32"
|
|
CREATE_NO_WINDOW = 0x08000000 if IS_WIN else 0
|
|
|
|
|
|
def normalize_server(addr: str) -> str:
|
|
"""Accept what users type - 'netdesk.space', 'https://x', 'x:8443',
|
|
'wss://x/ws' - and return the websocket URL of the server."""
|
|
a = (addr or "").strip().rstrip("/")
|
|
if not a:
|
|
return ""
|
|
if a.startswith("https://"):
|
|
a = "wss://" + a[8:]
|
|
elif a.startswith("http://"):
|
|
a = "ws://" + a[7:]
|
|
elif not a.startswith(("ws://", "wss://")):
|
|
host = a.split("/")[0]
|
|
is_local = (host.split(":")[0] in ("localhost", "127.0.0.1")
|
|
or host.startswith("192.168.") or host.startswith("10."))
|
|
a = ("ws://" if is_local else "wss://") + a
|
|
_, rest = a.split("://", 1)
|
|
if "/" not in rest:
|
|
a += "/ws"
|
|
return a
|
|
|
|
|
|
def exe_name_config() -> dict:
|
|
"""Configuration baked into the file name (RustDesk-style), e.g.
|
|
'NetDeskHost-host=desk.example.com,key=SECRET.exe' -> preset server + key.
|
|
Lets you hand out a pre-configured exe with no manual setup."""
|
|
if not getattr(sys, "frozen", False):
|
|
return {}
|
|
name = os.path.splitext(os.path.basename(sys.executable))[0]
|
|
out = {}
|
|
for part in name.replace("--", ",").split(","):
|
|
for key, field in (("host=", "server"), ("key=", "reg_key")):
|
|
idx = part.find(key)
|
|
if idx != -1:
|
|
out[field] = part[idx + len(key):].strip()
|
|
if out.get("server"):
|
|
out["server"] = normalize_server(out["server"])
|
|
return out
|
|
|
|
|
|
DEFAULT_SERVER = exe_name_config().get("server") or BUILTIN_SERVER
|
|
|
|
QUALITY_PRESETS = {
|
|
"low": dict(fps=15, bitrate=2000, scale=0.75),
|
|
"balanced": dict(fps=30, bitrate=6000, scale=1.0),
|
|
"high": dict(fps=30, bitrate=12000, scale=1.0),
|
|
"extreme": dict(fps=60, bitrate=20000, scale=1.0),
|
|
}
|
|
|
|
|
|
# ================================================================ config
|
|
|
|
def default_config_dir():
|
|
"""Config location. Frozen (exe) builds are portable: settings and logs
|
|
live next to the executable, so the folder can be carried on a USB stick
|
|
and moved freely. If that folder is not writable (network share, read-only
|
|
mount, Program Files), fall back to the user profile so the app still
|
|
works. Source (non-frozen) runs always use the user profile folder."""
|
|
if getattr(sys, "frozen", False):
|
|
exe_dir = Path(sys.executable).resolve().parent
|
|
try:
|
|
probe = exe_dir / ".nd_write_test"
|
|
probe.write_bytes(b"")
|
|
probe.unlink()
|
|
return exe_dir
|
|
except Exception:
|
|
pass
|
|
if IS_WIN:
|
|
d = Path(os.environ.get("APPDATA", Path.home())) / "NetDesk"
|
|
else:
|
|
d = Path.home() / ".netdesk"
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
return d
|
|
|
|
|
|
def programdata_dir():
|
|
d = Path(os.environ.get("ProgramData", r"C:\ProgramData")) / "NetDesk"
|
|
return d
|
|
|
|
|
|
_CONFIG_DIR = None
|
|
|
|
|
|
def appdata_dir():
|
|
global _CONFIG_DIR
|
|
d = _CONFIG_DIR or default_config_dir()
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
return d
|
|
|
|
|
|
def make_verifier(password: str, salt: bytes = None, iters: int = 200000):
|
|
salt = salt or secrets.token_bytes(16)
|
|
kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=iters)
|
|
verifier = kdf.derive(password.encode("utf-8"))
|
|
return salt, iters, verifier
|
|
|
|
|
|
# ---------------------------------------------------------------- TOTP (RFC 6238, stdlib only)
|
|
|
|
def totp_now(secret: str, t: float = None, step: int = 30, digits: int = 6) -> str:
|
|
if t is None:
|
|
t = time.time()
|
|
key = secret.upper().rstrip("=")
|
|
key += "=" * ((8 - len(key) % 8) % 8)
|
|
key = base64.b32decode(key)
|
|
counter = int(t // step).to_bytes(8, "big")
|
|
digest = hmac.new(key, counter, hashlib.sha1).digest()
|
|
off = digest[-1] & 0x0F
|
|
code = ((digest[off] & 0x7F) << 24 | digest[off + 1] << 16 |
|
|
digest[off + 2] << 8 | digest[off + 3]) % (10 ** digits)
|
|
return str(code).zfill(digits)
|
|
|
|
|
|
def totp_verify(secret: str, code: str, window: int = 1) -> bool:
|
|
code = (code or "").strip()
|
|
if not code.isdigit():
|
|
return False
|
|
now = time.time()
|
|
return any(hmac.compare_digest(totp_now(secret, now + i * 30), code)
|
|
for i in range(-window, window + 1))
|
|
|
|
|
|
def gen_password(n: int = 8) -> str:
|
|
alphabet = "abcdefghijkmnpqrstuvwxyz23456789" # no look-alikes (l/1, o/0)
|
|
return "".join(secrets.choice(alphabet) for _ in range(n))
|
|
|
|
|
|
class Config:
|
|
def __init__(self, path: Path):
|
|
self.path = path
|
|
self.data = {}
|
|
if path.exists():
|
|
try:
|
|
self.data = json.loads(path.read_text("utf-8"))
|
|
except Exception:
|
|
log.warning("config unreadable, regenerating")
|
|
if "id" not in self.data:
|
|
self.data["id"] = str(secrets.randbelow(900000000) + 100000000)
|
|
if "verifier" not in self.data:
|
|
pw = gen_password()
|
|
salt, iters, ver = make_verifier(pw)
|
|
self.data.update(salt=salt.hex(), iters=iters, verifier=ver.hex(), pw_plain=pw)
|
|
self._new_password = pw
|
|
else:
|
|
self._new_password = None
|
|
baked = exe_name_config()
|
|
if baked.get("server") and not self.data.get("server_locked_by_user"):
|
|
self.data["server"] = baked["server"]
|
|
if baked.get("reg_key") and not self.data.get("reg_key"):
|
|
self.data["reg_key"] = baked["reg_key"]
|
|
self.data.setdefault("server", DEFAULT_SERVER)
|
|
self.data.setdefault("name", os.environ.get("COMPUTERNAME", "host"))
|
|
self.data.setdefault("perms", {"input": True, "clipboard": True, "files": True, "audio": True})
|
|
self.data.setdefault("lock_on_end", False)
|
|
self.data.setdefault("privacy_on_session", False)
|
|
self.data.setdefault("allow_display_power", False)
|
|
# optional viewer IP allowlist: comma/space separated IPs or CIDRs
|
|
self.data.setdefault("allow_ips", "")
|
|
self.save()
|
|
|
|
def save(self):
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = self.path.with_suffix(".tmp")
|
|
tmp.write_text(json.dumps(self.data, indent=2), "utf-8")
|
|
os.replace(tmp, self.path)
|
|
|
|
def reload(self):
|
|
"""Pick up changes made by the desktop app (password, account, server)
|
|
without restarting the engine or service."""
|
|
try:
|
|
fresh = json.loads(self.path.read_text("utf-8"))
|
|
if isinstance(fresh, dict) and fresh.get("id") and fresh.get("verifier"):
|
|
self.data = fresh
|
|
except Exception:
|
|
pass
|
|
|
|
@property
|
|
def id(self): return self.data["id"]
|
|
|
|
def set_password(self, pw: str):
|
|
salt, iters, ver = make_verifier(pw)
|
|
self.data.update(salt=salt.hex(), iters=iters, verifier=ver.hex(), pw_plain=pw)
|
|
self.save()
|
|
|
|
def set_totp(self, secret: str):
|
|
self.data["totp"] = secret
|
|
self.save()
|
|
|
|
def check_proof(self, nonce: str, proof: str) -> bool:
|
|
want = hashlib.sha256((nonce + self.data["verifier"]).encode()).hexdigest()
|
|
return secrets.compare_digest(want, proof.lower())
|
|
|
|
def check_proof_bound(self, nonce: str, host_pub_b64: str, proof: str) -> bool:
|
|
"""Channel-bound proof (MITM hardening): hex(SHA256(nonce+verifier+host_pub)).
|
|
Binds the host's ephemeral ECDH pubkey so a relay that swapped it is
|
|
detected (it can't produce a matching proof without the verifier)."""
|
|
want = hashlib.sha256(
|
|
(nonce + self.data["verifier"] + host_pub_b64).encode()).hexdigest()
|
|
return secrets.compare_digest(want, proof.lower())
|
|
|
|
|
|
# ================================================================ Windows layer
|
|
|
|
if IS_WIN:
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
from ctypes.wintypes import DWORD, WORD, LONG
|
|
|
|
ULONG_PTR = ctypes.c_size_t
|
|
user32 = ctypes.windll.user32
|
|
kernel32 = ctypes.windll.kernel32
|
|
advapi32 = ctypes.windll.advapi32
|
|
|
|
def _set_dpi_awareness():
|
|
"""Make the process PER-MONITOR-DPI-AWARE (v2) BEFORE mss/capture/SendInput
|
|
run. Without this, screen-capture geometry, monitor origins and injected
|
|
mouse coordinates disagree on scaled or mixed-DPI multi-monitor setups,
|
|
so the remote cursor lands in the wrong place. RustDesk does the same."""
|
|
try:
|
|
user32.SetProcessDpiAwarenessContext(ctypes.c_void_p(-4)) # PER_MONITOR_AWARE_V2
|
|
return
|
|
except Exception:
|
|
pass
|
|
try:
|
|
ctypes.windll.shcore.SetProcessDpiAwareness(2) # PER_MONITOR_AWARE
|
|
return
|
|
except Exception:
|
|
pass
|
|
try:
|
|
user32.SetProcessDPIAware() # system-DPI-aware (last resort)
|
|
except Exception:
|
|
pass
|
|
|
|
_set_dpi_awareness()
|
|
|
|
class MOUSEINPUT(ctypes.Structure):
|
|
_fields_ = [("dx", LONG), ("dy", LONG), ("mouseData", DWORD),
|
|
("dwFlags", DWORD), ("time", DWORD), ("dwExtraInfo", ULONG_PTR)]
|
|
|
|
class KEYBDINPUT(ctypes.Structure):
|
|
_fields_ = [("wVk", WORD), ("wScan", WORD), ("dwFlags", DWORD),
|
|
("time", DWORD), ("dwExtraInfo", ULONG_PTR)]
|
|
|
|
class _U(ctypes.Union):
|
|
# NOTE: do NOT add a c_byte padding member here. An oversized union makes
|
|
# sizeof(INPUT) 48 instead of the 40 Windows x64 requires, and SendInput
|
|
# then rejects EVERY event (returns 0) - no mouse, no keyboard at all.
|
|
_fields_ = [("mi", MOUSEINPUT), ("ki", KEYBDINPUT)]
|
|
|
|
class INPUT(ctypes.Structure):
|
|
_fields_ = [("type", DWORD), ("u", _U)]
|
|
|
|
class PROCESS_INFORMATION(ctypes.Structure):
|
|
_fields_ = [("hProcess", wintypes.HANDLE), ("hThread", wintypes.HANDLE),
|
|
("dwProcessId", DWORD), ("dwThreadId", DWORD)]
|
|
|
|
class STARTUPINFOW(ctypes.Structure):
|
|
_fields_ = [("cb", DWORD), ("lpReserved", wintypes.LPWSTR), ("lpDesktop", wintypes.LPWSTR),
|
|
("lpTitle", wintypes.LPWSTR), ("dwX", DWORD), ("dwY", DWORD),
|
|
("dwXSize", DWORD), ("dwYSize", DWORD), ("dwXCountChars", DWORD),
|
|
("dwYCountChars", DWORD), ("dwFillAttribute", DWORD), ("dwFlags", DWORD),
|
|
("wShowWindow", WORD), ("cbReserved2", WORD), ("lpReserved2", wintypes.LPBYTE),
|
|
("hStdInput", wintypes.HANDLE), ("hStdOutput", wintypes.HANDLE),
|
|
("hStdError", wintypes.HANDLE)]
|
|
|
|
INPUT_MOUSE, INPUT_KEYBOARD = 0, 1
|
|
KEYEVENTF_KEYUP, KEYEVENTF_EXTENDEDKEY, KEYEVENTF_SCANCODE = 0x0002, 0x0001, 0x0008
|
|
MOUSEEVENTF_MOVE, MOUSEEVENTF_ABSOLUTE, MOUSEEVENTF_VIRTUALDESKTOP = 0x0001, 0x8000, 0x4000
|
|
MOUSEEVENTF_LEFTDOWN, MOUSEEVENTF_LEFTUP = 0x0002, 0x0004
|
|
MOUSEEVENTF_RIGHTDOWN, MOUSEEVENTF_RIGHTUP = 0x0008, 0x0010
|
|
MOUSEEVENTF_MIDDLEDOWN, MOUSEEVENTF_MIDDLEUP = 0x0020, 0x0040
|
|
MOUSEEVENTF_WHEEL, MOUSEEVENTF_HWHEEL = 0x0800, 0x1000
|
|
CF_UNICODETEXT, GMEM_MOVEABLE = 13, 0x0002
|
|
|
|
# 64-bit safety: handle-returning APIs must be declared, else ctypes
|
|
# truncates pointers to 32 bits and random access violations follow.
|
|
user32.GetClipboardData.restype = ctypes.c_void_p
|
|
user32.GetClipboardData.argtypes = [wintypes.UINT]
|
|
user32.SetClipboardData.restype = ctypes.c_void_p
|
|
user32.SetClipboardData.argtypes = [wintypes.UINT, wintypes.HGLOBAL]
|
|
kernel32.GlobalAlloc.restype = wintypes.HGLOBAL
|
|
kernel32.GlobalAlloc.argtypes = [wintypes.UINT, ctypes.c_size_t]
|
|
kernel32.GlobalLock.restype = ctypes.c_void_p
|
|
kernel32.GlobalLock.argtypes = [wintypes.HGLOBAL]
|
|
kernel32.GlobalUnlock.argtypes = [wintypes.HGLOBAL]
|
|
user32.MapVirtualKeyW.restype = wintypes.UINT
|
|
user32.MapVirtualKeyW.argtypes = [wintypes.UINT, wintypes.UINT]
|
|
user32.SendInput.argtypes = [wintypes.UINT, ctypes.POINTER(INPUT), ctypes.c_int]
|
|
|
|
# --- desktop following (logon screen / UAC secure desktop) ---
|
|
DESKTOP_ACCESS = 0x1 | 0x2 | 0x8 | 0x80 | 0x100 # read|createwin|hook|write|switch
|
|
user32.OpenInputDesktop.restype = wintypes.HANDLE
|
|
user32.OpenInputDesktop.argtypes = [DWORD, wintypes.BOOL, DWORD]
|
|
user32.CloseDesktop.argtypes = [wintypes.HANDLE]
|
|
user32.SetThreadDesktop.restype = wintypes.BOOL
|
|
user32.SetThreadDesktop.argtypes = [wintypes.HANDLE]
|
|
user32.GetUserObjectInformationW.argtypes = [wintypes.HANDLE, wintypes.INT,
|
|
ctypes.c_void_p, DWORD, ctypes.POINTER(DWORD)]
|
|
user32.GetThreadDesktop.restype = wintypes.HANDLE
|
|
user32.GetThreadDesktop.argtypes = [DWORD]
|
|
kernel32.GetCurrentThreadId.restype = DWORD
|
|
user32.LockWorkStation.restype = wintypes.BOOL
|
|
|
|
_desktop_tls = threading.local()
|
|
|
|
def _desktop_name(hdesk):
|
|
if not hdesk:
|
|
return None
|
|
buf = ctypes.create_unicode_buffer(64)
|
|
need = DWORD()
|
|
if user32.GetUserObjectInformationW(hdesk, 2, buf, 64, ctypes.byref(need)): # UOI_NAME
|
|
return buf.value
|
|
return None
|
|
|
|
def attach_desktop():
|
|
"""Follow the active input desktop so injection reaches the logon screen
|
|
and UAC secure desktop (service/SYSTEM mode). CRITICAL: only call
|
|
SetThreadDesktop when the input desktop is DIFFERENT from the desktop this
|
|
thread is already on. Switching to a fresh handle of the *same* desktop
|
|
silently breaks SendInput (mouse/keyboard stop working entirely) - which
|
|
is exactly what happens in normal user mode, where the thread already owns
|
|
the visible 'Default' desktop."""
|
|
if getattr(_desktop_tls, "name", None) is None:
|
|
# where this thread starts (usually 'Default' in a user session)
|
|
_desktop_tls.name = _desktop_name(
|
|
user32.GetThreadDesktop(kernel32.GetCurrentThreadId()))
|
|
h = user32.OpenInputDesktop(0, True, DESKTOP_ACCESS)
|
|
if not h:
|
|
return _desktop_tls.name
|
|
name = _desktop_name(h)
|
|
if not name or name == _desktop_tls.name:
|
|
user32.CloseDesktop(h) # already on the input desktop: do NOT switch
|
|
return _desktop_tls.name
|
|
# input desktop changed (user switch / logon screen / UAC): follow it
|
|
if user32.SetThreadDesktop(h):
|
|
old = getattr(_desktop_tls, "hdesk", None)
|
|
if old:
|
|
user32.CloseDesktop(old)
|
|
_desktop_tls.hdesk = h
|
|
_desktop_tls.name = name
|
|
else:
|
|
user32.CloseDesktop(h)
|
|
return _desktop_tls.name
|
|
|
|
def desktop_name():
|
|
return getattr(_desktop_tls, "name", None)
|
|
|
|
# --- input ---
|
|
VK = {
|
|
**{f"Key{c}": 0x41 + i for i, c in enumerate("ABCDEFGHIJKLMNOPQRSTUVWXYZ")},
|
|
**{f"Digit{c}": 0x30 + i for i, c in enumerate("0123456789")},
|
|
**{f"Numpad{c}": 0x60 + i for i, c in enumerate("0123456789")},
|
|
**{f"F{n}": 0x70 + n - 1 for n in range(1, 25)},
|
|
"Backquote": 0xC0, "Minus": 0xBD, "Equal": 0xBB, "BracketLeft": 0xDB,
|
|
"BracketRight": 0xDD, "Backslash": 0xDC, "Semicolon": 0xBA, "Quote": 0xDE,
|
|
"Comma": 0xBC, "Period": 0xBE, "Slash": 0xBF, "IntlBackslash": 0xE2,
|
|
"CapsLock": 0x14, "Space": 0x20, "Enter": 0x0D, "NumpadEnter": 0x0D,
|
|
"Backspace": 0x08, "Tab": 0x09, "Escape": 0x1B,
|
|
"Delete": 0x2E, "Insert": 0x2D, "Home": 0x24, "End": 0x23,
|
|
"PageUp": 0x21, "PageDown": 0x22,
|
|
"ArrowLeft": 0x25, "ArrowUp": 0x26, "ArrowRight": 0x27, "ArrowDown": 0x28,
|
|
"ControlLeft": 0xA2, "ControlRight": 0xA3, "AltLeft": 0xA4, "AltRight": 0xA5,
|
|
"ShiftLeft": 0xA0, "ShiftRight": 0xA1,
|
|
"MetaLeft": 0x5B, "MetaRight": 0x5C, "ContextMenu": 0x5D,
|
|
"PrintScreen": 0x2C, "ScrollLock": 0x91, "Pause": 0x13, "NumLock": 0x90,
|
|
"NumpadMultiply": 0x6A, "NumpadAdd": 0x6B, "NumpadSubtract": 0x6D,
|
|
"NumpadDecimal": 0x6E, "NumpadDivide": 0x6F,
|
|
}
|
|
EXTENDED_CODES = {
|
|
"ArrowRight", "ArrowLeft", "ArrowUp", "ArrowDown", "PageUp", "PageDown",
|
|
"Home", "End", "Insert", "Delete", "NumLock", "PrintScreen", "Pause",
|
|
"ControlRight", "AltRight", "MetaLeft", "MetaRight", "ContextMenu",
|
|
"NumpadEnter", "NumpadDivide", "ScrollLock",
|
|
}
|
|
|
|
user32.SendInput.restype = wintypes.UINT
|
|
|
|
def send_inputs(*inputs):
|
|
arr = (INPUT * len(inputs))(*inputs)
|
|
n = user32.SendInput(len(inputs), arr, ctypes.sizeof(INPUT))
|
|
if n != len(inputs):
|
|
# 0 here almost always means a bad cbSize/struct layout (see the _U
|
|
# union note) or input blocked by UIPI/secure desktop - make it loud
|
|
log.warning("SendInput injected %d/%d (sizeof(INPUT)=%d, err=%d)",
|
|
n, len(inputs), ctypes.sizeof(INPUT), ctypes.get_last_error())
|
|
|
|
def _mi(flags, data=0, dx=0, dy=0):
|
|
return INPUT(INPUT_MOUSE, _U(mi=MOUSEINPUT(dx, dy, data & 0xFFFFFFFF, flags, 0, 0)))
|
|
|
|
# SM_*VIRTUALSCREEN indices - the whole multi-monitor bounding box, read
|
|
# cheaply (no mss allocation per mouse move)
|
|
SM_XVIRTUALSCREEN, SM_YVIRTUALSCREEN = 76, 77
|
|
SM_CXVIRTUALSCREEN, SM_CYVIRTUALSCREEN = 78, 79
|
|
MOUSEEVENTF_XDOWN, MOUSEEVENTF_XUP = 0x0080, 0x0100
|
|
XBUTTON1, XBUTTON2 = 0x0001, 0x0002
|
|
|
|
def virtual_desktop_box():
|
|
return {"left": user32.GetSystemMetrics(SM_XVIRTUALSCREEN),
|
|
"top": user32.GetSystemMetrics(SM_YVIRTUALSCREEN),
|
|
"w": user32.GetSystemMetrics(SM_CXVIRTUALSCREEN),
|
|
"h": user32.GetSystemMetrics(SM_CYVIRTUALSCREEN)}
|
|
|
|
def inject_mouse_move(x, y, union_box=None):
|
|
u = union_box or virtual_desktop_box()
|
|
nx = int(round((x - u["left"]) * 65535 / max(1, u["w"] - 1)))
|
|
ny = int(round((y - u["top"]) * 65535 / max(1, u["h"] - 1)))
|
|
send_inputs(_mi(MOUSEEVENTF_MOVE | MOUSEEVENTF_ABSOLUTE | MOUSEEVENTF_VIRTUALDESKTOP, dx=nx, dy=ny))
|
|
|
|
def inject_mouse_buttons(new_b, old_b):
|
|
evts = []
|
|
for mask, down, up in (
|
|
(1, MOUSEEVENTF_LEFTDOWN, MOUSEEVENTF_LEFTUP),
|
|
(2, MOUSEEVENTF_RIGHTDOWN, MOUSEEVENTF_RIGHTUP),
|
|
(4, MOUSEEVENTF_MIDDLEDOWN, MOUSEEVENTF_MIDDLEUP),
|
|
):
|
|
if new_b & mask and not old_b & mask:
|
|
evts.append(_mi(down))
|
|
if not new_b & mask and old_b & mask:
|
|
evts.append(_mi(up))
|
|
# side buttons (back/forward): bit 8 = X1, bit 16 = X2
|
|
for mask, xb in ((8, XBUTTON1), (16, XBUTTON2)):
|
|
if new_b & mask and not old_b & mask:
|
|
evts.append(_mi(MOUSEEVENTF_XDOWN, data=xb))
|
|
if not new_b & mask and old_b & mask:
|
|
evts.append(_mi(MOUSEEVENTF_XUP, data=xb))
|
|
if evts:
|
|
send_inputs(*evts)
|
|
|
|
def inject_wheel(v, h):
|
|
evts = []
|
|
if v:
|
|
evts.append(_mi(MOUSEEVENTF_WHEEL, int(v * 120)))
|
|
if h:
|
|
evts.append(_mi(MOUSEEVENTF_HWHEEL, int(h * 120)))
|
|
if evts:
|
|
send_inputs(*evts)
|
|
|
|
_pressed_keys = set()
|
|
|
|
def inject_key(code, down):
|
|
vk = VK.get(code)
|
|
if not vk:
|
|
return
|
|
scan = user32.MapVirtualKeyW(vk, 0)
|
|
flags = KEYEVENTF_SCANCODE | (KEYEVENTF_EXTENDEDKEY if code in EXTENDED_CODES else 0)
|
|
if not down:
|
|
flags |= KEYEVENTF_KEYUP
|
|
if down:
|
|
_pressed_keys.add(code)
|
|
else:
|
|
_pressed_keys.discard(code)
|
|
if scan:
|
|
send_inputs(INPUT(INPUT_KEYBOARD, _U(ki=KEYBDINPUT(0, scan, flags, 0, 0))))
|
|
else:
|
|
send_inputs(INPUT(INPUT_KEYBOARD, _U(ki=KEYBDINPUT(vk, 0, flags, 0, 0))))
|
|
|
|
def release_all_keys():
|
|
for code in list(_pressed_keys):
|
|
inject_key(code, False)
|
|
_pressed_keys.clear()
|
|
|
|
class InputPump(threading.Thread):
|
|
"""Dedicated input thread that follows the active desktop before each
|
|
injection (works on logon screen / UAC when running as SYSTEM)."""
|
|
|
|
def __init__(self, host):
|
|
super().__init__(daemon=True)
|
|
self.host = host
|
|
self.q = queue.Queue(maxsize=2048)
|
|
self.buttons = 0
|
|
self._last_attach = 0.0
|
|
|
|
def put(self, item):
|
|
try:
|
|
self.q.put_nowait(item)
|
|
except queue.Full:
|
|
# never drop a button/key/wheel event (a lost "up" = a stuck
|
|
# key or button). Only mouse *moves* are safe to drop under load;
|
|
# evict the oldest queued move to make room for this event.
|
|
if item[0] != "move":
|
|
try:
|
|
for _ in range(len(self.q.queue)):
|
|
old = self.q.get_nowait()
|
|
if old[0] == "move":
|
|
break
|
|
except Exception:
|
|
pass
|
|
try:
|
|
self.q.put_nowait(item)
|
|
except Exception:
|
|
pass
|
|
|
|
def run(self):
|
|
while True:
|
|
item = self.q.get()
|
|
try:
|
|
# latest-wins for motion: if moves piled up (slow link burst,
|
|
# injector stalled), only the newest position matters - never
|
|
# skip past a button/key/wheel event though
|
|
if item[0] == "move":
|
|
while True:
|
|
try:
|
|
nxt = self.q.queue[0]
|
|
except IndexError:
|
|
break
|
|
if nxt[0] != "move":
|
|
break
|
|
item = self.q.get_nowait()
|
|
# follow the active input desktop (logon screen / UAC) but not
|
|
# on every single event - and INSIDE the try so a transient
|
|
# OpenInputDesktop failure can never kill the input thread
|
|
tnow = time.time()
|
|
if item[0] == "key" or tnow - self._last_attach > 0.2:
|
|
attach_desktop()
|
|
self._last_attach = tnow
|
|
if item[0] == "move":
|
|
# cheap virtual-desktop bounds (no per-move mss allocation)
|
|
inject_mouse_move(item[1], item[2], virtual_desktop_box())
|
|
elif item[0] == "buttons":
|
|
inject_mouse_buttons(item[1], self.buttons)
|
|
self.buttons = item[1]
|
|
elif item[0] == "wheel":
|
|
inject_wheel(item[1], item[2])
|
|
elif item[0] == "key":
|
|
if self.host.cfg.data["perms"].get("input", True):
|
|
inject_key(item[1], item[2])
|
|
elif item[0] == "wake":
|
|
send_inputs(_mi(MOUSEEVENTF_MOVE, dx=1, dy=0),
|
|
_mi(MOUSEEVENTF_MOVE, dx=-1, dy=0))
|
|
except Exception as e:
|
|
log.debug("input pump error: %s", e)
|
|
|
|
# --- remote cursor shape (RustDesk-style overlay: sent separately from video
|
|
# so idle screens still update the cursor and no re-encode is forced) ---
|
|
gdi32 = ctypes.windll.gdi32
|
|
|
|
class CURSORINFO(ctypes.Structure):
|
|
_fields_ = [("cbSize", DWORD), ("flags", DWORD), ("hCursor", wintypes.HANDLE),
|
|
("ptScreenPos", wintypes.POINT)]
|
|
|
|
class ICONINFO(ctypes.Structure):
|
|
_fields_ = [("fIcon", wintypes.BOOL), ("xHotspot", DWORD), ("yHotspot", DWORD),
|
|
("hbmMask", wintypes.HANDLE), ("hbmColor", wintypes.HANDLE)]
|
|
|
|
class BITMAP(ctypes.Structure):
|
|
_fields_ = [("bmType", LONG), ("bmWidth", LONG), ("bmHeight", LONG),
|
|
("bmWidthBytes", LONG), ("bmPlanes", WORD), ("bmBitsPixel", WORD),
|
|
("bmBits", ctypes.c_void_p)]
|
|
|
|
class BITMAPINFOHEADER(ctypes.Structure):
|
|
_fields_ = [("biSize", DWORD), ("biWidth", LONG), ("biHeight", LONG),
|
|
("biPlanes", WORD), ("biBitCount", WORD), ("biCompression", DWORD),
|
|
("biSizeImage", DWORD), ("biXPelsPerMeter", LONG), ("biYPelsPerMeter", LONG),
|
|
("biClrUsed", DWORD), ("biClrImportant", DWORD)]
|
|
|
|
user32.GetCursorInfo.argtypes = [ctypes.POINTER(CURSORINFO)]
|
|
user32.GetIconInfo.argtypes = [wintypes.HANDLE, ctypes.POINTER(ICONINFO)]
|
|
CURSOR_SHOWING = 0x0001
|
|
_cursor_png_cache = {}
|
|
# shared system cursors -> CSS name. The viewer then uses the OS-native pointer
|
|
# (crisp at any DPI), and it still works when GetIconInfo refuses a handle that
|
|
# belongs to another session (engine running as SYSTEM in session 0).
|
|
_STD_CURSORS = {}
|
|
user32.LoadCursorW.restype = ctypes.c_void_p
|
|
for _idc, _css in ((32512, "default"), (32513, "text"), (32514, "wait"),
|
|
(32515, "crosshair"), (32516, "default"), (32642, "nwse-resize"),
|
|
(32643, "nesw-resize"), (32644, "ew-resize"), (32645, "ns-resize"),
|
|
(32646, "move"), (32648, "not-allowed"), (32649, "pointer"),
|
|
(32650, "progress"), (32651, "help")):
|
|
try:
|
|
_h = user32.LoadCursorW(None, ctypes.c_void_p(_idc)) # MAKEINTRESOURCE
|
|
if _h:
|
|
_STD_CURSORS[int(_h)] = _css
|
|
except Exception:
|
|
pass
|
|
|
|
def cursor_state():
|
|
"""(hCursor, global_x, global_y) or None. hCursor==0 means hidden."""
|
|
ci = CURSORINFO()
|
|
ci.cbSize = ctypes.sizeof(CURSORINFO)
|
|
if not user32.GetCursorInfo(ctypes.byref(ci)):
|
|
return None
|
|
showing = bool(ci.flags & CURSOR_SHOWING) and ci.hCursor
|
|
return (int(ci.hCursor) if showing else 0, ci.ptScreenPos.x, ci.ptScreenPos.y)
|
|
|
|
def cursor_png(hcursor):
|
|
"""RGBA PNG (base64) + hotspot for a cursor handle, cached by handle."""
|
|
if hcursor in _cursor_png_cache:
|
|
return _cursor_png_cache[hcursor]
|
|
import numpy as _np
|
|
from PIL import Image as _Img
|
|
import io as _io
|
|
ii = ICONINFO()
|
|
if not user32.GetIconInfo(hcursor, ctypes.byref(ii)):
|
|
return None
|
|
try:
|
|
bm = BITMAP()
|
|
src = ii.hbmColor if ii.hbmColor else ii.hbmMask
|
|
gdi32.GetObjectW(src, ctypes.sizeof(BITMAP), ctypes.byref(bm))
|
|
w, h = bm.bmWidth, bm.bmHeight
|
|
if not ii.hbmColor:
|
|
h //= 2
|
|
hdc = user32.GetDC(0)
|
|
bi = BITMAPINFOHEADER()
|
|
bi.biSize = ctypes.sizeof(BITMAPINFOHEADER)
|
|
bi.biWidth = w; bi.biPlanes = 1; bi.biBitCount = 32; bi.biCompression = 0
|
|
|
|
def dib(hbm, rows):
|
|
bi.biHeight = -abs(rows)
|
|
buf = (ctypes.c_char * (w * abs(rows) * 4))()
|
|
gdi32.GetDIBits(hdc, hbm, 0, abs(rows), buf, ctypes.byref(bi), 0)
|
|
return _np.frombuffer(bytes(buf), dtype=_np.uint8).reshape(abs(rows), w, 4)
|
|
|
|
if ii.hbmColor:
|
|
col = dib(ii.hbmColor, h)
|
|
if col[:, :, 3].any():
|
|
rgba = col[:, :, [2, 1, 0, 3]].copy()
|
|
else:
|
|
andm = dib(ii.hbmMask, h)[:, :, 0]
|
|
alpha = _np.where(andm == 0, 255, 0).astype(_np.uint8)
|
|
rgba = _np.dstack([col[:, :, 2], col[:, :, 1], col[:, :, 0], alpha])
|
|
else:
|
|
both = dib(ii.hbmMask, h * 2)[:, :, 0]
|
|
AND, XOR = both[:h], both[h:]
|
|
alpha = _np.where(AND == 0, 255, 0).astype(_np.uint8)
|
|
lum = _np.where((AND == 0) & (XOR > 0), 255, 0).astype(_np.uint8)
|
|
rgba = _np.dstack([lum, lum, lum, alpha])
|
|
user32.ReleaseDC(0, hdc)
|
|
b = _io.BytesIO()
|
|
_Img.fromarray(rgba, "RGBA").save(b, "PNG")
|
|
out = {"w": int(w), "h": int(h), "hx": int(ii.xHotspot), "hy": int(ii.yHotspot),
|
|
"png": base64.b64encode(b.getvalue()).decode()}
|
|
if len(_cursor_png_cache) < 64:
|
|
_cursor_png_cache[hcursor] = out
|
|
return out
|
|
except Exception:
|
|
return None
|
|
finally:
|
|
try:
|
|
if ii.hbmColor:
|
|
gdi32.DeleteObject(ii.hbmColor)
|
|
if ii.hbmMask:
|
|
gdi32.DeleteObject(ii.hbmMask)
|
|
except Exception:
|
|
pass
|
|
|
|
# --- clipboard ---
|
|
def clip_get():
|
|
if not user32.OpenClipboard(None):
|
|
return None
|
|
try:
|
|
h = user32.GetClipboardData(CF_UNICODETEXT)
|
|
if not h:
|
|
return None
|
|
ptr = kernel32.GlobalLock(h)
|
|
if not ptr:
|
|
return None
|
|
try:
|
|
return ctypes.wstring_at(ptr)
|
|
finally:
|
|
kernel32.GlobalUnlock(h)
|
|
finally:
|
|
user32.CloseClipboard()
|
|
|
|
def clip_set(text):
|
|
data = text.encode("utf-16-le") + b"\x00\x00"
|
|
if not user32.OpenClipboard(None):
|
|
return False
|
|
try:
|
|
user32.EmptyClipboard()
|
|
h = kernel32.GlobalAlloc(GMEM_MOVEABLE, len(data))
|
|
if not h:
|
|
return False
|
|
ptr = kernel32.GlobalLock(h)
|
|
if not ptr:
|
|
return False
|
|
ctypes.memmove(ptr, data, len(data))
|
|
kernel32.GlobalUnlock(h)
|
|
user32.SetClipboardData(CF_UNICODETEXT, h)
|
|
return True
|
|
finally:
|
|
user32.CloseClipboard()
|
|
|
|
# --- display power (privacy mode) ---
|
|
HWND_BROADCAST, WM_SYSCOMMAND, SC_MONITORPOWER = 0xFFFF, 0x0112, 0xF170
|
|
SMTO_ABORTIFHUNG = 0x0002
|
|
|
|
def monitor_power(state): # -1 on, 2 off
|
|
res = DWORD()
|
|
user32.SendMessageTimeoutW(HWND_BROADCAST, WM_SYSCOMMAND, SC_MONITORPOWER,
|
|
state, SMTO_ABORTIFHUNG, 500, ctypes.byref(res))
|
|
|
|
# --- process manager ---
|
|
def process_list():
|
|
out = subprocess.run(["tasklist", "/FO", "CSV", "/NH"], capture_output=True,
|
|
text=True, creationflags=0x08000000, timeout=15).stdout
|
|
items = []
|
|
for line in out.splitlines():
|
|
parts = [p.strip('"') for p in line.split('","')]
|
|
if len(parts) >= 5:
|
|
try:
|
|
mem = int(parts[4].replace("K", "").replace(",", "").strip())
|
|
except Exception:
|
|
mem = 0
|
|
items.append({"name": parts[0], "pid": int(parts[1]), "mem": mem})
|
|
items.sort(key=lambda i: -i["mem"])
|
|
return items[:300]
|
|
|
|
def process_kill(pid):
|
|
r = subprocess.run(["taskkill", "/F", "/PID", str(pid)], capture_output=True,
|
|
text=True, creationflags=0x08000000, timeout=10)
|
|
return r.returncode == 0, (r.stdout or r.stderr).strip()
|
|
|
|
def print_file(path):
|
|
"""Print a document on this machine's default printer (print verb)."""
|
|
if not os.path.isfile(path):
|
|
return False, "file not found"
|
|
try:
|
|
r = ctypes.windll.shell32.ShellExecuteW(0, "print", path, None, None, 0)
|
|
if r > 32:
|
|
return True, "sent to default printer"
|
|
return False, f"no print association (code {r})"
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
# --- power / commands ---
|
|
def sys_command(action, cmd=None):
|
|
CREATE_NO_WINDOW = 0x08000000
|
|
try:
|
|
if action == "lock":
|
|
user32.LockWorkStation()
|
|
return True, "locked"
|
|
if action == "logoff":
|
|
subprocess.Popen("shutdown /l", shell=True, creationflags=CREATE_NO_WINDOW)
|
|
return True, "logging off"
|
|
if action == "reboot":
|
|
subprocess.Popen("shutdown /r /t 3", shell=True, creationflags=CREATE_NO_WINDOW)
|
|
return True, "rebooting"
|
|
if action == "shutdown":
|
|
subprocess.Popen("shutdown /s /t 3", shell=True, creationflags=CREATE_NO_WINDOW)
|
|
return True, "shutting down"
|
|
if action == "run" and cmd:
|
|
subprocess.Popen(cmd, shell=True, creationflags=CREATE_NO_WINDOW)
|
|
return True, "started"
|
|
return False, "unknown command"
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
# --- session/process helpers for the service broker ---
|
|
kernel32.WTSGetActiveConsoleSessionId.restype = DWORD
|
|
kernel32.WTSGetActiveConsoleSessionId.argtypes = []
|
|
kernel32.ProcessIdToSessionId.restype = wintypes.BOOL
|
|
kernel32.ProcessIdToSessionId.argtypes = [DWORD, ctypes.POINTER(DWORD)]
|
|
kernel32.GetCurrentProcess.restype = wintypes.HANDLE
|
|
kernel32.GetCurrentProcess.argtypes = []
|
|
kernel32.TerminateProcess.argtypes = [wintypes.HANDLE, wintypes.UINT]
|
|
|
|
TOKEN_ALL_ACCESS = 0xF01FF
|
|
TokenSessionId = 12
|
|
SecurityImpersonation, TokenPrimary = 2, 1
|
|
DETACHED_PROCESS, CREATE_NEW_PROCESS_GROUP = 0x00000008, 0x00000200
|
|
CREATE_UNICODE_ENVIRONMENT, CREATE_BREAKAWAY_FROM_JOB = 0x00000400, 0x01000000
|
|
|
|
advapi32.OpenProcessToken.argtypes = [wintypes.HANDLE, DWORD, ctypes.POINTER(wintypes.HANDLE)]
|
|
advapi32.DuplicateTokenEx.argtypes = [wintypes.HANDLE, DWORD, wintypes.LPVOID,
|
|
wintypes.INT, wintypes.INT,
|
|
ctypes.POINTER(wintypes.HANDLE)]
|
|
advapi32.SetTokenInformation.argtypes = [wintypes.HANDLE, wintypes.INT,
|
|
wintypes.LPVOID, DWORD]
|
|
advapi32.CreateProcessAsUserW.argtypes = [wintypes.HANDLE, wintypes.LPCWSTR,
|
|
wintypes.LPWSTR, wintypes.LPVOID,
|
|
wintypes.LPVOID, wintypes.BOOL, DWORD,
|
|
wintypes.LPVOID, wintypes.LPCWSTR,
|
|
ctypes.POINTER(STARTUPINFOW),
|
|
ctypes.POINTER(PROCESS_INFORMATION)]
|
|
userenv = ctypes.windll.userenv
|
|
userenv.CreateEnvironmentBlock.argtypes = [ctypes.POINTER(wintypes.LPVOID),
|
|
wintypes.HANDLE, wintypes.BOOL]
|
|
userenv.DestroyEnvironmentBlock.argtypes = [wintypes.LPVOID]
|
|
|
|
def console_session_id():
|
|
return kernel32.WTSGetActiveConsoleSessionId()
|
|
|
|
def my_session_id():
|
|
sid = DWORD(0)
|
|
kernel32.ProcessIdToSessionId(kernel32.GetCurrentProcessId(), ctypes.byref(sid))
|
|
return sid.value
|
|
|
|
def spawn_in_console_session(cmdline, session_id, desktop="Winsta0\\Default"):
|
|
"""Spawn a detached SYSTEM process inside the given console session
|
|
(needs SeTcbPrivilege, i.e. the caller runs as SYSTEM)."""
|
|
htok = wintypes.HANDLE()
|
|
if not advapi32.OpenProcessToken(kernel32.GetCurrentProcess(),
|
|
TOKEN_ALL_ACCESS, ctypes.byref(htok)):
|
|
raise OSError("OpenProcessToken failed")
|
|
hdup = wintypes.HANDLE()
|
|
if not advapi32.DuplicateTokenEx(htok, TOKEN_ALL_ACCESS, None,
|
|
SecurityImpersonation, TokenPrimary,
|
|
ctypes.byref(hdup)):
|
|
raise OSError("DuplicateTokenEx failed")
|
|
sid = DWORD(session_id)
|
|
if not advapi32.SetTokenInformation(hdup, TokenSessionId,
|
|
ctypes.byref(sid), ctypes.sizeof(sid)):
|
|
raise OSError("SetTokenInformation(TokenSessionId) failed")
|
|
env = wintypes.LPVOID()
|
|
have_env = userenv.CreateEnvironmentBlock(ctypes.byref(env), hdup, False)
|
|
si = STARTUPINFOW()
|
|
si.cb = ctypes.sizeof(si)
|
|
si.lpDesktop = desktop
|
|
pi = PROCESS_INFORMATION()
|
|
flags = DETACHED_PROCESS | CREATE_UNICODE_ENVIRONMENT | CREATE_BREAKAWAY_FROM_JOB
|
|
cmd = ctypes.create_unicode_buffer(cmdline)
|
|
if cmdline.startswith('"'):
|
|
app = cmdline[1:cmdline.index('"', 1)]
|
|
else:
|
|
app = cmdline.split(" ", 1)[0]
|
|
ok = advapi32.CreateProcessAsUserW(
|
|
hdup, app, cmd, None, None, False, flags,
|
|
env.value if have_env else None, None, ctypes.byref(si), ctypes.byref(pi))
|
|
if have_env:
|
|
userenv.DestroyEnvironmentBlock(env)
|
|
if not ok:
|
|
raise OSError(f"CreateProcessAsUser failed: {ctypes.GetLastError()}")
|
|
return pi.hProcess, pi.dwProcessId
|
|
|
|
def is_elevated():
|
|
try:
|
|
return bool(ctypes.windll.shell32.IsUserAnAdmin())
|
|
except Exception:
|
|
return False
|
|
|
|
elif sys.platform == "darwin":
|
|
# ---------------- macOS (Apple Silicon native, pure ctypes) ----------------
|
|
import ctypes.util
|
|
import platform
|
|
|
|
_cg = ctypes.cdll.LoadLibrary(ctypes.util.find_library("CoreGraphics"))
|
|
_cf = ctypes.cdll.LoadLibrary(ctypes.util.find_library("CoreFoundation"))
|
|
|
|
class CGPoint(ctypes.Structure):
|
|
_fields_ = [("x", ctypes.c_double), ("y", ctypes.c_double)]
|
|
|
|
class _Size(ctypes.Structure):
|
|
_fields_ = [("width", ctypes.c_double), ("height", ctypes.c_double)]
|
|
|
|
class CGRect2(ctypes.Structure):
|
|
_fields_ = [("origin", CGPoint), ("size", _Size)]
|
|
|
|
_cg.CGEventCreateMouseEvent.restype = ctypes.c_void_p
|
|
_cg.CGEventCreateMouseEvent.argtypes = [ctypes.c_void_p, ctypes.c_uint32, CGPoint, ctypes.c_uint32]
|
|
_cg.CGEventCreateScrollWheelEvent.restype = ctypes.c_void_p
|
|
_cg.CGEventCreateScrollWheelEvent.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_int32]
|
|
_cg.CGEventCreateKeyboardEvent.restype = ctypes.c_void_p
|
|
_cg.CGEventCreateKeyboardEvent.argtypes = [ctypes.c_void_p, ctypes.c_uint16, ctypes.c_bool]
|
|
_cg.CGEventPost.argtypes = [ctypes.c_uint32, ctypes.c_void_p]
|
|
_cg.CGEventSetLocation.argtypes = [ctypes.c_void_p, CGPoint]
|
|
_cg.CFRelease.argtypes = [ctypes.c_void_p]
|
|
_cg.CGMainDisplayID.restype = ctypes.c_void_p
|
|
_cg.CGDisplayBounds.restype = CGRect2
|
|
_cg.CGDisplayBounds.argtypes = [ctypes.c_void_p]
|
|
_cg.CGDisplayPixelsWide.restype = ctypes.c_uint
|
|
_cg.CGDisplayPixelsWide.argtypes = [ctypes.c_void_p]
|
|
|
|
kCGHIDEventTap = 0
|
|
kCGEventMouseMoved, kCGEventLeftDown, kCGEventLeftUp = 5, 1, 2
|
|
kCGEventRightDown, kCGEventRightUp = 3, 4
|
|
kCGEventOtherDown, kCGEventOtherUp = 25, 26
|
|
kCGScrollEventUnitLine = 1
|
|
# US-ANSI virtual keycodes for our JS KeyboardEvent.code names
|
|
VK_MAC = {
|
|
**{f"Key{c}": v for c, v in zip(
|
|
"ABCDEFGHIJKLMNOPQRSTUVWXYZ",
|
|
[0x00, 0x0B, 0x08, 0x02, 0x0E, 0x03, 0x05, 0x04, 0x22, 0x26,
|
|
0x28, 0x25, 0x2E, 0x2D, 0x1F, 0x23, 0x0C, 0x0F, 0x01, 0x11,
|
|
0x20, 0x09, 0x0D, 0x07, 0x10, 0x06])},
|
|
**{f"Digit{c}": v for c, v in zip(
|
|
"1234567890", [0x12, 0x13, 0x14, 0x15, 0x17, 0x16, 0x1A, 0x1C, 0x19, 0x1D])},
|
|
**{f"F{n}": v for n, v in zip(range(1, 13),
|
|
[0x7A, 0x78, 0x63, 0x76, 0x60, 0x61, 0x62, 0x64, 0x65, 0x6D, 0x67, 0x6F])},
|
|
"Backquote": 0x32, "Minus": 0x1B, "Equal": 0x18, "BracketLeft": 0x21,
|
|
"BracketRight": 0x1E, "Backslash": 0x2A, "Semicolon": 0x29, "Quote": 0x27,
|
|
"Comma": 0x2B, "Period": 0x2F, "Slash": 0x2C,
|
|
"CapsLock": 0x39, "Space": 0x31, "Enter": 0x24, "NumpadEnter": 0x24,
|
|
"Backspace": 0x33, "Tab": 0x30, "Escape": 0x35,
|
|
"Delete": 0x75, "Insert": 0, "Home": 0x73, "End": 0x77,
|
|
"PageUp": 0x74, "PageDown": 0x79,
|
|
"ArrowLeft": 0x7B, "ArrowUp": 0x7E, "ArrowRight": 0x7C, "ArrowDown": 0x7D,
|
|
"ControlLeft": 0x3B, "ControlRight": 0x3E, "AltLeft": 0x3A, "AltRight": 0x3D,
|
|
"ShiftLeft": 0x38, "ShiftRight": 0x3C,
|
|
"MetaLeft": 0x37, "MetaRight": 0x37, "ContextMenu": 0,
|
|
"Numpad0": 0x52, "Numpad1": 0x53, "Numpad2": 0x54, "Numpad3": 0x55,
|
|
"Numpad4": 0x56, "Numpad5": 0x57, "Numpad6": 0x58, "Numpad7": 0x59,
|
|
"Numpad8": 0x5B, "Numpad9": 0x5C,
|
|
}
|
|
|
|
_mac_tls = threading.local()
|
|
|
|
def _display_info():
|
|
"""[(scale, origin_x, origin_y, point_w, point_h)] per active display."""
|
|
if getattr(_mac_tls, "displays", None):
|
|
return _mac_tls.displays
|
|
infos = []
|
|
count = 32
|
|
arr = (ctypes.c_void_p * count)()
|
|
actual = ctypes.c_uint32()
|
|
if _cg.CGGetActiveDisplayList(count, arr, ctypes.byref(actual)) == 0:
|
|
for i in range(actual.value):
|
|
did = arr[i]
|
|
bounds = _cg.CGDisplayBounds(did)
|
|
pw = _cg.CGDisplayPixelsWide(did)
|
|
scale = (pw / bounds.size.width) if bounds.size.width else 1.0
|
|
infos.append((scale, bounds.origin.x, bounds.origin.y,
|
|
bounds.size.width, bounds.size.height))
|
|
_mac_tls.displays = infos or [(1.0, 0.0, 0.0, 1920.0, 1080.0)]
|
|
return _mac_tls.displays
|
|
|
|
def pixel_to_global(x, y):
|
|
"""Map a physical-pixel coordinate on the virtual desktop to global points."""
|
|
for scale, ox, oy, pw, ph in _display_info():
|
|
if x < pw * scale and y < ph * scale:
|
|
return CGPoint(ox + x / scale, oy + y / scale)
|
|
scale, ox, oy, pw, ph = _display_info()[0]
|
|
return CGPoint(ox + x / scale, oy + y / scale)
|
|
|
|
def attach_desktop():
|
|
return "Default"
|
|
|
|
def desktop_name():
|
|
return "Default"
|
|
|
|
def _post_event(ev):
|
|
_cg.CGEventPost(kCGHIDEventTap, ev)
|
|
_cg.CFRelease(ev)
|
|
|
|
def _mouse_event(etype, button, x, y):
|
|
return _cg.CGEventCreateMouseEvent(None, etype, pixel_to_global(x, y), button)
|
|
|
|
def send_inputs(*inputs): # parity stub (mac uses discrete events)
|
|
pass
|
|
|
|
class InputPump(threading.Thread):
|
|
def __init__(self, host):
|
|
super().__init__(daemon=True)
|
|
self.host = host
|
|
self.q = queue.Queue(maxsize=512)
|
|
self.buttons = 0
|
|
|
|
def put(self, item):
|
|
try:
|
|
self.q.put_nowait(item)
|
|
except Exception:
|
|
pass
|
|
|
|
def run(self):
|
|
while True:
|
|
item = self.q.get()
|
|
try:
|
|
if item[0] == "move":
|
|
_post_event(_mouse_event(kCGEventMouseMoved, 0, item[1], item[2]))
|
|
elif item[0] == "buttons":
|
|
b = item[1]
|
|
x, y = self._pos()
|
|
for mask, (down_t, up_t, btn) in (
|
|
(1, (kCGEventLeftDown, kCGEventLeftUp, 0)),
|
|
(2, (kCGEventRightDown, kCGEventRightUp, 1)),
|
|
(4, (kCGEventOtherDown, kCGEventOtherUp, 2))):
|
|
if b & mask and not self.buttons & mask:
|
|
_post_event(_mouse_event(down_t, btn, x, y))
|
|
if not b & mask and self.buttons & mask:
|
|
_post_event(_mouse_event(up_t, btn, x, y))
|
|
self.buttons = b
|
|
elif item[0] == "wheel":
|
|
v, h = item[1], item[2]
|
|
ev = _cg.CGEventCreateScrollWheelEvent(None, kCGScrollEventUnitLine, 2, int(v), int(h))
|
|
_post_event(ev)
|
|
elif item[0] == "key":
|
|
if self.host.cfg.data["perms"].get("input", True):
|
|
kc = VK_MAC.get(item[1])
|
|
if kc:
|
|
_post_event(_cg.CGEventCreateKeyboardEvent(None, kc, item[2]))
|
|
elif item[0] == "wake":
|
|
x, y = self._pos()
|
|
_post_event(_mouse_event(kCGEventMouseMoved, 0, x + 1, y))
|
|
except Exception as e:
|
|
log.debug("mac input error: %s", e)
|
|
|
|
def _pos(self):
|
|
return getattr(self.host, "_last_xy", (0, 0))
|
|
|
|
def clip_get():
|
|
r = subprocess.run(["/usr/bin/pbpaste"], capture_output=True, text=True, timeout=5)
|
|
return r.stdout if r.returncode == 0 else None
|
|
|
|
def clip_set(text):
|
|
r = subprocess.run(["/usr/bin/pbcopy"], input=text, text=True, timeout=5)
|
|
return r.returncode == 0
|
|
|
|
def monitor_power(state): # -1 ignored; 2 = display sleep
|
|
if state == 2:
|
|
subprocess.Popen(["/usr/bin/pmset", "displaysleepnow"])
|
|
|
|
def process_list():
|
|
r = subprocess.run(["ps", "-axo", "pid=,rss=,comm="], capture_output=True, text=True, timeout=15)
|
|
items = []
|
|
for line in r.stdout.splitlines():
|
|
parts = line.strip().split(None, 2)
|
|
if len(parts) == 3 and parts[0].isdigit():
|
|
items.append({"pid": int(parts[0]), "mem": int(parts[1]) * 1024,
|
|
"name": os.path.basename(parts[2])})
|
|
items.sort(key=lambda i: -i["mem"])
|
|
return items[:300]
|
|
|
|
def process_kill(pid):
|
|
try:
|
|
os.kill(pid, 9)
|
|
return True, "killed"
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
def print_file(path):
|
|
if not os.path.isfile(path):
|
|
return False, "file not found"
|
|
r = subprocess.run(["/usr/bin/lp", path], capture_output=True, text=True, timeout=10)
|
|
return r.returncode == 0, (r.stdout or r.stderr).strip() or "queued"
|
|
|
|
def sys_command(action, cmd=None):
|
|
try:
|
|
if action == "lock":
|
|
subprocess.run(["osascript", "-e",
|
|
'tell application "System Events" to keystroke "q" using {command down, control down}'],
|
|
timeout=5)
|
|
return True, "locked"
|
|
if action == "logoff":
|
|
return False, "use lock or reboot"
|
|
if action == "reboot":
|
|
subprocess.Popen(["osascript", "-e", 'tell app "System Events" to restart'])
|
|
return True, "rebooting"
|
|
if action == "shutdown":
|
|
subprocess.Popen(["osascript", "-e", 'tell app "System Events" to shut down'])
|
|
return True, "shutting down"
|
|
if action == "run" and cmd:
|
|
subprocess.Popen(cmd, shell=True)
|
|
return True, "started"
|
|
return False, "unknown command"
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
def console_session_id():
|
|
return 1
|
|
|
|
def my_session_id():
|
|
return 1
|
|
|
|
def spawn_in_console_session(cmdline, sid, desktop=None):
|
|
raise OSError("macOS uses LaunchAgent, not a session broker")
|
|
|
|
def is_elevated():
|
|
return os.geteuid() == 0
|
|
|
|
else: # ---- other non-Windows stubs ----
|
|
def attach_desktop(): return "Default"
|
|
def desktop_name(): return "Default"
|
|
def clip_get(): return None
|
|
def clip_set(t): return False
|
|
def release_all_keys(): pass
|
|
def monitor_power(state): pass
|
|
def process_list(): return []
|
|
def process_kill(pid): return False, "unsupported"
|
|
def print_file(path): return False, "unsupported"
|
|
def sys_command(a, c=None): return False, "Windows only"
|
|
def console_session_id(): return 0
|
|
def my_session_id(): return 0
|
|
def spawn_in_console_session(cmdline, sid, desktop=None): raise OSError("Windows only")
|
|
def is_elevated(): return os.geteuid() == 0
|
|
|
|
|
|
# ================================================================ session / host
|
|
|
|
class Session:
|
|
def __init__(self, host, sid):
|
|
self.host = host
|
|
self.sid = sid
|
|
self.active = False
|
|
self.e2e = None
|
|
self.worker = None
|
|
self.audio = None
|
|
self.stop_evt = threading.Event()
|
|
self.mon = 1
|
|
self.fps = 30
|
|
self.bitrate = 6000
|
|
self.scale = 1.0
|
|
self.apply_evt = threading.Event()
|
|
self.uploads = {}
|
|
self.monitors = []
|
|
|
|
|
|
class Host:
|
|
"""The desktop worker. Talks over exactly one of:
|
|
link='server' direct outbound websocket to the NetDesk server (simple mode)
|
|
link='broker' websocket to the local service broker (service mode)
|
|
link='direct' a direct viewer websocket (LAN / port-forward)"""
|
|
|
|
def __init__(self, cfg: Config, args, link="server"):
|
|
self.cfg = cfg
|
|
self.args = args
|
|
self.link = link
|
|
self.loop = None
|
|
self.session: Session | None = None
|
|
self._current_sid = None
|
|
# WebRTC P2P fast path (video+input peer-to-peer; relay = fallback)
|
|
self.rtc_pc = None
|
|
self.rtc_media = None
|
|
self.rtc_ctl = None
|
|
self.rtc_active = False
|
|
self.ws_server = None # server link socket (simple mode)
|
|
self.ws_broker = None # broker socket (service mode)
|
|
self.ws_direct = None # direct viewer socket
|
|
self._pending = None
|
|
self._mouse_b = 0
|
|
self.clip_last_sent = None
|
|
self.clip_thread = None
|
|
self.pump = None
|
|
self.privacy = False
|
|
self.reg_error = ""
|
|
self.account = ""
|
|
self._tunnels = {} # tid -> (reader, writer)
|
|
self.audio_available = False
|
|
try:
|
|
import pyaudiowpatch # noqa: F401
|
|
self.audio_available = (not getattr(args, "no_audio", False)
|
|
and cfg.data["perms"].get("audio", True))
|
|
except Exception:
|
|
log.info("PyAudioWPatch not installed - audio disabled (pip install PyAudioWPatch)")
|
|
if IS_WIN:
|
|
self.pump = InputPump(self)
|
|
self.pump.start()
|
|
|
|
# ---------------- outgoing (thread-safe, non-blocking) ----------------
|
|
def _post(self, coro):
|
|
def _drain(fut):
|
|
try:
|
|
fut.exception()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
asyncio.run_coroutine_threadsafe(coro, self.loop).add_done_callback(_drain)
|
|
except RuntimeError:
|
|
pass
|
|
|
|
def _sock(self):
|
|
if self.link == "direct":
|
|
return self.ws_direct
|
|
if self.link == "broker":
|
|
return self.ws_broker
|
|
return self.ws_server
|
|
|
|
async def _send_str(self, ws, data):
|
|
if ws is not None and not ws.closed:
|
|
await ws.send_str(data)
|
|
|
|
# ---------------- WebRTC P2P (video + input take the direct UDP path) ----
|
|
async def rtc_answer(self, m):
|
|
"""Viewer sent an SDP offer over the relay: answer it. On success the
|
|
media/ctl DataChannels open and video+input flow peer-to-peer."""
|
|
try:
|
|
await self.rtc_close()
|
|
# ICE servers come with the offer (server-minted, incl. the nearest
|
|
# TURN relay) so the host never needs TURN credentials itself.
|
|
ice = []
|
|
for srv in (m.get("ice") or []):
|
|
try:
|
|
urls = srv.get("urls")
|
|
urls = urls if isinstance(urls, list) else [urls]
|
|
# aiortc uses ONE turn url: prefer UDP (lowest latency)
|
|
urls.sort(key=lambda u: (0 if "transport=udp" in u else 1))
|
|
ice.append(RTCIceServer(urls=urls, username=srv.get("username"),
|
|
credential=srv.get("credential")))
|
|
except Exception:
|
|
pass
|
|
if not ice:
|
|
ice = [RTCIceServer(urls=["stun:stun.cloudflare.com:3478"])]
|
|
pc = RTCPeerConnection(RTCConfiguration(iceServers=ice))
|
|
self.rtc_pc = pc
|
|
self._rtc_turn = any("turn" in str(s.urls) for s in ice)
|
|
|
|
@pc.on("datachannel")
|
|
def on_dc(ch):
|
|
if ch.label == "media":
|
|
self.rtc_media = ch
|
|
elif ch.label in ("ctl", "move"):
|
|
# "ctl" = ordered (clicks/keys); "move" = unordered+unreliable
|
|
# mouse motion so a lost packet never stalls the pointer
|
|
if ch.label == "ctl":
|
|
self.rtc_ctl = ch
|
|
else:
|
|
# tell the viewer we understand the motion channel; older
|
|
# engines stay silent and the viewer keeps motion on "ctl"
|
|
try:
|
|
ch.send(json.dumps({"t": "move.ok"}))
|
|
except Exception:
|
|
pass
|
|
|
|
@ch.on("message")
|
|
def on_ctl(msg):
|
|
# input JSON arrives here instead of the relay socket
|
|
try:
|
|
if isinstance(msg, str):
|
|
self.on_message(json.loads(msg))
|
|
except Exception:
|
|
pass
|
|
|
|
@ch.on("close")
|
|
def on_close():
|
|
if self.rtc_active:
|
|
log.info("rtc: channel closed - back to relay")
|
|
self.rtc_active = False
|
|
self._rtc_check_open()
|
|
|
|
@ch.on("open")
|
|
def on_open():
|
|
self._rtc_check_open()
|
|
|
|
@pc.on("connectionstatechange")
|
|
async def on_state():
|
|
if pc.connectionState in ("failed", "closed", "disconnected"):
|
|
if self.rtc_active:
|
|
log.info("rtc: %s - back to relay", pc.connectionState)
|
|
self.rtc_active = False
|
|
|
|
await pc.setRemoteDescription(
|
|
RTCSessionDescription(sdp=m.get("sdp", ""), type="offer"))
|
|
answer = await pc.createAnswer()
|
|
await pc.setLocalDescription(answer) # gathers ICE (non-trickle)
|
|
sid = self.session.sid if self.session else None
|
|
self.send_json({"t": "rtc.answer", "sid": sid,
|
|
"sdp": pc.localDescription.sdp})
|
|
log.info("rtc: sent answer, waiting for P2P channels")
|
|
except Exception as e:
|
|
log.warning("rtc setup failed: %s", e)
|
|
try:
|
|
self.send_json({"t": "rtc.unavailable",
|
|
"sid": self.session.sid if self.session else None,
|
|
"msg": str(e)[:120]})
|
|
except Exception:
|
|
pass
|
|
|
|
def _rtc_check_open(self):
|
|
mch, cch = self.rtc_media, self.rtc_ctl
|
|
if (mch is not None and cch is not None
|
|
and mch.readyState == "open" and cch.readyState == "open"
|
|
and not self.rtc_active):
|
|
self.rtc_active = True
|
|
route = "unknown"
|
|
try: # which tier won: direct P2P or the TURN relay?
|
|
pair = self.rtc_pc.sctp.transport.transport._connection._nominated
|
|
for p in pair.values():
|
|
route = ("TURN relay" if p.local_candidate.type == "relay"
|
|
or p.remote_candidate.type == "relay" else "DIRECT P2P")
|
|
except Exception:
|
|
pass
|
|
log.info("rtc: ACTIVE via %s - video and input bypass the WSS relay", route)
|
|
|
|
async def rtc_close(self):
|
|
self.rtc_active = False
|
|
pc, self.rtc_pc = self.rtc_pc, None
|
|
self.rtc_media = self.rtc_ctl = None
|
|
if pc is not None:
|
|
try:
|
|
await pc.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def send_json(self, m: dict):
|
|
if self._current_sid and "sid" not in m:
|
|
m["sid"] = self._current_sid
|
|
self._post(self._send_str(self._sock(), json.dumps(m)))
|
|
|
|
def send_binary(self, payload: bytes):
|
|
sess = self.session
|
|
if not sess:
|
|
return
|
|
self._post(self._send_bin(sess, payload))
|
|
|
|
def send_raw(self, payload: bytes):
|
|
"""Send without session gating/E2E (used by TCP tunnels: the payload
|
|
is TLS-protected to the server, which relays it to the viewer)."""
|
|
ws = self._sock()
|
|
async def _go():
|
|
if ws is not None and not ws.closed:
|
|
await ws.send_bytes(payload)
|
|
self._post(_go())
|
|
|
|
async def _send_bin(self, sess: Session, payload: bytes):
|
|
if sess.e2e:
|
|
nonce = secrets.token_bytes(12)
|
|
payload = nonce + sess.e2e.encrypt(nonce, payload, None)
|
|
# P2P fast path: same encrypted bytes, direct UDP instead of the relay.
|
|
# If the channel is congested we DROP the frame (fresh beats late);
|
|
# if it's gone we fall through to the relay socket seamlessly.
|
|
if self.rtc_active:
|
|
ch = self.rtc_media
|
|
try:
|
|
if ch is not None and ch.readyState == "open":
|
|
if ch.bufferedAmount < 262144:
|
|
ch.send(payload)
|
|
return
|
|
except Exception:
|
|
self.rtc_active = False
|
|
ws = self._sock()
|
|
if ws is not None and not ws.closed:
|
|
await ws.send_bytes(payload)
|
|
|
|
# ---------------- auth ----------------
|
|
def make_challenge(self, sid):
|
|
nonce = secrets.token_hex(16)
|
|
self._current_sid = sid
|
|
self._pending = {"sid": sid, "nonce": nonce,
|
|
"priv": ec.generate_private_key(ec.SECP256R1())}
|
|
pub_b = self._pending["priv"].public_key().public_bytes(
|
|
serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint)
|
|
pub_b64 = base64.b64encode(pub_b).decode()
|
|
self._pending["pub_b64"] = pub_b64 # for channel-bound proof verification
|
|
return {
|
|
"t": "session.challenge", "sid": sid,
|
|
"salt": base64.b64encode(bytes.fromhex(self.cfg.data["salt"])).decode(),
|
|
"iters": self.cfg.data["iters"],
|
|
"nonce": nonce, "e2e": True,
|
|
"totp": bool(self.cfg.data.get("totp")),
|
|
"pub": pub_b64,
|
|
}
|
|
|
|
def finish_auth(self, m):
|
|
p = getattr(self, "_pending", None)
|
|
sid = m.get("sid")
|
|
if not p or p["sid"] != sid:
|
|
self.send_json({"t": "session.deny", "reason": "no pending session"})
|
|
return
|
|
totp_secret = self.cfg.data.get("totp")
|
|
if totp_secret and not totp_verify(totp_secret, m.get("totp", "")):
|
|
log.warning("bad 2FA code (sid=%s)", sid)
|
|
time.sleep(0.4)
|
|
self.send_json({"t": "session.deny", "reason": "bad 2FA code"})
|
|
self._pending = None
|
|
self._current_sid = None
|
|
return
|
|
# Channel-bound proof (MITM hardening): if the viewer sent proof_bind,
|
|
# it must verify against OUR ephemeral pubkey — a relay that swapped the
|
|
# pub in the challenge can't satisfy it. Absent → legacy viewer, fall
|
|
# back to the plain proof (no regression).
|
|
pbind = m.get("proof_bind")
|
|
if pbind and not self.cfg.check_proof_bound(p["nonce"], p.get("pub_b64", ""), pbind):
|
|
log.warning("channel-bound proof MISMATCH — possible relay key-swap (sid=%s)", sid)
|
|
time.sleep(0.4)
|
|
self.send_json({"t": "session.deny", "reason": "key verification failed"})
|
|
self._pending = None
|
|
self._current_sid = None
|
|
return
|
|
if not self.cfg.check_proof(p["nonce"], m.get("proof", "")):
|
|
log.warning("bad password attempt (sid=%s)", sid)
|
|
time.sleep(0.4)
|
|
self.send_json({"t": "session.deny", "reason": "wrong password"})
|
|
self._pending = None
|
|
self._current_sid = None
|
|
return
|
|
aesgcm = None
|
|
viewer_pub_b = m.get("pub")
|
|
if viewer_pub_b:
|
|
try:
|
|
peer = ec.EllipticCurvePublicKey.from_encoded_point(
|
|
ec.SECP256R1(), base64.b64decode(viewer_pub_b))
|
|
shared = p["priv"].exchange(ec.ECDH(), peer)
|
|
aesgcm = AESGCM(hashlib.sha256(shared + b"e2e").digest())
|
|
except Exception as e:
|
|
log.warning("E2E negotiation failed (%s); transport-only", e)
|
|
self._pending = None
|
|
|
|
sess = Session(self, sid)
|
|
sess.e2e = aesgcm
|
|
try:
|
|
with mss.mss() as sct:
|
|
sess.monitors = [
|
|
{"i": i, "l": mon["left"], "t": mon["top"], "w": mon["width"], "h": mon["height"]}
|
|
for i, mon in enumerate(sct.monitors[1:], start=1)
|
|
]
|
|
except Exception:
|
|
sess.monitors = []
|
|
mon = next((mm for mm in sess.monitors if mm["i"] == sess.mon), None)
|
|
info = {
|
|
"mon": sess.mon, "monitors": sess.monitors,
|
|
"w": mon["w"] if mon else 0, "h": mon["h"] if mon else 0,
|
|
"fps": sess.fps, "bitrate": sess.bitrate,
|
|
"perms": self.cfg.data["perms"], "lock_on_end": self.cfg.data.get("lock_on_end", False),
|
|
"e2e": aesgcm is not None,
|
|
"audio": self.audio_available,
|
|
"host_name": self.cfg.data["name"],
|
|
"version": VERSION, "os": sys.platform,
|
|
"logon_screen": (desktop_name() or "").lower().startswith("winlogon"),
|
|
}
|
|
self.session = sess
|
|
sess.active = True
|
|
if self.cfg.data.get("privacy_on_session") and self.cfg.data.get("allow_display_power"):
|
|
# silent/privacy session: blank the physical monitor instead of
|
|
# waking the displays - nothing appears on the host's screen
|
|
self.privacy = True
|
|
self._display_power(2)
|
|
elif self.pump:
|
|
self.pump.put(("wake",)) # wake the displays for the viewer
|
|
info["privacy_active"] = self.privacy
|
|
self.send_json({"t": "session.accept", "info": info})
|
|
sess.worker = threading.Thread(target=capture_worker, args=(sess,), daemon=True)
|
|
sess.worker.start()
|
|
if self.audio_available:
|
|
sess.audio = threading.Thread(target=audio_worker, args=(sess,), daemon=True)
|
|
sess.audio.start()
|
|
if self.clip_thread is None or not self.clip_thread.is_alive():
|
|
self.clip_thread = threading.Thread(target=clipboard_worker, args=(self,), daemon=True)
|
|
self.clip_thread.start()
|
|
log.info("session %s accepted (e2e=%s)", sid, aesgcm is not None)
|
|
|
|
def end_session(self, reason):
|
|
sess = self.session
|
|
if not sess:
|
|
return
|
|
sess.active = False
|
|
sess.stop_evt.set()
|
|
for fh in sess.uploads.values():
|
|
try: fh.close()
|
|
except Exception: pass
|
|
sess.uploads.clear()
|
|
release_all_keys()
|
|
if RTC_AVAILABLE and self.rtc_pc is not None:
|
|
try:
|
|
if self.loop is not None:
|
|
asyncio.run_coroutine_threadsafe(self.rtc_close(), self.loop)
|
|
except Exception:
|
|
pass
|
|
if self.privacy:
|
|
monitor_power(-1)
|
|
self.privacy = False
|
|
if getattr(self, "_wp_hidden", False):
|
|
# never leave the desk without its wallpaper
|
|
try:
|
|
wallpaper_set(getattr(self, "_wp_orig", "") or "")
|
|
except Exception:
|
|
pass
|
|
self._wp_hidden = False
|
|
if getattr(self, "_input_blocked", False):
|
|
try:
|
|
ctypes.windll.user32.BlockInput(False)
|
|
except Exception:
|
|
pass
|
|
self._input_blocked = False
|
|
mr = getattr(self, "_mode_restore", None)
|
|
if mr:
|
|
# never leave the remote PC in a resolution the viewer picked
|
|
try:
|
|
set_display_mode(mr["name"], mr["w"], mr["h"])
|
|
except Exception:
|
|
pass
|
|
self._mode_restore = None
|
|
if self.cfg.data.get("lock_on_end"):
|
|
try:
|
|
sys_command("lock")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
self.send_json({"t": "session.end", "sid": sess.sid, "reason": reason})
|
|
except Exception:
|
|
pass
|
|
self.session = None
|
|
self._current_sid = None
|
|
log.info("session ended: %s", reason)
|
|
|
|
def _display_power(self, state):
|
|
"""Hardware display control. Sleeping the display requires the user's
|
|
explicit opt-in (allow_display_power); waking it is always permitted."""
|
|
if state == 2 and not self.cfg.data.get("allow_display_power"):
|
|
log.info("display sleep requested but allow_display_power is off - ignored")
|
|
return
|
|
threading.Thread(target=monitor_power, args=(state,), daemon=True).start()
|
|
if state == 2:
|
|
self._start_privacy_watchdog()
|
|
|
|
def _privacy_watchdog(self):
|
|
"""Safety net: if privacy (monitor off) is ever engaged with no active
|
|
session - crash, race, stale config - force the monitor back on."""
|
|
while self.privacy:
|
|
time.sleep(2)
|
|
sess = self.session
|
|
if self.privacy and (sess is None or not sess.active):
|
|
self.privacy = False
|
|
monitor_power(-1)
|
|
log.info("privacy watchdog: restored monitor (no active session)")
|
|
|
|
def _start_privacy_watchdog(self):
|
|
threading.Thread(target=self._privacy_watchdog, daemon=True).start()
|
|
|
|
# ---------------- messages ----------------
|
|
def on_message(self, m: dict):
|
|
t = m.get("t")
|
|
sess = self.session
|
|
|
|
if t == "wol":
|
|
# another computer on this account asked us to wake a LAN neighbour
|
|
threading.Thread(target=send_wol, args=(str(m.get("mac", "")),),
|
|
daemon=True).start()
|
|
return
|
|
|
|
if t == "session.begin":
|
|
if sess and sess.active:
|
|
self.send_json({"t": "session.deny", "sid": m.get("sid"), "reason": "host busy"})
|
|
return
|
|
self.cfg.reload() # password / 2FA changed in the app apply immediately
|
|
allow = (self.cfg.data.get("allow_ips") or "").strip()
|
|
peer = str(m.get("peer") or "")
|
|
if allow and peer and not ip_allowed(peer, allow):
|
|
log.warning("session denied: viewer %s is not on the IP allowlist", peer)
|
|
self.send_json({"t": "session.deny", "sid": m.get("sid"),
|
|
"reason": "this computer only accepts connections from allowed IPs"})
|
|
return
|
|
self.send_json(self.make_challenge(m.get("sid") or uuid.uuid4().hex))
|
|
return
|
|
if t == "session.auth":
|
|
self.finish_auth(m)
|
|
return
|
|
|
|
# TCP tunnels live outside the session (server-gated, portal-authed)
|
|
if t == "tunnel.open":
|
|
self.tunnel_open(int(m.get("tid", 0)), m.get("thost", "127.0.0.1"),
|
|
int(m.get("port", 0)))
|
|
return
|
|
if t == "tunnel.close":
|
|
self.tunnel_close(int(m.get("tid", 0)))
|
|
return
|
|
|
|
if not sess or m.get("sid") != sess.sid:
|
|
return
|
|
|
|
if t == "session.end":
|
|
self.end_session("viewer left")
|
|
elif t == "ping":
|
|
self.send_json({"t": "pong", "c": m.get("c")})
|
|
elif t == "set.quality":
|
|
sess.fps = int(m.get("fps", 30))
|
|
sess.bitrate = int(m.get("bitrate", 6000))
|
|
sess.scale = float(m.get("scale", 1.0))
|
|
sess.apply_evt.set()
|
|
elif t == "monitor.switch":
|
|
try:
|
|
sess.mon = int(m.get("mon", 1))
|
|
sess.apply_evt.set()
|
|
except Exception:
|
|
pass
|
|
elif t == "keyframe":
|
|
sess.force_idr = True
|
|
elif t == "rtc.offer":
|
|
if RTC_AVAILABLE:
|
|
asyncio.ensure_future(self.rtc_answer(m))
|
|
else:
|
|
self.send_json({"t": "rtc.unavailable", "sid": sess.sid})
|
|
elif t == "i.mouse":
|
|
self.do_mouse(m)
|
|
elif t == "i.key":
|
|
if self.pump:
|
|
self.pump.put(("key", m.get("c", ""), bool(m.get("d"))))
|
|
elif t == "clip.set":
|
|
if self.cfg.data["perms"].get("clipboard", True) and m.get("origin") == "viewer":
|
|
clip_set(m.get("text", "")[:1024 * 1024])
|
|
self.clip_last_sent = m.get("text", "")
|
|
elif t == "chat":
|
|
log.info("chat from viewer: %s", m.get("text", "")[:200])
|
|
elif t == "set.cfg":
|
|
for key in ("lock_on_end", "privacy_on_session", "allow_display_power"):
|
|
if key in m:
|
|
self.cfg.data[key] = bool(m[key])
|
|
self.cfg.save()
|
|
self.send_json({"t": "sys.res", "ok": True, "msg": "config saved"})
|
|
elif t == "file.roots":
|
|
self.file_roots()
|
|
elif t == "file.list":
|
|
self.file_list(m.get("path", ""))
|
|
elif t == "file.get":
|
|
self.file_get(m.get("path"), int(m.get("fid", 0)))
|
|
elif t == "file.put":
|
|
self.file_put_begin(m.get("path"), int(m.get("fid", 0)), int(m.get("size", 0)))
|
|
elif t == "file.mkdir":
|
|
self.file_simple(m.get("path"), lambda p: os.makedirs(p, exist_ok=True))
|
|
elif t == "file.del":
|
|
self.file_simple(m.get("path"),
|
|
(lambda p: shutil.rmtree(p)) if m.get("dir") else (lambda p: os.remove(p)))
|
|
elif t == "file.ren":
|
|
self.file_rename(m.get("src"), m.get("dst"))
|
|
elif t == "sys.cmd":
|
|
self.do_sys(m)
|
|
|
|
def _session_device(self):
|
|
"""GDI display device of the monitor the current session is viewing.
|
|
Uses the geometry the capture worker already cached on the session -
|
|
NEVER calls mss here (a second mss in the event-loop thread crashes
|
|
the process on Windows via a GDI conflict with the capture thread)."""
|
|
sess = self.session
|
|
mi = sess.mon if sess else 1
|
|
geo = None
|
|
for mm in (getattr(sess, "monitors", None) or []):
|
|
if mm["i"] == mi:
|
|
geo = {"left": mm["l"], "top": mm["t"]}
|
|
break
|
|
if geo is None:
|
|
# fall back to the primary display device
|
|
devs = _display_devices()
|
|
if not devs:
|
|
raise RuntimeError("no display device found")
|
|
return devs[0]
|
|
dev = _device_for_monitor(geo)
|
|
if not dev:
|
|
raise RuntimeError("no display device found")
|
|
return dev
|
|
|
|
def do_sys(self, m):
|
|
# A system command must NEVER take down the engine or the session.
|
|
try:
|
|
self._do_sys(m)
|
|
except Exception as e:
|
|
log.warning("sys command %r failed: %s", m.get("a"), e)
|
|
try:
|
|
self.send_json({"t": "sys.res", "ok": False,
|
|
"kind": m.get("a"), "msg": str(e)[:200]})
|
|
except Exception:
|
|
pass
|
|
|
|
def _do_sys(self, m):
|
|
a = m.get("a", "")
|
|
if a == "processes":
|
|
try:
|
|
self.send_json({"t": "sys.res", "ok": True, "kind": "processes",
|
|
"items": process_list()})
|
|
except Exception as e:
|
|
self.send_json({"t": "sys.res", "ok": False, "msg": str(e)})
|
|
elif a == "kill":
|
|
ok, msg = process_kill(int(m.get("pid", 0)))
|
|
self.send_json({"t": "sys.res", "ok": ok, "msg": msg})
|
|
elif a == "privacy":
|
|
if not self.cfg.data.get("allow_display_power"):
|
|
self.send_json({"t": "sys.res", "ok": False, "kind": "privacy",
|
|
"msg": "display control disabled (enable in the NetDesk app)"})
|
|
return
|
|
self.privacy = not self.privacy
|
|
self._display_power(2 if self.privacy else -1)
|
|
self.send_json({"t": "sys.res", "ok": True, "kind": "privacy", "on": self.privacy,
|
|
"msg": "monitor off" if self.privacy else "monitor on"})
|
|
elif a == "run":
|
|
ok, msg = sys_command("run", m.get("cmd"))
|
|
self.send_json({"t": "sys.res", "ok": ok, "msg": msg})
|
|
elif a == "print":
|
|
path = m.get("path", "")
|
|
ok, msg = print_file(path)
|
|
self.send_json({"t": "sys.res", "ok": ok, "msg": msg})
|
|
elif a == "modes":
|
|
try:
|
|
dev = self._session_device()
|
|
self.send_json({"t": "sys.res", "ok": True, "kind": "modes",
|
|
"items": display_modes(dev["name"]),
|
|
"current": [dev["w"], dev["h"]]})
|
|
except Exception as e:
|
|
self.send_json({"t": "sys.res", "ok": False, "kind": "modes",
|
|
"msg": str(e)[:200]})
|
|
elif a == "setmode":
|
|
try:
|
|
dev = self._session_device()
|
|
# remember the ORIGINAL resolution so we restore it at
|
|
# session end - never leave the remote PC in a changed mode
|
|
if not getattr(self, "_mode_restore", None):
|
|
self._mode_restore = {"name": dev["name"], "w": dev["w"], "h": dev["h"]}
|
|
ok, msg = set_display_mode(dev["name"], int(m.get("w", 0)), int(m.get("h", 0)))
|
|
if ok and self.session:
|
|
self.session.force_idr = True
|
|
self.send_json({"t": "sys.res", "ok": ok, "kind": "setmode", "msg": msg})
|
|
except Exception as e:
|
|
self.send_json({"t": "sys.res", "ok": False, "kind": "setmode",
|
|
"msg": str(e)[:200]})
|
|
elif a == "wallpaper":
|
|
try:
|
|
if getattr(self, "_wp_hidden", False):
|
|
wallpaper_set(getattr(self, "_wp_orig", "") or "")
|
|
self._wp_hidden = False
|
|
else:
|
|
self._wp_orig = wallpaper_get()
|
|
wallpaper_set("")
|
|
self._wp_hidden = True
|
|
self.send_json({"t": "sys.res", "ok": True, "kind": "wallpaper",
|
|
"on": self._wp_hidden,
|
|
"msg": "wallpaper hidden" if self._wp_hidden else "wallpaper restored"})
|
|
except Exception as e:
|
|
self.send_json({"t": "sys.res", "ok": False, "kind": "wallpaper",
|
|
"msg": str(e)[:200]})
|
|
elif a == "cad":
|
|
ok, msg = send_cad()
|
|
self.send_json({"t": "sys.res", "ok": ok, "kind": "cad", "msg": msg})
|
|
elif a == "block_input":
|
|
# RustDesk parity: freeze the remote user's own mouse/keyboard while
|
|
# you control. Needs an elevated agent; auto-releases at session end
|
|
# (and Windows releases it if the process ever dies).
|
|
want = not getattr(self, "_input_blocked", False)
|
|
ok = bool(ctypes.windll.user32.BlockInput(want))
|
|
if ok:
|
|
self._input_blocked = want
|
|
msg = "remote user's input BLOCKED" if want else "remote user's input released"
|
|
else:
|
|
msg = ("blocking needs the elevated service agent "
|
|
"(install the boot-time service)")
|
|
self.send_json({"t": "sys.res", "ok": ok, "kind": "block_input",
|
|
"on": getattr(self, "_input_blocked", False), "msg": msg})
|
|
else:
|
|
ok, msg = sys_command(a)
|
|
self.send_json({"t": "sys.res", "ok": ok, "msg": msg})
|
|
|
|
def on_binary(self, data: bytes):
|
|
if data and data[0] == 0xF3 and len(data) >= 5:
|
|
# tunnel payload from the server (not session-E2E encrypted)
|
|
tid = int.from_bytes(data[1:5], "little")
|
|
self.on_tunnel_data(tid, data[5:])
|
|
return
|
|
sess = self.session
|
|
if not sess:
|
|
return
|
|
if sess.e2e:
|
|
try:
|
|
data = sess.e2e.decrypt(data[:12], data[12:], None)
|
|
except Exception:
|
|
return
|
|
if not data or data[0] != 0xF2:
|
|
return
|
|
fid = int.from_bytes(data[1:5], "little")
|
|
off = int.from_bytes(data[5:13], "little")
|
|
chunk = data[13:]
|
|
fh = sess.uploads.get(fid)
|
|
if fh:
|
|
fh.seek(off)
|
|
fh.write(chunk)
|
|
left = getattr(fh, "_nd_left", None)
|
|
if left is not None:
|
|
fh._nd_left = left - len(chunk)
|
|
if fh._nd_left <= 0:
|
|
final = fh._nd_final
|
|
fh.close()
|
|
sess.uploads.pop(fid, None)
|
|
try:
|
|
os.replace(final + ".netdesk.part", final)
|
|
self.send_json({"t": "file.ok", "msg": "upload complete", "refresh": True})
|
|
except Exception as e:
|
|
self.send_json({"t": "file.err", "msg": str(e)})
|
|
|
|
# ---------------- input ----------------
|
|
def do_mouse(self, m):
|
|
if not self.cfg.data["perms"].get("input", True) or not self.pump:
|
|
return
|
|
x, y = int(m.get("x", 0)), int(m.get("y", 0))
|
|
# viewer sends coords in the ENCODED image space; convert to physical
|
|
scale = getattr(self.session, "scale", 1.0) if self.session else 1.0
|
|
if scale and abs(scale - 1.0) > 0.01:
|
|
x = int(x / scale)
|
|
y = int(y / scale)
|
|
# viewer coords are relative to the selected monitor; the injector uses
|
|
# virtual-desktop absolute coords, so add the monitor's origin
|
|
sess = self.session
|
|
if sess and sess.monitors:
|
|
mon = next((mm for mm in sess.monitors if mm["i"] == sess.mon), None)
|
|
if mon:
|
|
x += mon["l"]
|
|
y += mon["t"]
|
|
b = int(m.get("b", 0))
|
|
w, hw = float(m.get("w", 0) or 0), float(m.get("hw", 0) or 0)
|
|
if (x, y) != getattr(self, "_last_xy", None):
|
|
self._last_xy = (x, y)
|
|
self.pump.put(("move", x, y))
|
|
if b != self._mouse_b:
|
|
self.pump.put(("buttons", b))
|
|
self._mouse_b = b
|
|
if w or hw:
|
|
self.pump.put(("wheel", w, hw)) # fractional notches preserved
|
|
|
|
# ---------------- TCP tunnels ----------------
|
|
def tunnel_open(self, tid: int, thost: str, port: int):
|
|
if tid in self._tunnels:
|
|
return
|
|
async def _go():
|
|
try:
|
|
r, w = await asyncio.wait_for(asyncio.open_connection(thost, port), 6)
|
|
except Exception as e:
|
|
self.send_json({"t": "tunnel.err", "tid": tid, "msg": str(e)})
|
|
return
|
|
self._tunnels[tid] = (r, w)
|
|
self.send_json({"t": "tunnel.ok", "tid": tid})
|
|
try:
|
|
while True:
|
|
data = await r.read(65536)
|
|
if not data:
|
|
break
|
|
self.send_raw(b"\xf3" + tid.to_bytes(4, "little") + data)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self.tunnel_close(tid)
|
|
self.send_json({"t": "tunnel.closed", "tid": tid})
|
|
self._post(_go())
|
|
|
|
def tunnel_close(self, tid: int):
|
|
conn = self._tunnels.pop(tid, None)
|
|
if conn:
|
|
try:
|
|
conn[1].close()
|
|
except Exception:
|
|
pass
|
|
|
|
def on_tunnel_data(self, tid: int, payload: bytes):
|
|
conn = self._tunnels.get(tid)
|
|
if conn:
|
|
try:
|
|
conn[1].write(payload)
|
|
except Exception:
|
|
self.tunnel_close(tid)
|
|
|
|
# ---------------- files ----------------
|
|
def file_roots(self):
|
|
items = []
|
|
if IS_WIN:
|
|
for letter in string.ascii_uppercase:
|
|
drv = f"{letter}:\\"
|
|
if os.path.exists(drv):
|
|
items.append({"n": drv, "t": "d", "s": 0, "m": 0})
|
|
else:
|
|
items.append({"n": "/", "t": "d", "s": 0, "m": 0})
|
|
self.send_json({"t": "file.roots.res", "path": "", "items": items})
|
|
|
|
def file_list(self, path):
|
|
try:
|
|
if not path:
|
|
return self.file_roots()
|
|
items = []
|
|
with os.scandir(path) as it:
|
|
for e in it:
|
|
try:
|
|
st = e.stat()
|
|
items.append({"n": e.name, "t": "d" if e.is_dir() else "f",
|
|
"s": 0 if e.is_dir() else st.st_size, "m": int(st.st_mtime)})
|
|
except Exception:
|
|
continue
|
|
items.sort(key=lambda i: (i["t"] != "d", i["n"].lower()))
|
|
self.send_json({"t": "file.list.res", "path": path, "items": items[:2000]})
|
|
except Exception as e:
|
|
self.send_json({"t": "file.err", "msg": f"cannot list {path}: {e}"})
|
|
|
|
def file_get(self, path, fid):
|
|
def worker():
|
|
try:
|
|
sent = 0
|
|
with open(path, "rb") as f:
|
|
while True:
|
|
chunk = f.read(128 * 1024)
|
|
if not chunk:
|
|
break
|
|
self.send_binary(0xF1.to_bytes(1, "little") + fid.to_bytes(4, "little")
|
|
+ sent.to_bytes(8, "little") + chunk)
|
|
sent += len(chunk)
|
|
self.send_json({"t": "file.done", "fid": fid, "path": path, "size": sent})
|
|
except Exception as e:
|
|
self.send_json({"t": "file.err", "msg": str(e)})
|
|
threading.Thread(target=worker, daemon=True).start()
|
|
|
|
@staticmethod
|
|
def _desktop_dir():
|
|
"""The console user's real Desktop (handles OneDrive redirection)."""
|
|
try:
|
|
buf = ctypes.c_wchar_p()
|
|
# SHGetKnownFolderPath(FOLDERID_Desktop)
|
|
import uuid as _u
|
|
fid = (ctypes.c_ubyte * 16).from_buffer_copy(
|
|
_u.UUID("B4BFCC3A-DB2C-424C-B029-7FE99A87C641").bytes_le)
|
|
if ctypes.windll.shell32.SHGetKnownFolderPath(
|
|
ctypes.byref(fid), 0, None, ctypes.byref(buf)) == 0:
|
|
p = buf.value
|
|
ctypes.windll.ole32.CoTaskMemFree(buf)
|
|
if p and os.path.isdir(p):
|
|
return p
|
|
except Exception:
|
|
pass
|
|
p = os.path.join(os.path.expanduser("~"), "Desktop")
|
|
return p if os.path.isdir(p) else os.path.expanduser("~")
|
|
|
|
def file_put_begin(self, path, fid, size):
|
|
sess = self.session
|
|
if not sess or not self.cfg.data["perms"].get("files", True):
|
|
return
|
|
try:
|
|
# TeamViewer-style: a drop on the screen (relative path, no target
|
|
# folder chosen) lands on the remote DESKTOP, never in our CWD.
|
|
if not os.path.isabs(path):
|
|
rel = os.path.normpath(path).replace("/", "\\")
|
|
if rel.startswith("..") or ":" in rel:
|
|
rel = os.path.basename(rel)
|
|
path = os.path.join(self._desktop_dir(), rel)
|
|
if fid in sess.uploads:
|
|
sess.uploads[fid].close()
|
|
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
|
fh = open(path + ".netdesk.part", "wb")
|
|
fh._nd_final = path
|
|
fh._nd_left = size
|
|
sess.uploads[fid] = fh
|
|
except Exception as e:
|
|
self.send_json({"t": "file.err", "msg": str(e)})
|
|
|
|
def file_simple(self, path, fn):
|
|
try:
|
|
fn(path)
|
|
self.send_json({"t": "file.ok", "msg": "ok", "refresh": True})
|
|
except Exception as e:
|
|
self.send_json({"t": "file.err", "msg": str(e)})
|
|
|
|
def file_rename(self, src, dst):
|
|
self.file_simple(src, lambda p: os.replace(p, dst))
|
|
|
|
|
|
# ================================================================ capture + encode
|
|
|
|
# H.264 encoder selection. Prefer a HARDWARE encoder (near-zero CPU, constant
|
|
# low latency regardless of screen complexity) and fall back to software
|
|
# libx264, so NetDesk runs smoothly on ANY PC + GPU:
|
|
# NVIDIA -> h264_nvenc, Intel/QuickSync -> h264_qsv, AMD -> h264_amf,
|
|
# macOS -> videotoolbox, Linux -> vaapi, else -> libx264 (software).
|
|
# The winner is probed ONCE (opening a HW encoder with no capable GPU raises)
|
|
# and cached; if a driver rejects the tuned options we retry the same encoder
|
|
# with bare options before moving on, so quirky hardware still gets accel.
|
|
_VIDEO_ENCODER = None # cached ffmpeg encoder name after the first successful probe
|
|
|
|
|
|
def _encoder_chain(pref="auto"):
|
|
"""Ranked encoder candidates for a user preference. "auto" prefers a
|
|
DISCRETE-GPU hardware encoder (NVENC/AMF - genuinely lower latency and
|
|
offloads the CPU) and otherwise uses software libx264, which on a capable
|
|
CPU is actually faster than Intel iGPU QuickSync for live screen content
|
|
(measured) - so QSV is only chosen when explicitly requested. Any HW encoder
|
|
stays available via "hardware" or a specific name ("qsv"/"nvenc"/"amf")."""
|
|
pref = (pref or "auto").lower()
|
|
if pref == "software":
|
|
return ["libx264"]
|
|
if pref in ("nvenc", "qsv", "amf", "mf"):
|
|
return [f"h264_{pref}", "libx264"]
|
|
if IS_WIN:
|
|
if pref == "hardware":
|
|
return ["h264_nvenc", "h264_qsv", "h264_amf", "libx264"]
|
|
return ["h264_nvenc", "h264_amf", "libx264"] # auto
|
|
if IS_MAC:
|
|
return ["libx264"] if pref == "software" else ["h264_videotoolbox", "libx264"]
|
|
return ["h264_vaapi", "libx264"]
|
|
|
|
|
|
def _encoder_opts(name, fps, bitrate_kbps):
|
|
"""Return (input_pix_fmt, options dict) tuned for low latency.
|
|
HW encoders take nv12; libx264 takes yuv420p. `forced_idr`/`forced-idr`
|
|
makes pict_type=I produce a real IDR on demand (monitor switch / recovery)."""
|
|
maxrate = bitrate_kbps
|
|
bufsize = max(bitrate_kbps * 2 // max(fps, 1), 150) # ~250ms VBV
|
|
if name == "libx264":
|
|
# TCP smoothness: a raw keyframe is a 10-30x burst that stalls the pipe
|
|
# for ~an RTT on long links. Cap instantaneous rate with a tight VBV,
|
|
# stretch keyint and kill scene-cut bursts; force_idr still gives
|
|
# on-demand keyframes for monitor switches / decoder recovery.
|
|
return "yuv420p", {
|
|
"preset": "ultrafast", "tune": "zerolatency",
|
|
"x264opts": (f"vbv-maxrate={maxrate}:vbv-bufsize={bufsize}:"
|
|
f"keyint-min={fps * 4}:scenecut=0"),
|
|
}
|
|
if name == "h264_nvenc":
|
|
return "nv12", {"preset": "p1", "tune": "ull", "rc": "cbr",
|
|
"forced-idr": "1"}
|
|
if name == "h264_qsv":
|
|
return "nv12", {"preset": "veryfast", "low_power": "1",
|
|
"async_depth": "1", "look_ahead": "0",
|
|
"forced_idr": "1"}
|
|
if name == "h264_amf":
|
|
return "nv12", {"usage": "ultralowlatency", "rc": "cbr",
|
|
"quality": "speed", "forced_idr": "1"}
|
|
if name == "h264_videotoolbox":
|
|
return "nv12", {"realtime": "1"}
|
|
if name == "h264_vaapi":
|
|
return "nv12", {}
|
|
return "yuv420p", {}
|
|
|
|
|
|
def make_encoder(w, h, fps, bitrate_kbps, pref="auto"):
|
|
"""Build the best available H.264 encoder for this machine, honoring the
|
|
`video_encoder` preference ("auto"/"software"/"hardware"/name).
|
|
Returns (CodecContext, input_pix_fmt)."""
|
|
global _VIDEO_ENCODER
|
|
order = []
|
|
if _VIDEO_ENCODER:
|
|
order.append(_VIDEO_ENCODER)
|
|
for n in _encoder_chain(pref):
|
|
if n not in order:
|
|
order.append(n)
|
|
err = None
|
|
for name in order:
|
|
fmt, opts = _encoder_opts(name, fps, bitrate_kbps)
|
|
# tuned options first, then bare options (for quirky HW drivers)
|
|
attempts = [opts] if name == "libx264" else [opts, {}]
|
|
for att in attempts:
|
|
try:
|
|
c = av.CodecContext.create(name, "w")
|
|
c.width, c.height, c.pix_fmt = w, h, fmt
|
|
c.time_base = Fraction(1, 1000)
|
|
c.framerate = Fraction(fps, 1)
|
|
c.bit_rate = bitrate_kbps * 1000
|
|
c.gop_size = fps * 4
|
|
c.max_b_frames = 0
|
|
if att:
|
|
c.options = att
|
|
c.open()
|
|
if _VIDEO_ENCODER != name:
|
|
_VIDEO_ENCODER = name
|
|
log.info("video encoder: %s (%s)", name,
|
|
"hardware" if name != "libx264" else "software")
|
|
return c, fmt
|
|
except Exception as e:
|
|
err = e
|
|
if _VIDEO_ENCODER == name: # cached winner stopped working - re-probe
|
|
_VIDEO_ENCODER = None
|
|
raise RuntimeError(f"no usable H.264 encoder: {err}")
|
|
|
|
|
|
def _enum_monitor_positions():
|
|
"""Map each Windows display device name -> (left, top, width, height,
|
|
primary) via EnumDisplayMonitors. Used to pair a dxcam output (which knows
|
|
its own \\\\.\\DISPLAYn name) with the mss monitor at the same position,
|
|
because dxcam's output order does NOT match mss's monitor order."""
|
|
out = {}
|
|
if not IS_WIN:
|
|
return out
|
|
try:
|
|
from ctypes import wintypes # noqa
|
|
|
|
class _RECT(ctypes.Structure):
|
|
_fields_ = [("left", ctypes.c_long), ("top", ctypes.c_long),
|
|
("right", ctypes.c_long), ("bottom", ctypes.c_long)]
|
|
|
|
class _MI(ctypes.Structure):
|
|
_fields_ = [("cbSize", ctypes.c_ulong), ("rcMonitor", _RECT),
|
|
("rcWork", _RECT), ("dwFlags", ctypes.c_ulong),
|
|
("szDevice", ctypes.c_wchar * 32)]
|
|
|
|
user32 = ctypes.windll.user32
|
|
proc = ctypes.WINFUNCTYPE(ctypes.c_int, ctypes.c_ulong, ctypes.c_ulong,
|
|
ctypes.POINTER(_RECT), ctypes.c_double)
|
|
|
|
def _cb(hmon, hdc, lprc, data):
|
|
mi = _MI()
|
|
mi.cbSize = ctypes.sizeof(_MI)
|
|
if user32.GetMonitorInfoW(hmon, ctypes.byref(mi)):
|
|
r = mi.rcMonitor
|
|
out[mi.szDevice] = (r.left, r.top, r.right - r.left,
|
|
r.bottom - r.top, bool(mi.dwFlags & 1))
|
|
return 1
|
|
|
|
user32.EnumDisplayMonitors(0, 0, proc(_cb), 0)
|
|
except Exception:
|
|
pass
|
|
return out
|
|
|
|
|
|
def _dxcam_output_count():
|
|
try:
|
|
return max(1, _dxcam.output_info().count("Output["))
|
|
except Exception:
|
|
return 1
|
|
|
|
|
|
def capture_worker(sess: Session):
|
|
host = sess.host
|
|
w = h = None
|
|
enc = None
|
|
enc_fmt = "yuv420p" # pixel format the chosen encoder wants (yuv420p / nv12)
|
|
last_frame = None
|
|
idx = 0
|
|
force_idr = True
|
|
last_desktop = None
|
|
sct = None
|
|
cam = None # dxcam camera (DXGI fast path), or None -> mss
|
|
cam_mon = None # the sess.mon the current camera was built for
|
|
cam_first = None # the camera's priming frame, sent as the first paint
|
|
last_cur_h = None
|
|
last_cur_pos = None
|
|
cur_mon_seen = None # detect monitor switches to force a keyframe
|
|
enc_pref = host.cfg.data.get("video_encoder", "auto")
|
|
use_dx = _DX_AVAILABLE and str(
|
|
host.cfg.data.get("capture_backend", "auto")).lower() != "mss"
|
|
log.info("capture worker started (dxgi=%s encoder=%s)", use_dx, enc_pref)
|
|
|
|
def release_cam():
|
|
nonlocal cam, cam_mon, cam_first
|
|
try:
|
|
if cam is not None:
|
|
cam.release()
|
|
except Exception:
|
|
pass
|
|
cam = None
|
|
cam_mon = None
|
|
cam_first = None
|
|
|
|
def ensure_cam(mon_index, mon_geo):
|
|
"""(Re)create the dxcam camera for the mss monitor at mon_geo, mapping by
|
|
device position (dxcam output order != mss order). Returns True when the
|
|
DXGI fast path is live, False to use mss. Any failure -> mss."""
|
|
nonlocal cam, cam_mon, cam_first
|
|
if not use_dx or mon_index < 1:
|
|
release_cam()
|
|
return False
|
|
if cam is not None and cam_mon == mon_index:
|
|
return True
|
|
release_cam()
|
|
target = (mon_geo["left"], mon_geo["top"])
|
|
names = _enum_monitor_positions()
|
|
target_name = next((nm for nm, g in names.items()
|
|
if (g[0], g[1]) == target), None)
|
|
for idx_o in range(_dxcam_output_count()):
|
|
try:
|
|
c = _dxcam.create(output_idx=idx_o, output_color="BGRA")
|
|
except Exception:
|
|
continue
|
|
if c is None:
|
|
continue
|
|
dn = getattr(getattr(c, "_output", None), "devicename", "")
|
|
if target_name is not None and dn != target_name:
|
|
try: c.release()
|
|
except Exception: pass
|
|
continue
|
|
probe = c.grab()
|
|
for _ in range(4):
|
|
if probe is not None:
|
|
break
|
|
time.sleep(0.01)
|
|
probe = c.grab()
|
|
if (probe is not None and probe.shape[1] == mon_geo["width"]
|
|
and probe.shape[0] == mon_geo["height"]):
|
|
cam = c
|
|
cam_mon = mon_index
|
|
cam_first = probe
|
|
log.info("dxcam output %d (%s) -> monitor %d (%dx%d)",
|
|
idx_o, dn, mon_index, probe.shape[1], probe.shape[0])
|
|
return True
|
|
try: c.release()
|
|
except Exception: pass
|
|
log.info("dxcam unavailable for monitor %d; using mss", mon_index)
|
|
return False
|
|
|
|
def update_cursor(mon_geo, scale):
|
|
nonlocal last_cur_h, last_cur_pos
|
|
if not IS_WIN:
|
|
return
|
|
try:
|
|
cs = cursor_state()
|
|
if cs is None:
|
|
return
|
|
hcur, gx, gy = cs
|
|
cx = int((gx - mon_geo["left"]) * scale)
|
|
cy = int((gy - mon_geo["top"]) * scale)
|
|
if hcur != last_cur_h:
|
|
last_cur_h = hcur
|
|
if hcur == 0:
|
|
host.send_json({"t": "cursor", "hidden": True})
|
|
else:
|
|
shp = cursor_png(hcur)
|
|
css = _STD_CURSORS.get(hcur)
|
|
if shp or css:
|
|
msg = {"t": "cursor.shape", "id": hcur}
|
|
if shp:
|
|
msg.update(shp)
|
|
if css:
|
|
msg["css"] = css
|
|
host.send_json(msg)
|
|
if hcur != 0 and (cx, cy) != last_cur_pos:
|
|
last_cur_pos = (cx, cy)
|
|
host.send_json({"t": "cursor", "id": hcur, "x": cx, "y": cy})
|
|
except Exception:
|
|
pass
|
|
|
|
def rebuild():
|
|
nonlocal sct, enc, last_frame, w, h, force_idr
|
|
try:
|
|
if sct:
|
|
sct.close()
|
|
except Exception:
|
|
pass
|
|
sct = mss.mss()
|
|
release_cam()
|
|
try:
|
|
if enc is not None:
|
|
enc.close()
|
|
except Exception:
|
|
pass
|
|
enc = None
|
|
last_frame = None
|
|
w = h = None
|
|
force_idr = True
|
|
|
|
try:
|
|
while sess.active and not sess.stop_evt.is_set():
|
|
t0 = time.time()
|
|
# follow the active desktop (logon screen / UAC switches)
|
|
name = attach_desktop()
|
|
if name != last_desktop:
|
|
log.info("input desktop switched to %s", name)
|
|
last_desktop = name
|
|
rebuild()
|
|
host.send_json({"t": "session.info", "info": {
|
|
"mon": sess.mon, "monitors": sess.monitors,
|
|
"w": w or 0, "h": h or 0, "fps": sess.fps, "bitrate": sess.bitrate,
|
|
"logon_screen": bool(name and name.lower().startswith("winlogon")),
|
|
"host_name": host.cfg.data["name"],
|
|
"e2e": sess.e2e is not None, "audio": host.audio_available,
|
|
"perms": host.cfg.data["perms"], "lock_on_end": host.cfg.data.get("lock_on_end", False),
|
|
"privacy_active": host.privacy}})
|
|
if not name:
|
|
time.sleep(0.25)
|
|
continue
|
|
|
|
mon = sct.monitors[min(sess.mon, len(sct.monitors) - 1)]
|
|
# A monitor switch (even between same-size displays) invalidates the
|
|
# encoder's reference frames - force a keyframe so the new screen is
|
|
# not predicted from the old one.
|
|
if cur_mon_seen is not None and cur_mon_seen != sess.mon:
|
|
force_idr = True
|
|
last_frame = None
|
|
cur_mon_seen = sess.mon
|
|
ow, oh = mon["width"] & ~1, mon["height"] & ~1
|
|
scale = max(0.25, min(sess.scale, 1.0))
|
|
tw = max(2, int(ow * scale) & ~1)
|
|
th = max(2, int(oh * scale) & ~1)
|
|
if enc is None or (w, h) != (tw, th) or sess.apply_evt.is_set():
|
|
sess.apply_evt.clear()
|
|
if (w, h) != (tw, th):
|
|
w, h = tw, th
|
|
force_idr = True
|
|
last_frame = None
|
|
if enc is not None:
|
|
try: enc.close()
|
|
except Exception: pass
|
|
enc, enc_fmt = make_encoder(w, h, sess.fps, sess.bitrate, enc_pref)
|
|
host.send_json({"t": "session.info", "info": {
|
|
"mon": sess.mon, "monitors": sess.monitors, "w": w, "h": h,
|
|
"fps": sess.fps, "bitrate": sess.bitrate,
|
|
"logon_screen": str(last_desktop or "").lower().startswith("winlogon"),
|
|
"host_name": host.cfg.data["name"], "e2e": sess.e2e is not None,
|
|
"audio": host.audio_available, "perms": host.cfg.data["perms"], "lock_on_end": host.cfg.data.get("lock_on_end", False),
|
|
"privacy_active": host.privacy}})
|
|
need_kf = bool(force_idr or getattr(sess, "force_idr", False))
|
|
dx_active = ensure_cam(sess.mon, mon)
|
|
if dx_active:
|
|
# DXGI Desktop Duplication fast path (~4ms vs ~45ms for mss).
|
|
if cam_first is not None:
|
|
full, cam_first = cam_first, None # priming frame
|
|
else:
|
|
try:
|
|
full = cam.grab() # None == unchanged
|
|
except Exception as e:
|
|
log.debug("dxcam grab failed (%s); using mss", e)
|
|
release_cam()
|
|
continue
|
|
if full is None:
|
|
# Screen unchanged: keep the cursor live; only re-encode the
|
|
# last frame when a keyframe is actually due.
|
|
update_cursor(mon, scale)
|
|
if not need_kf or last_frame is None:
|
|
time.sleep(max(0.0, 1.0 / sess.fps - (time.time() - t0)))
|
|
continue
|
|
arr = last_frame
|
|
else:
|
|
arr = full[:oh, :ow, :]
|
|
last_frame = arr
|
|
update_cursor(mon, scale)
|
|
else:
|
|
# Fallback: mss GDI capture + full-frame change detection.
|
|
try:
|
|
img = sct.grab(mon)
|
|
except Exception as e:
|
|
log.debug("grab failed (%s)", e)
|
|
time.sleep(0.3)
|
|
rebuild()
|
|
continue
|
|
arr = np.frombuffer(img.raw, dtype=np.uint8).reshape(
|
|
img.size.height, img.size.width, 4)
|
|
arr = arr[:oh, :ow, :]
|
|
update_cursor(mon, scale)
|
|
unchanged = (last_frame is not None and arr.shape == last_frame.shape
|
|
and np.array_equal(arr, last_frame) and not need_kf)
|
|
if unchanged:
|
|
time.sleep(max(0.0, 1.0 / sess.fps - (time.time() - t0)))
|
|
continue
|
|
last_frame = arr
|
|
frame = av.VideoFrame.from_ndarray(arr, format="bgra")
|
|
frame = frame.reformat(width=w, height=h, format=enc_fmt)
|
|
frame.pts = idx * 1000 // max(1, sess.fps)
|
|
frame.time_base = Fraction(1, 1000)
|
|
idx += 1
|
|
if force_idr or getattr(sess, "force_idr", False):
|
|
try:
|
|
frame.pict_type = av.video.frame.PictureType.I
|
|
except Exception:
|
|
pass
|
|
sess.force_idr = False
|
|
force_idr = False
|
|
for pkt in enc.encode(frame):
|
|
host.send_binary(b"\x01" + bytes([1 if pkt.is_keyframe else 0]) + bytes(pkt))
|
|
delay = 1.0 / sess.fps - (time.time() - t0)
|
|
if delay > 0:
|
|
time.sleep(delay)
|
|
except Exception as e:
|
|
log.exception("capture worker crashed: %s", e)
|
|
finally:
|
|
release_cam()
|
|
try:
|
|
if enc is not None:
|
|
enc.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
if sct is not None:
|
|
sct.close()
|
|
except Exception:
|
|
pass
|
|
log.info("capture worker stopped")
|
|
|
|
|
|
# ================================================================ audio / clipboard workers
|
|
|
|
def audio_worker(sess: Session):
|
|
host = sess.host
|
|
try:
|
|
import pyaudiowpatch as pyaudio
|
|
import numpy as _np
|
|
except Exception:
|
|
return
|
|
ctx = None
|
|
stream = None
|
|
p = None
|
|
try:
|
|
p = pyaudio.PyAudio()
|
|
try:
|
|
dev = p.get_default_wasapi_loopback()
|
|
except Exception:
|
|
return
|
|
stream = p.open(format=pyaudio.paInt16, channels=2, rate=48000, input=True,
|
|
frames_per_buffer=960, input_device_index=dev["index"])
|
|
ctx = av.CodecContext.create("libopus", "w")
|
|
ctx.sample_rate = 48000
|
|
ctx.layout = "stereo"
|
|
ctx.format = "s16"
|
|
ctx.bit_rate = 96000
|
|
ctx.open()
|
|
log.info("audio capture running")
|
|
while sess.active and not sess.stop_evt.is_set():
|
|
data = stream.read(960, exception_on_overflow=False)
|
|
samples = _np.frombuffer(data, dtype=_np.int16)
|
|
if not samples.any():
|
|
continue
|
|
frame = av.AudioFrame.from_ndarray(samples.reshape(1, -1), format="s16", layout="stereo")
|
|
frame.sample_rate = 48000
|
|
frame.pts = None
|
|
for pkt in ctx.encode(frame):
|
|
host.send_binary(b"\xa1" + bytes(pkt))
|
|
except Exception as e:
|
|
log.info("audio disabled: %s", e)
|
|
finally:
|
|
try:
|
|
if stream is not None:
|
|
stream.stop_stream()
|
|
stream.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
if p is not None:
|
|
p.terminate()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
if ctx is not None:
|
|
ctx.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def clipboard_worker(host: Host):
|
|
while True:
|
|
time.sleep(0.5)
|
|
sess = host.session
|
|
if not sess or not sess.active:
|
|
continue
|
|
if not host.cfg.data["perms"].get("clipboard", True):
|
|
continue
|
|
txt = clip_get()
|
|
if txt is not None and txt != host.clip_last_sent:
|
|
host.clip_last_sent = txt
|
|
host.send_json({"t": "clip.set", "text": txt[:1024 * 1024], "origin": "host"})
|
|
|
|
|
|
# ================================================================ transports (host side)
|
|
|
|
from contextlib import asynccontextmanager
|
|
|
|
@asynccontextmanager
|
|
async def ws_connect_tolerant(http, url, **kw):
|
|
"""Async-CM ws_connect that falls back to unverified TLS for self-signed
|
|
certificates. Acceptable for personal deployments with a self-signed cert;
|
|
use a real certificate in production."""
|
|
try:
|
|
async with http.ws_connect(url, **kw) as ws:
|
|
yield ws
|
|
except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError) as e:
|
|
log.warning("self-signed certificate on %s - accepting (INSECURE)", url)
|
|
async with http.ws_connect(url, ssl=False, **kw) as ws:
|
|
yield ws
|
|
|
|
|
|
# ---------------- display modes / wallpaper / SAS (TeamViewer-style View & Actions menus)
|
|
class _DEVMODEW(ctypes.Structure):
|
|
_fields_ = [
|
|
("dmDeviceName", ctypes.c_wchar * 32),
|
|
("dmSpecVersion", ctypes.c_ushort), ("dmDriverVersion", ctypes.c_ushort),
|
|
("dmSize", ctypes.c_ushort), ("dmDriverExtra", ctypes.c_ushort),
|
|
("dmFields", ctypes.c_ulong),
|
|
("dmPositionX", ctypes.c_long), ("dmPositionY", ctypes.c_long),
|
|
("dmDisplayOrientation", ctypes.c_ulong), ("dmDisplayFixedOutput", ctypes.c_ulong),
|
|
("dmColor", ctypes.c_short), ("dmDuplex", ctypes.c_short),
|
|
("dmYResolution", ctypes.c_short), ("dmTTOption", ctypes.c_short),
|
|
("dmCollate", ctypes.c_short),
|
|
("dmFormName", ctypes.c_wchar * 32),
|
|
("dmLogPixels", ctypes.c_ushort),
|
|
("dmBitsPerPel", ctypes.c_ulong),
|
|
("dmPelsWidth", ctypes.c_ulong), ("dmPelsHeight", ctypes.c_ulong),
|
|
("dmDisplayFlags", ctypes.c_ulong), ("dmDisplayFrequency", ctypes.c_ulong),
|
|
("dmICMMethod", ctypes.c_ulong), ("dmICMIntent", ctypes.c_ulong),
|
|
("dmMediaType", ctypes.c_ulong), ("dmDitherType", ctypes.c_ulong),
|
|
("dmReserved1", ctypes.c_ulong), ("dmReserved2", ctypes.c_ulong),
|
|
("dmPanningWidth", ctypes.c_ulong), ("dmPanningHeight", ctypes.c_ulong),
|
|
]
|
|
|
|
|
|
class _DISPLAY_DEVICEW(ctypes.Structure):
|
|
_fields_ = [("cb", ctypes.c_ulong), ("DeviceName", ctypes.c_wchar * 32),
|
|
("DeviceString", ctypes.c_wchar * 128), ("StateFlags", ctypes.c_ulong),
|
|
("DeviceID", ctypes.c_wchar * 128), ("DeviceKey", ctypes.c_wchar * 128)]
|
|
|
|
|
|
def _display_devices():
|
|
"""Attached display devices with their current desktop position + size."""
|
|
u = ctypes.windll.user32
|
|
out, i = [], 0
|
|
while True:
|
|
dd = _DISPLAY_DEVICEW()
|
|
dd.cb = ctypes.sizeof(dd)
|
|
if not u.EnumDisplayDevicesW(None, i, ctypes.byref(dd), 0):
|
|
break
|
|
i += 1
|
|
if not (dd.StateFlags & 0x1): # DISPLAY_DEVICE_ATTACHED_TO_DESKTOP
|
|
continue
|
|
dm = _DEVMODEW()
|
|
dm.dmSize = ctypes.sizeof(dm)
|
|
if u.EnumDisplaySettingsW(dd.DeviceName, -1, ctypes.byref(dm)): # current
|
|
out.append({"name": dd.DeviceName, "x": dm.dmPositionX, "y": dm.dmPositionY,
|
|
"w": int(dm.dmPelsWidth), "h": int(dm.dmPelsHeight)})
|
|
return out
|
|
|
|
|
|
def _device_for_monitor(mon_geo):
|
|
"""Match an mss monitor dict (left/top) to its GDI display device."""
|
|
devs = _display_devices()
|
|
for d in devs:
|
|
if d["x"] == mon_geo.get("left") and d["y"] == mon_geo.get("top"):
|
|
return d
|
|
return devs[0] if devs else None
|
|
|
|
|
|
def display_modes(device_name):
|
|
"""All (w, h) modes the device supports, largest first."""
|
|
u = ctypes.windll.user32
|
|
modes, i = set(), 0
|
|
while True:
|
|
dm = _DEVMODEW()
|
|
dm.dmSize = ctypes.sizeof(dm)
|
|
if not u.EnumDisplaySettingsW(device_name, i, ctypes.byref(dm)):
|
|
break
|
|
i += 1
|
|
if dm.dmPelsWidth >= 800 and dm.dmPelsHeight >= 600:
|
|
modes.add((int(dm.dmPelsWidth), int(dm.dmPelsHeight)))
|
|
return sorted(modes, reverse=True)
|
|
|
|
|
|
def set_display_mode(device_name, w, h):
|
|
"""Temporarily change the display resolution (not persisted to registry)."""
|
|
u = ctypes.windll.user32
|
|
dm = _DEVMODEW()
|
|
dm.dmSize = ctypes.sizeof(dm)
|
|
if not u.EnumDisplaySettingsW(device_name, -1, ctypes.byref(dm)):
|
|
return False, "cannot read current mode"
|
|
dm.dmPelsWidth, dm.dmPelsHeight = int(w), int(h)
|
|
dm.dmFields = 0x80000 | 0x100000 # DM_PELSWIDTH | DM_PELSHEIGHT
|
|
r = u.ChangeDisplaySettingsExW(device_name, ctypes.byref(dm), None, 0, None)
|
|
return (r == 0), ("resolution changed" if r == 0
|
|
else f"display rejected the mode (code {r})")
|
|
|
|
|
|
def wallpaper_get():
|
|
buf = ctypes.create_unicode_buffer(520)
|
|
ctypes.windll.user32.SystemParametersInfoW(0x73, 512, buf, 0) # SPI_GETDESKWALLPAPER
|
|
return buf.value
|
|
|
|
|
|
def wallpaper_set(path):
|
|
# SPI_SETDESKWALLPAPER with SPIF_UPDATEINIFILE | SPIF_SENDCHANGE
|
|
return bool(ctypes.windll.user32.SystemParametersInfoW(0x14, 0, path or "", 3))
|
|
|
|
|
|
def send_cad():
|
|
"""Secure Attention Sequence (Ctrl+Alt+Del). Only takes effect when the
|
|
agent runs as SYSTEM (boot-time service) and Windows SAS policy allows it."""
|
|
try:
|
|
sas = ctypes.WinDLL("sas.dll")
|
|
sas.SendSAS(0) # as service
|
|
sas.SendSAS(1) # as user (whichever the policy permits)
|
|
return True, "Ctrl+Alt+Del sent (takes effect when SAS policy allows it)"
|
|
except Exception as e:
|
|
return False, f"Ctrl+Alt+Del unavailable: {str(e)[:120]}"
|
|
|
|
|
|
def _mac_str() -> str:
|
|
"""Primary NIC MAC for Wake-on-LAN (empty if only a random/virtual one)."""
|
|
try:
|
|
n = uuid.getnode()
|
|
if (n >> 40) & 0x02: # locally-administered bit: random, not a real NIC
|
|
return ""
|
|
return ":".join(f"{(n >> s) & 0xFF:02x}" for s in range(40, -1, -8))
|
|
except Exception:
|
|
return ""
|
|
|
|
|
|
def send_wol(mac: str):
|
|
"""Broadcast a Wake-on-LAN magic packet on the local network."""
|
|
import socket as _sk
|
|
clean = "".join(ch for ch in mac if ch in "0123456789abcdefABCDEF")
|
|
if len(clean) != 12:
|
|
return
|
|
payload = b"\xff" * 6 + bytes.fromhex(clean) * 16
|
|
try:
|
|
s = _sk.socket(_sk.AF_INET, _sk.SOCK_DGRAM)
|
|
s.setsockopt(_sk.SOL_SOCKET, _sk.SO_BROADCAST, 1)
|
|
for port in (9, 7):
|
|
s.sendto(payload, ("255.255.255.255", port))
|
|
s.close()
|
|
log.info("WoL magic packet broadcast for %s", mac)
|
|
except Exception as e:
|
|
log.warning("WoL send failed: %s", e)
|
|
|
|
|
|
def ip_allowed(ip: str, allow_csv: str) -> bool:
|
|
"""True if ip matches any entry (IP or CIDR) in the allowlist string."""
|
|
import ipaddress
|
|
try:
|
|
a = ipaddress.ip_address(ip)
|
|
except ValueError:
|
|
return True # unparseable/missing peer info: never lock the owner out
|
|
for part in allow_csv.replace(",", " ").split():
|
|
try:
|
|
if a in ipaddress.ip_network(part, strict=False):
|
|
return True
|
|
except ValueError:
|
|
continue
|
|
return False
|
|
|
|
|
|
def reg_message(cfg: Config, args, takeover=False) -> dict:
|
|
return {"t": "reg.host", "id": cfg.id,
|
|
"key": getattr(args, "reg_key", None) or cfg.data.get("reg_key", ""),
|
|
"name": cfg.data.get("name", "host"), "ver": VERSION,
|
|
# desktop-app account token: binds this computer to "My computers"
|
|
"token": cfg.data.get("account_token", ""),
|
|
# MAC lets another computer on the account wake this one (WoL)
|
|
"mac": _mac_str(),
|
|
# on our first registration only, take over any idle stuck engine
|
|
# holding this ID (never displaces an active session server-side)
|
|
"takeover": bool(takeover)}
|
|
|
|
|
|
def write_status(cfg: Config, state: str, server: str, in_session=False, account=""):
|
|
"""Heartbeat file the desktop app reads (written by engine and broker).
|
|
|
|
A stale duplicate engine (rejected with "already connected elsewhere")
|
|
must never stomp the status of the engine that actually holds the ID -
|
|
that ping-pong made the app flicker Online/Stopped every few seconds."""
|
|
try:
|
|
if state != "online":
|
|
try:
|
|
d = json.loads((cfg.path.parent / "gui_status.json").read_text("utf-8"))
|
|
if (d.get("state") == "online" and d.get("pid") != os.getpid()
|
|
and time.time() - d.get("ts", 0) < 8):
|
|
return # another live engine of ours is online - keep its status
|
|
except Exception:
|
|
pass
|
|
(cfg.path.parent / "gui_status.json").write_text(json.dumps({
|
|
"state": state, "id": cfg.id, "name": cfg.data.get("name", ""),
|
|
"server": server, "in_session": bool(in_session), "account": account,
|
|
"ts": int(time.time()), "pid": os.getpid()}), "utf-8")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def pump_socket_to_host(host: Host, ws, label):
|
|
"""Feed one inbound websocket into the host (used by simple/direct/agent modes)."""
|
|
async for msg in ws:
|
|
if msg.type == WSMsgType.TEXT:
|
|
try:
|
|
m = json.loads(msg.data)
|
|
except Exception:
|
|
continue
|
|
if m.get("t") == "host_status":
|
|
if not m.get("ok"):
|
|
log.error("server rejected registration: %s", m.get("msg"))
|
|
host.reg_error = m.get("msg") or "rejected"
|
|
else:
|
|
host.reg_error = ""
|
|
host.account = m.get("account", "")
|
|
continue
|
|
host.on_message(m)
|
|
elif msg.type == WSMsgType.BINARY:
|
|
host.on_binary(msg.data)
|
|
elif msg.type == WSMsgType.ERROR:
|
|
break
|
|
|
|
|
|
async def run_simple_mode(host: Host, args):
|
|
"""Original mode: this process talks to the server itself + optional direct listener."""
|
|
async def direct_listen():
|
|
if not args.listen:
|
|
return
|
|
|
|
async def handler(req):
|
|
ws = web.WebSocketResponse(heartbeat=20.0, max_msg_size=64 * 1024 * 1024)
|
|
await ws.prepare(req)
|
|
pending_sid = None
|
|
try:
|
|
async for msg in ws:
|
|
if msg.type == WSMsgType.TEXT:
|
|
try:
|
|
m = json.loads(msg.data)
|
|
except Exception:
|
|
continue
|
|
t = m.get("t")
|
|
if t == "session.begin":
|
|
if host.session and host.session.active:
|
|
await ws.send_json({"t": "session.deny", "sid": None, "reason": "host busy"})
|
|
continue
|
|
pending_sid = uuid.uuid4().hex
|
|
host.link = "direct"
|
|
host.ws_direct = ws
|
|
host.on_message({"t": "session.begin", "sid": pending_sid,
|
|
"peer": req.remote or ""})
|
|
else:
|
|
if pending_sid and not m.get("sid"):
|
|
m["sid"] = pending_sid
|
|
host.on_message(m)
|
|
elif msg.type == WSMsgType.BINARY:
|
|
host.on_binary(msg.data)
|
|
finally:
|
|
if host.link == "direct":
|
|
host.end_session("direct viewer left")
|
|
host.link = "server"
|
|
host.ws_direct = None
|
|
return ws
|
|
|
|
app = web.Application()
|
|
app.router.add_get("/ws", handler)
|
|
runner = web.AppRunner(app, access_log=None)
|
|
await runner.setup()
|
|
await web.TCPSite(runner, "0.0.0.0", args.listen).start()
|
|
log.info("direct listener on port %d", args.listen)
|
|
|
|
if args.listen:
|
|
await direct_listen()
|
|
|
|
url = args.server or host.cfg.data.get("server", "") or DEFAULT_SERVER
|
|
if not url:
|
|
log.info("no --server given; direct mode only")
|
|
return
|
|
|
|
async def gui_status_writer():
|
|
"""Heartbeat for the desktop GUI (gui_status.json next to the config)."""
|
|
while True:
|
|
online = host.ws_server is not None and not host.ws_server.closed
|
|
err = getattr(host, "reg_error", "")
|
|
state = ("error: " + err) if err else ("online" if online else "connecting")
|
|
write_status(host.cfg, state, url, bool(host.session), getattr(host, "account", ""))
|
|
await asyncio.sleep(3)
|
|
|
|
status_task = asyncio.ensure_future(gui_status_writer())
|
|
host._status_task = status_task
|
|
backoff = 1
|
|
while True:
|
|
try:
|
|
timeout = aiohttp.ClientTimeout(total=None, connect=15)
|
|
async with aiohttp.ClientSession(timeout=timeout) as http:
|
|
async with ws_connect_tolerant(http, url, heartbeat=20, max_msg_size=64 * 1024 * 1024) as ws:
|
|
host.ws_server = ws
|
|
host.link = "server"
|
|
backoff = 1
|
|
host.cfg.reload()
|
|
# take over on our first attempt, AND whenever the last try
|
|
# was bounced because a STALE idle engine still holds the ID
|
|
# (server never lets takeover kill an active session, so this
|
|
# only ever reclaims a dead/ghost holder - the real fix for
|
|
# "reconnect says the ID is already connected elsewhere")
|
|
prev_err = (getattr(host, "reg_error", "") or "").lower()
|
|
want_takeover = (not getattr(host, "_did_register", False)
|
|
or "already connected" in prev_err)
|
|
host._did_register = True
|
|
await ws.send_json(reg_message(host.cfg, args, takeover=want_takeover))
|
|
log.info("connected to server %s as ID %s (takeover=%s)",
|
|
url, host.cfg.id, want_takeover)
|
|
await pump_socket_to_host(host, ws, "server")
|
|
except Exception as e:
|
|
log.warning("server link lost (%s); retrying in %ss", e, backoff)
|
|
host.ws_server = None
|
|
if host.session and host.link == "server":
|
|
host.end_session("server link lost")
|
|
await asyncio.sleep(backoff)
|
|
backoff = min(backoff * 2, 5)
|
|
|
|
|
|
# ================================================================ service broker
|
|
|
|
class Broker:
|
|
"""Boot-time SYSTEM service core: holds the server link, spawns a desktop
|
|
agent into the console session, and relays traffic. Runs in session 0."""
|
|
|
|
def __init__(self, cfg: Config, args):
|
|
self.cfg = cfg
|
|
self.args = args
|
|
self.token = secrets.token_hex(16)
|
|
self.port = 0
|
|
self.agent_ws = None
|
|
self.agent_proc = None # (handle|None, pid, session)
|
|
self.server_ws = None
|
|
self.direct_ws = None # direct viewer socket
|
|
self.direct_sid = None
|
|
self.session_origin = "server"
|
|
self.reg_error = ""
|
|
self.account = ""
|
|
|
|
# ---------- spawn / supervise agent ----------
|
|
def agent_cmdline(self):
|
|
if getattr(sys, "frozen", False):
|
|
prog = f'"{sys.executable}"'
|
|
else:
|
|
prog = f'"{sys.executable}" "{os.path.abspath(__file__)}"'
|
|
return (f'{prog} --session-agent --config "{self.cfg.path.parent}" '
|
|
f'--broker-url ws://127.0.0.1:{self.port}/agent --token {self.token}')
|
|
|
|
def spawn_agent(self, session_id):
|
|
cmd = self.agent_cmdline()
|
|
try:
|
|
handle, pid = spawn_in_console_session(cmd, session_id)
|
|
log.info("spawned SYSTEM agent pid=%s session=%s", pid, session_id)
|
|
return (handle, pid, None)
|
|
except Exception as e:
|
|
log.warning("SYSTEM spawn failed (%s); spawning plain child", e)
|
|
proc = subprocess.Popen(cmd, creationflags=0x00000008 | 0x00000200) # DETACHED|NEW_GROUP
|
|
return (None, proc.pid, proc)
|
|
|
|
def kill_agent(self):
|
|
if self.agent_ws is not None:
|
|
self.kill_agent_ws_only()
|
|
info = self.agent_proc
|
|
if info:
|
|
handle, pid, popen = info
|
|
if popen is not None:
|
|
popen.kill()
|
|
elif handle:
|
|
kernel32.TerminateProcess(handle, 1)
|
|
kernel32.CloseHandle(handle)
|
|
else:
|
|
subprocess.run(["taskkill", "/F", "/PID", str(pid)],
|
|
capture_output=True, creationflags=0x08000000)
|
|
self.agent_proc = None
|
|
|
|
async def lifecycle(self):
|
|
while True:
|
|
try:
|
|
con = console_session_id()
|
|
info = self.agent_proc
|
|
alive = False
|
|
if info:
|
|
handle, pid, popen = info
|
|
if popen is not None:
|
|
alive = popen.poll() is None
|
|
elif handle:
|
|
alive = kernel32.WaitForSingleObject(handle, 0) == 0x102 # STILL_ACTIVE
|
|
if alive:
|
|
sid = DWORD(0)
|
|
if kernel32.ProcessIdToSessionId(pid, ctypes.byref(sid)) and sid.value != con:
|
|
alive = False # console session changed (user switch)
|
|
if not alive:
|
|
if info:
|
|
log.info("respawning agent (old pid=%s)", info[1])
|
|
self.kill_agent()
|
|
self.agent_proc = self.spawn_agent(con)
|
|
except Exception as e:
|
|
log.debug("lifecycle: %s", e)
|
|
await asyncio.sleep(2)
|
|
|
|
# ---------- sockets ----------
|
|
async def server_link(self):
|
|
url = self.args.server or self.cfg.data.get("server", "") or DEFAULT_SERVER
|
|
if not url:
|
|
log.warning("broker: no server URL configured (--server)")
|
|
return
|
|
backoff = 1
|
|
while True:
|
|
try:
|
|
timeout = aiohttp.ClientTimeout(total=None, connect=15)
|
|
async with aiohttp.ClientSession(timeout=timeout) as http:
|
|
async with ws_connect_tolerant(http, url, heartbeat=20, max_msg_size=64 * 1024 * 1024) as ws:
|
|
self.server_ws = ws
|
|
backoff = 1
|
|
self.cfg.reload()
|
|
prev_err = (getattr(self, "reg_error", "") or "").lower()
|
|
want_takeover = (not getattr(self, "_did_register", False)
|
|
or "already connected" in prev_err)
|
|
self._did_register = True
|
|
await ws.send_json(reg_message(self.cfg, self.args, takeover=want_takeover))
|
|
log.info("broker connected to %s as ID %s (takeover=%s)",
|
|
url, self.cfg.id, want_takeover)
|
|
async for msg in ws:
|
|
if msg.type == WSMsgType.TEXT:
|
|
m = json.loads(msg.data) if msg.data.startswith("{") else {}
|
|
if m.get("t") == "host_status":
|
|
if not m.get("ok"):
|
|
log.error("registration rejected: %s", m.get("msg"))
|
|
self.reg_error = m.get("msg") or "rejected"
|
|
else:
|
|
self.reg_error = ""
|
|
self.account = m.get("account", "")
|
|
continue
|
|
if m.get("t") == "session.begin" and self.agent_ws is None:
|
|
await ws.send_json({"t": "session.deny", "sid": m.get("sid"),
|
|
"reason": "host starting, try again shortly"})
|
|
continue
|
|
if self.agent_ws and not self.agent_ws.closed:
|
|
await self.agent_ws.send_str(msg.data)
|
|
elif msg.type == WSMsgType.BINARY:
|
|
if self.agent_ws and not self.agent_ws.closed:
|
|
await self.agent_ws.send_bytes(msg.data)
|
|
except Exception as e:
|
|
log.warning("broker server link lost (%s); retry in %ss", e, backoff)
|
|
self.server_ws = None
|
|
await asyncio.sleep(backoff)
|
|
backoff = min(backoff * 2, 5)
|
|
|
|
def route_from_agent(self, data: str):
|
|
"""Agent output goes to whoever started the session (direct client or server)."""
|
|
if self.session_origin == "direct" and self.direct_ws and not self.direct_ws.closed:
|
|
asyncio.ensure_future(self.direct_ws.send_str(data))
|
|
elif self.server_ws and not self.server_ws.closed:
|
|
asyncio.ensure_future(self.server_ws.send_str(data))
|
|
|
|
async def agent_handler(self, req):
|
|
if req.query.get("token") != self.token:
|
|
return web.Response(status=403, text="bad token")
|
|
ws = web.WebSocketResponse(heartbeat=15.0, max_msg_size=64 * 1024 * 1024)
|
|
await ws.prepare(req)
|
|
self.kill_agent_ws_only()
|
|
self.agent_ws = ws
|
|
log.info("desktop agent attached (session %s)", my_session_id())
|
|
try:
|
|
async for msg in ws:
|
|
if msg.type == WSMsgType.TEXT:
|
|
self.route_from_agent(msg.data)
|
|
elif msg.type == WSMsgType.BINARY:
|
|
if self.session_origin == "direct" and self.direct_ws and not self.direct_ws.closed:
|
|
asyncio.ensure_future(self.direct_ws.send_bytes(msg.data))
|
|
elif self.server_ws and not self.server_ws.closed:
|
|
asyncio.ensure_future(self.server_ws.send_bytes(msg.data))
|
|
finally:
|
|
if self.agent_ws is ws:
|
|
self.agent_ws = None
|
|
log.info("desktop agent detached")
|
|
return ws
|
|
|
|
def kill_agent_ws_only(self):
|
|
if self.agent_ws is not None:
|
|
try:
|
|
asyncio.ensure_future(self.agent_ws.close())
|
|
except Exception:
|
|
pass
|
|
self.agent_ws = None
|
|
|
|
async def direct_handler(self, req):
|
|
ws = web.WebSocketResponse(heartbeat=20.0, max_msg_size=64 * 1024 * 1024)
|
|
await ws.prepare(req)
|
|
old = self.direct_ws
|
|
if old is not None and not old.closed:
|
|
try: await old.close()
|
|
except Exception: pass
|
|
self.direct_ws = ws
|
|
try:
|
|
async for msg in ws:
|
|
if msg.type == WSMsgType.TEXT:
|
|
m = json.loads(msg.data) if msg.data.startswith("{") else {}
|
|
if m.get("t") == "session.begin":
|
|
self.direct_sid = uuid.uuid4().hex
|
|
m["sid"] = self.direct_sid
|
|
self.session_origin = "direct"
|
|
if self.agent_ws and not self.agent_ws.closed:
|
|
await self.agent_ws.send_str(json.dumps(m))
|
|
continue
|
|
if self.agent_ws and not self.agent_ws.closed:
|
|
await self.agent_ws.send_str(msg.data)
|
|
elif msg.type == WSMsgType.BINARY:
|
|
if self.agent_ws and not self.agent_ws.closed:
|
|
await self.agent_ws.send_bytes(msg.data)
|
|
finally:
|
|
if self.direct_ws is ws:
|
|
self.direct_ws = None
|
|
self.session_origin = "server"
|
|
return ws
|
|
|
|
async def run(self):
|
|
# agent socket
|
|
app = web.Application()
|
|
app.router.add_get("/agent", self.agent_handler)
|
|
runner = web.AppRunner(app, access_log=None)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, "127.0.0.1", 0)
|
|
await site.start()
|
|
self.port = site._server.sockets[0].getsockname()[1]
|
|
(self.cfg.path.parent / "broker.json").write_text(
|
|
json.dumps({"port": self.port, "token": self.token}), "utf-8")
|
|
log.info("broker socket on 127.0.0.1:%d", self.port)
|
|
|
|
# optional direct listener
|
|
if self.args.listen:
|
|
dapp = web.Application()
|
|
dapp.router.add_get("/ws", self.direct_handler)
|
|
drunner = web.AppRunner(dapp, access_log=None)
|
|
await drunner.setup()
|
|
await web.TCPSite(drunner, "0.0.0.0", self.args.listen).start()
|
|
log.info("direct listener on port %d", self.args.listen)
|
|
|
|
async def status():
|
|
url = self.args.server or self.cfg.data.get("server", "") or DEFAULT_SERVER
|
|
while True:
|
|
up = self.server_ws is not None and not self.server_ws.closed
|
|
err = getattr(self, "reg_error", "")
|
|
state = ("error: " + err) if err else ("online" if up else "connecting")
|
|
write_status(self.cfg, state, url, False, getattr(self, "account", ""))
|
|
await asyncio.sleep(3)
|
|
|
|
await asyncio.gather(self.server_link(), self.lifecycle(), status())
|
|
|
|
|
|
def kill_other_engines():
|
|
"""Windows: make sure only THIS engine instance runs (frozen exe)."""
|
|
if not IS_WIN or not getattr(sys, "frozen", False):
|
|
return
|
|
me = os.getpid()
|
|
subprocess.run(
|
|
["powershell", "-NoProfile", "-Command",
|
|
f"Get-CimInstance Win32_Process -Filter \"Name='{os.path.basename(sys.executable)}'\" | "
|
|
f"Where-Object {{ $_.ProcessId -ne {me} -and $_.CommandLine -like '*--engine*' }} | "
|
|
"ForEach-Object { Stop-Process -Id $_.ProcessId -Force }"],
|
|
capture_output=True, creationflags=CREATE_NO_WINDOW)
|
|
|
|
|
|
async def run_agent_mode(args):
|
|
"""Desktop worker spawned by the broker (runs inside the console session)."""
|
|
cfg = Config(Path(args.config) / "config.json" if args.config else appdata_dir() / "config.json")
|
|
host = Host(cfg, args, link="broker")
|
|
host.loop = asyncio.get_running_loop()
|
|
kill_other_engines()
|
|
backoff = 1
|
|
timeout = aiohttp.ClientTimeout(total=None, connect=5, sock_read=30)
|
|
while True:
|
|
try:
|
|
async with aiohttp.ClientSession(timeout=timeout) as http:
|
|
async with http.ws_connect(f"{args.broker_url}?token={args.token}",
|
|
heartbeat=30, max_msg_size=64 * 1024 * 1024) as ws:
|
|
host.ws_broker = ws
|
|
backoff = 1
|
|
log.info("agent linked to broker")
|
|
await pump_socket_to_host(host, ws, "broker")
|
|
except Exception as e:
|
|
log.warning("agent/broker link error: %s", e)
|
|
host.ws_broker = None
|
|
if host.session:
|
|
host.end_session("broker link lost")
|
|
await asyncio.sleep(backoff)
|
|
backoff = min(backoff * 2, 10)
|
|
# exit if we are no longer in the console session (the broker respawns us)
|
|
if IS_WIN and my_session_id() != console_session_id():
|
|
log.info("session changed; agent exiting")
|
|
return
|
|
|
|
|
|
# ================================================================ install / uninstall
|
|
|
|
SERVICE_TASK = "NetDesk"
|
|
|
|
|
|
def install_service(args):
|
|
if not IS_WIN:
|
|
print("service install is Windows-only")
|
|
return
|
|
if not is_elevated():
|
|
print("Please run --install from an elevated (Administrator) prompt once.")
|
|
return
|
|
cfg_dir = appdata_dir() # the portable folder (or APPDATA) the APP writes to
|
|
cfg = Config(cfg_dir / "config.json")
|
|
# Point the SYSTEM broker at the SAME config file the app writes, so any
|
|
# password / 2FA / settings change made in the app applies to the boot
|
|
# service automatically (cfg.reload() at each connection) - no re-install,
|
|
# no extra elevation. Also mirror it into ProgramData as a fallback the
|
|
# service can always read even if the portable folder becomes unavailable.
|
|
pd = programdata_dir()
|
|
pd.mkdir(parents=True, exist_ok=True)
|
|
(pd / "config.json").write_text(json.dumps(cfg.data, indent=2), "utf-8")
|
|
svc_cfg = cfg_dir if os.access(cfg_dir, os.W_OK) else pd
|
|
if getattr(sys, "frozen", False):
|
|
# packaged exe: run it directly with --broker (no python, no script path)
|
|
tr = f'"{sys.executable}" --broker --config "{svc_cfg}"'
|
|
else:
|
|
py = Path(sys.executable).with_name("pythonw.exe")
|
|
if not py.exists():
|
|
py = Path(sys.executable)
|
|
tr = f'"{py}" "{os.path.abspath(__file__)}" --broker --config "{svc_cfg}"'
|
|
server = args.server or cfg.data.get("server", "")
|
|
if server:
|
|
tr += f' --server "{server}"'
|
|
reg = args.reg_key or cfg.data.get("reg_key", "")
|
|
if reg:
|
|
tr += f' --reg-key "{reg}"'
|
|
if args.listen:
|
|
tr += f" --listen {args.listen}"
|
|
r = subprocess.run(["schtasks", "/Create", "/TN", SERVICE_TASK, "/SC", "ONSTART",
|
|
"/RU", "SYSTEM", "/RL", "HIGHEST", "/TR", tr, "/F"],
|
|
capture_output=True, text=True)
|
|
result_path = programdata_dir() / "install_result.json"
|
|
if r.returncode != 0:
|
|
result_path.write_text(json.dumps(
|
|
{"ok": False, "error": (r.stderr or r.stdout or "schtasks failed").strip()}), "utf-8")
|
|
print("schtasks failed:", r.stderr or r.stdout)
|
|
return
|
|
# handover: stop only the user-run ENGINE so the service owns the connection.
|
|
# Do NOT blanket-kill NetDeskHost.exe - that also kills the GUI (which is
|
|
# polling for success) and this very installer, so the install looked like it
|
|
# "never finished" and the toggle stayed off.
|
|
me = os.getpid()
|
|
subprocess.run(
|
|
["powershell", "-NoProfile", "-Command",
|
|
"Get-CimInstance Win32_Process | Where-Object { "
|
|
"($_.Name -eq 'NetDeskHost.exe' -or $_.CommandLine -like '*netdesk_host.py*') "
|
|
f"-and $_.CommandLine -like '*--engine*' -and $_.ProcessId -ne {me} }} | "
|
|
"ForEach-Object { Stop-Process -Id $_.ProcessId -Force }"],
|
|
capture_output=True, creationflags=0x08000000)
|
|
subprocess.run(["schtasks", "/Run", "/TN", SERVICE_TASK], capture_output=True)
|
|
result_path.write_text(json.dumps({"ok": True, "id": cfg.id}), "utf-8")
|
|
print("=" * 62)
|
|
print(" NetDesk service installed and started (boot-time, SYSTEM).")
|
|
print(f" Host ID : {cfg.id} (service reads config from {svc_cfg})")
|
|
print(" Reachable at the logon screen, across reboots, no auto-logon.")
|
|
print(f" Remove: {Path(__file__).name} --uninstall")
|
|
print("=" * 62)
|
|
|
|
|
|
def uninstall_service():
|
|
if not IS_WIN:
|
|
return
|
|
subprocess.run(["schtasks", "/End", "/TN", SERVICE_TASK], capture_output=True)
|
|
r = subprocess.run(["schtasks", "/Delete", "/TN", SERVICE_TASK, "/F"],
|
|
capture_output=True, text=True)
|
|
if r.returncode != 0:
|
|
print("schtasks delete failed:", r.stderr or r.stdout)
|
|
return
|
|
# stop the broker/agent processes this service started (targeted, by command line)
|
|
subprocess.run(
|
|
["powershell", "-NoProfile", "-Command",
|
|
"Get-CimInstance Win32_Process | Where-Object { $_.CommandLine -like '*--broker*' "
|
|
"-or $_.CommandLine -like '*--session-agent*' } | "
|
|
"ForEach-Object { Stop-Process -Id $_.ProcessId -Force }"],
|
|
capture_output=True, creationflags=0x08000000)
|
|
print("NetDesk service removed.")
|
|
|
|
|
|
# ================================================================ main
|
|
|
|
def main():
|
|
# frozen/CLI behavior: double-click (no args) opens the desktop app,
|
|
# explicit flags keep the classic engine/service CLI.
|
|
gui_wanted = ("--gui" in sys.argv) or (
|
|
len(sys.argv) == 1 and getattr(sys, "frozen", False))
|
|
if "--gui-selftest" in sys.argv:
|
|
import netdesk_gui
|
|
netdesk_gui.selftest()
|
|
return
|
|
if gui_wanted:
|
|
try:
|
|
import netdesk_gui
|
|
except Exception as e:
|
|
print("GUI unavailable (%s); engine CLI only." % e)
|
|
import traceback; traceback.print_exc()
|
|
return
|
|
netdesk_gui.run()
|
|
return
|
|
|
|
ap = argparse.ArgumentParser(description="NetDesk host agent")
|
|
ap.add_argument("--server", help="NetDesk server ws(s):// URL")
|
|
ap.add_argument("--reg-key", help="server registration key")
|
|
ap.add_argument("--id", help="override host ID (9 digits)")
|
|
ap.add_argument("--set-password", dest="newpw", help="set a new permanent password")
|
|
ap.add_argument("--set-totp", dest="newtotp", nargs="?", const="auto",
|
|
help="enable TOTP 2FA (optional secret, base32; 'auto' generates one)")
|
|
ap.add_argument("--name", help="friendly name shown to viewers")
|
|
ap.add_argument("--listen", type=int, metavar="PORT", help="direct-access listener port")
|
|
ap.add_argument("--no-audio", action="store_true")
|
|
ap.add_argument("--config", help="config directory (service mode)")
|
|
ap.add_argument("--broker", action="store_true", help=argparse.SUPPRESS)
|
|
ap.add_argument("--broker-url", dest="broker_url", help=argparse.SUPPRESS)
|
|
ap.add_argument("--session-agent", action="store_true", help=argparse.SUPPRESS)
|
|
ap.add_argument("--token", help=argparse.SUPPRESS)
|
|
ap.add_argument("--engine", action="store_true", help=argparse.SUPPRESS)
|
|
ap.add_argument("--gui", action="store_true", help="open the desktop control panel")
|
|
ap.add_argument("--install", action="store_true", help="install boot-time service (admin once)")
|
|
ap.add_argument("--uninstall", action="store_true", help="remove the service")
|
|
ap.add_argument("--verbose", action="store_true")
|
|
args = ap.parse_args()
|
|
|
|
global _CONFIG_DIR
|
|
if args.config:
|
|
_CONFIG_DIR = Path(args.config)
|
|
|
|
logging.basicConfig(
|
|
level=logging.DEBUG if args.verbose else logging.INFO,
|
|
format="%(asctime)s %(levelname)-7s %(message)s",
|
|
handlers=[logging.StreamHandler(),
|
|
logging.FileHandler(appdata_dir() / "host.log", encoding="utf-8")],
|
|
)
|
|
|
|
if args.install:
|
|
install_service(args)
|
|
return
|
|
if args.uninstall:
|
|
uninstall_service()
|
|
return
|
|
|
|
cfg = Config(appdata_dir() / "config.json")
|
|
if args.id and args.id.isdigit():
|
|
cfg.data["id"] = args.id
|
|
if args.name:
|
|
cfg.data["name"] = args.name
|
|
if args.server:
|
|
cfg.data["server"] = args.server
|
|
cfg.save()
|
|
|
|
if args.newpw:
|
|
cfg.set_password(args.newpw)
|
|
print("[netdesk] permanent password updated.")
|
|
return
|
|
if args.newtotp is not None:
|
|
if args.newtotp == "auto":
|
|
secret = "".join(secrets.choice(string.ascii_uppercase + "234567") for _ in range(32))
|
|
else:
|
|
secret = args.newtotp
|
|
cfg.set_totp(secret)
|
|
print("[netdesk] TOTP 2FA ENABLED.")
|
|
print(f" Secret (base32, add to your authenticator): {secret}")
|
|
print(" Viewer will ask for the 6-digit code at connect.")
|
|
return
|
|
|
|
if args.broker:
|
|
if not IS_WIN:
|
|
print("[netdesk] the session broker is Windows-only; on macOS use a LaunchAgent (see scripts/deploy)")
|
|
return
|
|
broker = Broker(cfg, args)
|
|
print(f"[netdesk] broker starting (ID {cfg.id}, session {my_session_id()})")
|
|
try:
|
|
asyncio.run(broker.run())
|
|
except KeyboardInterrupt:
|
|
pass
|
|
return
|
|
if args.session_agent:
|
|
if not args.broker_url or not args.token:
|
|
print("--session-agent needs --broker-url and --token")
|
|
return
|
|
print(f"[netdesk] desktop agent starting (ID {cfg.id}, session {my_session_id()})")
|
|
try:
|
|
asyncio.run(run_agent_mode(args))
|
|
except KeyboardInterrupt:
|
|
pass
|
|
return
|
|
|
|
if not args.engine:
|
|
print("=" * 60)
|
|
print(f" NetDesk host v{VERSION}")
|
|
print(f" Host ID : {cfg.id}")
|
|
print(f" Name : {cfg.data['name']}")
|
|
print(f" Server : {args.server or cfg.data.get('server') or '(none - direct mode only)'}")
|
|
print(f" Config : {cfg.path}")
|
|
if cfg.data.get("totp"):
|
|
print(" 2FA : TOTP enabled")
|
|
if cfg._new_password and not args.engine:
|
|
print(f" NEW PASSWORD (save it): {cfg._new_password}")
|
|
if not args.engine:
|
|
print("[netdesk] running - press Ctrl+C to stop")
|
|
print("=" * 60)
|
|
|
|
host = Host(cfg, args, link="server")
|
|
|
|
async def amain():
|
|
host.loop = asyncio.get_running_loop()
|
|
await run_simple_mode(host, args)
|
|
|
|
try:
|
|
asyncio.run(amain())
|
|
except KeyboardInterrupt:
|
|
host.end_session("host stopping")
|
|
print("\r[netdesk] stopped")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|