Add TLS transport for Remote Desktop host and viewer

RemoteDesktopHost and RemoteDesktopViewer now accept an ssl.SSLContext;
when provided, the host wraps each accepted connection server-side and
the viewer wraps the connect socket client-side. Failed handshakes on
the host are logged and the raw socket is closed before the client
handler is registered, so a TLS-only host can be hit by plain TCP
viewers without leaking entries into the connected_clients counter.

Tests use a self-signed loopback certificate generated with
cryptography to cover: full TLS round-trip with both a trusting and an
insecure client context, plain viewer rejected against a TLS host,
TLS-only viewer rejected against a plain host, and confirmation that
the wrapped socket is an SSLSocket after connect.

Author: JE-Chen

je_auto_control/utils/remote_desktop/host.py

@@ -1,6 +1,7 @@
 """TCP host that streams JPEG frames and applies viewer input."""
 import json
 import socket
+import ssl
 import threading
 import time
 from io import BytesIO
@@ -214,6 +215,7 @@ def __init__(self, token: str,
                  frame_provider: Optional[FrameProvider] = None,
                  input_dispatcher: Optional[InputDispatcher] = None,
                  host_id: Optional[str] = None,
+                 ssl_context: Optional[ssl.SSLContext] = None,
                  ) -> None:
         if not isinstance(token, str) or not token:
             raise ValueError("token must be a non-empty string")
@@ -224,6 +226,7 @@ def __init__(self, token: str,
         self._host_id = (validate_host_id(host_id) if host_id
                          else load_or_create_host_id())
         self._token = token
+        self._ssl_context = ssl_context
         self._bind = bind
         self._requested_port = int(port)
         self._period = 1.0 / float(fps)
@@ -332,7 +335,10 @@ def _accept_loop(self) -> None:
                 continue
             except OSError:
                 return
-            handler = _ClientHandler(self, client_sock, address)
+            wrapped = self._maybe_wrap_tls(client_sock, address)
+            if wrapped is None:
+                continue
+            handler = _ClientHandler(self, wrapped, address)
             with self._clients_lock:
                 if len(self._clients) >= self._max_clients:
                     autocontrol_logger.info(
@@ -345,6 +351,29 @@ def _accept_loop(self) -> None:
             handler.start()
             self._reap_dead_clients()
 
+    def _maybe_wrap_tls(self, client_sock: socket.socket,
+                        address) -> Optional[socket.socket]:
+        """Return a TLS-wrapped socket when an ssl_context is configured."""
+        if self._ssl_context is None:
+            return client_sock
+        try:
+            client_sock.settimeout(_AUTH_TIMEOUT_S)
+            wrapped = self._ssl_context.wrap_socket(
+                client_sock, server_side=True,
+            )
+            wrapped.settimeout(None)
+            return wrapped
+        except (ssl.SSLError, OSError) as error:
+            autocontrol_logger.info(
+                "remote_desktop TLS handshake from %s failed: %r",
+                address, error,
+            )
+            try:
+                client_sock.close()
+            except OSError:
+                pass
+            return None
+
     def _capture_loop(self) -> None:
         next_tick = time.monotonic()
         while not self._shutdown.is_set():

je_auto_control/utils/remote_desktop/registry.py

@@ -5,6 +5,7 @@
 references here keeps :mod:`action_executor` thin and avoids circular
 imports between the executor and the host/viewer classes.
 """
+import ssl
 from typing import Any, Callable, Dict, Optional, Sequence
 
 from je_auto_control.utils.remote_desktop.host import RemoteDesktopHost
@@ -36,14 +37,16 @@ def start_host(self, token: str,
                    quality: int = 70,
                    region: Optional[Sequence[int]] = None,
                    max_clients: int = 4,
-                   host_id: Optional[str] = None) -> Dict[str, Any]:
+                   host_id: Optional[str] = None,
+                   ssl_context: Optional[ssl.SSLContext] = 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_id=host_id, ssl_context=ssl_context,
         )
         host.start()
         self._host = host
@@ -75,19 +78,24 @@ def connect_viewer(self, host: str, port: int, token: str,
                        on_frame: Optional[FrameCallback] = None,
                        on_error: Optional[ErrorCallback] = None,
                        expected_host_id: Optional[str] = None,
+                       ssl_context: Optional[ssl.SSLContext] = None,
+                       server_hostname: 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. When ``expected_host_id`` is provided
         the handshake is rejected if the server reports a different ID.
+        Pass an ``ssl_context`` to upgrade the connection to TLS.
         """
         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,
+            ssl_context=ssl_context,
+            server_hostname=server_hostname,
         )
         viewer.connect(timeout=float(timeout))
         self._viewer = viewer

je_auto_control/utils/remote_desktop/viewer.py

@@ -1,6 +1,7 @@
 """TCP viewer that receives JPEG frames and forwards input messages."""
 import json
 import socket
+import ssl
 import threading
 from typing import Any, Callable, Mapping, Optional
 
@@ -43,6 +44,8 @@ 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,
+                 ssl_context: Optional[ssl.SSLContext] = None,
+                 server_hostname: Optional[str] = None,
                  ) -> None:
         if not isinstance(host, str) or not host:
             raise ValueError("host must be a non-empty string")
@@ -56,6 +59,8 @@ def __init__(self, host: str, port: int, token: str,
         self._expected_host_id = (validate_host_id(expected_host_id)
                                   if expected_host_id else None)
         self._remote_host_id: Optional[str] = None
+        self._ssl_context = ssl_context
+        self._server_hostname = server_hostname
         self._sock: Optional[socket.socket] = None
         self._send_lock = threading.Lock()
         self._shutdown = threading.Event()
@@ -72,22 +77,23 @@ def remote_host_id(self) -> Optional[str]:
         return self._remote_host_id
 
     def connect(self, timeout: float = _DEFAULT_CONNECT_TIMEOUT_S) -> None:
-        """Open the TCP connection and complete the auth handshake.
+        """Open the (optionally TLS) connection and complete the auth handshake.
 
         Spawns a receiver thread on success. Raises
         :class:`AuthenticationError` if the handshake fails.
         """
         if self._connected:
             return
-        sock = socket.create_connection(
+        raw_sock = socket.create_connection(
             (self._host, self._port), timeout=timeout,
         )
-        sock.settimeout(_DEFAULT_AUTH_TIMEOUT_S)
+        raw_sock.settimeout(_DEFAULT_AUTH_TIMEOUT_S)
         try:
+            sock = self._maybe_wrap_tls(raw_sock)
             self._handshake(sock)
-        except (AuthenticationError, ProtocolError, OSError):
+        except (AuthenticationError, ProtocolError, OSError, ssl.SSLError):
             try:
-                sock.close()
+                raw_sock.close()
             except OSError:
                 pass
             raise
@@ -100,6 +106,19 @@ def connect(self, timeout: float = _DEFAULT_CONNECT_TIMEOUT_S) -> None:
         )
         self._receiver.start()
 
+    def _maybe_wrap_tls(self, raw_sock: socket.socket) -> socket.socket:
+        """Return a TLS-wrapped socket when an ssl_context was configured."""
+        if self._ssl_context is None:
+            return raw_sock
+        hostname = self._server_hostname or self._host
+        if (self._ssl_context.check_hostname is False
+                and self._ssl_context.verify_mode == ssl.CERT_NONE):
+            # ``wrap_socket`` rejects server_hostname when verification is off.
+            hostname = None
+        return self._ssl_context.wrap_socket(
+            raw_sock, server_hostname=hostname,
+        )
+
     def disconnect(self, timeout: float = 2.0) -> None:
         """Close the connection and join the receiver thread."""
         self._shutdown.set()

test/unit_test/headless/test_remote_desktop_tls.py

@@ -0,0 +1,188 @@
+"""End-to-end TLS tests using a self-signed loopback certificate."""
+import datetime
+import ipaddress
+import socket
+import ssl
+import time
+from pathlib import Path
+from typing import Tuple
+
+import pytest
+
+cryptography = pytest.importorskip("cryptography")
+
+from cryptography import x509  # noqa: E402
+from cryptography.hazmat.primitives import hashes, serialization  # noqa: E402
+from cryptography.hazmat.primitives.asymmetric import rsa  # noqa: E402
+from cryptography.x509.oid import NameOID  # noqa: E402
+
+from je_auto_control.utils.remote_desktop import (
+    RemoteDesktopHost, RemoteDesktopViewer,
+)
+from je_auto_control.utils.remote_desktop.protocol import AuthenticationError
+
+
+def _wait_until(predicate, timeout: float = 2.0,
+                interval: float = 0.02) -> bool:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        time.sleep(interval)
+    return predicate()
+
+
+def _generate_self_signed(tmp_path: Path) -> Tuple[Path, Path]:
+    """Write a self-signed cert + key for ``127.0.0.1`` to ``tmp_path``."""
+    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    name = x509.Name([
+        x509.NameAttribute(NameOID.COMMON_NAME, "remote-desktop-test"),
+    ])
+    now = datetime.datetime.now(datetime.timezone.utc)
+    cert = (
+        x509.CertificateBuilder()
+        .subject_name(name)
+        .issuer_name(name)
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(now - datetime.timedelta(minutes=1))
+        .not_valid_after(now + datetime.timedelta(days=1))
+        .add_extension(
+            x509.SubjectAlternativeName([
+                x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
+                x509.DNSName("localhost"),
+            ]),
+            critical=False,
+        )
+        .sign(private_key=key, algorithm=hashes.SHA256())
+    )
+    cert_path = tmp_path / "cert.pem"
+    key_path = tmp_path / "key.pem"
+    cert_path.write_bytes(
+        cert.public_bytes(serialization.Encoding.PEM)
+    )
+    key_path.write_bytes(
+        key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+    )
+    return cert_path, key_path
+
+
+def _server_context(cert_path: Path, key_path: Path) -> ssl.SSLContext:
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+    ctx.load_cert_chain(certfile=str(cert_path), keyfile=str(key_path))
+    return ctx
+
+
+def _trusting_client_context(ca_path: Path) -> ssl.SSLContext:
+    """Verifying client context that trusts only the supplied test CA cert."""
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    ctx.load_verify_locations(cafile=str(ca_path))
+    ctx.check_hostname = True
+    ctx.verify_mode = ssl.CERT_REQUIRED
+    return ctx
+
+
+def _insecure_client_context() -> ssl.SSLContext:
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    ctx.check_hostname = False
+    ctx.verify_mode = ssl.CERT_NONE
+    return ctx
+
+
+def _start_tls_host(tmp_path: Path) -> Tuple[RemoteDesktopHost, Path, Path]:
+    cert_path, key_path = _generate_self_signed(tmp_path)
+    server_ctx = _server_context(cert_path, key_path)
+    host = RemoteDesktopHost(
+        token="tok", bind="127.0.0.1", port=0, fps=50.0,
+        frame_provider=lambda: b"tls-frame",
+        input_dispatcher=lambda *_a, **_k: None,
+        host_id="111111111", ssl_context=server_ctx,
+    )
+    host.start()
+    return host, cert_path, key_path
+
+
+def test_tls_round_trip_with_trusting_client(tmp_path):
+    host, cert_path, _ = _start_tls_host(tmp_path)
+    try:
+        client_ctx = _trusting_client_context(cert_path)
+        received = []
+        viewer = RemoteDesktopViewer(
+            host="127.0.0.1", port=host.port, token="tok",
+            on_frame=received.append,
+            ssl_context=client_ctx,
+        )
+        viewer.connect(timeout=2.0)
+        assert _wait_until(lambda: len(received) >= 1, timeout=2.0)
+        assert all(frame == b"tls-frame" for frame in received)
+        viewer.disconnect()
+    finally:
+        host.stop(timeout=1.0)
+
+
+def test_tls_round_trip_with_insecure_client(tmp_path):
+    host, _, _ = _start_tls_host(tmp_path)
+    try:
+        viewer = RemoteDesktopViewer(
+            host="127.0.0.1", port=host.port, token="tok",
+            ssl_context=_insecure_client_context(),
+        )
+        viewer.connect(timeout=2.0)
+        assert viewer.connected
+        viewer.disconnect()
+    finally:
+        host.stop(timeout=1.0)
+
+
+def test_plain_viewer_against_tls_host_fails(tmp_path):
+    """A non-TLS viewer cannot finish the handshake against a TLS host."""
+    host, _, _ = _start_tls_host(tmp_path)
+    try:
+        viewer = RemoteDesktopViewer(
+            host="127.0.0.1", port=host.port, token="tok",
+        )
+        with pytest.raises((OSError, AuthenticationError)):
+            viewer.connect(timeout=2.0)
+        # Host should refuse to count an incomplete handshake as connected.
+        assert _wait_until(lambda: host.connected_clients == 0, timeout=2.0)
+    finally:
+        host.stop(timeout=1.0)
+
+
+def test_tls_client_against_plain_host_fails():
+    """A TLS-only viewer cannot speak to a plain TCP host."""
+    host = RemoteDesktopHost(
+        token="tok", bind="127.0.0.1", port=0, fps=50.0,
+        frame_provider=lambda: b"plain",
+        input_dispatcher=lambda *_a, **_k: None,
+        host_id="222222222",
+    )
+    host.start()
+    try:
+        viewer = RemoteDesktopViewer(
+            host="127.0.0.1", port=host.port, token="tok",
+            ssl_context=_insecure_client_context(),
+        )
+        with pytest.raises((OSError, ssl.SSLError, AuthenticationError)):
+            viewer.connect(timeout=2.0)
+    finally:
+        host.stop(timeout=1.0)
+
+
+def test_tls_uses_socket_class_after_wrap(tmp_path):
+    """After connect, the viewer's socket should be an SSLSocket."""
+    host, cert_path, _ = _start_tls_host(tmp_path)
+    try:
+        viewer = RemoteDesktopViewer(
+            host="127.0.0.1", port=host.port, token="tok",
+            ssl_context=_trusting_client_context(cert_path),
+        )
+        viewer.connect(timeout=2.0)
+        assert isinstance(viewer._sock, ssl.SSLSocket)  # noqa: SLF001
+        viewer.disconnect()
+    finally:
+        host.stop(timeout=1.0)