wip

terminus/app/consumers.py

@@ -1,31 +1,74 @@
 import asyncio
+import os
+import ssl
+import websockets
 from channels.generic.websocket import AsyncWebsocketConsumer
+from django.conf import settings
+from urllib.parse import quote
+
+
+class GatewayConnection:
+    _ssl_context: ssl.SSLContext = None
+
+    def __init__(self, host: str, port: int):
+        if settings.CONNECTION_GATEWAY_AUTH_KEY and not GatewayConnection._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
+            GatewayConnection._ssl_context = ctx
+
+        proto = 'wss' if GatewayConnection._ssl_context else 'ws'
+        self.url = f'{proto}://localhost:9000/connect/{quote(host)}:{quote(str(port))}'
+
+    async def connect(self):
+        self.context = websockets.connect(self.url, ssl=GatewayConnection._ssl_context)
+        self.socket = await self.context.__aenter__()
+
+    async def send(self, data):
+        await self.socket.send(data)
+
+    def recv(self, timeout=None):
+        return asyncio.wait_for(self.socket.recv(), timeout)
+
+    async def close(self):
+        await self.socket.close()
+        await self.context.__aexit__(None, None, None)
 
 
 class TCPConsumer(AsyncWebsocketConsumer):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.loop = asyncio.get_event_loop()
-
     async def connect(self):
-        self.reader, self.writer = await asyncio.open_connection('192.168.78.233', 22)
-        self._socket_reader = self.loop.create_task(self.socket_reader())
+        self.closed = False
+        self.conn = GatewayConnection(
+            self.scope['url_route']['kwargs']['host'],
+            int(self.scope['url_route']['kwargs']['port']),
+        )
+        await self.conn.connect()
         await self.accept()
+        self.reader = asyncio.get_event_loop().create_task(self.socket_reader())
 
     async def disconnect(self, close_code):
-        await self.writer.drain()
-        self.writer.close()
-        await self._socket_reader
+        self.closed = True
+        await self.conn.close()
 
     async def receive(self, bytes_data):
-        self.writer.write(bytes_data)
+        await self.conn.send(bytes_data)
 
     async def socket_reader(self):
         while True:
-            await self.reader._wait_for_data('read')
-            data = bytes(self.reader._buffer.copy())
-            if not data:
+            if self.closed:
+                return
+            try:
+                data = await self.conn.recv(timeout=10)
+            except asyncio.TimeoutError:
+                continue
+            except websockets.exceptions.ConnectionClosed:
                 await self.close()
                 return
-            del self.reader._buffer[:]
             await self.send(bytes_data=data)

terminus/app/urls.py

@@ -1,4 +1,4 @@
-from django.urls import path, include
+from django.urls import path, re_path, include
 from rest_framework import routers
 
 from . import api
@@ -22,5 +22,5 @@ urlpatterns = [
 ]
 
 websocket_urlpatterns = [
-    path('api/1/gateway/tcp', consumers.TCPConsumer.as_asgi()),
+    re_path(r'^api/1/gateway/tcp/(?P<host>[^/]+):(?P<port>\d+)$', consumers.TCPConsumer.as_asgi()),
 ]