"""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.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 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 Returns agent_id and token for subsequent API calls. """ 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=None, # Agents register without tenant initially ) 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")