""" Superadmin tunnel management API endpoints. """ from typing import Annotated, Optional from fastapi import APIRouter, Depends, HTTPException, status from sqlalchemy.ext.asyncio import AsyncSession from pydantic import BaseModel from sqlalchemy import select from app.api.deps import get_current_superadmin from app.core.database import get_db from app.models.device import Device from app.models.user import User from app.services.tunnel_service import tunnel_service router = APIRouter(prefix="/tunnels", tags=["superadmin-tunnels"]) class TunnelEnableResponse(BaseModel): """Response when enabling tunnel""" session_uuid: str device_id: str tunnel_type: str status: str class TunnelStatusResponse(BaseModel): """Tunnel session status response""" session_uuid: str status: str # "waiting" | "ready" | "failed" device_tunnel_port: Optional[int] = None ttyd_port: Optional[int] = None tunnel_url: Optional[str] = None @router.post("/devices/{device_id}/{tunnel_type}") async def enable_tunnel( device_id: int, tunnel_type: str, db: Annotated[AsyncSession, Depends(get_db)], current_user: Annotated[User, Depends(get_current_superadmin)] ) -> TunnelEnableResponse: """ Enable SSH or Dashboard tunnel for device. Creates session and triggers device to connect tunnel. Args: device_id: Device numeric ID tunnel_type: "ssh" or "dashboard" Returns: session_uuid: UUID for polling status """ if tunnel_type not in ["ssh", "dashboard"]: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="tunnel_type must be 'ssh' or 'dashboard'" ) # Get device result = await db.execute( select(Device).where(Device.id == device_id) ) device = result.scalar_one_or_none() if not device: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Device not found" ) # Create tunnel session session = tunnel_service.create_session( device_id=device.mac_address, admin_user=current_user.email, tunnel_type=tunnel_type ) # Update device config to enable tunnel if device.config is None: device.config = {} tunnel_key = f"{tunnel_type}_tunnel" if tunnel_key not in device.config: device.config[tunnel_key] = {} device.config[tunnel_key]["enabled"] = True # Mark as modified (SQLAlchemy JSON field) from sqlalchemy.orm import attributes attributes.flag_modified(device, "config") await db.commit() return TunnelEnableResponse( session_uuid=session.uuid, device_id=device.mac_address, tunnel_type=tunnel_type, status="waiting" ) @router.get("/sessions/{session_uuid}/status") async def get_tunnel_status( session_uuid: str, current_user: Annotated[User, Depends(get_current_superadmin)] ) -> TunnelStatusResponse: """ Poll tunnel session status. Returns: - waiting: Device not yet connected - ready: Tunnel connected, ttyd spawned - failed: Session expired or failed """ session = tunnel_service.get_session(session_uuid) if not session: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Session not found or expired" ) # Build response response = TunnelStatusResponse( session_uuid=session.uuid, status=session.status, device_tunnel_port=session.device_tunnel_port, ttyd_port=session.ttyd_port ) # If ready, spawn ttyd if not already spawned if session.status == "ready" and session.device_tunnel_port: if not session.ttyd_port: try: # Spawn ttyd ttyd_port = tunnel_service.spawn_ttyd( session_uuid=session.uuid, device_tunnel_port=session.device_tunnel_port ) response.ttyd_port = ttyd_port response.tunnel_url = f"/admin/{session.tunnel_type}/{session.uuid}" except Exception as e: print(f"[tunnel] Failed to spawn ttyd: {e}") session.status = "failed" response.status = "failed" else: response.tunnel_url = f"/admin/{session.tunnel_type}/{session.uuid}" return response @router.post("/sessions/{session_uuid}/heartbeat") async def session_heartbeat( session_uuid: str, current_user: Annotated[User, Depends(get_current_superadmin)] ): """ Browser sends heartbeat every 30 seconds to keep session alive. """ success = tunnel_service.update_heartbeat(session_uuid) if not success: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Session not found" ) return {"success": True}