Support for WS

app.py

@@ -7,7 +7,10 @@ import time
 from datetime import datetime, timezone
 from socketserver import ThreadingMixIn
 import os
-
+import asyncio
+import websockets
+import hashlib
+import base64
 # Define the target server to proxy requests to
 class ProxyHandler(http.server.BaseHTTPRequestHandler):
     def __init__(self, configuration, docker_client):
@@ -73,6 +76,15 @@ class ProxyHandler(http.server.BaseHTTPRequestHandler):
             return
         
         activity[proxy_host_configuration['domain']] = datetime.now(timezone.utc)
+
+        # Check if this is a WebSocket request
+        if self.headers.get("Upgrade", "").lower() == "websocket":
+            print("Request is WS connecting to {0}".format(container['container_name']))
+            print("Request is WS connecting to {0}:{1}".format(proxy_host_configuration['proxy_host'],proxy_host_configuration['proxy_port']))
+            activity[proxy_host_configuration['domain']] = True
+            self.upgrade_to_websocket(proxy_host_configuration['proxy_host'], proxy_host_configuration['proxy_port'])
+            return
+        
         # Open a connection to the target server
         conn = http.client.HTTPConnection(proxy_host_configuration['proxy_host'], proxy_host_configuration['proxy_port'])
         conn.request(method, self.path, headers=self.headers)
@@ -122,6 +134,56 @@ class ProxyHandler(http.server.BaseHTTPRequestHandler):
 
         self.wfile.flush()
 
+    async def websocket_proxy(self, target_host, target_port):
+        server_ws = None
+        try:
+            client_connection = self.connection
+            # Establish server connection to backend
+            server_ws = await websockets.connect(f"ws://{target_host}:{target_port}")
+            
+            # Bridge function to handle message forwarding
+            async def bridge_websockets():
+                try:
+                    while True:
+                        # Wait for a message from the client
+                        client_message = await client_connection.recv()
+                        # Send it to the server
+                        await server_ws.send(client_message)
+                        # Wait for a message from the server
+                        server_message = await server_ws.recv()
+                        # Send it to the client
+                        await client_connection.send(server_message)
+                except websockets.exceptions.ConnectionClosed as e:
+                    print(f"WebSocket connection closed: {e}")
+                except Exception as e:
+                    print(f"Error during WebSocket communication: {e}")
+
+            # Run the bridge coroutine
+            await bridge_websockets()
+
+        except Exception as e:
+            print(f"WebSocket proxy encountered an error: {e}")
+        finally:
+            if server_ws:
+                await server_ws.close()
+
+    def upgrade_to_websocket(self, target_host, target_port):
+        """
+        Handles WebSocket upgrade requests and spawns an asyncio WebSocket proxy.
+        """
+        key = self.headers['Sec-WebSocket-Key']
+        accept_val = base64.b64encode(hashlib.sha1((key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode('utf-8')).digest()).decode('utf-8')
+
+
+        self.send_response(101)  # Switching Protocols
+        self.send_header("Upgrade", "websocket")
+        self.send_header("Connection", "Upgrade")
+        self.send_header("Sec-WebSocket-Accept", accept_val)
+        self.end_headers()
+
+        loop = asyncio.new_event_loop()
+        threading.Thread(target=loop.run_until_complete, args=(self.websocket_proxy(target_host, target_port),)).start()
+
 class ThreadedHTTPServer(ThreadingMixIn, http.server.HTTPServer):
     """Handle requests in a separate thread."""
 
@@ -148,6 +210,9 @@ class BackgroundTasks(threading.Thread):
                             if (apps['domain'] in activity):
                                 dt = activity[apps['domain']]
 
+                            if (dt == True):
+                                continue
+
                             diff_seconds = (datetime.now(timezone.utc) - dt).total_seconds()
                             if(diff_seconds > apps['proxy_timeout_seconds']):
                                 print("stopping container: {0} ({1}) after {2}s".format(container['container_name'], container_object.id, diff_seconds))
@@ -157,6 +222,30 @@ class BackgroundTasks(threading.Thread):
             
             time.sleep(sleep_time)
 
+async def websocket_proxy(client_ws, target_host, target_port):
+    """
+    Forwards WebSocket messages between the client and the target container.
+    """
+    try:
+        async with websockets.connect(f"ws://{target_host}:{target_port}") as server_ws:
+            # Create tasks to read from both directions
+            async def forward_client_to_server():
+                async for message in client_ws:
+                    await server_ws.send(message)
+
+            async def forward_server_to_client():
+                async for message in server_ws:
+                    await client_ws.send(message)
+
+            await asyncio.gather(forward_client_to_server(), forward_server_to_client())
+
+    except Exception as e:
+        print(f"WebSocket error: {e}")
+    finally:
+        await client_ws.close()
+
+# MAIN #
+
 if __name__ == '__main__':
     activity = {}
 
@@ -173,4 +262,5 @@ if __name__ == '__main__':
     proxy_handler = ProxyHandler(configuration, docker_client)
     httpd = ThreadedHTTPServer(server_address, proxy_handler)
     print('Reverse proxy server running on port {0}...'.format(configuration['proxy_port']))
-    httpd.serve_forever()
+    httpd.serve_forever()
+