feat: add secure agent enrollment
This commit is contained in:
lussierj committed
1 parent
82309db79a
commit
d9f0da694b
26 files changed
+2942
-76
No files matched your search
@@ -0,0 +1 @@
|
||||
"""Secure one-time Agent enrollment."""
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Dependency wiring for secure Agent enrollment."""
|
||||
|
||||
from app.agent_enrollment.repository import AgentEnrollmentRepository
|
||||
from app.agent_enrollment.service import AgentEnrollmentService
|
||||
|
||||
|
||||
def get_agent_enrollment_service() -> AgentEnrollmentService:
|
||||
return AgentEnrollmentService(AgentEnrollmentRepository())
|
||||
@@ -0,0 +1,21 @@
|
||||
"""Domain exceptions for secure Agent enrollment."""
|
||||
|
||||
|
||||
class AgentEnrollmentError(Exception):
|
||||
"""Base Agent enrollment error."""
|
||||
|
||||
|
||||
class AgentEnrollmentNotFound(AgentEnrollmentError):
|
||||
"""The requested Site or Agent Zone does not exist."""
|
||||
|
||||
|
||||
class AgentEnrollmentConflict(AgentEnrollmentError):
|
||||
"""The requested enrollment conflicts with current Agent routing state."""
|
||||
|
||||
|
||||
class AgentEnrollmentTokenInvalid(AgentEnrollmentError):
|
||||
"""The supplied bootstrap token cannot be redeemed."""
|
||||
|
||||
|
||||
class AgentEnrollmentCredentialCollision(AgentEnrollmentError):
|
||||
"""A generated machine credential collided with an existing credential."""
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Typed contracts for secure Agent enrollment."""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Literal
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
|
||||
class AgentEnrollmentCreateRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
site_id: UUID
|
||||
agent_zone_id: UUID
|
||||
display_name: str = Field(min_length=1, max_length=120)
|
||||
|
||||
@field_validator("display_name")
|
||||
@classmethod
|
||||
def normalize_display_name(cls, value: str) -> str:
|
||||
normalized = value.strip()
|
||||
if not normalized:
|
||||
raise ValueError("display_name cannot be blank")
|
||||
return normalized
|
||||
|
||||
|
||||
class AgentEnrollmentCreated(BaseModel):
|
||||
id: UUID
|
||||
site_id: UUID
|
||||
site_name: str
|
||||
agent_zone_id: UUID
|
||||
agent_zone_name: str
|
||||
display_name: str
|
||||
created_at: datetime
|
||||
expires_at: datetime
|
||||
enrollment_token: str
|
||||
|
||||
|
||||
class AgentEnrollmentCreateResponse(BaseModel):
|
||||
data: AgentEnrollmentCreated
|
||||
meta: dict[str, object] = Field(default_factory=dict)
|
||||
error: None = None
|
||||
|
||||
|
||||
class AgentEnrollmentSummary(BaseModel):
|
||||
id: UUID
|
||||
site_id: UUID
|
||||
site_name: str
|
||||
agent_zone_id: UUID
|
||||
agent_zone_name: str
|
||||
display_name: str
|
||||
created_by_user_id: UUID | None = None
|
||||
created_at: datetime
|
||||
expires_at: datetime
|
||||
used_at: datetime | None = None
|
||||
revoked_at: datetime | None = None
|
||||
enrolled_agent_id: str | None = None
|
||||
status: Literal["waiting", "used", "revoked", "expired"]
|
||||
|
||||
|
||||
class AgentEnrollmentListResponse(BaseModel):
|
||||
data: list[AgentEnrollmentSummary]
|
||||
meta: dict[str, object] = Field(default_factory=dict)
|
||||
error: None = None
|
||||
|
||||
|
||||
class AgentEnrollRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
token: str = Field(min_length=32, max_length=256)
|
||||
|
||||
@field_validator("token")
|
||||
@classmethod
|
||||
def normalize_token(cls, value: str) -> str:
|
||||
normalized = value.strip()
|
||||
if not normalized:
|
||||
raise ValueError("token cannot be blank")
|
||||
return normalized
|
||||
|
||||
|
||||
class AgentEnrollCredentials(BaseModel):
|
||||
agent_id: str
|
||||
agent_secret: str
|
||||
site_id: UUID
|
||||
site_name: str
|
||||
agent_zone_id: UUID
|
||||
agent_zone_name: str
|
||||
display_name: str
|
||||
ip_ranges: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentEnrollResponse(BaseModel):
|
||||
data: AgentEnrollCredentials
|
||||
meta: dict[str, object] = Field(default_factory=dict)
|
||||
error: None = None
|
||||
@@ -0,0 +1,506 @@
|
||||
"""PostgreSQL persistence for secure Agent enrollment."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from uuid import UUID
|
||||
|
||||
import psycopg
|
||||
from psycopg.errors import UniqueViolation
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
from app.agent_enrollment.exceptions import (
|
||||
AgentEnrollmentConflict,
|
||||
AgentEnrollmentCredentialCollision,
|
||||
AgentEnrollmentNotFound,
|
||||
AgentEnrollmentTokenInvalid,
|
||||
)
|
||||
from app.config import settings
|
||||
|
||||
|
||||
ConnectionFactory = Callable[[], psycopg.Connection]
|
||||
|
||||
TOKEN_HASH_CONSTRAINT = "uq_agent_enrollments_token_hash"
|
||||
OPEN_ZONE_CONSTRAINT = "uq_agent_enrollments_open_zone"
|
||||
AGENT_SECRET_HASH_CONSTRAINT = "uq_minerops_site_agents_agent_secret_hash"
|
||||
AGENT_ZONE_CONSTRAINT = "uq_minerops_site_agents_agent_zone_id"
|
||||
AGENT_ID_CONSTRAINT = "minerops_site_agents_pkey"
|
||||
|
||||
|
||||
def default_connection_factory() -> psycopg.Connection:
|
||||
return psycopg.connect(
|
||||
host=settings.postgres_host,
|
||||
port=settings.postgres_port,
|
||||
dbname=settings.postgres_db,
|
||||
user=settings.postgres_user,
|
||||
password=settings.postgres_password,
|
||||
connect_timeout=5,
|
||||
)
|
||||
|
||||
|
||||
class AgentEnrollmentRepository:
|
||||
def __init__(self, connection_factory: ConnectionFactory | None = None):
|
||||
self._connection_factory = connection_factory or default_connection_factory
|
||||
|
||||
def create_enrollment(
|
||||
self,
|
||||
*,
|
||||
token_hash: str,
|
||||
site_id: UUID,
|
||||
agent_zone_id: UUID,
|
||||
display_name: str,
|
||||
actor_user_id: UUID,
|
||||
) -> dict:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT zone.id,
|
||||
zone.site_id,
|
||||
zone.name AS agent_zone_name,
|
||||
zone.is_enabled,
|
||||
site.name AS site_name,
|
||||
agent.agent_id AS assigned_agent_id
|
||||
FROM agent_zones AS zone
|
||||
JOIN sites AS site ON site.id = zone.site_id
|
||||
LEFT JOIN minerops_site_agents AS agent
|
||||
ON agent.agent_zone_id = zone.id
|
||||
WHERE zone.id = %s
|
||||
AND zone.site_id = %s
|
||||
FOR UPDATE OF zone
|
||||
""",
|
||||
(agent_zone_id, site_id),
|
||||
)
|
||||
zone = cur.fetchone()
|
||||
|
||||
if not zone:
|
||||
raise AgentEnrollmentNotFound(
|
||||
"Agent Zone was not found for the selected Site"
|
||||
)
|
||||
|
||||
if zone["is_enabled"] is not True:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Agent Zone is disabled"
|
||||
)
|
||||
|
||||
if zone["assigned_agent_id"]:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Agent Zone already has an assigned Agent"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"SELECT clock_timestamp() AS operation_now"
|
||||
)
|
||||
operation_now = cur.fetchone()["operation_now"]
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE agent_enrollments
|
||||
SET revoked_at = %s
|
||||
WHERE agent_zone_id = %s
|
||||
AND used_at IS NULL
|
||||
AND revoked_at IS NULL
|
||||
RETURNING id
|
||||
""",
|
||||
(operation_now, agent_zone_id),
|
||||
)
|
||||
replaced_ids = [
|
||||
str(row["id"]) for row in cur.fetchall()
|
||||
]
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_enrollments (
|
||||
token_hash,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
display_name,
|
||||
created_by_user_id,
|
||||
created_at,
|
||||
expires_at
|
||||
)
|
||||
VALUES (
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s + interval '30 minutes'
|
||||
)
|
||||
RETURNING
|
||||
id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
display_name,
|
||||
created_at,
|
||||
expires_at
|
||||
""",
|
||||
(
|
||||
token_hash,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
display_name,
|
||||
actor_user_id,
|
||||
operation_now,
|
||||
operation_now,
|
||||
),
|
||||
)
|
||||
enrollment = dict(cur.fetchone())
|
||||
|
||||
enrollment["site_name"] = zone["site_name"]
|
||||
enrollment["agent_zone_name"] = zone["agent_zone_name"]
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO audit_logs (
|
||||
actor_user_id,
|
||||
action,
|
||||
entity_type,
|
||||
entity_id,
|
||||
new_value_json,
|
||||
created_at
|
||||
)
|
||||
VALUES (
|
||||
%s,
|
||||
'agent.enrollment.created',
|
||||
'agent_enrollment',
|
||||
%s,
|
||||
%s,
|
||||
%s
|
||||
)
|
||||
""",
|
||||
(
|
||||
actor_user_id,
|
||||
enrollment["id"],
|
||||
Jsonb(
|
||||
{
|
||||
"site_id": str(site_id),
|
||||
"agent_zone_id": str(agent_zone_id),
|
||||
"display_name": display_name,
|
||||
"expires_at": enrollment[
|
||||
"expires_at"
|
||||
].isoformat(),
|
||||
"replaced_enrollment_ids": replaced_ids,
|
||||
}
|
||||
),
|
||||
operation_now,
|
||||
),
|
||||
)
|
||||
|
||||
connection.commit()
|
||||
return enrollment
|
||||
|
||||
except UniqueViolation as exc:
|
||||
connection.rollback()
|
||||
constraint = exc.diag.constraint_name
|
||||
|
||||
if constraint == TOKEN_HASH_CONSTRAINT:
|
||||
raise AgentEnrollmentCredentialCollision(
|
||||
"Generated enrollment token collided"
|
||||
) from exc
|
||||
|
||||
if constraint == OPEN_ZONE_CONSTRAINT:
|
||||
raise AgentEnrollmentConflict(
|
||||
"An active enrollment already exists for this Agent Zone"
|
||||
) from exc
|
||||
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def list_enrollments(
|
||||
self,
|
||||
*,
|
||||
site_id: UUID | None = None,
|
||||
agent_zone_id: UUID | None = None,
|
||||
limit: int = 100,
|
||||
) -> list[dict]:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
enrollment.id,
|
||||
enrollment.site_id,
|
||||
site.name AS site_name,
|
||||
enrollment.agent_zone_id,
|
||||
zone.name AS agent_zone_name,
|
||||
enrollment.display_name,
|
||||
enrollment.created_by_user_id,
|
||||
enrollment.created_at,
|
||||
enrollment.expires_at,
|
||||
enrollment.used_at,
|
||||
enrollment.revoked_at,
|
||||
enrollment.enrolled_agent_id,
|
||||
CASE
|
||||
WHEN enrollment.used_at IS NOT NULL
|
||||
THEN 'used'
|
||||
WHEN enrollment.revoked_at IS NOT NULL
|
||||
THEN 'revoked'
|
||||
WHEN enrollment.expires_at <= now()
|
||||
THEN 'expired'
|
||||
ELSE 'waiting'
|
||||
END AS status
|
||||
FROM agent_enrollments AS enrollment
|
||||
JOIN sites AS site
|
||||
ON site.id = enrollment.site_id
|
||||
JOIN agent_zones AS zone
|
||||
ON zone.id = enrollment.agent_zone_id
|
||||
WHERE (%s::uuid IS NULL OR enrollment.site_id = %s::uuid)
|
||||
AND (
|
||||
%s::uuid IS NULL
|
||||
OR enrollment.agent_zone_id = %s::uuid
|
||||
)
|
||||
ORDER BY enrollment.created_at DESC, enrollment.id DESC
|
||||
LIMIT %s
|
||||
""",
|
||||
(
|
||||
site_id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
agent_zone_id,
|
||||
limit,
|
||||
),
|
||||
)
|
||||
return [dict(row) for row in cur.fetchall()]
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def redeem_enrollment(
|
||||
self,
|
||||
*,
|
||||
token_hash: str,
|
||||
agent_id: str,
|
||||
agent_secret_hash: str,
|
||||
) -> dict:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
display_name,
|
||||
expires_at,
|
||||
used_at,
|
||||
revoked_at
|
||||
FROM agent_enrollments
|
||||
WHERE token_hash = %s
|
||||
FOR UPDATE
|
||||
""",
|
||||
(token_hash,),
|
||||
)
|
||||
enrollment = cur.fetchone()
|
||||
|
||||
if not enrollment:
|
||||
raise AgentEnrollmentTokenInvalid(
|
||||
"Invalid or expired enrollment token"
|
||||
)
|
||||
|
||||
if (
|
||||
enrollment["used_at"] is not None
|
||||
or enrollment["revoked_at"] is not None
|
||||
):
|
||||
raise AgentEnrollmentTokenInvalid(
|
||||
"Invalid or expired enrollment token"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
zone.id,
|
||||
zone.site_id,
|
||||
zone.name AS agent_zone_name,
|
||||
zone.is_enabled,
|
||||
site.name AS site_name
|
||||
FROM agent_zones AS zone
|
||||
JOIN sites AS site ON site.id = zone.site_id
|
||||
WHERE zone.id = %s
|
||||
AND zone.site_id = %s
|
||||
FOR UPDATE OF zone
|
||||
""",
|
||||
(
|
||||
enrollment["agent_zone_id"],
|
||||
enrollment["site_id"],
|
||||
),
|
||||
)
|
||||
zone = cur.fetchone()
|
||||
|
||||
if not zone:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Enrollment Agent Zone no longer exists"
|
||||
)
|
||||
|
||||
if zone["is_enabled"] is not True:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Agent Zone is disabled"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT agent_id
|
||||
FROM minerops_site_agents
|
||||
WHERE agent_zone_id = %s
|
||||
FOR UPDATE
|
||||
""",
|
||||
(enrollment["agent_zone_id"],),
|
||||
)
|
||||
existing_agent = cur.fetchone()
|
||||
|
||||
if existing_agent:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Agent Zone already has an assigned Agent"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"SELECT clock_timestamp() AS operation_now"
|
||||
)
|
||||
operation_now = cur.fetchone()["operation_now"]
|
||||
|
||||
if enrollment["expires_at"] <= operation_now:
|
||||
raise AgentEnrollmentTokenInvalid(
|
||||
"Invalid or expired enrollment token"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO minerops_site_agents (
|
||||
agent_id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
site_name,
|
||||
display_name,
|
||||
agent_secret_hash,
|
||||
credential_created_at
|
||||
)
|
||||
VALUES (
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s
|
||||
)
|
||||
""",
|
||||
(
|
||||
agent_id,
|
||||
enrollment["site_id"],
|
||||
enrollment["agent_zone_id"],
|
||||
zone["site_name"],
|
||||
enrollment["display_name"],
|
||||
agent_secret_hash,
|
||||
operation_now,
|
||||
),
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE agent_enrollments
|
||||
SET used_at = %s,
|
||||
enrolled_agent_id = %s
|
||||
WHERE id = %s
|
||||
RETURNING used_at
|
||||
""",
|
||||
(operation_now, agent_id, enrollment["id"]),
|
||||
)
|
||||
used_at = cur.fetchone()["used_at"]
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT network::text AS network
|
||||
FROM agent_zone_network_ranges
|
||||
WHERE agent_zone_id = %s
|
||||
ORDER BY network, id
|
||||
""",
|
||||
(enrollment["agent_zone_id"],),
|
||||
)
|
||||
networks = [
|
||||
row["network"] for row in cur.fetchall()
|
||||
]
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO audit_logs (
|
||||
actor_user_id,
|
||||
action,
|
||||
entity_type,
|
||||
entity_id,
|
||||
new_value_json,
|
||||
created_at
|
||||
)
|
||||
VALUES (
|
||||
NULL,
|
||||
'agent.enrollment.redeemed',
|
||||
'agent_enrollment',
|
||||
%s,
|
||||
%s,
|
||||
%s
|
||||
)
|
||||
""",
|
||||
(
|
||||
enrollment["id"],
|
||||
Jsonb(
|
||||
{
|
||||
"agent_id": agent_id,
|
||||
"site_id": str(enrollment["site_id"]),
|
||||
"agent_zone_id": str(
|
||||
enrollment["agent_zone_id"]
|
||||
),
|
||||
"display_name": enrollment[
|
||||
"display_name"
|
||||
],
|
||||
"used_at": used_at.isoformat(),
|
||||
}
|
||||
),
|
||||
operation_now,
|
||||
),
|
||||
)
|
||||
|
||||
result = {
|
||||
"agent_id": agent_id,
|
||||
"site_id": enrollment["site_id"],
|
||||
"site_name": zone["site_name"],
|
||||
"agent_zone_id": enrollment["agent_zone_id"],
|
||||
"agent_zone_name": zone["agent_zone_name"],
|
||||
"display_name": enrollment["display_name"],
|
||||
"ip_ranges": networks,
|
||||
}
|
||||
|
||||
connection.commit()
|
||||
return result
|
||||
|
||||
except UniqueViolation as exc:
|
||||
connection.rollback()
|
||||
constraint = exc.diag.constraint_name
|
||||
|
||||
if constraint in {
|
||||
AGENT_ID_CONSTRAINT,
|
||||
AGENT_SECRET_HASH_CONSTRAINT,
|
||||
}:
|
||||
raise AgentEnrollmentCredentialCollision(
|
||||
"Generated Agent credential collided"
|
||||
) from exc
|
||||
|
||||
if constraint == AGENT_ZONE_CONSTRAINT:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Agent Zone already has an assigned Agent"
|
||||
) from exc
|
||||
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
|
||||
finally:
|
||||
connection.close()
|
||||
@@ -0,0 +1,108 @@
|
||||
"""HTTP routes for Agent enrollment and bootstrap."""
|
||||
|
||||
from typing import Annotated
|
||||
from uuid import UUID
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, status
|
||||
|
||||
from app.agent_enrollment.dependencies import get_agent_enrollment_service
|
||||
from app.agent_enrollment.exceptions import (
|
||||
AgentEnrollmentConflict,
|
||||
AgentEnrollmentNotFound,
|
||||
AgentEnrollmentTokenInvalid,
|
||||
)
|
||||
from app.agent_enrollment.models import (
|
||||
AgentEnrollmentCreateRequest,
|
||||
AgentEnrollmentCreateResponse,
|
||||
AgentEnrollmentListResponse,
|
||||
AgentEnrollRequest,
|
||||
AgentEnrollResponse,
|
||||
)
|
||||
from app.agent_enrollment.service import AgentEnrollmentService
|
||||
from app.identity import require_permission
|
||||
|
||||
|
||||
router = APIRouter(tags=["agent-enrollment"])
|
||||
|
||||
EnrollmentService = Annotated[
|
||||
AgentEnrollmentService,
|
||||
Depends(get_agent_enrollment_service),
|
||||
]
|
||||
EnrollmentViewer = Annotated[
|
||||
dict,
|
||||
Depends(require_permission("agents.view")),
|
||||
]
|
||||
EnrollmentManager = Annotated[
|
||||
dict,
|
||||
Depends(require_permission("agents.manage")),
|
||||
]
|
||||
|
||||
|
||||
@router.post(
|
||||
"/api/agent-enrollments",
|
||||
response_model=AgentEnrollmentCreateResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
def create_agent_enrollment(
|
||||
request: AgentEnrollmentCreateRequest,
|
||||
response: Response,
|
||||
current_user: EnrollmentManager,
|
||||
service: EnrollmentService,
|
||||
) -> AgentEnrollmentCreateResponse:
|
||||
response.headers["Cache-Control"] = "no-store"
|
||||
|
||||
try:
|
||||
enrollment = service.create_enrollment(
|
||||
request,
|
||||
actor_user_id=UUID(str(current_user["id"])),
|
||||
)
|
||||
return AgentEnrollmentCreateResponse(data=enrollment)
|
||||
except AgentEnrollmentNotFound as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
except AgentEnrollmentConflict as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
|
||||
|
||||
@router.get(
|
||||
"/api/agent-enrollments",
|
||||
response_model=AgentEnrollmentListResponse,
|
||||
)
|
||||
def list_agent_enrollments(
|
||||
current_user: EnrollmentViewer,
|
||||
service: EnrollmentService,
|
||||
site_id: UUID | None = Query(default=None),
|
||||
agent_zone_id: UUID | None = Query(default=None),
|
||||
limit: int = Query(default=100, ge=1, le=500),
|
||||
) -> AgentEnrollmentListResponse:
|
||||
del current_user
|
||||
|
||||
return AgentEnrollmentListResponse(
|
||||
data=service.list_enrollments(
|
||||
site_id=site_id,
|
||||
agent_zone_id=agent_zone_id,
|
||||
limit=limit,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/api/agent/enroll",
|
||||
response_model=AgentEnrollResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
def enroll_agent(
|
||||
request: AgentEnrollRequest,
|
||||
response: Response,
|
||||
service: EnrollmentService,
|
||||
) -> AgentEnrollResponse:
|
||||
response.headers["Cache-Control"] = "no-store"
|
||||
|
||||
try:
|
||||
return AgentEnrollResponse(data=service.redeem(request))
|
||||
except AgentEnrollmentTokenInvalid as exc:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid or expired enrollment token",
|
||||
) from exc
|
||||
except AgentEnrollmentConflict as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Application service for secure one-time Agent enrollment."""
|
||||
|
||||
import secrets
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from app.agent_auth.service import hash_agent_secret
|
||||
from app.agent_enrollment.exceptions import (
|
||||
AgentEnrollmentConflict,
|
||||
AgentEnrollmentCredentialCollision,
|
||||
)
|
||||
from app.agent_enrollment.models import (
|
||||
AgentEnrollmentCreateRequest,
|
||||
AgentEnrollRequest,
|
||||
)
|
||||
|
||||
|
||||
MAX_GENERATION_ATTEMPTS = 3
|
||||
|
||||
|
||||
def generate_enrollment_token() -> str:
|
||||
return secrets.token_urlsafe(32)
|
||||
|
||||
|
||||
def generate_agent_secret() -> str:
|
||||
return secrets.token_urlsafe(32)
|
||||
|
||||
|
||||
def generate_agent_id() -> str:
|
||||
return f"agent-{uuid4().hex}"
|
||||
|
||||
|
||||
class AgentEnrollmentService:
|
||||
def __init__(self, repository):
|
||||
self.repository = repository
|
||||
|
||||
def create_enrollment(
|
||||
self,
|
||||
request: AgentEnrollmentCreateRequest,
|
||||
*,
|
||||
actor_user_id: UUID,
|
||||
) -> dict:
|
||||
for _ in range(MAX_GENERATION_ATTEMPTS):
|
||||
token = generate_enrollment_token()
|
||||
try:
|
||||
enrollment = self.repository.create_enrollment(
|
||||
token_hash=hash_agent_secret(token),
|
||||
site_id=request.site_id,
|
||||
agent_zone_id=request.agent_zone_id,
|
||||
display_name=request.display_name,
|
||||
actor_user_id=actor_user_id,
|
||||
)
|
||||
return {
|
||||
**enrollment,
|
||||
"enrollment_token": token,
|
||||
}
|
||||
except AgentEnrollmentCredentialCollision:
|
||||
continue
|
||||
|
||||
raise AgentEnrollmentConflict(
|
||||
"Unable to generate a unique enrollment token"
|
||||
)
|
||||
|
||||
def list_enrollments(
|
||||
self,
|
||||
*,
|
||||
site_id: UUID | None = None,
|
||||
agent_zone_id: UUID | None = None,
|
||||
limit: int = 100,
|
||||
) -> list[dict]:
|
||||
return self.repository.list_enrollments(
|
||||
site_id=site_id,
|
||||
agent_zone_id=agent_zone_id,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
def redeem(self, request: AgentEnrollRequest) -> dict:
|
||||
token_hash = hash_agent_secret(request.token)
|
||||
|
||||
for _ in range(MAX_GENERATION_ATTEMPTS):
|
||||
agent_id = generate_agent_id()
|
||||
agent_secret = generate_agent_secret()
|
||||
|
||||
try:
|
||||
result = self.repository.redeem_enrollment(
|
||||
token_hash=token_hash,
|
||||
agent_id=agent_id,
|
||||
agent_secret_hash=hash_agent_secret(agent_secret),
|
||||
)
|
||||
return {
|
||||
**result,
|
||||
"agent_secret": agent_secret,
|
||||
}
|
||||
except AgentEnrollmentCredentialCollision:
|
||||
continue
|
||||
|
||||
raise AgentEnrollmentConflict(
|
||||
"Unable to generate unique Agent credentials"
|
||||
)
|
||||
Reference in new issue
Block a user