1 files changed, 8 insertions, 15 deletions

diff --git a/ttun_server/__init__.py b/ttun_server/__init__.py
index 6c77858..d54227a 100644
--- a/ttun_server/__init__.py
+++ b/ttun_server/__init__.py
@@ -2,30 +2,23 @@ import logging
 import os
 
 from fastapi import FastAPI
-from starlette.routing import Host, Route, Router, WebSocketRoute
+from starlette.routing import Host, Route, WebSocketRoute
 
-from ttun_server.endpoints import health, proxy
-from .websockets import WebsocketProxy, Tunnel
+from proxy import app as proxy_app
+from ttun_server.endpoints import health, base_endpoints
+from ttun_server.websockets import tunnel
 
 logging.basicConfig(level=getattr(logging, os.environ.get('LOG_LEVEL', 'INFO')))
 
-base_router = Router(routes=[
-    Route('/health/', health),
-    WebSocketRoute('/tunnel/', Tunnel)
-])
-
-server = FastAPI(
+app = FastAPI(
     debug=True,
     routes=[
-        Host(os.environ['TUNNEL_DOMAIN'], base_router, 'base'),
-        Route('/{path:path}', proxy),
-        WebSocketRoute('/{path:path}', WebsocketProxy)
+        WebSocketRoute('/tunnel/', endpoint=tunnel),
+        Route('/health/', endpoint=health),
+        Host(f'{{subdomain}}.{os.environ['TUNNEL_DOMAIN']}', app=proxy_app)
     ]
 )
 
-server.post()
-
-
 try:
     from ._version import version
     __version__ = version