342 lines
12 KiB
Python
342 lines
12 KiB
Python
from datetime import UTC, datetime
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from sqlalchemy import and_, delete, or_, select
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ..cursors import decode, encode
|
|
from ..db import get_session
|
|
from ..models import Agent as AgentModel
|
|
from ..models import AgentBlacklist, Event, Incident, NotificationOutbox
|
|
from ..schemas import (
|
|
ActionCountResponse,
|
|
Agent,
|
|
AgentBlacklistItem,
|
|
AgentBlacklistPage,
|
|
AgentsPage,
|
|
)
|
|
from ..timeutil import to_utc
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _is_pending(a: AgentModel) -> bool:
|
|
return a.pending_key_hash is not None
|
|
|
|
|
|
def _to_schema(a: AgentModel) -> Agent:
|
|
pending = _is_pending(a)
|
|
hostname = a.pending_hostname if pending and a.pending_hostname is not None else a.hostname
|
|
last_seen_at = (
|
|
a.pending_seen_at if pending and a.pending_seen_at is not None else a.last_seen_at
|
|
)
|
|
features = a.pending_features if pending and a.pending_features is not None else a.features
|
|
labels = a.pending_labels if pending and a.pending_labels is not None else a.labels
|
|
return Agent(
|
|
agent_id=a.agent_id,
|
|
hostname=hostname,
|
|
status=a.status,
|
|
admission_state="pending" if pending or a.accepted_key_hash is None else "accepted",
|
|
rotation_pending=pending and a.accepted_key_hash is not None,
|
|
key_fingerprint=(a.pending_key_fingerprint if pending else a.accepted_key_fingerprint),
|
|
last_seen_at=to_utc(last_seen_at),
|
|
features=features or {"push": False, "metrics": False},
|
|
labels=labels or {},
|
|
)
|
|
|
|
|
|
def _blacklist_schema(row: AgentBlacklist) -> AgentBlacklistItem:
|
|
return AgentBlacklistItem(
|
|
id=row.key_hash,
|
|
key_fingerprint=row.key_fingerprint,
|
|
agent_id=row.agent_id,
|
|
hostname=row.hostname,
|
|
created_at=to_utc(row.created_at),
|
|
last_seen_at=to_utc(row.last_seen_at),
|
|
)
|
|
|
|
|
|
async def _get_agent(session: AsyncSession, agent_id: str) -> AgentModel:
|
|
res = await session.execute(select(AgentModel).where(AgentModel.agent_id == agent_id))
|
|
a = res.scalar_one_or_none()
|
|
if a is None:
|
|
raise HTTPException(status_code=404, detail="agent not found")
|
|
return a
|
|
|
|
|
|
def _clear_pending(a: AgentModel) -> None:
|
|
a.pending_key_hash = None
|
|
a.pending_key_fingerprint = None
|
|
a.pending_hostname = None
|
|
a.pending_seen_at = None
|
|
a.pending_features = None
|
|
a.pending_labels = None
|
|
|
|
|
|
def _accept_pending(a: AgentModel) -> bool:
|
|
if not _is_pending(a):
|
|
return False
|
|
now = datetime.now(UTC)
|
|
a.accepted_key_hash = a.pending_key_hash
|
|
a.accepted_key_fingerprint = a.pending_key_fingerprint
|
|
a.accepted_at = now
|
|
a.hostname = a.pending_hostname or a.hostname
|
|
a.last_seen_at = a.pending_seen_at or a.last_seen_at
|
|
if a.pending_features is not None:
|
|
a.features = a.pending_features
|
|
if a.pending_labels is not None:
|
|
a.labels = a.pending_labels
|
|
a.status = "alive"
|
|
_clear_pending(a)
|
|
return True
|
|
|
|
|
|
async def _block_pending(session: AsyncSession, a: AgentModel) -> bool:
|
|
if not _is_pending(a):
|
|
return False
|
|
now = datetime.now(UTC)
|
|
stmt = pg_insert(AgentBlacklist).values(
|
|
key_hash=a.pending_key_hash,
|
|
key_fingerprint=a.pending_key_fingerprint,
|
|
agent_id=a.agent_id,
|
|
hostname=a.pending_hostname or a.hostname,
|
|
last_seen_at=a.pending_seen_at or now,
|
|
)
|
|
stmt = stmt.on_conflict_do_update(
|
|
index_elements=[AgentBlacklist.key_hash],
|
|
set_={
|
|
"agent_id": stmt.excluded.agent_id,
|
|
"hostname": stmt.excluded.hostname,
|
|
"last_seen_at": stmt.excluded.last_seen_at,
|
|
},
|
|
)
|
|
await session.execute(stmt)
|
|
_clear_pending(a)
|
|
if a.accepted_key_hash is None:
|
|
await _delete_agent_data(session, a.agent_id)
|
|
return True
|
|
|
|
|
|
async def _delete_agent_data(session: AsyncSession, agent_id: str) -> None:
|
|
incident_ids = select(Incident.id).where(Incident.agent_id == agent_id)
|
|
await session.execute(
|
|
delete(NotificationOutbox).where(NotificationOutbox.incident_id.in_(incident_ids))
|
|
)
|
|
await session.execute(delete(Incident).where(Incident.agent_id == agent_id))
|
|
await session.execute(delete(Event).where(Event.agent_id == agent_id))
|
|
await session.execute(delete(AgentModel).where(AgentModel.agent_id == agent_id))
|
|
|
|
|
|
async def _remove_pending_or_agent(session: AsyncSession, a: AgentModel) -> None:
|
|
if _is_pending(a) and a.accepted_key_hash is not None:
|
|
_clear_pending(a)
|
|
return
|
|
await _delete_agent_data(session, a.agent_id)
|
|
|
|
|
|
def _agent_rank(a: AgentModel) -> int:
|
|
return a.admission_rank
|
|
|
|
|
|
def _decode_agent_cursor(cursor: str) -> tuple[int, str]:
|
|
rank_raw, agent_id = decode(cursor, 2)
|
|
try:
|
|
rank = int(rank_raw)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=400, detail="invalid cursor") from exc
|
|
if rank not in (0, 1):
|
|
raise HTTPException(status_code=400, detail="invalid cursor")
|
|
return rank, agent_id
|
|
|
|
|
|
@router.get("/agents", response_model=AgentsPage)
|
|
async def list_agents(
|
|
cursor: str | None = Query(default=None, max_length=512),
|
|
back: bool = Query(default=False),
|
|
limit: int = Query(default=100, ge=1, le=500),
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> AgentsPage:
|
|
if back:
|
|
stmt = select(AgentModel).order_by(
|
|
AgentModel.admission_rank.desc(), AgentModel.agent_id.desc()
|
|
)
|
|
else:
|
|
stmt = select(AgentModel).order_by(
|
|
AgentModel.admission_rank.asc(), AgentModel.agent_id.asc()
|
|
)
|
|
if cursor is not None:
|
|
last_rank, last_id = _decode_agent_cursor(cursor)
|
|
if back:
|
|
stmt = stmt.where(
|
|
or_(
|
|
AgentModel.admission_rank < last_rank,
|
|
and_(
|
|
AgentModel.admission_rank == last_rank,
|
|
AgentModel.agent_id < last_id,
|
|
),
|
|
)
|
|
)
|
|
else:
|
|
stmt = stmt.where(
|
|
or_(
|
|
AgentModel.admission_rank > last_rank,
|
|
and_(
|
|
AgentModel.admission_rank == last_rank,
|
|
AgentModel.agent_id > last_id,
|
|
),
|
|
)
|
|
)
|
|
rows = (await session.execute(stmt.limit(limit + 1))).scalars().all()
|
|
has_more = len(rows) > limit
|
|
rows = rows[:limit]
|
|
if back:
|
|
rows = list(reversed(rows))
|
|
next_cursor: str | None = None
|
|
prev_cursor: str | None = None
|
|
if rows:
|
|
if back:
|
|
next_cursor = encode(_agent_rank(rows[-1]), rows[-1].agent_id)
|
|
if has_more:
|
|
prev_cursor = encode(_agent_rank(rows[0]), rows[0].agent_id)
|
|
else:
|
|
if has_more:
|
|
next_cursor = encode(_agent_rank(rows[-1]), rows[-1].agent_id)
|
|
if cursor is not None:
|
|
prev_cursor = encode(_agent_rank(rows[0]), rows[0].agent_id)
|
|
return AgentsPage(
|
|
items=[_to_schema(a) for a in rows],
|
|
next_cursor=next_cursor,
|
|
prev_cursor=prev_cursor,
|
|
)
|
|
|
|
|
|
@router.get("/agents/admission/blacklist", response_model=AgentBlacklistPage)
|
|
async def list_agent_blacklist(
|
|
cursor: str | None = Query(default=None, max_length=512),
|
|
limit: int = Query(default=100, ge=1, le=500),
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> AgentBlacklistPage:
|
|
stmt = select(AgentBlacklist).order_by(AgentBlacklist.key_fingerprint, AgentBlacklist.key_hash)
|
|
if cursor is not None:
|
|
last_fingerprint, last_hash = decode(cursor, 2)
|
|
stmt = stmt.where(
|
|
or_(
|
|
AgentBlacklist.key_fingerprint > last_fingerprint,
|
|
and_(
|
|
AgentBlacklist.key_fingerprint == last_fingerprint,
|
|
AgentBlacklist.key_hash > last_hash,
|
|
),
|
|
)
|
|
)
|
|
rows = (await session.execute(stmt.limit(limit + 1))).scalars().all()
|
|
has_more = len(rows) > limit
|
|
rows = rows[:limit]
|
|
return AgentBlacklistPage(
|
|
items=[_blacklist_schema(r) for r in rows],
|
|
next_cursor=encode(rows[-1].key_fingerprint, rows[-1].key_hash)
|
|
if rows and has_more
|
|
else None,
|
|
prev_cursor=None,
|
|
)
|
|
|
|
|
|
@router.delete("/agents/admission/blacklist", response_model=ActionCountResponse)
|
|
async def clear_agent_blacklist(
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> ActionCountResponse:
|
|
res = await session.execute(delete(AgentBlacklist))
|
|
await session.commit()
|
|
return ActionCountResponse(count=res.rowcount or 0)
|
|
|
|
|
|
@router.delete("/agents/admission/blacklist/{key_id}", response_model=ActionCountResponse)
|
|
async def delete_agent_blacklist_item(
|
|
key_id: str,
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> ActionCountResponse:
|
|
res = await session.execute(delete(AgentBlacklist).where(AgentBlacklist.key_hash == key_id))
|
|
await session.commit()
|
|
return ActionCountResponse(count=res.rowcount or 0)
|
|
|
|
|
|
@router.post("/agents/admission/accept-all", response_model=ActionCountResponse)
|
|
async def accept_all_pending_agents(
|
|
include_rotation: bool = Query(default=False),
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> ActionCountResponse:
|
|
stmt = select(AgentModel).where(AgentModel.pending_key_hash.is_not(None))
|
|
if not include_rotation:
|
|
stmt = stmt.where(AgentModel.accepted_key_hash.is_(None))
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
count = sum(1 for a in rows if _accept_pending(a))
|
|
await session.commit()
|
|
return ActionCountResponse(count=count)
|
|
|
|
|
|
@router.post("/agents/admission/block-all", response_model=ActionCountResponse)
|
|
async def block_all_pending_agents(
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> ActionCountResponse:
|
|
rows = (
|
|
(await session.execute(select(AgentModel).where(AgentModel.pending_key_hash.is_not(None))))
|
|
.scalars()
|
|
.all()
|
|
)
|
|
count = 0
|
|
for a in rows:
|
|
if await _block_pending(session, a):
|
|
count += 1
|
|
await session.commit()
|
|
return ActionCountResponse(count=count)
|
|
|
|
|
|
@router.post("/agents/admission/remove-all", response_model=ActionCountResponse)
|
|
async def remove_all_pending_agents(
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> ActionCountResponse:
|
|
rows = (
|
|
(await session.execute(select(AgentModel).where(AgentModel.pending_key_hash.is_not(None))))
|
|
.scalars()
|
|
.all()
|
|
)
|
|
for agent in rows:
|
|
await _remove_pending_or_agent(session, agent)
|
|
await session.commit()
|
|
return ActionCountResponse(count=len(rows))
|
|
|
|
|
|
@router.post("/agents/{agent_id}/accept", response_model=ActionCountResponse)
|
|
async def accept_agent(
|
|
agent_id: str, session: AsyncSession = Depends(get_session)
|
|
) -> ActionCountResponse:
|
|
agent = await _get_agent(session, agent_id)
|
|
count = 1 if _accept_pending(agent) else 0
|
|
await session.commit()
|
|
return ActionCountResponse(count=count)
|
|
|
|
|
|
@router.post("/agents/{agent_id}/block", response_model=ActionCountResponse)
|
|
async def block_agent(
|
|
agent_id: str, session: AsyncSession = Depends(get_session)
|
|
) -> ActionCountResponse:
|
|
agent = await _get_agent(session, agent_id)
|
|
count = 1 if await _block_pending(session, agent) else 0
|
|
await session.commit()
|
|
return ActionCountResponse(count=count)
|
|
|
|
|
|
@router.delete("/agents/{agent_id}", response_model=ActionCountResponse)
|
|
async def delete_agent(
|
|
agent_id: str, session: AsyncSession = Depends(get_session)
|
|
) -> ActionCountResponse:
|
|
agent = await _get_agent(session, agent_id)
|
|
await _remove_pending_or_agent(session, agent)
|
|
await session.commit()
|
|
return ActionCountResponse(count=1)
|
|
|
|
|
|
@router.get("/agents/{agent_id}", response_model=Agent)
|
|
async def get_agent(agent_id: str, session: AsyncSession = Depends(get_session)) -> Agent:
|
|
return _to_schema(await _get_agent(session, agent_id))
|