feat: add TLS support

Author: robgee86

src/arduino/app_peripherals/camera/websocket_camera.py

@@ -19,7 +19,7 @@
 from arduino.app_utils import Logger
 
 from .base_camera import BaseCamera
-from .errors import CameraOpenError
+from .errors import CameraConfigError, CameraOpenError
 
 logger = Logger("WebSocketCamera")
 
@@ -41,16 +41,18 @@ class WebSocketCamera(BaseCamera):
 
     Secure communication with the WebSocket server is supported in three security modes:
     - Security disabled (empty secret)
-    - Authenticated (secret + enable_encryption=False) - HMAC-SHA256
-    - Authenticated + Encrypted (secret + enable_encryption=True) - ChaCha20-Poly1305
+    - Authenticated (secret + encrypt=False) - HMAC-SHA256
+    - Authenticated + Encrypted (secret + encrypt=True) - ChaCha20-Poly1305
     """
 
     def __init__(
         self,
         port: int = 8080,
         timeout: int = 3,
+        certs_dir_path: str = "/app/certs",
+        use_tls: bool = False,
         secret: str = "",
-        enable_encryption: bool = False,
+        encrypt: bool = False,
         resolution: tuple[int, int] = (640, 480),
         fps: int = 10,
         adjustments: Callable[[np.ndarray], np.ndarray] | None = None,
@@ -62,26 +64,54 @@ def __init__(
         Args:
             port (int): Port to bind the server to
             timeout (int): Connection timeout in seconds
+            certs_dir_path (str): Path to the directory containing TLS certificates
+            use_tls (bool): Enable TLS for secure connections. If True, 'encrypt' will
+                be ignored. Use this for transport-level security with clients that can
+                accept self-signed certificates or when supplying your own certificates.
             secret (str): Secret key for authentication/encryption (empty = security disabled)
-            enable_encryption (bool): Enable encryption (only effective if secret is provided)
+            encrypt (bool): Enable encryption (only effective if secret is provided)
             resolution (tuple[int, int]): Resolution as (width, height)
             fps (int): Frames per second to capture
             adjustments (Callable[[np.ndarray], np.ndarray] | None): Function to adjust frames
             auto_reconnect (bool): Enable automatic reconnection on failure
         """
         super().__init__(resolution, fps, adjustments, auto_reconnect)
 
-        self.protocol = "ws"
-        self.port = port
-        self.timeout = timeout
-        self.codec = BPPCodec(secret, enable_encryption)
+        if use_tls and encrypt:
+            logger.warning("Encryption is redundant over TLS connections, disabling encryption.")
+            encrypt = False
+
+        self.codec = BPPCodec(secret, encrypt)
         self.secret = secret
-        self.enable_encryption = enable_encryption
+        self.encrypt = encrypt
         self.logger = logger
+        self.name = self.__class__.__name__
 
-        host_ip = os.getenv("HOST_IP")
+        # Address and port configuration
+        self.use_tls = use_tls
+        self.protocol = "wss" if use_tls else "ws"
         self._bind_ip = "0.0.0.0"
+        host_ip = os.getenv("HOST_IP")
         self.ip = host_ip if host_ip is not None else self._bind_ip
+        if port < 0 or port > 65535:
+            raise CameraConfigError(f"Invalid port number: {port}")
+        self.port = port
+        if timeout <= 0:
+            raise CameraConfigError(f"Invalid timeout value: {timeout}")
+        self.timeout = timeout
+
+        # TLS configuration
+        if self.use_tls:
+            import ssl
+            from arduino.app_utils.tls_cert_manager import TLSCertificateManager
+
+            try:
+                cert_path, key_path = TLSCertificateManager.get_or_create_certificates(certs_dir=certs_dir_path)
+                self._ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+                self._ssl_context.load_cert_chain(cert_path, key_path)
+                logger.info(f"SSL context created with certificate: {cert_path}")
+            except Exception as e:
+                raise RuntimeError("Failed to configure TLS certificate. Please check certificates and the certs directory.") from e
 
         self._frame_queue = queue.Queue(1)
         self._server = None
@@ -101,7 +131,7 @@ def security_mode(self) -> str:
         """Return current security mode for logging/debugging."""
         if not self.secret:
             return "none"
-        elif self.enable_encryption:
+        elif self.encrypt:
             return "encrypted (ChaCha20-Poly1305)"
         else:
             return "authenticated (HMAC-SHA256)"
@@ -136,12 +166,11 @@ def _start_server_thread(self) -> None:
         finally:
             if self._loop and not self._loop.is_closed():
                 self._loop.close()
+                self._loop = None
 
     async def _start_server(self) -> None:
         """Start the WebSocket server."""
         try:
-            self._stop_event.clear()
-
             self._server = await asyncio.wait_for(
                 websockets.serve(
                     self._ws_handler,
@@ -152,6 +181,7 @@ async def _start_server(self) -> None:
                     close_timeout=self.timeout,
                     ping_interval=20,
                     max_size=5 * 1024 * 1024,  # Limit max message size for security
+                    ssl=self._ssl_context if self.use_tls else None,
                 ),
                 timeout=self.timeout,
             )
@@ -161,7 +191,7 @@ async def _start_server(self) -> None:
                 server_socket = list(self._server.sockets)[0]
                 self.port = server_socket.getsockname()[1]
 
-            await self._stop_event.wait()
+            await self._server.wait_closed()
 
         except TimeoutError as e:
             self.logger.error(f"Failed to start WebSocket server within {self.timeout}s: {e}")
@@ -170,10 +200,7 @@ async def _start_server(self) -> None:
             self.logger.error(f"Failed to start WebSocket server: {e}")
             raise
         finally:
-            if self._server:
-                self._server.close()
-                await self._server.wait_closed()
-                self._server = None
+            self._server = None
 
     async def _ws_handler(self, conn: websockets.ServerConnection) -> None:
         """Handle a connected WebSocket client. Only one client allowed at a time."""
@@ -200,14 +227,14 @@ async def _ws_handler(self, conn: websockets.ServerConnection) -> None:
         try:
             # Send welcome message
             try:
-                welcome = json.dumps({
+                welcome = {
                     "status": "connected",
                     "message": "Connected to camera server",
                     "security_mode": self.security_mode,
                     "resolution": self.resolution,
                     "fps": self.fps,
-                })
-                await self._send_to_client(welcome)
+                }
+                await self._send_to_client(json.dumps(welcome))
             except Exception as e:
                 self.logger.warning(f"Failed to send welcome message: {e}")
 
@@ -241,6 +268,7 @@ async def _ws_handler(self, conn: websockets.ServerConnection) -> None:
                     self.logger.debug(f"Client removed: {client_addr}")
 
     def _parse_message(self, message: websockets.Data) -> np.ndarray | None:
+        """Parse WebSocket message to extract a video frame."""
         if isinstance(message, str):
             try:
                 message = base64.b64decode(message)
@@ -261,7 +289,7 @@ def _close_camera(self):
         """Stop the WebSocket server."""
         if self._loop and not self._loop.is_closed() and self._loop.is_running():
             try:
-                future = asyncio.run_coroutine_threadsafe(self._stop_and_disconnect_client(), self._loop)
+                future = asyncio.run_coroutine_threadsafe(self._disconnect_and_stop(), self._loop)
                 future.result(1.0)
             except CancelledError:
                 self.logger.debug(f"Error stopping WebSocket server: CancelledError")
@@ -281,26 +309,23 @@ def _close_camera(self):
         except queue.Empty:
             pass
 
-        # Reset state
-        self._server = None
-        self._loop = None
-        self._client = None
-
-    async def _stop_and_disconnect_client(self):
-        """Cleanly disconnect client with goodbye message."""
-        if self._client:
-            try:
-                self.logger.debug("Disconnecting client...")
-                goodbye = json.dumps({"status": "disconnecting", "message": "Server is shutting down"})
-                await self._send_to_client(goodbye)
-            except Exception as e:
-                self.logger.warning(f"Failed to send goodbye message: {e}")
-            finally:
-                if self._client:
-                    await self._client.close()
-                    self.logger.debug("Client connection closed")
-
-        self._stop_event.set()
+    async def _disconnect_and_stop(self):
+        """Cleanly disconnect client with goodbye message and stop the server."""
+        async with self._client_lock:
+            if self._client:
+                try:
+                    self.logger.debug("Disconnecting client...")
+                    goodbye = json.dumps({"status": "disconnecting", "message": "Server is shutting down"})
+                    await self._send_to_client(goodbye)
+                except Exception as e:
+                    self.logger.warning(f"Failed to send goodbye message: {e}")
+                finally:
+                    if self._client:
+                        await self._client.close()
+                        self.logger.debug("Client connection closed")
+
+        if self._server:
+            self._server.close()
 
     def _read_frame(self) -> np.ndarray | None:
         """Read a single frame from the queue."""