"""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 = 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")