letsbe-orchestrator/app/routes/agents.py

"""Agent management endpoints."""

import secrets
import uuid

from fastapi import APIRouter, Header, HTTPException, status
from sqlalchemy import select

from app.db import AsyncSessionDep
from app.models.agent import Agent, AgentStatus
from app.models.base import utc_now
from app.models.tenant import Tenant
from app.schemas.agent import (
    AgentHeartbeatResponse,
    AgentRegisterRequest,
    AgentRegisterResponse,
)

router = APIRouter(prefix="/agents", tags=["Agents"])


# --- Helper functions (embryonic service layer) ---


async def get_agent_by_id(db: AsyncSessionDep, agent_id: uuid.UUID) -> Agent | None:
    """Retrieve an agent by ID."""
    result = await db.execute(select(Agent).where(Agent.id == agent_id))
    return result.scalar_one_or_none()


async def get_tenant_by_id(db: AsyncSessionDep, tenant_id: uuid.UUID) -> Tenant | None:
    """Retrieve a tenant by ID."""
    result = await db.execute(select(Tenant).where(Tenant.id == tenant_id))
    return result.scalar_one_or_none()


async def validate_agent_token(
    db: AsyncSessionDep,
    agent_id: uuid.UUID,
    authorization: str | None,
) -> Agent:
    """
    Validate agent exists and token matches.

    Args:
        db: Database session
        agent_id: Agent UUID
        authorization: Authorization header value

    Returns:
        Agent if valid

    Raises:
        HTTPException: 401 if invalid
    """
    if authorization is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing Authorization header",
        )

    # Parse Bearer token
    parts = authorization.split(" ", 1)
    if len(parts) != 2 or parts[0].lower() != "bearer":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid Authorization header format. Expected: Bearer <token>",
        )

    token = parts[1]

    # Find and validate agent
    agent = await get_agent_by_id(db, agent_id)
    if agent is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid agent credentials",
        )

    # Use secrets.compare_digest for timing-attack-safe comparison
    if not secrets.compare_digest(agent.token, token):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid agent credentials",
        )

    return agent


# --- Route handlers (thin controllers) ---


@router.post(
    "/register",
    response_model=AgentRegisterResponse,
    status_code=status.HTTP_201_CREATED,
)
async def register_agent(
    request: AgentRegisterRequest,
    db: AsyncSessionDep,
) -> AgentRegisterResponse:
    """
    Register a new SysAdmin agent.

    - **hostname**: Agent hostname (will be used as name)
    - **version**: Agent software version
    - **metadata**: Optional JSON metadata
    - **tenant_id**: Optional tenant UUID to associate the agent with

    Returns agent_id and token for subsequent API calls.

    If tenant_id is provided but invalid, returns 404 Not Found.
    """
    # Validate tenant exists if provided
    if request.tenant_id is not None:
        tenant = await get_tenant_by_id(db, request.tenant_id)
        if tenant is None:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"Tenant {request.tenant_id} not found",
            )

    agent_id = uuid.uuid4()
    token = secrets.token_hex(32)

    agent = Agent(
        id=agent_id,
        name=request.hostname,
        version=request.version,
        status=AgentStatus.ONLINE.value,
        last_heartbeat=utc_now(),
        token=token,
        tenant_id=request.tenant_id,
    )

    db.add(agent)
    await db.commit()

    return AgentRegisterResponse(agent_id=agent_id, token=token)


@router.post(
    "/{agent_id}/heartbeat",
    response_model=AgentHeartbeatResponse,
)
async def agent_heartbeat(
    agent_id: uuid.UUID,
    db: AsyncSessionDep,
    authorization: str | None = Header(None),
) -> AgentHeartbeatResponse:
    """
    Send heartbeat from agent.

    Updates last_heartbeat timestamp and sets status to online.
    Requires Bearer token authentication.
    """
    agent = await validate_agent_token(db, agent_id, authorization)

    # Update heartbeat
    agent.last_heartbeat = utc_now()
    agent.status = AgentStatus.ONLINE.value

    await db.commit()

    return AgentHeartbeatResponse(status="ok")