mavutil: added wss client support

Author: tridge

mavutil.py

@@ -13,6 +13,7 @@
 import re
 import platform
 from pymavlink import mavexpression
+import ssl
 
 # We want to re-export x25crc here
 from pymavlink.generator.mavcrc import x25crc as x25crc
@@ -1853,8 +1854,9 @@ def write(self, buf):
                 self.close_port()
             pass
 
+
 class mavwebsocket_client(mavfile):
-    '''client using WebSocket over TCP'''
+    '''client using WebSocket over TCP with WS and WSS support'''
     def __init__(self,
                  device,
                  source_system=255,
@@ -1863,15 +1865,17 @@ def __init__(self,
                  use_native=default_native):
         self.resource = "/"
         a = device.split(':')
-        if len(a) < 2:
-            raise ValueError("TCP ports must be specified as host:port")
-        self.host = a[0]
-        self.port = int(a[1])
-        if len(a) > 2:
-            self.resource = a[2]
+        protocol = a[0]
+        if len(a) < 3:
+            raise ValueError("WebSocket ports must be specified as protocol:host:port")
+        self.host = a[1]
+        self.port = int(a[2])
+        if len(a) > 3:
+            self.resource = a[3]
         self.sock = None
+        self.use_ssl = protocol.lower() == 'wss'
         self.connect()
-        mavfile.__init__(self, self.sock.fileno(), "ws:" + device, source_system=source_system, source_component=source_component, use_native=use_native)
+        mavfile.__init__(self, self.sock.fileno(), protocol + device, source_system=source_system, source_component=source_component, use_native=use_native)
 
     def connect(self):
         self.close()
@@ -1883,20 +1887,42 @@ def connect(self):
             BytesMessage,
         )
         try:
-            self.sock = socket.create_connection((self.host, self.port))
+            # Create basic socket connection
+            raw_sock = socket.create_connection((self.host, self.port))
+            
+            # Wrap with SSL if using WSS
+            if self.use_ssl:
+                context = ssl.create_default_context()
+                # Optional: For testing with self-signed certificates, uncomment:
+                # context.check_hostname = False
+                # context.verify_mode = ssl.CERT_NONE
+                self.sock = context.wrap_socket(raw_sock, server_hostname=self.host)
+            else:
+                self.sock = raw_sock
+                
         except socket.error as e:
-            if e.errno in [ errno.ECONNREFUSED, errno.EHOSTUNREACH ]:
+            if e.errno in [errno.ECONNREFUSED, errno.EHOSTUNREACH]:
                 return
             raise
+        except ssl.SSLError as e:
+            print(f"SSL Error: {e}")
+            raise
+            
         self.sock.setblocking(1)
         self.ws = WSConnection(ConnectionType.CLIENT)
         b = self.ws.send(Request(host=self.host, target=self.resource))
         self.sock.send(b)
         self.buffer = b''
-
+        
         # wait for handshake response
         while True:
-            data = self.sock.recv(4096)
+            try:
+                data = self.sock.recv(4096)
+            except ssl.SSLError as e:
+                raise RuntimeError(f"WebSocket SSL handshake failed: {e}")
+            except socket.error as e:
+                raise RuntimeError(f"WebSocket handshake failed: {e}")
+                
             if not data:
                 raise RuntimeError("WebSocket handshake failed")
             self.ws.receive_data(data)
@@ -1905,7 +1931,7 @@ def connect(self):
                     self.sock.setblocking(0)
                     return
 
-    def recv(self,n=None):
+    def recv(self, n=None):
         from wsproto.events import (
             BytesMessage,
             CloseConnection
@@ -1918,10 +1944,22 @@ def recv(self,n=None):
             return out
         try:
             data = self.sock.recv(n)
+        except ssl.SSLError as e:
+            # Handle SSL-specific errors
+            if e.errno == ssl.SSL_ERROR_WANT_READ:
+                return b""  # Need more data, try again later
+            elif e.errno == ssl.SSL_ERROR_WANT_WRITE:
+                return b""  # SSL needs to write, try again later
+            elif e.errno in [errno.EAGAIN, errno.EWOULDBLOCK]:
+                return b""
+            else:
+                # Real SSL error, reconnect
+                self.connect()
+                return b''
         except socket.error as e:
-            if e.errno in [ errno.EAGAIN, errno.EWOULDBLOCK ]:
+            if e.errno in [errno.EAGAIN, errno.EWOULDBLOCK]:
                 return b""
-            if e.errno in [ errno.ECONNRESET, errno.EPIPE ]:
+            if e.errno in [errno.ECONNRESET, errno.EPIPE]:
                 self.connect()
                 return b''
             raise
@@ -1944,16 +1982,26 @@ def write(self, data):
         b = self.ws.send(BytesMessage(data=data))
         try:
             self.sock.send(b)
+        except ssl.SSLError as e:
+            # Handle SSL-specific errors
+            if e.errno == ssl.SSL_ERROR_WANT_READ:
+                return  # SSL needs to read, try again later
+            elif e.errno == ssl.SSL_ERROR_WANT_WRITE:
+                return  # SSL needs to write, try again later
+            else:
+                # Real SSL error, reconnect
+                self.connect()
         except socket.error as e:
-            if e.errno in [ errno.EPIPE ]:
+            if e.errno in [errno.EPIPE]:
                 self.connect()
             pass
 
     def close(self):
         if self.sock:
             self.sock.close()
             self.sock = None
-        
+
+
 def mavlink_connection(device, baud=115200, source_system=255, source_component=0,
                        planner_format=None, write=False, append=False,
                        robust_parsing=True, notimestamps=False, input=True,
@@ -1987,8 +2035,8 @@ def mavlink_connection(device, baud=115200, source_system=255, source_component=
         return mavudp(device[9:], input=False, source_system=source_system, source_component=source_component, use_native=use_native, broadcast=True)
     if device.startswith('wsserver:'):
         return mavwebsocket(device[9:], source_system=source_system, source_component=source_component, use_native=use_native)
-    if device.startswith('ws:'):
-        return mavwebsocket_client(device[3:], source_system=source_system, source_component=source_component, use_native=use_native)
+    if device.startswith('ws:') or device.startswith('wss:'):
+        return mavwebsocket_client(device, source_system=source_system, source_component=source_component, use_native=use_native)
     # For legacy purposes we accept the following syntax and let the caller to specify direction
     if device.startswith('udp:'):
         return mavudp(device[4:], input=input, source_system=source_system, source_component=source_component, use_native=use_native)