mybeacon-backend/backend/app/api/v1/tunnel_port.py

"""
Device API endpoint for reporting tunnel port allocation.
"""

from typing import Annotated, Optional

from fastapi import APIRouter, Depends, Header, HTTPException, status
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import get_db
from app.models.device import Device
from app.services.tunnel_service import tunnel_service


router = APIRouter(tags=["device-api"])


class TunnelPortRequest(BaseModel):
    """Device reports tunnel port"""
    tunnel_type: str  # "ssh" | "dashboard"
    port: Optional[int] = None
    status: str  # "connected" | "disconnected"


async def _auth_device_token(
    authorization: str | None, db: AsyncSession
) -> Device:
    """Authenticate device by token from Authorization header."""
    if not authorization or not authorization.lower().startswith("bearer "):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing token",
        )

    token = authorization.split(None, 1)[1]

    result = await db.execute(select(Device).where(Device.device_token == token))
    device = result.scalar_one_or_none()

    if not device:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid token",
        )

    return device


@router.post("/tunnel-port")
async def report_tunnel_port(
    req: TunnelPortRequest,
    db: Annotated[AsyncSession, Depends(get_db)],
    authorization: Annotated[str | None, Header()] = None,
):
    """
    Device reports tunnel port allocation after establishing reverse SSH tunnel.

    Args:
        tunnel_type: "ssh" or "dashboard"
        port: Allocated port number (if connected)
        status: "connected" or "disconnected"

    Returns:
        success: True
    """
    device = await _auth_device_token(authorization, db)

    if req.tunnel_type not in ["ssh", "dashboard"]:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="tunnel_type must be 'ssh' or 'dashboard'"
        )

    # Update tunnel status
    tunnel_service.report_device_port(
        device_id=device.mac_address,
        tunnel_type=req.tunnel_type,
        port=req.port,
        status=req.status
    )

    print(f"[tunnel] Device {device.mac_address} reported {req.tunnel_type} "
          f"tunnel {req.status}" + (f" on port {req.port}" if req.port else ""))

    return {"success": True}