implementing ssl for server and client;

README.md

@@ -58,7 +58,25 @@ client = TranscriptionClient(
   9090,
   lang="en",
   translate=False,
-  model="small"
+  model="small",
+  secure_websocket=False,
+  sslopt={},
+)
+
+client("tests/jfk.wav")
+```
+
+- To transcribe an audio file with ssl:
+```python
+from whisper_live.client import TranscriptionClient
+client = TranscriptionClient(
+  "localhost",
+  9090,
+  lang="en",
+  translate=False,
+  model="small",
+  secure_websocket=True,
+  sslopt={"ca_certs": "cert.pem"},
 )
 
 client("tests/jfk.wav")
@@ -73,7 +91,9 @@ client = TranscriptionClient(
   9090,
   lang="hi",
   translate=True,
-  model="small"
+  model="small",
+  secure_websocket=False,
+  sslopt={},
 )
 client()
 ```

run_server.py

@@ -1,4 +1,5 @@
 import argparse
+import ssl
 from whisper_live.server import TranscriptionServer
 
 if __name__ == "__main__":
@@ -21,18 +22,26 @@
     parser.add_argument('--trt_multilingual', '-m',
                         action="store_true",
                         help='Boolean only for TensorRT model. True if multilingual.')
+    parser.add_argument('--ssl_cert_path', '-ssl',
+                        type=str,
+                        default=None,
+                        help='Path to cert.pem and key.pem if ssl should be used.')
     args = parser.parse_args()
 
     if args.backend == "tensorrt":
         if args.trt_model_path is None:
             raise ValueError("Please Provide a valid tensorrt model path")
-
+    ssl_context = None
+    if args.ssl_cert_path is not None:
+        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        ssl_context.load_cert_chain(certfile=f"{args.ssl_cert_path}/cert.pem", keyfile=f"{args.ssl_cert_path}/key.pem")
     server = TranscriptionServer()
     server.run(
         "0.0.0.0",
         port=args.port, 
         backend=args.backend,
         faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
         whisper_tensorrt_path=args.trt_model_path,
-        trt_multilingual=args.trt_multilingual
-    )
+        trt_multilingual=args.trt_multilingual,
+        ssl_context=ssl_context
+    )

whisper_live/client.py

@@ -79,7 +79,9 @@ def __init__(
         lang=None,
         translate=False,
         model="small",
-        srt_file_path="output.srt"
+        srt_file_path="output.srt",
+        secure_websocket=False,
+        sslopt={},
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -125,7 +127,8 @@ def __init__(
         )
 
         if host is not None and port is not None:
-            socket_url = f"ws://{host}:{port}"
+            socket_protocol = "wss" if secure_websocket else "ws"
+            socket_url = f"{socket_protocol}://{host}:{port}"
             self.client_socket = websocket.WebSocketApp(
                 socket_url,
                 on_open=lambda ws: self.on_open(ws),
@@ -142,7 +145,7 @@ def __init__(
         Client.INSTANCES[self.uid] = self
 
         # start websocket client in a thread
-        self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
+        self.ws_thread = threading.Thread(target=self.client_socket.run_forever, kwargs={"sslopt": sslopt})
         self.ws_thread.setDaemon(True)
         self.ws_thread.start()
 
@@ -559,8 +562,10 @@ def __init__(self,
         lang=None,
         translate=False,
         model="small",
+        secure_websocket=False,
+        sslopt={},
     ):
-        self.client = Client(host, port, lang, translate, model)
+        self.client = Client(host, port, lang, translate, model, secure_websocket, sslopt)
 
     def __call__(self, audio=None, hls_url=None):
         """

whisper_live/server.py

@@ -218,7 +218,8 @@ def run(self,
             backend="tensorrt", 
             faster_whisper_custom_model_path=None,
             whisper_tensorrt_path=None, 
-            trt_multilingual=False
+            trt_multilingual=False,
+            ssl_context=None
         ):
         """
         Run the transcription server.
@@ -236,7 +237,8 @@ def run(self,
                 trt_multilingual=trt_multilingual
             ),
             host,
-            port
+            port,
+            ssl_context=ssl_context
         ) as server:
             server.serve_forever()