Adom Desktop - RustDesk Bridge

#!/usr/bin/env python3
"""Adom Desktop - RustDesk Bridge (Phase 1: client side).

Drives the CONTROLLING side of RustDesk from cloud Claude via adom-desktop:
launch a connection out to a remote ID/IP, list active sessions, close them,
and ensure the client binary is present (download is NON-BLOCKING).

HTTP contract (adom-desktop bridge SDK):
  GET  /status                  -> {ok, bridge, version, uptimeSec, led, summary, tooltip}
  POST /command {command,args}  -> {success, output|error, _hint, ...}

The runtime spawns this as:  python server.py --port <ephemeral>
Every verb returns a rich `_hint`. Nothing blocks on a heavy install.
"""
import csv
import glob
import io
import json
import os
import re
import subprocess
import sys
import threading
import time
import urllib.request
import ctypes
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

BRIDGE_NAME = "rustdesk"
BRIDGE_VERSION = "0.3.2"
DEFAULT_PORT = 8810
RUSTDESK_DL_URL = "https://github.com/rustdesk/rustdesk/releases/download/1.4.8/rustdesk-1.4.8-x86_64.exe"

_CREATE_FLAGS = 0x00000008 | 0x00000200  # DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP
_NO_WINDOW = 0x08000000  # CREATE_NO_WINDOW for tasklist/taskkill helpers
_START = time.time()

# ---- non-blocking client-download state (so no verb blocks on the ~24 MB fetch) ----
_dl_lock = threading.Lock()
_dl = {"state": "idle", "pct": 0, "path": None, "error": None}

# ---- safeguard: auto-disconnect a session when THIS (client) machine goes idle ----
# A live RustDesk client session keeps the local display awake (like media playback).
# So if the user walks away (no input on this laptop) we drop the session -> the screen
# can sleep/screensaver normally. Active use resets the timer (mouse/keys into the
# session ARE local input). The timer is the OS's own last-input clock (self-refreshing:
# any activity zeroes it), NOT a countdown from connect. The HOST is untouched. 0 disables.
# Default 60 min: long enough not to interrupt passive watching, short enough to kill the
# overnight/walked-away case (burn-in needs hours of static content; 60 min never gets there).
IDLE_DISCONNECT_SECS = int(os.environ.get("RUSTDESK_IDLE_DISCONNECT_SECS", "3600"))


# ---------------------------------------------------------------- helpers
def _candidate_exes():
    la = os.environ.get("LOCALAPPDATA", "")
    pf = os.environ.get("ProgramFiles", r"C:\Program Files")
    up = os.environ.get("USERPROFILE", "")
    out = [
        os.path.join(la, "rustdesk", "rustdesk.exe"),
        os.path.join(pf, "RustDesk", "rustdesk.exe"),
    ]
    if up:
        out += sorted(glob.glob(os.path.join(up, "Downloads", "rustdesk-*.exe")), reverse=True)
    return out


def find_rustdesk():
    for p in _candidate_exes():
        if p and os.path.isfile(p):
            return p
    from shutil import which
    return which("rustdesk")


def _tasklist_rustdesk():
    try:
        out = subprocess.run(
            ["tasklist", "/FI", "IMAGENAME eq rustdesk.exe", "/V", "/FO", "CSV", "/NH"],
            capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
        ).stdout
    except Exception:
        return []
    rows = []
    for r in csv.reader(io.StringIO(out)):
        if len(r) >= 9 and r[0].lower().startswith("rustdesk"):
            try:
                pid = int(r[1])
            except ValueError:
                continue
            rows.append({"pid": pid, "title": r[-1]})
    return rows


_SESSION_RE = re.compile(r"^\s*(\S.*?)\s*-\s*Remote Desktop\s*-\s*RustDesk", re.I)


def _sessions():
    out = []
    for p in _tasklist_rustdesk():
        m = _SESSION_RE.match(p.get("title") or "")
        if m:
            out.append({"pid": p["pid"], "peer": m.group(1).strip(), "title": p["title"]})
    return out


def _norm(s):
    return re.sub(r"\s+", "", s or "")


def _download_client():
    up = os.environ.get("USERPROFILE", "")
    dest = os.path.join(up, "Downloads", "rustdesk-1.4.8-x86_64.exe")

    def hook(blocks, bs, total):
        if total and total > 0:
            with _dl_lock:
                _dl["pct"] = min(99, int(blocks * bs * 100 / total))

    try:
        urllib.request.urlretrieve(RUSTDESK_DL_URL, dest, hook)
        with _dl_lock:
            _dl.update(state="done", pct=100, path=dest, error=None)
    except Exception as e:
        with _dl_lock:
            _dl.update(state="error", error=str(e))


def _user_idle_secs():
    """Seconds since the last user input on THIS (client) machine. 0 on non-Windows/error."""
    try:
        class _LII(ctypes.Structure):
            _fields_ = [("cbSize", ctypes.c_uint), ("dwTime", ctypes.c_uint)]
        lii = _LII()
        lii.cbSize = ctypes.sizeof(lii)
        if ctypes.windll.user32.GetLastInputInfo(ctypes.byref(lii)):
            tick = ctypes.windll.kernel32.GetTickCount()
            return max(0, (int(tick) - int(lii.dwTime)) // 1000)
    except Exception:
        pass
    return 0


def _idle_monitor():
    """Drop active sessions once the local user has been idle past IDLE_DISCONNECT_SECS,
    so the client's screen can sleep. Never touches the host side."""
    while True:
        time.sleep(30)
        limit = IDLE_DISCONNECT_SECS
        if limit <= 0:
            continue
        sess = _sessions()
        if sess and _user_idle_secs() >= limit:
            for s in sess:
                try:
                    subprocess.run(["taskkill", "/PID", str(s["pid"]), "/F"],
                                   capture_output=True, timeout=10, creationflags=_NO_WINDOW)
                except Exception:
                    pass
            sys.stderr.write("rustdesk bridge: auto-disconnected %d idle session(s) after %ds idle\n"
                             % (len(sess), limit))
            sys.stderr.flush()


# ---------------------------------------------------------------- verbs
def v_connect(args):
    to = (args.get("to") or args.get("id") or args.get("ip") or "").strip()
    if not to:
        return {"success": False, "error": "missing 'to' (RustDesk ID or IP)",
                "_hint": "Pass {\"to\":\"<id-or-ip>\",\"password\":\"<pw>\"}. Get the host's ID/password from its RustDesk window or the host-side bridge."}
    exe = find_rustdesk()
    if not exe:
        return {"success": False, "error": "rustdesk client not found",
                "_hint": "The client isn't installed. Call rustdesk_ensure_client (non-blocking download), poll rustdesk_client_status until installed:true, then retry."}
    cmd = [exe, "--connect", to]
    pw = args.get("password")
    if pw:
        cmd += ["--password", str(pw)]
    try:
        p = subprocess.Popen(cmd, creationflags=_CREATE_FLAGS,
                             stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception as e:
        return {"success": False, "error": "launch failed: %s" % e,
                "_hint": "The client exe was found but wouldn't launch. Check rustdesk_client_status; the exe path may be stale."}
    return {"success": True, "output": "connecting to %s" % to,
            "to": to, "pid": p.pid, "withPassword": bool(pw), "exe": exe,
            "_hint": "Session window opens as '%s - Remote Desktop - RustDesk'. NOTE: while open this keeps the CONTROLLING machine's display awake (RustDesk treats it like live playback) - the bridge auto-disconnects after %ds of local idle so the screen can still sleep; close sooner with rustdesk_disconnect. Confirm with rustdesk_client_status; without a password it waits at a prompt on the controlling machine." % (to, IDLE_DISCONNECT_SECS)}


def v_disconnect(args):
    to = (args.get("to") or args.get("id") or "").strip()
    killed = []
    for s in _sessions():
        if not to or _norm(s["peer"]) == _norm(to):
            try:
                subprocess.run(["taskkill", "/PID", str(s["pid"]), "/F"],
                               capture_output=True, timeout=10, creationflags=_NO_WINDOW)
                killed.append(s["peer"])
            except Exception:
                pass
    hint = ("Closed %d session window(s). " % len(killed)) + (
        "Pass {\"to\":\"<id>\"} to close just one; omit it to close all." if not to
        else "Run rustdesk_client_status to confirm none remain.")
    return {"success": True, "output": "disconnected %d session(s)" % len(killed),
            "disconnected": killed, "count": len(killed), "_hint": hint}


def v_client_status(args):
    exe = find_rustdesk()
    sess = _sessions()
    with _dl_lock:
        dl = dict(_dl)
    if not exe and dl["state"] == "downloading":
        hint = "Client still downloading (%d%%). Keep polling; when downloadState is 'done', rustdesk_connect will work." % dl["pct"]
    elif not exe:
        hint = "Client not installed. Call rustdesk_ensure_client (non-blocking) to fetch the portable build, then poll here."
    elif sess:
        hint = "%d active session(s): %s. Close with rustdesk_disconnect." % (len(sess), ", ".join(s["peer"] for s in sess))
    else:
        hint = "Client installed, no active sessions. rustdesk_connect {to,password} to start one."
    return {"success": True,
            "output": "%d active session(s)" % len(sess),
            "installed": bool(exe), "path": exe, "sessions": sess,
            "downloadState": dl["state"], "downloadPct": dl["pct"], "downloadError": dl["error"],
            "bridgeVersion": BRIDGE_VERSION, "_hint": hint}


def v_ensure_client(args):
    exe = find_rustdesk()
    if exe:
        return {"success": True, "output": "client present", "present": True, "path": exe,
                "_hint": "RustDesk client already installed — rustdesk_connect is ready."}
    with _dl_lock:
        if _dl["state"] == "downloading":
            return {"success": True, "output": "download already in progress", "downloading": True,
                    "pct": _dl["pct"],
                    "_hint": "A background download is already running (%d%%). Poll rustdesk_client_status (downloadPct)." % _dl["pct"]}
        _dl.update(state="downloading", pct=0, path=None, error=None)
    threading.Thread(target=_download_client, daemon=True).start()
    return {"success": True, "output": "started background download of the portable client",
            "downloading": True, "pct": 0,
            "_hint": "NON-BLOCKING: the ~24 MB portable client is downloading in the background (returns instantly). Poll rustdesk_client_status — when downloadState becomes 'done'/installed:true, rustdesk_connect works."}


def v_readiness(args):
    """READ-ONLY cold-start probe. Never downloads. {ready, installing, installProgressPct, lastError}."""
    exe = find_rustdesk()
    with _dl_lock:
        dl = dict(_dl)
    ready = bool(exe)
    installing = (not ready) and dl["state"] == "downloading"
    if ready:
        hint = "Ready - rustdesk_connect {to,password} works now."
    elif installing:
        hint = "Client downloading in the background (%d%%). Re-poll rustdesk_readiness; never block on it." % dl["pct"]
    else:
        hint = "Client not present. Call rustdesk_prewarm (non-blocking) to fetch it, then poll here."
    return {"success": True, "ready": ready, "installing": installing,
            "installProgressPct": (dl["pct"] if installing else (100 if ready else 0)),
            "lastError": dl["error"], "_hint": hint}


def v_prewarm(args):
    """Kick off the client download in the background; supports {\"wait\":true} to block until ready."""
    if find_rustdesk():
        return {"success": True, "ready": True, "installing": False, "output": "already warm",
                "_hint": "Client already present - nothing to prewarm; rustdesk_connect is ready."}
    started = False
    with _dl_lock:
        if _dl["state"] != "downloading":
            _dl.update(state="downloading", pct=0, path=None, error=None)
            started = True
    if started:
        threading.Thread(target=_download_client, daemon=True).start()
    if args.get("wait"):
        deadline = time.time() + 50  # stay under AD's ~60s verb cap
        while time.time() < deadline:
            if find_rustdesk():
                return {"success": True, "ready": True, "installing": False, "output": "warmed",
                        "_hint": "Client downloaded - rustdesk_connect is ready."}
            with _dl_lock:
                err = _dl["error"]
            if err:
                return {"success": False, "ready": False, "error": err,
                        "_hint": "Download failed: %s. Retry rustdesk_prewarm." % err}
            time.sleep(1.5)
        with _dl_lock:
            pct = _dl["pct"]
        return {"success": True, "ready": False, "installing": True, "installProgressPct": pct,
                "_hint": "Still downloading (%d%%) after the wait window; poll rustdesk_readiness." % pct}
    with _dl_lock:
        pct = _dl["pct"]
    return {"success": True, "ready": False, "installing": True, "installProgressPct": pct,
            "_hint": "NON-BLOCKING prewarm started; poll rustdesk_readiness until ready:true (pass {\"wait\":true} to block)."}


_CATALOG = [
    {"name": "rustdesk_connect", "summary": "Open a remote-control session to a host by RustDesk ID or IP.",
     "input": {"to": "ID or IP (required)", "password": "optional, auto-authenticates"},
     "output": "{to, pid, withPassword}", "example": {"to": "123456789", "password": "<pw>"}},
    {"name": "rustdesk_disconnect", "summary": "Close a session window (or all sessions).",
     "input": {"to": "optional ID; omit = close all"}, "output": "{disconnected[], count}"},
    {"name": "rustdesk_client_status", "summary": "Installed path + active sessions + download state (may spawn the bridge).",
     "input": {}, "output": "{installed, path, sessions[], downloadState, downloadPct}"},
    {"name": "rustdesk_readiness", "summary": "READ-ONLY cold-start probe; never downloads.",
     "input": {}, "output": "{ready, installing, installProgressPct, lastError}"},
    {"name": "rustdesk_prewarm", "summary": "Kick off the client download in the background; {wait:true} blocks until ready.",
     "input": {"wait": "optional bool"}, "output": "{ready, installing, installProgressPct}"},
    {"name": "rustdesk_ensure_client", "summary": "Ensure the client binary is present (non-blocking download).",
     "input": {}, "output": "{present|downloading, path}"},
    {"name": "rustdesk_describe", "summary": "This machine-readable verb catalog.", "input": {}, "output": "{verbs[]}"},
]


def v_describe(args):
    return {"success": True, "verbs": _CATALOG, "version": BRIDGE_VERSION,
            "_hint": "80%% path: rustdesk_connect / rustdesk_client_status / rustdesk_disconnect. Cold start: rustdesk_readiness (read-only) + rustdesk_prewarm. Connect with {to,password}."}


VERBS = {
    "rustdesk_connect": v_connect,
    "rustdesk_disconnect": v_disconnect,
    "rustdesk_client_status": v_client_status,
    "rustdesk_ensure_client": v_ensure_client,
    "rustdesk_readiness": v_readiness,
    "rustdesk_prewarm": v_prewarm,
    "rustdesk_describe": v_describe,
}


# ---------------------------------------------------------------- status chip
def status_payload():
    exe = find_rustdesk()
    sess = _sessions()
    with _dl_lock:
        dl = dict(_dl)
    if not exe and dl["state"] == "downloading":
        led, summary = "yellow", "client %d%%" % dl["pct"]
    elif not exe:
        led, summary = "yellow", "client missing"
    elif sess:
        led, summary = "green", "%d session%s" % (len(sess), "" if len(sess) == 1 else "s")
    else:
        led, summary = "green", "ready"
    tooltip = "RustDesk client bridge v%s\n%s" % (
        BRIDGE_VERSION,
        ("Sessions: " + ", ".join(s["peer"] for s in sess)) if sess else "No active sessions",
    )
    return {"ok": True, "bridge": BRIDGE_NAME, "version": BRIDGE_VERSION,
            "uptimeSec": int(time.time() - _START),
            "led": led, "summary": summary, "tooltip": tooltip}


# ---------------------------------------------------------------- http
class Handler(BaseHTTPRequestHandler):
    def _send(self, code, obj):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *a):
        pass

    def do_GET(self):
        if self.path.rstrip("/") in ("/status", "/health", ""):
            self._send(200, status_payload())
        else:
            self._send(404, {"ok": False, "error": "not found"})

    def do_POST(self):
        if self.path.rstrip("/") != "/command":
            self._send(404, {"success": False, "error": "not found"})
            return
        try:
            n = int(self.headers.get("Content-Length", 0) or 0)
            body = json.loads(self.rfile.read(n) or b"{}")
        except Exception as e:
            self._send(400, {"success": False, "error": "bad json: %s" % e})
            return
        verb = body.get("command") or body.get("verb") or ""
        args = body.get("args") or {}
        fn = VERBS.get(verb)
        if not fn:
            self._send(404, {"success": False, "error": "unknown verb: %s" % verb,
                             "known": list(VERBS),
                             "_hint": "Valid verbs: " + ", ".join(VERBS)})
            return
        try:
            self._send(200, fn(args))
        except Exception as e:
            self._send(200, {"success": False, "error": "%s: %s" % (type(e).__name__, e),
                             "_hint": "Unexpected bridge error; retry. If it persists, bridge_list to check instanceCount, or reinstall."})


def _parse_port(argv):
    for i, a in enumerate(argv):
        if a == "--port" and i + 1 < len(argv):
            try:
                return int(argv[i + 1])
            except ValueError:
                pass
        if a.startswith("--port="):
            try:
                return int(a.split("=", 1)[1])
            except ValueError:
                pass
    env = os.environ.get("BRIDGE_PORT") or os.environ.get("ADOM_BRIDGE_PORT")
    return int(env) if env else DEFAULT_PORT


def main():
    port = _parse_port(sys.argv[1:])
    srv = ThreadingHTTPServer(("127.0.0.1", port), Handler)
    sys.stderr.write("rustdesk bridge v%s listening on 127.0.0.1:%d (idle-disconnect %ds)\n"
                     % (BRIDGE_VERSION, port, IDLE_DISCONNECT_SECS))
    sys.stderr.flush()
    threading.Thread(target=_idle_monitor, daemon=True).start()
    srv.serve_forever()


if __name__ == "__main__":
    main()