"""PostgreSQL persistence for secure Agent enrollment.""" import hmac 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" OPEN_TARGET_CONSTRAINT = "uq_agent_enrollments_open_target_agent" 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 create_credential_transition( self, *, token_hash: str, target_agent_id: str, actor_user_id: UUID, ) -> dict: connection = self._connection_factory() try: with connection.cursor(row_factory=dict_row) as cur: cur.execute( """ SELECT agent.agent_id, agent.site_id, agent.agent_zone_id, agent.display_name, agent.agent_secret_hash, agent.credential_created_at, agent.credential_revoked_at, agent.legacy_auth_disabled_at, site.name AS site_name, zone.name AS agent_zone_name, zone.is_enabled AS agent_zone_enabled, retirement.retired_at FROM minerops_site_agents AS agent JOIN sites AS site ON site.id = agent.site_id JOIN agent_zones AS zone ON zone.id = agent.agent_zone_id AND zone.site_id = agent.site_id LEFT JOIN agent_legacy_auth_retirements AS retirement ON retirement.agent_id = agent.agent_id WHERE agent.agent_id = %s FOR UPDATE OF agent, zone """, (target_agent_id,), ) agent = cur.fetchone() if not agent: raise AgentEnrollmentNotFound( "Existing Agent with durable Site and Zone routing was not found" ) if agent["agent_zone_enabled"] is not True: raise AgentEnrollmentConflict("Agent Zone is disabled") if any( agent[field] is not None for field in ( "agent_secret_hash", "credential_created_at", "credential_revoked_at", "legacy_auth_disabled_at", "retired_at", ) ): raise AgentEnrollmentConflict( "Existing Agent is not eligible for legacy credential transition" ) 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["agent_zone_id"]), ) replaced_ids = [str(row["id"]) for row in cur.fetchall()] display_name = ( str(agent["display_name"] or agent["agent_id"]).strip()[:120] ) cur.execute( """ INSERT INTO agent_enrollments ( token_hash, site_id, agent_zone_id, display_name, created_by_user_id, created_at, expires_at, target_agent_id ) VALUES (%s, %s, %s, %s, %s, %s, %s + interval '30 minutes', %s) RETURNING id, site_id, agent_zone_id, display_name, created_at, expires_at, target_agent_id """, ( token_hash, agent["site_id"], agent["agent_zone_id"], display_name, actor_user_id, operation_now, operation_now, agent["agent_id"], ), ) transition = dict(cur.fetchone()) transition["site_name"] = agent["site_name"] transition["agent_zone_name"] = agent["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.credential_transition.created', 'agent_enrollment', %s, %s, %s ) """, ( actor_user_id, transition["id"], Jsonb( { "target_agent_id": agent["agent_id"], "site_id": str(agent["site_id"]), "agent_zone_id": str(agent["agent_zone_id"]), "expires_at": transition["expires_at"].isoformat(), "replaced_enrollment_ids": replaced_ids, } ), operation_now, ), ) connection.commit() return transition except UniqueViolation as exc: connection.rollback() constraint = exc.diag.constraint_name if constraint == TOKEN_HASH_CONSTRAINT: raise AgentEnrollmentCredentialCollision( "Generated transition token collided" ) from exc if constraint in {OPEN_ZONE_CONSTRAINT, OPEN_TARGET_CONSTRAINT}: raise AgentEnrollmentConflict( "An active credential transition already exists for this Agent" ) 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.target_agent_id, enrollment.enrolled_agent_id, CASE WHEN enrollment.target_agent_id IS NULL THEN 'new_agent' ELSE 'existing_agent_credential_transition' END AS purpose, 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, target_agent_id 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 or enrollment["target_agent_id"] 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" ) # Retire the identity before storing its credential. Both # writes use this transaction, so a failed issuance cannot # leave a new misleading tombstone. ON CONFLICT preserves an # earlier retirement if this identity has credential history. cur.execute( """ INSERT INTO agent_legacy_auth_retirements ( agent_id, retired_at ) VALUES (%s, %s) ON CONFLICT (agent_id) DO NOTHING """, (agent_id, operation_now), ) cur.execute( """ INSERT INTO minerops_site_agents ( agent_id, site_id, agent_zone_id, site_name, display_name, agent_secret_hash, credential_created_at, legacy_auth_disabled_at ) VALUES ( %s, %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, 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() def redeem_credential_transition( self, *, token_hash: 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, target_agent_id, enrolled_agent_id FROM agent_enrollments WHERE token_hash = %s FOR UPDATE """, (token_hash,), ) enrollment = cur.fetchone() if ( not enrollment or enrollment["target_agent_id"] is None or enrollment["revoked_at"] is not None ): raise AgentEnrollmentTokenInvalid( "Invalid or expired credential transition token" ) cur.execute( """ SELECT agent.agent_id, agent.site_id, agent.agent_zone_id, agent.display_name, agent.agent_secret_hash, agent.credential_created_at, agent.credential_revoked_at, agent.legacy_auth_disabled_at, site.name AS site_name, zone.name AS agent_zone_name, zone.is_enabled AS agent_zone_enabled, retirement.retired_at FROM minerops_site_agents AS agent JOIN sites AS site ON site.id = agent.site_id JOIN agent_zones AS zone ON zone.id = agent.agent_zone_id AND zone.site_id = agent.site_id LEFT JOIN agent_legacy_auth_retirements AS retirement ON retirement.agent_id = agent.agent_id WHERE agent.agent_id = %s FOR UPDATE OF agent, zone """, (enrollment["target_agent_id"],), ) agent = cur.fetchone() if not agent: raise AgentEnrollmentConflict( "Credential transition target no longer exists with durable routing" ) if ( agent["site_id"] != enrollment["site_id"] or agent["agent_zone_id"] != enrollment["agent_zone_id"] ): raise AgentEnrollmentConflict( "Existing Agent routing changed after transition authorization" ) if agent["agent_zone_enabled"] is not True: raise AgentEnrollmentConflict("Agent Zone is disabled") def transition_result() -> dict: 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"],), ) return { "agent_id": agent["agent_id"], "site_id": agent["site_id"], "site_name": agent["site_name"], "agent_zone_id": agent["agent_zone_id"], "agent_zone_name": agent["agent_zone_name"], "display_name": str( agent["display_name"] or enrollment["display_name"] ), "ip_ranges": [row["network"] for row in cur.fetchall()], } # A retry after an ambiguous network response is successful only # for the exact locally generated credential that won the first # transaction. It never rotates or reveals a credential. if enrollment["used_at"] is not None: stored_hash = str(agent["agent_secret_hash"] or "") if ( enrollment["enrolled_agent_id"] == enrollment["target_agent_id"] and stored_hash and hmac.compare_digest(stored_hash, agent_secret_hash) and agent["credential_created_at"] is not None and agent["credential_revoked_at"] is None and agent["legacy_auth_disabled_at"] is not None and agent["retired_at"] is not None ): result = transition_result() connection.commit() return result raise AgentEnrollmentTokenInvalid( "Credential transition token has already been consumed" ) 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 credential transition token" ) if any( agent[field] is not None for field in ( "agent_secret_hash", "credential_created_at", "credential_revoked_at", "legacy_auth_disabled_at", "retired_at", ) ): raise AgentEnrollmentConflict( "Existing Agent is no longer eligible for credential transition" ) # The durable retirement, Agent credential, enrollment terminal # state, and audit record commit together. Any later failure # rolls all of them back. cur.execute( """ INSERT INTO agent_legacy_auth_retirements ( agent_id, retired_at ) VALUES (%s, %s) """, (agent["agent_id"], operation_now), ) cur.execute( """ UPDATE minerops_site_agents SET agent_secret_hash = %s, credential_created_at = %s, legacy_auth_disabled_at = %s WHERE agent_id = %s RETURNING agent_id """, ( agent_secret_hash, operation_now, operation_now, agent["agent_id"], ), ) if not cur.fetchone(): raise AgentEnrollmentConflict( "Credential transition target disappeared" ) cur.execute( """ UPDATE agent_enrollments SET used_at = %s, enrolled_agent_id = target_agent_id WHERE id = %s AND used_at IS NULL AND revoked_at IS NULL RETURNING used_at, enrolled_agent_id """, (operation_now, enrollment["id"]), ) completed = cur.fetchone() if not completed: raise AgentEnrollmentConflict( "Credential transition state changed concurrently" ) cur.execute( """ INSERT INTO audit_logs ( actor_user_id, action, entity_type, entity_id, new_value_json, created_at ) VALUES ( NULL, 'agent.credential_transition.redeemed', 'agent_enrollment', %s, %s, %s ) """, ( enrollment["id"], Jsonb( { "agent_id": agent["agent_id"], "site_id": str(agent["site_id"]), "agent_zone_id": str(agent["agent_zone_id"]), "used_at": completed["used_at"].isoformat(), } ), operation_now, ), ) result = transition_result() connection.commit() return result except UniqueViolation as exc: connection.rollback() if exc.diag.constraint_name == AGENT_SECRET_HASH_CONSTRAINT: raise AgentEnrollmentCredentialCollision( "Agent credential is already in use" ) from exc raise except Exception: connection.rollback() raise finally: connection.close()