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,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()
|
||||
Reference in new issue
Block a user