166 lines
4.4 KiB
Python
166 lines
4.4 KiB
Python
"""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")
|