Add persistent host ID handshake for Remote Desktop

Each host now exposes a stable 9-digit numeric ID — short enough to
read aloud, persisted at ~/.je_auto_control/remote_host_id so it stays
the same across restarts. The ID is announced inside AUTH_OK as JSON
so only authenticated viewers see it. Viewers that pass
expected_host_id raise AuthenticationError when the announced ID does
not match, defending against TCP-level impersonation by a different
process listening on the same address.

The ID is *not* a substitute for the auth token — token-based HMAC
gates the actual session; the ID is meant to be shared (token + ID
together identify a host).

Author: JE-Chen

je_auto_control/utils/remote_desktop/__init__.py

@@ -10,6 +10,10 @@
 front-end.
 """
 from je_auto_control.utils.remote_desktop.host import RemoteDesktopHost
+from je_auto_control.utils.remote_desktop.host_id import (
+    HostIdError, format_host_id, generate_host_id, load_or_create_host_id,
+    parse_host_id, validate_host_id,
+)
 from je_auto_control.utils.remote_desktop.input_dispatch import (
     InputDispatchError, dispatch_input,
 )
@@ -25,4 +29,6 @@
     "InputDispatchError", "AuthenticationError", "ProtocolError",
     "MessageType", "encode_frame", "decode_frame_header",
     "dispatch_input", "registry",
+    "HostIdError", "format_host_id", "generate_host_id",
+    "load_or_create_host_id", "parse_host_id", "validate_host_id",
 ]

je_auto_control/utils/remote_desktop/host.py

@@ -10,6 +10,9 @@
 from je_auto_control.utils.remote_desktop.auth import (
     NONCE_BYTES, make_nonce, verify_response,
 )
+from je_auto_control.utils.remote_desktop.host_id import (
+    load_or_create_host_id, validate_host_id,
+)
 from je_auto_control.utils.remote_desktop.input_dispatch import (
     InputDispatchError, dispatch_input,
 )
@@ -102,7 +105,10 @@ def _authenticate(self) -> None:
         if not verify_response(self._host._token, nonce, payload):
             self._send(MessageType.AUTH_FAIL, b"bad token")
             raise AuthenticationError("bad token")
-        self._send(MessageType.AUTH_OK, b"")
+        ok_payload = json.dumps(
+            {"host_id": self._host.host_id}, ensure_ascii=False,
+        ).encode("utf-8")
+        self._send(MessageType.AUTH_OK, ok_payload)
         self._sock.settimeout(None)
 
     def _send(self, message_type: MessageType, payload: bytes) -> None:
@@ -207,13 +213,16 @@ def __init__(self, token: str,
                  max_clients: int = 4,
                  frame_provider: Optional[FrameProvider] = None,
                  input_dispatcher: Optional[InputDispatcher] = None,
+                 host_id: Optional[str] = None,
                  ) -> None:
         if not isinstance(token, str) or not token:
             raise ValueError("token must be a non-empty string")
         if fps <= 0:
             raise ValueError("fps must be positive")
         if not 1 <= int(quality) <= 95:
             raise ValueError("quality must be in [1, 95]")
+        self._host_id = (validate_host_id(host_id) if host_id
+                         else load_or_create_host_id())
         self._token = token
         self._bind = bind
         self._requested_port = int(port)
@@ -236,6 +245,11 @@ def __init__(self, token: str,
 
     # public API ----------------------------------------------------------
 
+    @property
+    def host_id(self) -> str:
+        """The 9-digit numeric ID viewers use to verify this host."""
+        return self._host_id
+
     @property
     def port(self) -> int:
         return self._port

je_auto_control/utils/remote_desktop/host_id.py

@@ -0,0 +1,81 @@
+"""Stable, persistent host identifier exposed during the auth handshake.
+
+Each host has a 9-digit numeric ID — short enough to read aloud, long
+enough to be hard to guess by chance. The ID is generated on first use
+and cached at ``~/.je_auto_control/remote_host_id`` so it stays the same
+across restarts; users hand the ID + token + address to the people they
+want to connect, and viewers can verify ``expected_host_id`` after auth
+to defend against TCP-level impersonation.
+
+The ID is *not* a substitute for the auth token — it is broadcast in
+plain text inside ``AUTH_OK`` and is meant to be shared. Token-based
+HMAC auth gates the actual session.
+"""
+import os
+import re
+import secrets
+from pathlib import Path
+from typing import Optional
+
+_HOST_ID_DIGITS = 9
+_DEFAULT_PATH_RELATIVE = ".je_auto_control/remote_host_id"
+_HOST_ID_PATTERN = re.compile(r"^\d{9}$")
+
+
+class HostIdError(ValueError):
+    """Raised when a host ID is malformed."""
+
+
+def generate_host_id() -> str:
+    """Return a fresh random 9-digit host ID (zero-padded)."""
+    return f"{secrets.randbelow(10 ** _HOST_ID_DIGITS):0{_HOST_ID_DIGITS}d}"
+
+
+def validate_host_id(value: str) -> str:
+    """Return ``value`` unchanged if it is a valid 9-digit host ID."""
+    if not isinstance(value, str) or _HOST_ID_PATTERN.fullmatch(value) is None:
+        raise HostIdError(
+            f"host_id must be {_HOST_ID_DIGITS} numeric digits, got {value!r}"
+        )
+    return value
+
+
+def format_host_id(value: str) -> str:
+    """Render a host ID with grouping for display (e.g. ``123 456 789``)."""
+    digits = validate_host_id(value)
+    return f"{digits[:3]} {digits[3:6]} {digits[6:]}"
+
+
+def parse_host_id(value: str) -> str:
+    """Strip whitespace / separators from user input and validate."""
+    if not isinstance(value, str):
+        raise HostIdError(f"host_id must be a string, got {type(value).__name__}")
+    cleaned = re.sub(r"[\s\-_]", "", value)
+    return validate_host_id(cleaned)
+
+
+def default_host_id_path() -> Path:
+    """Return the on-disk path used to persist the host ID."""
+    home = Path(os.path.expanduser("~"))
+    return home / _DEFAULT_PATH_RELATIVE
+
+
+def load_or_create_host_id(path: Optional[Path] = None) -> str:
+    """Return the persisted host ID, creating one on first call."""
+    target = Path(path) if path is not None else default_host_id_path()
+    if target.exists():
+        try:
+            existing = target.read_text(encoding="utf-8").strip()
+            return validate_host_id(existing)
+        except (OSError, HostIdError):
+            # Corrupt / unreadable — regenerate rather than fail the host.
+            pass
+    new_id = generate_host_id()
+    try:
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(new_id, encoding="utf-8")
+    except OSError:
+        # Persisting is best-effort; an in-memory ID still works for the
+        # current process even if the home directory is read-only.
+        pass
+    return new_id

je_auto_control/utils/remote_desktop/registry.py

@@ -35,13 +35,15 @@ def start_host(self, token: str,
                    fps: float = 10.0,
                    quality: int = 70,
                    region: Optional[Sequence[int]] = None,
-                   max_clients: int = 4) -> Dict[str, Any]:
+                   max_clients: int = 4,
+                   host_id: Optional[str] = None) -> Dict[str, Any]:
         """Stop any existing host, then start a fresh one with the given config."""
         self.stop_host()
         host = RemoteDesktopHost(
             token=token, bind=bind, port=int(port),
             fps=float(fps), quality=int(quality),
             region=region, max_clients=int(max_clients),
+            host_id=host_id,
         )
         host.start()
         self._host = host
@@ -57,28 +59,35 @@ def stop_host(self, timeout: float = 2.0) -> Dict[str, Any]:
     def host_status(self) -> Dict[str, Any]:
         host = self._host
         if host is None:
-            return {"running": False, "port": 0, "connected_clients": 0}
+            return {
+                "running": False, "port": 0, "connected_clients": 0,
+                "host_id": None,
+            }
         return {
             "running": host.is_running,
             "port": host.port,
             "connected_clients": host.connected_clients,
+            "host_id": host.host_id,
         }
 
     def connect_viewer(self, host: str, port: int, token: str,
                        timeout: float = 5.0,
                        on_frame: Optional[FrameCallback] = None,
                        on_error: Optional[ErrorCallback] = None,
+                       expected_host_id: Optional[str] = None,
                        ) -> Dict[str, Any]:
         """Disconnect any existing viewer, then connect a fresh one.
 
         ``on_frame`` and ``on_error`` are wired before the receiver
         thread starts, so no frame can arrive while the GUI is still
-        attaching its callbacks.
+        attaching its callbacks. When ``expected_host_id`` is provided
+        the handshake is rejected if the server reports a different ID.
         """
         self.disconnect_viewer()
         viewer = RemoteDesktopViewer(
             host=host, port=int(port), token=token,
             on_frame=on_frame, on_error=on_error,
+            expected_host_id=expected_host_id,
         )
         viewer.connect(timeout=float(timeout))
         self._viewer = viewer
@@ -94,8 +103,11 @@ def disconnect_viewer(self, timeout: float = 2.0) -> Dict[str, Any]:
     def viewer_status(self) -> Dict[str, Any]:
         viewer = self._viewer
         if viewer is None:
-            return {"connected": False}
-        return {"connected": viewer.connected}
+            return {"connected": False, "host_id": None}
+        return {
+            "connected": viewer.connected,
+            "host_id": viewer.remote_host_id,
+        }
 
     def send_input(self, action: Dict[str, Any]) -> Dict[str, Any]:
         """Forward ``action`` through the connected viewer, raise if offline."""

je_auto_control/utils/remote_desktop/viewer.py

@@ -6,6 +6,7 @@
 
 from je_auto_control.utils.logging.logging_instance import autocontrol_logger
 from je_auto_control.utils.remote_desktop.auth import compute_response
+from je_auto_control.utils.remote_desktop.host_id import validate_host_id
 from je_auto_control.utils.remote_desktop.protocol import (
     AuthenticationError, MessageType, ProtocolError,
     encode_frame, read_message,
@@ -18,6 +19,18 @@
 _DEFAULT_CONNECT_TIMEOUT_S = 5.0
 
 
+def _extract_host_id(payload: bytes) -> Optional[str]:
+    """Pull ``host_id`` out of an AUTH_OK payload (JSON or empty)."""
+    if not payload:
+        return None
+    try:
+        body = json.loads(payload.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return None
+    value = body.get("host_id") if isinstance(body, dict) else None
+    return value if isinstance(value, str) else None
+
+
 class RemoteDesktopViewer:
     """Connect to a :class:`RemoteDesktopHost` and stream frames + input.
 
@@ -29,6 +42,7 @@ class RemoteDesktopViewer:
     def __init__(self, host: str, port: int, token: str,
                  on_frame: Optional[FrameCallback] = None,
                  on_error: Optional[ErrorCallback] = None,
+                 expected_host_id: Optional[str] = None,
                  ) -> None:
         if not isinstance(host, str) or not host:
             raise ValueError("host must be a non-empty string")
@@ -39,6 +53,9 @@ def __init__(self, host: str, port: int, token: str,
         self._token = token
         self._on_frame = on_frame
         self._on_error = on_error
+        self._expected_host_id = (validate_host_id(expected_host_id)
+                                  if expected_host_id else None)
+        self._remote_host_id: Optional[str] = None
         self._sock: Optional[socket.socket] = None
         self._send_lock = threading.Lock()
         self._shutdown = threading.Event()
@@ -49,6 +66,11 @@ def __init__(self, host: str, port: int, token: str,
     def connected(self) -> bool:
         return self._connected and not self._shutdown.is_set()
 
+    @property
+    def remote_host_id(self) -> Optional[str]:
+        """The host ID announced in AUTH_OK; ``None`` until handshake completes."""
+        return self._remote_host_id
+
     def connect(self, timeout: float = _DEFAULT_CONNECT_TIMEOUT_S) -> None:
         """Open the TCP connection and complete the auth handshake.
 
@@ -138,6 +160,8 @@ def _handshake(self, sock: socket.socket) -> None:
         sock.sendall(encode_frame(MessageType.AUTH_RESPONSE, response))
         msg_type, payload = read_message(sock)
         if msg_type is MessageType.AUTH_OK:
+            self._remote_host_id = _extract_host_id(payload)
+            self._verify_host_id(self._remote_host_id)
             return
         if msg_type is MessageType.AUTH_FAIL:
             raise AuthenticationError(
@@ -147,6 +171,16 @@ def _handshake(self, sock: socket.socket) -> None:
             f"unexpected handshake reply {msg_type.name}"
         )
 
+    def _verify_host_id(self, announced: Optional[str]) -> None:
+        """Reject the connection when the server's ID does not match expectation."""
+        if self._expected_host_id is None:
+            return
+        if announced != self._expected_host_id:
+            raise AuthenticationError(
+                f"host_id mismatch: expected {self._expected_host_id}, "
+                f"got {announced!r}"
+            )
+
     def _recv_loop(self) -> None:
         sock = self._sock
         if sock is None: