4 files changed, 138 insertions, 0 deletions

diff --git a/ttun_server/__init__.py b/ttun_server/__init__.py
new file mode 100644
index 0000000..b8fd114
--- /dev/null
+++ b/ttun_server/__init__.py
@@ -0,0 +1,12 @@
+from starlette.applications import Starlette
+from starlette.routing import Route, WebSocketRoute
+
+from ttun_server.endpoints import Proxy, Tunnel
+
+server = Starlette(
+    debug=True,
+    routes=[
+        Route('/{path:path}', Proxy),
+        WebSocketRoute('/tunnel/', Tunnel)
+    ]
+)
diff --git a/ttun_server/connections.py b/ttun_server/connections.py
new file mode 100644
index 0000000..a8dabcf
--- /dev/null
+++ b/ttun_server/connections.py
@@ -0,0 +1,3 @@
+from ttun_server.types import Connection
+
+connections: dict[str, Connection] = {}
diff --git a/ttun_server/endpoints.py b/ttun_server/endpoints.py
new file mode 100644
index 0000000..d59cb7c
--- /dev/null
+++ b/ttun_server/endpoints.py
@@ -0,0 +1,98 @@
+import asyncio
+import os
+from asyncio import Queue
+from base64 import b64decode, b64encode
+from typing import Optional, Any
+from uuid import uuid4
+
+from starlette.endpoints import HTTPEndpoint, WebSocketEndpoint
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import Scope, Receive, Send
+from starlette.websockets import WebSocket
+
+from ttun_server.types import Connection, RequestData, Config, ResponseData
+
+from ttun_server.connections import connections
+
+
+class Proxy(HTTPEndpoint):
+    async def dispatch(self) -> None:
+        request = Request(self.scope, self.receive)
+
+        [subdomain, *_] = request.headers['host'].split('.')
+        response = Response(content='Not Found', status_code=404)
+
+        if subdomain in connections:
+            connection = connections[subdomain]
+
+            await connection['requests'].put(RequestData(
+                method=request.method,
+                path=str(request.url).replace(str(request.base_url), '/'),
+                headers=dict(request.headers),
+                cookies=dict(request.cookies),
+                body=b64encode(await request.body()).decode()
+            ))
+
+            _response = await connection['responses'].get()
+            response = Response(
+                status_code=_response['status'],
+                headers=_response['headers'],
+                content=b64decode(_response['body'].encode())
+            )
+
+        await response(self.scope, self.receive, self.send)
+
+
+class Tunnel(WebSocketEndpoint):
+    encoding = 'json'
+
+    def __init__(self, scope: Scope, receive: Receive, send: Send):
+        super().__init__(scope, receive, send)
+        self.request_task = None
+        self.config: Optional[Config] = None
+
+    @property
+    def requests(self) -> Queue[RequestData]:
+        return connections[self.config['subdomain']]['requests']
+
+    @property
+    def responses(self) -> Queue[ResponseData]:
+        return connections[self.config['subdomain']]['responses']
+
+    async def handle_requests(self, websocket: WebSocket):
+        while request := await self.requests.get():
+            await websocket.send_json(request)
+
+    async def on_connect(self, websocket: WebSocket) -> None:
+        await websocket.accept()
+        self.config = await websocket.receive_json()
+
+        if self.config['subdomain'] is None \
+                or self.config['subdomain'] in connections:
+            self.config['subdomain'] = uuid4().hex
+
+
+        connections[self.config['subdomain']] = Connection(
+            requests=Queue(),
+            responses=Queue(),
+        )
+
+        hostname = os.environ.get("TUNNEL_DOMAIN")
+        protocol = "https" if os.environ.get("SECURE", False) else "http"
+
+        await websocket.send_json({
+            'url': f'{protocol}://{self.config["subdomain"]}.{hostname}'
+        })
+
+        self.request_task = asyncio.create_task(self.handle_requests(websocket))
+
+    async def on_receive(self, websocket: WebSocket, data: Any) -> None:
+        await self.responses.put(data)
+
+    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:
+        if self.config is not None and self.config['subdomain'] in connections:
+            del connections[self.config['subdomain']]
+
+        if self.request_task is not None:
+            self.request_task.cancel()
diff --git a/ttun_server/types.py b/ttun_server/types.py
new file mode 100644
index 0000000..0b2fb87
--- /dev/null
+++ b/ttun_server/types.py
@@ -0,0 +1,25 @@
+from asyncio import Queue
+from typing import TypedDict, Optional
+
+
+class Config(TypedDict):
+    subdomain: str
+
+class RequestData(TypedDict):
+    method: str
+    path: str
+    headers: dict
+    cookies: dict
+    body: Optional[str]
+
+
+class ResponseData(TypedDict):
+    status: int
+    headers: dict
+    cookies: dict
+    body: Optional[str]
+
+
+class Connection(TypedDict):
+    requests: Queue[RequestData]
+    responses: Queue[ResponseData]