hub: Add websocket capability

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>

diff --git a/src/pakfire/hub.py b/src/pakfire/hub.py
index 95d2ad002fd5164e068df0452a2fd612516eed86..be5a509e3ed547cf8c820246d01480c171bc8a83 100644 (file)
--- a/src/pakfire/hub.py
+++ b/src/pakfire/hub.py
@@ -27,6 +27,7 @@ import logging
 import os.path
 import psutil
 import tornado.httpclient
+import tornado.websocket
 import urllib.parse
 
 from . import _pakfire
@@ -55,10 +56,15 @@ class Hub(object):
 
                # XXX support proxies
 
-       async def _request(self, method, path, body=None, body_producer=None, **kwargs):
+       async def _request(self, method, path, websocket=False,
+                       body=None, body_producer=None, on_message_callback=None, **kwargs):
                # Make absolute URL
                url = urllib.parse.urljoin(self.url, path)
 
+               # Change scheme for websocket
+               if websocket and url.startswith("https://"):
+                       url = url.replace("https://", "wss://")
+
                # Convert any query arguments
                query_args = urllib.parse.urlencode(kwargs)
 
@@ -82,6 +88,13 @@ class Hub(object):
                        body_producer=body_producer,
                )
 
+               # Is this a web socket request?
+               if websocket:
+                       return await tornado.websocket.websocket_connect(
+                               req,
+                               on_message_callback=on_message_callback,
+                       )
+
                # Send the request and wait for a response
                res = await self.client.fetch(req)