wip

tabby/app/consumers.py

@@ -1,11 +1,15 @@
 import asyncio
+import json
 import os
+import secrets
 import ssl
 import websockets
 from channels.generic.websocket import AsyncWebsocketConsumer
 from django.conf import settings
 from urllib.parse import quote
 
+from .models import Gateway
+
 
 class GatewayConnection:
     _ssl_context: ssl.SSLContext = None
@@ -42,6 +46,47 @@ class GatewayConnection:
         await self.context.__aexit__(None, None, None)
 
 
+class GatewayAdminConnection:
+    _ssl_context: ssl.SSLContext = None
+
+    def __init__(self, gateway: Gateway):
+        if not settings.CONNECTION_GATEWAY_AUTH_KEY:
+            raise RuntimeError('CONNECTION_GATEWAY_AUTH_KEY is required to manage connection gateways')
+        if not GatewayAdminConnection._ssl_context:
+            ctx = ssl.create_default_context(purpose=ssl.Purpose.CLIENT_AUTH)
+            ctx.load_cert_chain(
+                os.path.realpath(settings.CONNECTION_GATEWAY_AUTH_CERTIFICATE),
+                os.path.realpath(settings.CONNECTION_GATEWAY_AUTH_KEY),
+            )
+            if settings.CONNECTION_GATEWAY_AUTH_CA:
+                ctx.load_verify_locations(
+                    cafile=os.path.realpath(settings.CONNECTION_GATEWAY_AUTH_CA),
+                )
+                ctx.verify_mode = ssl.CERT_REQUIRED
+            GatewayAdminConnection._ssl_context = ctx
+
+        self.url = f'wss://{gateway.host}:{gateway.admin_port}'
+
+    async def connect(self):
+        self.context = websockets.connect(self.url, ssl=GatewayAdminConnection._ssl_context)
+        self.socket = await self.context.__aenter__()
+
+    async def authorize_client(self) -> str:
+        token = secrets.token_hex(32)
+        await self.send(json.dumps({
+            '_': 'authorize-client',
+            'token': token,
+        }))
+        return token
+
+    async def send(self, data):
+        await self.socket.send(data)
+
+    async def close(self):
+        await self.socket.close()
+        await self.context.__aexit__(None, None, None)
+
+
 class TCPConsumer(AsyncWebsocketConsumer):
     async def connect(self):
         self.closed = False