507 lines
17 KiB
Python
507 lines
17 KiB
Python
"""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()
|