Added websocket support

5 files changed, 234 insertions, 59 deletions

diff --git a/ttun/__main__.py b/ttun/__main__.py
index 4e693fb..dd83e53 100644
--- a/ttun/__main__.py
+++ b/ttun/__main__.py
@@ -1,21 +1,23 @@
 import asyncio
+import logging
+import os
 import re
-import time
 from argparse import ArgumentDefaultsHelpFormatter
 from argparse import ArgumentParser
 from asyncio import FIRST_EXCEPTION
 from asyncio.exceptions import CancelledError
 from asyncio.exceptions import TimeoutError
-from typing import Dict
 from typing import Tuple
 
-from websockets.exceptions import ConnectionClosedError
-
 from ttun.client import Client
 from ttun.inspect_server import Server
 from ttun.settings import SERVER_HOSTNAME
 from ttun.settings import SERVER_USING_SSL
 
+logging.basicConfig(encoding="utf-8")
+logging.getLogger("asyncio").setLevel(os.environ.get("LOGGING_LEVEL", "NOTSET"))
+logging.getLogger("websockets").setLevel(os.environ.get("LOGGING_LEVEL", "NOTSET"))
+
 inspect_queue = asyncio.Queue()
 
 
@@ -73,7 +75,6 @@ def main():
         headers=args.header,
     )
 
-
     try:
         loop = asyncio.get_running_loop()
     except RuntimeError:
diff --git a/ttun/client.py b/ttun/client.py
index a75c882..b19bb47 100644
--- a/ttun/client.py
+++ b/ttun/client.py
@@ -1,5 +1,8 @@
 import asyncio
 import json
+import logging
+import os
+import sys
 from asyncio import get_running_loop
 from base64 import b64decode
 from base64 import b64encode
@@ -24,10 +27,18 @@ from websockets.exceptions import ConnectionClosed
 from ttun import __version__
 from ttun.pubsub import PubSub
 from ttun.types import Config
+from ttun.types import HttpMessage
+from ttun.types import HttpMessageType
+from ttun.types import HttpRequestData
+from ttun.types import HttpResponseData
 from ttun.types import Message
 from ttun.types import MessageType
-from ttun.types import RequestData
-from ttun.types import ResponseData
+from ttun.types import WebsocketMessage
+from ttun.types import WebsocketMessageData
+from ttun.types import WebsocketMessageType
+
+logger = logging.getLogger(__name__)
+logger.setLevel(os.environ.get("LOGGING_LEVEL", "NOTSET"))
 
 
 class Client:
@@ -48,9 +59,12 @@ class Client:
         self.connection: WebSocketClientProtocol = None
 
         self.proxy_origin = f'{"https" if https else "http"}://{to}:{port}'
+        self.ws_proxy_origin = f'{"wss" if https else "ws"}://{to}:{port}'
 
         self.headers = [] if headers is None else headers
 
+        self.websocket_connections = {}
+
     async def send(self, data: dict):
         await self.connection.send(json.dumps(data))
 
@@ -86,63 +100,181 @@ class Client:
     def session(self):
         return ClientSession(base_url=self.proxy_origin, cookie_jar=DummyCookieJar())
 
+    async def handle_request(self, message: HttpMessage, session: ClientSession = None):
+        if session is None:
+            session = self.session()
+
+        request: HttpRequestData = message["payload"]
+
+        request["headers"] = [
+            *request["headers"],
+            *self.headers,
+        ]
+
+        async def response_handler(
+            response: HttpResponseData, identifier=message["identifier"]
+        ):
+            await self.send(
+                HttpMessage(
+                    type=HttpMessageType.response.value,
+                    identifier=identifier,
+                    payload=response,
+                )
+            )
+
+        await self.proxy_request(
+            session=session,
+            request=request,
+            on_response=response_handler,
+        )
+
+    async def receive_websocket_message(self, message: str, idenitfier: str):
+        message_data = WebsocketMessage(
+            identifier=idenitfier,
+            type=WebsocketMessageType.message.value,
+            payload=WebsocketMessageData(body=b64encode(message.encode()).decode()),
+        )
+        await self.send(message_data)
+
+        await PubSub.publish(
+            {
+                "type": "websocket_outbound",
+                "payload": {
+                    "id": message_data["identifier"],
+                    "timestamp": datetime.now().isoformat(),
+                    **message_data["payload"],
+                },
+            }
+        )
+
+    async def connect_websocket(self, message: WebsocketMessage):
+        assert not message["identifier"] in self.websocket_connections
+
+        start = perf_counter()
+        await PubSub.publish(
+            {
+                "type": "websocket_connect",
+                "payload": {
+                    "id": message["identifier"],
+                    "timestamp": datetime.now().isoformat(),
+                    **message["payload"],
+                },
+            }
+        )
+
+        async with websockets.connect(
+            f'{self.ws_proxy_origin}/{message["payload"]["path"]}'
+        ) as connection:
+            end = perf_counter()
+            self.websocket_connections[message["identifier"]] = connection
+
+            await self.send(
+                WebsocketMessage(
+                    identifier=message["identifier"],
+                    type=WebsocketMessageType.ack.value,
+                    payload=None,
+                )
+            )
+
+            await PubSub.publish(
+                {
+                    "type": "websocket_connected",
+                    "payload": {
+                        "id": message["identifier"],
+                        "timing": end - start,
+                    },
+                }
+            )
+
+            async for m in connection:
+                await self.receive_websocket_message(m, message["identifier"])
+
+    async def send_websocket_message(self, message: WebsocketMessage):
+        assert message["identifier"] in self.websocket_connections
+        await self.websocket_connections[message["identifier"]].send(
+            b64decode(message["payload"]["body"]).decode()
+        )
+
+        await PubSub.publish(
+            {
+                "type": "websocket_inbound",
+                "payload": {
+                    "id": message["identifier"],
+                    "timestamp": datetime.now().isoformat(),
+                    **message["payload"],
+                },
+            }
+        )
+
+    async def disconnect_websocket(self, message: WebsocketMessage):
+        assert message["identifier"] in self.websocket_connections
+
+        await self.websocket_connections[message["identifier"]].close()
+
+        self.websocket_connections[message["identifier"]] = None
+        await PubSub.publish(
+            {
+                "type": "websocket_disconnect",
+                "payload": {
+                    "id": message["identifier"],
+                    "timestamp": datetime.now().isoformat(),
+                    **message["payload"],
+                },
+            }
+        )
+
     async def handle_messages(self):
         loop = get_running_loop()
+        tasks = set()
+
         async with self.session() as session:
             while True:
                 try:
                     message: Message = await self.receive()
+                    logger.debug(message)
 
-                    try:
+                    ma