hermes-hub/app/main.py

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, Response
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pathlib import Path
from typing import Optional, Tuple
from contextlib import asynccontextmanager
import time

from app.config import settings, UNIVERSES, BASE_DIR
from app.routers import api, proxy, chat
from app.proxy import close_http_client, proxy_request, proxy_websocket
from app.personas import ensure_dirs
from app.db import init_db

APP_VERSION = f"2.1.{int(time.time())}"

def get_subdomain_target(host: str) -> Optional[Tuple[str, str]]:
    """
    Inspects Host header and returns (target_backend_url, universe_id) if it matches a universe hostname.
    
    Hostnames (Dash-based 1st level):
      tt-hub.yesminedor.tn       -> (HERMES_TT_URL, 'tt')
      nyora-hub.yesminedor.tn    -> (HERMES_NYORA_URL, 'nyora')
      perso-hub.yesminedor.tn    -> (HERMES_PERSO_URL, 'perso')
      nabil-hub.yesminedor.tn    -> (HERMES_NABIL_URL, 'nabil')
      dsh-hub.yesminedor.tn      -> (HERMES_DSH_URL, 'dsh')
      files-hub.yesminedor.tn    -> (DSH_FILEBROWSER_URL, 'dsh')
    
    Apex / Hub UI:
      hub.yesminedor.tn          -> None (Serves index.html Workspace Switcher)
    """
    if not host:
        return None
    
    hostname = host.split(":")[0].lower()
    
    # The main hub switcher UI should NOT be proxied
    if hostname in ("hub.yesminedor.tn", "localhost", "127.0.0.1"):
        return None
    
    sub = None
    # 1. New 1st level dash-based production hostnames (*-hub.yesminedor.tn)
    if hostname.endswith("-hub.yesminedor.tn"):
        sub = hostname.rsplit("-hub.yesminedor.tn", 1)[0]
    # 2. Backward compatibility (*.hub.yesminedor.tn)
    elif hostname.endswith(".hub.yesminedor.tn"):
        sub = hostname.rsplit(".hub.yesminedor.tn", 1)[0]
    # 3. Localhost dash development (*-hub.localhost)
    elif hostname.endswith("-hub.localhost"):
        sub = hostname.rsplit("-hub.localhost", 1)[0]
    # 4. Localhost dot development (*.localhost)
    elif hostname.endswith(".localhost"):
        sub = hostname.rsplit(".localhost", 1)[0]
    
    if not sub or sub in ("hub", "www"):
        return None
    
    if sub in ("files", "dsh-files"):
        return (settings.dsh_filebrowser_url, "dsh")
    
    if sub in UNIVERSES:
        return (UNIVERSES[sub].backend_url, sub)
        
    return None

@asynccontextmanager
async def lifespan(app: FastAPI):
    ensure_dirs()
    init_db()
    yield
    await close_http_client()

app = FastAPI(
    title=settings.app_name,
    version="2.1.0",
    lifespan=lifespan
)

# Cache-Control & Anti-Stale Middleware for Hub UI & Statics
@app.middleware("http")
async def cache_control_middleware(request: Request, call_next):
    response: Response = await call_next(request)
    path = request.url.path
    if path.startswith("/static/") or path == "/" or path.endswith(".html"):
        response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate, max-age=0"
        response.headers["Pragma"] = "no-cache"
        response.headers["Expires"] = "0"
    return response

# Subdomain Routing Middleware for HTTP
@app.middleware("http")
async def subdomain_routing_middleware(request: Request, call_next):
    host = request.headers.get("host", "")
    target = get_subdomain_target(host)
    if target:
        backend_url, universe_id = target
        path = request.url.path
        return await proxy_request(
            request=request,
            backend_url=backend_url,
            path=path,
            universe_id=universe_id
        )
    return await call_next(request)

# Mount static files
static_dir = BASE_DIR / "static"
static_dir.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")

# Templates
templates_dir = BASE_DIR / "app" / "templates"
templates = Jinja2Templates(directory=str(templates_dir))

# Include Routers
app.include_router(api.router)
app.include_router(proxy.router)
app.include_router(chat.router)

@app.get("/", response_class=HTMLResponse)
@app.head("/", response_class=HTMLResponse)
async def index_view(request: Request):
    return templates.TemplateResponse(
        request=request,
        name="index.html",
        context={
            "universes": UNIVERSES,
            "app_name": settings.app_name,
            "version": APP_VERSION
        }
    )

# WebSocket Proxy Route (intercepts any WebSocket connection across subdomains)
@app.websocket("/{path:path}")
@app.websocket("")
async def websocket_proxy_endpoint(websocket: WebSocket, path: str = ""):
    host = websocket.headers.get("host", "")
    target = get_subdomain_target(host)
    
    if target:
        backend_url, universe_id = target
        await proxy_websocket(
            client_ws=websocket,
            backend_url=backend_url,
            path=path,
            universe_id=universe_id
        )
    elif path.startswith("u/"):
        parts = path.split("/", 2)
        if len(parts) >= 2 and parts[1] in UNIVERSES:
            u_id = parts[1]
            subpath = parts[2] if len(parts) > 2 else ""
            await proxy_websocket(
                client_ws=websocket,
                backend_url=UNIVERSES[u_id].backend_url,
                path=subpath,
                universe_id=u_id
            )
        else:
            await websocket.close(code=1008)
    else:
        await websocket.close(code=1008)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app.main:app", host=settings.host, port=settings.port, reload=settings.debug)